Add SQLAgentComponent build method to create SQL agent with SQLDatabaseToolkit

This commit is contained in:
anovazzi1 2024-01-15 20:47:06 -03:00
commit 6fe02442dd

View file

@ -3,7 +3,10 @@ from langflow import CustomComponent
from typing import Union, Callable
from langchain.agents import AgentExecutor
from langflow.field_typing import BaseLanguageModel
from langchain_community.agent_toolkits.sql.base import create_sql_agent
from langchain.sql_database import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
class SQLAgentComponent(CustomComponent):
display_name = "SQLAgent"
description = "Construct an SQL agent from an LLM and tools."
@ -12,14 +15,15 @@ class SQLAgentComponent(CustomComponent):
return {
"llm": {"display_name": "LLM"},
"database_uri": {"display_name": "Database URI"},
"verbose": {"display_name": "Verbose", "value": False,"advanced": True},
}
def build(
self,
llm: BaseLanguageModel,
database_uri: str,
verbose: bool = False,
) -> Union[AgentExecutor, Callable]:
# Assuming there is a constructor for SQLAgent that takes these parameters
# Since the actual implementation is not provided, this is a placeholder
# Replace SQLAgent with the actual class name if different
return AgentExecutor(llm=llm, database_uri=database_uri)
db = SQLDatabase.from_uri(database_uri)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
return create_sql_agent(llm=llm, toolkit=toolkit)