From 4f5f8258f3c5cbfa6032efdac1eef11d87976f26 Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 16 Jan 2024 18:13:16 -0300 Subject: [PATCH] Add typing imports and update return type in LLMMathChainComponent --- src/backend/langflow/components/chains/LLMMathChain.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/components/chains/LLMMathChain.py b/src/backend/langflow/components/chains/LLMMathChain.py index 406143418..1a6138a6e 100644 --- a/src/backend/langflow/components/chains/LLMMathChain.py +++ b/src/backend/langflow/components/chains/LLMMathChain.py @@ -1,10 +1,11 @@ - +from typing import Callable, Optional, Union from langflow import CustomComponent from langchain.chains import LLMChain from typing import Optional from langflow.field_typing import ( BaseLanguageModel, BaseMemory, + Chain ) class LLMMathChainComponent(CustomComponent): @@ -28,5 +29,5 @@ class LLMMathChainComponent(CustomComponent): input_key: str, output_key: str, memory: Optional[BaseMemory] = None, - ) -> LLMChain: + ) -> Union[Chain, Callable]: return LLMChain(llm=llm, prompt=llm_chain, input_key=input_key, output_key=output_key, memory=memory)