From ad68e3e41f7decff24196fe7d6cf007fe1b34c15 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 26 Feb 2024 09:50:39 -0300 Subject: [PATCH] Refactor LLMChainComponent to return the result of invoking the runnable --- .../langflow/components/chains/LLMChain.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/components/chains/LLMChain.py b/src/backend/langflow/components/chains/LLMChain.py index 8596dd0c9..5d660cd4b 100644 --- a/src/backend/langflow/components/chains/LLMChain.py +++ b/src/backend/langflow/components/chains/LLMChain.py @@ -1,8 +1,14 @@ -from typing import Callable, Optional, Union +from typing import Optional from langchain.chains import LLMChain + 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): @@ -22,5 +28,10 @@ class LLMChainComponent(CustomComponent): prompt: BasePromptTemplate, llm: BaseLanguageModel, memory: Optional[BaseMemory] = None, - ) -> Union[Chain, Callable, Text]: - return LLMChain(prompt=prompt, llm=llm, memory=memory) + ) -> Text: + 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