Add SQLAgentComponent build method to create SQL agent with SQLDatabaseToolkit
This commit is contained in:
parent
fbd64033dc
commit
6fe02442dd
1 changed files with 9 additions and 5 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue