From c70a227ee4417deabd5799888e6be64d7d1f447b Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 22 Feb 2024 10:40:46 -0300 Subject: [PATCH] Refactor ConversationChainComponent build method --- .../components/chains/ConversationChain.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/backend/langflow/components/chains/ConversationChain.py b/src/backend/langflow/components/chains/ConversationChain.py index f3a302fe8..43f71e67b 100644 --- a/src/backend/langflow/components/chains/ConversationChain.py +++ b/src/backend/langflow/components/chains/ConversationChain.py @@ -23,19 +23,25 @@ class ConversationChainComponent(CustomComponent): def build( self, + inputs: str, llm: BaseLanguageModel, memory: Optional[BaseMemory] = None, - inputs: dict = {}, ) -> Union[Chain, Callable, Text]: if memory is None: chain = ConversationChain(llm=llm) - chain = ConversationChain(llm=llm, memory=memory) + else: + chain = ConversationChain(llm=llm, memory=memory) result = chain.invoke(inputs) # result is an AIMessage which is a subclass of BaseMessage # We need to check if it is a string or a BaseMessage if hasattr(result, "content") and isinstance(result.content, str): - return result.content + self.status = "is message" + result = result.content elif isinstance(result, str): - return result - - return str(result) + self.status = "is_string" + result = result + else: + # is dict + result = result.get("response") + self.status = result + return result