Refactor LLMChainComponent to return the result of invoking the runnable
This commit is contained in:
parent
49df8e993d
commit
ad68e3e41f
1 changed files with 15 additions and 4 deletions
|
|
@ -1,8 +1,14 @@
|
||||||
from typing import Callable, Optional, Union
|
from typing import Optional
|
||||||
|
|
||||||
from langchain.chains import LLMChain
|
from langchain.chains import LLMChain
|
||||||
|
|
||||||
from langflow import CustomComponent
|
from langflow import CustomComponent
|
||||||
from langflow.field_typing import BaseLanguageModel, BaseMemory, BasePromptTemplate, Chain, Text
|
from langflow.field_typing import (
|
||||||
|
BaseLanguageModel,
|
||||||
|
BaseMemory,
|
||||||
|
BasePromptTemplate,
|
||||||
|
Text,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class LLMChainComponent(CustomComponent):
|
class LLMChainComponent(CustomComponent):
|
||||||
|
|
@ -22,5 +28,10 @@ class LLMChainComponent(CustomComponent):
|
||||||
prompt: BasePromptTemplate,
|
prompt: BasePromptTemplate,
|
||||||
llm: BaseLanguageModel,
|
llm: BaseLanguageModel,
|
||||||
memory: Optional[BaseMemory] = None,
|
memory: Optional[BaseMemory] = None,
|
||||||
) -> Union[Chain, Callable, Text]:
|
) -> Text:
|
||||||
return LLMChain(prompt=prompt, llm=llm, memory=memory)
|
runnable = LLMChain(prompt=prompt, llm=llm, memory=memory)
|
||||||
|
result_dict = runnable.invoke({})
|
||||||
|
output_key = runnable.output_key
|
||||||
|
result = result_dict[output_key]
|
||||||
|
self.status = result
|
||||||
|
return result
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue