feat(custom.py): pass llm parameter to SQLDatabaseToolkit constructor

This commit is contained in:
Gabriel Almeida 2023-04-27 12:25:13 -03:00 committed by Gabriel Luiz Freitas Almeida
commit f00420ecbe

View file

@ -174,7 +174,7 @@ class SQLAgent(AgentExecutor):
def from_toolkit_and_llm(cls, llm: BaseLLM, database_uri: str, **kwargs: Any):
"""Construct a sql agent from an LLM and tools."""
db = SQLDatabase.from_uri(database_uri)
toolkit = SQLDatabaseToolkit(db=db)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
# The right code should be this, but there is a problem with tools = toolkit.get_tools()
# related to `OPENAI_API_KEY`