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

Update SQLDatabaseChainComponent return type

Browse source
This commit is contained in:
anovazzi1 2024-01-16 19:17:52 -03:00
commit 8c61654bda
1 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/backend/langflow/components/chains/SQLDatabaseChain.py
View file

@ -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)