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

Add SQLAgentComponent build method to create SQL agent with SQLDatabaseToolkit

Browse source
This commit is contained in:
anovazzi1 2024-01-15 20:47:06 -03:00
commit 6fe02442dd
1 changed files with 9 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

12
src/backend/langflow/components/agents/SQLAgent.py
View file

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