From b56cca506d8f10d54f227b1246ad032a8e0503d4 Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 16 Jan 2024 18:22:26 -0300 Subject: [PATCH] Update LLMMathChainComponent build method signature --- .../langflow/components/chains/LLMMathChain.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/backend/langflow/components/chains/LLMMathChain.py b/src/backend/langflow/components/chains/LLMMathChain.py index 1a6138a6e..73f0eddc0 100644 --- a/src/backend/langflow/components/chains/LLMMathChain.py +++ b/src/backend/langflow/components/chains/LLMMathChain.py @@ -5,7 +5,8 @@ from typing import Optional from langflow.field_typing import ( BaseLanguageModel, BaseMemory, - Chain + Chain, + BasePromptTemplate ) class LLMMathChainComponent(CustomComponent): @@ -16,18 +17,16 @@ class LLMMathChainComponent(CustomComponent): def build_config(self): return { "llm": {"display_name": "LLM"}, - "llm_chain": {"display_name": "LLM Chain"}, + "prompt": {"display_name": "Prompt"}, "memory": {"display_name": "Memory"}, - "input_key": {"display_name": "Input Key"}, "output_key": {"display_name": "Output Key"}, } def build( self, llm: BaseLanguageModel, - llm_chain: LLMChain, - input_key: str, - output_key: str, + prompt: BasePromptTemplate, + output_key: str="text", memory: Optional[BaseMemory] = None, ) -> Union[Chain, Callable]: - return LLMChain(llm=llm, prompt=llm_chain, input_key=input_key, output_key=output_key, memory=memory) + return LLMChain(llm=llm, prompt=prompt, output_key=output_key, memory=memory)