diff --git a/src/backend/langflow/components/agents/SQLAgent.py b/src/backend/langflow/components/agents/SQLAgent.py index 8b5781dd9..7b1865637 100644 --- a/src/backend/langflow/components/agents/SQLAgent.py +++ b/src/backend/langflow/components/agents/SQLAgent.py @@ -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)