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

Update LLMMathChainComponent to use LLMMathChain

Browse source
This commit is contained in:
anovazzi1 2024-01-16 18:27:42 -03:00
commit 9bd1c755c4
1 changed files with 11 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1,12 +1,11 @@
from typing import Callable, Optional, Union
from langflow import CustomComponent from langflow import CustomComponent
from langchain.chains import LLMChain from langchain.chains import LLMChain,LLMMathChain
from typing import Optional from typing import Callable, Optional, Union
from langflow.field_typing import ( from langflow.field_typing import (
BaseLanguageModel, BaseLanguageModel,
BaseMemory, BaseMemory,
Chain, Chain
BasePromptTemplate
) )
class LLMMathChainComponent(CustomComponent): class LLMMathChainComponent(CustomComponent):
@ -17,16 +16,18 @@ class LLMMathChainComponent(CustomComponent):
def build_config(self): def build_config(self):
return { return {
"llm": {"display_name": "LLM"}, "llm": {"display_name": "LLM"},
"prompt": {"display_name": "Prompt"}, "llm_chain": {"display_name": "LLM Chain"},
"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,
prompt: BasePromptTemplate, llm_chain: LLMChain,
output_key: str="text", input_key: Optional[str]="question",
output_key: Optional[str]="answer",
memory: Optional[BaseMemory] = None, memory: Optional[BaseMemory] = None,
) -> Union[Chain, Callable]: ) -> Union[LLMMathChain, Callable,Chain]:
return LLMChain(llm=llm, prompt=prompt, output_key=output_key, memory=memory) return LLMMathChain(llm=llm, llm_chain=llm_chain, input_key=input_key, output_key=output_key, memory=memory)