From 8c61654bdafbee95974227d8bd2bdd34053f9cf9 Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Tue, 16 Jan 2024 19:17:52 -0300 Subject: [PATCH] Update SQLDatabaseChainComponent return type --- src/backend/langflow/components/chains/SQLDatabaseChain.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/components/chains/SQLDatabaseChain.py b/src/backend/langflow/components/chains/SQLDatabaseChain.py index a86555b35..80be78e72 100644 --- a/src/backend/langflow/components/chains/SQLDatabaseChain.py +++ b/src/backend/langflow/components/chains/SQLDatabaseChain.py @@ -7,6 +7,7 @@ from langflow.field_typing import ( Chain ) from langchain_community.utilities.sql_database import SQLDatabase +from langchain_experimental.sql.base import SQLDatabaseChain class SQLDatabaseChainComponent(CustomComponent): display_name = "SQLDatabaseChain" @@ -24,5 +25,5 @@ class SQLDatabaseChainComponent(CustomComponent): db: SQLDatabase, llm: BaseLanguageModel, prompt: BasePromptTemplate, - ) -> Union[Chain, Callable]: - return Chain(db=db, llm=llm, prompt=prompt) + ) -> Union[Chain, Callable,SQLDatabaseChain]: + return SQLDatabaseChain.from_llm(llm=llm, db=db, prompt=prompt)