feat: add SQL agent

src/backend/langflow/config.yaml

@@ -14,6 +14,7 @@ agents:
   - initialize_agent
   - VectorStoreAgent
   - VectorStoreRouterAgent
+  - SQLAgent
 
 prompts:
   - PromptTemplate

src/backend/langflow/custom/customs.py

@@ -10,6 +10,7 @@ CUSTOM_NODES = {
         "initialize_agent": nodes.InitializeAgentNode(),
         "VectorStoreAgent": nodes.VectorStoreAgentNode(),
         "VectorStoreRouterAgent": nodes.VectorStoreRouterAgentNode(),
+        "SQLAgent": nodes.SQLAgentNode(),
     },
 }
 

src/backend/langflow/graph/nodes.py

@@ -33,8 +33,8 @@ class AgentNode(Node):
 
             self._build()
 
-        #! Cannot deepcopy VectorStore
-        if self.node_type in ["VectorStoreAgent", "VectorStoreRouterAgent"]:
+        #! Cannot deepcopy VectorStore, VectorStoreRouter, or SQL agents
+        if self.node_type in ["VectorStoreAgent", "VectorStoreRouterAgent", "SQLAgent"]:
             return self._built_object
         return deepcopy(self._built_object)
 

src/backend/langflow/interface/agents/custom.py

@@ -1,8 +1,14 @@
 from typing import Any, List, Optional
 
 from langchain import LLMChain
-from langchain.agents import AgentExecutor, Tool, ZeroShotAgent, initialize_agent
+from langchain.agents import (
+    AgentExecutor,
+    Tool,
+    ZeroShotAgent,
+    initialize_agent,
+)
 from langchain.agents.agent_toolkits import (
+    SQLDatabaseToolkit,
     VectorStoreInfo,
     VectorStoreRouterToolkit,
     VectorStoreToolkit,
@@ -11,6 +17,7 @@ from langchain.agents.agent_toolkits.json.prompt import JSON_PREFIX, JSON_SUFFIX
 from langchain.agents.agent_toolkits.json.toolkit import JsonToolkit
 from langchain.agents.agent_toolkits.pandas.prompt import PREFIX as PANDAS_PREFIX
 from langchain.agents.agent_toolkits.pandas.prompt import SUFFIX as PANDAS_SUFFIX
+from langchain.agents.agent_toolkits.sql.prompt import SQL_PREFIX, SQL_SUFFIX
 from langchain.agents.agent_toolkits.vectorstore.prompt import (
     PREFIX as VECTORSTORE_PREFIX,
 )
@@ -18,10 +25,13 @@ from langchain.agents.agent_toolkits.vectorstore.prompt import (
     ROUTER_PREFIX as VECTORSTORE_ROUTER_PREFIX,
 )
 from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
+from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS as SQL_FORMAT_INSTRUCTIONS
 from langchain.llms.base import BaseLLM
 from langchain.memory.chat_memory import BaseChatMemory
 from langchain.schema import BaseLanguageModel
+from langchain.sql_database import SQLDatabase
 from langchain.tools.python.tool import PythonAstREPLTool
+from langchain.tools.sql_database.prompt import QUERY_CHECKER
 
 
 class JsonAgent(AgentExecutor):
@@ -146,6 +156,76 @@ class VectorStoreAgent(AgentExecutor):
         return super().run(*args, **kwargs)
 
 
+class SQLAgent(AgentExecutor):
+    """SQL agent"""
+
+    @staticmethod
+    def function_name():
+        return "SQLAgent"
+
+    @classmethod
+    def initialize(cls, *args, **kwargs):
+        return cls.from_toolkit_and_llm(*args, **kwargs)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    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)
+
+        # The right code should be this, but there is a problem with tools = toolkit.get_tools()
+        # related to `OPENAI_API_KEY`
+        # return create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)
+        from langchain.prompts import PromptTemplate
+        from langchain.tools.sql_database.tool import (
+            InfoSQLDatabaseTool,
+            ListSQLDatabaseTool,
+            QueryCheckerTool,
+            QuerySQLDataBaseTool,
+        )
+
+        llmchain = LLMChain(
+            llm=llm,
+            prompt=PromptTemplate(
+                template=QUERY_CHECKER, input_variables=["query", "dialect"]
+            ),
+        )
+
+        tools = [
+            QuerySQLDataBaseTool(db=db),  # type: ignore
+            InfoSQLDatabaseTool(db=db),  # type: ignore
+            ListSQLDatabaseTool(db=db),  # type: ignore
+            QueryCheckerTool(db=db, llm_chain=llmchain),  # type: ignore
+        ]
+
+        prefix = SQL_PREFIX.format(dialect=toolkit.dialect, top_k=10)
+        prompt = ZeroShotAgent.create_prompt(
+            tools=tools,  # type: ignore
+            prefix=prefix,
+            suffix=SQL_SUFFIX,
+            format_instructions=SQL_FORMAT_INSTRUCTIONS,
+        )
+        llm_chain = LLMChain(
+            llm=llm,
+            prompt=prompt,
+        )
+        tool_names = [tool.name for tool in tools]  # type: ignore
+        agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
+        return AgentExecutor.from_agent_and_tools(
+            agent=agent,
+            tools=tools,  # type: ignore
+            verbose=True,
+            max_iterations=15,
+            early_stopping_method="force",
+        )
+
+    def run(self, *args, **kwargs):
+        return super().run(*args, **kwargs)
+
+
 class VectorStoreRouterAgent(AgentExecutor):
     """Vector Store Router Agent"""
 
@@ -218,4 +298,5 @@ CUSTOM_AGENTS = {
     "initialize_agent": InitializeAgent,
     "VectorStoreAgent": VectorStoreAgent,
     "VectorStoreRouterAgent": VectorStoreRouterAgent,
+    "SQLAgent": SQLAgent,
 }

src/backend/langflow/template/nodes.py

@@ -304,6 +304,37 @@ class VectorStoreRouterAgentNode(FrontendNode):
         return super().to_dict()
 
 
+class SQLAgentNode(FrontendNode):
+    name: str = "SQLAgent"
+    template: Template = Template(
+        type_name="sql_agent",
+        fields=[
+            TemplateField(
+                field_type="str",
+                required=True,
+                placeholder="",
+                is_list=False,
+                show=True,
+                multiline=False,
+                value="",
+                name="database_uri",
+            ),
+            TemplateField(
+                field_type="BaseLanguageModel",
+                required=True,
+                show=True,
+                name="llm",
+                display_name="LLM",
+            ),
+        ],
+    )
+    description: str = """Construct an agent from a Vector Store Router."""
+    base_classes: list[str] = ["AgentExecutor"]
+
+    def to_dict(self):
+        return super().to_dict()
+
+
 class PromptFrontendNode(FrontendNode):
     @staticmethod
     def format_field(field: TemplateField, name: Optional[str] = None) -> None: