Update SQLDatabaseChainComponent return type
This commit is contained in:
parent
6fb0800d52
commit
8c61654bda
1 changed files with 3 additions and 2 deletions
|
|
@ -7,6 +7,7 @@ from langflow.field_typing import (
|
||||||
Chain
|
Chain
|
||||||
)
|
)
|
||||||
from langchain_community.utilities.sql_database import SQLDatabase
|
from langchain_community.utilities.sql_database import SQLDatabase
|
||||||
|
from langchain_experimental.sql.base import SQLDatabaseChain
|
||||||
|
|
||||||
class SQLDatabaseChainComponent(CustomComponent):
|
class SQLDatabaseChainComponent(CustomComponent):
|
||||||
display_name = "SQLDatabaseChain"
|
display_name = "SQLDatabaseChain"
|
||||||
|
|
@ -24,5 +25,5 @@ class SQLDatabaseChainComponent(CustomComponent):
|
||||||
db: SQLDatabase,
|
db: SQLDatabase,
|
||||||
llm: BaseLanguageModel,
|
llm: BaseLanguageModel,
|
||||||
prompt: BasePromptTemplate,
|
prompt: BasePromptTemplate,
|
||||||
) -> Union[Chain, Callable]:
|
) -> Union[Chain, Callable,SQLDatabaseChain]:
|
||||||
return Chain(db=db, llm=llm, prompt=prompt)
|
return SQLDatabaseChain.from_llm(llm=llm, db=db, prompt=prompt)
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue