jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Update LLMMathChainComponent build method signature

Browse source
This commit is contained in:
anovazzi1 2024-01-16 18:22:26 -03:00
commit b56cca506d
1 changed files with 6 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

13
src/backend/langflow/components/chains/LLMMathChain.py
View file

@ -5,7 +5,8 @@ from typing import Optional
from langflow.field_typing import ( from langflow.field_typing import (
BaseLanguageModel, BaseLanguageModel,
BaseMemory, BaseMemory,
Chain Chain,
BasePromptTemplate
) )
class LLMMathChainComponent(CustomComponent): class LLMMathChainComponent(CustomComponent):
@ -16,18 +17,16 @@ class LLMMathChainComponent(CustomComponent):
def build_config(self): def build_config(self):
return { return {
"llm": {"display_name": "LLM"}, "llm": {"display_name": "LLM"},
"llm_chain": {"display_name": "LLM Chain"}, "prompt": {"display_name": "Prompt"},
"memory": {"display_name": "Memory"}, "memory": {"display_name": "Memory"},
"input_key": {"display_name": "Input Key"},
"output_key": {"display_name": "Output Key"}, "output_key": {"display_name": "Output Key"},
} }
def build( def build(
self, self,
llm: BaseLanguageModel, llm: BaseLanguageModel,
llm_chain: LLMChain, prompt: BasePromptTemplate,
input_key: str, output_key: str="text",
output_key: str,
memory: Optional[BaseMemory] = None, memory: Optional[BaseMemory] = None,
) -> Union[Chain, Callable]: ) -> 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)