"""SQL agent.""" from __future__ import annotations from typing import ( TYPE_CHECKING, Any, Dict, List, Literal, Optional, Sequence, Union, cast, ) from langchain_core.messages import AIMessage, SystemMessage from langchain_core.prompts import BasePromptTemplate, PromptTemplate from langchain_core.prompts.chat import ( ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder, ) from langchain_community.agent_toolkits.sql.prompt import ( SQL_FUNCTIONS_SUFFIX, SQL_PREFIX, SQL_SUFFIX, ) from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit from langchain_community.tools.sql_database.tool import ( InfoSQLDatabaseTool, ListSQLDatabaseTool, ) if TYPE_CHECKING: from langchain_classic.agents.agent import AgentExecutor from langchain_classic.agents.agent_types import AgentType from langchain_core.callbacks import BaseCallbackManager from langchain_core.language_models import BaseLanguageModel from langchain_core.tools import BaseTool from langchain_community.utilities.sql_database import SQLDatabase def create_sql_agent( llm: BaseLanguageModel, toolkit: Optional[SQLDatabaseToolkit] = None, agent_type: Optional[ Union[AgentType, Literal["openai-tools", "tool-calling"]] ] = None, callback_manager: Optional[BaseCallbackManager] = None, prefix: Optional[str] = None, suffix: Optional[str] = None, format_instructions: Optional[str] = None, input_variables: Optional[List[str]] = None, top_k: int = 10, max_iterations: Optional[int] = 15, max_execution_time: Optional[float] = None, early_stopping_method: str = "force", verbose: bool = False, agent_executor_kwargs: Optional[Dict[str, Any]] = None, extra_tools: Sequence[BaseTool] = (), *, db: Optional[SQLDatabase] = None, prompt: Optional[BasePromptTemplate] = None, **kwargs: Any, ) -> AgentExecutor: """Construct a SQL agent from an LLM and toolkit or database. !!! warning This agent can execute arbitrary SQL against your database. By default, the agent is allowed to generate SQL strings and run them via the database connection. This is powerful, but it also means the agent can generate expensive or dangerous queries (e.g., long-running queries, large scans/joins, locking queries, or unintended writes depending on your database permissions). ``create_sql_agent`` returns a ``langchain_classic`` ``AgentExecutor``. ``AgentExecutor`` is an agent abstraction that has long been considered legacy and is not actively supported as the recommended foundation for new production applications. For production-grade agent development, prefer building with Deep Agents: https://github.com/langchain-ai/deepagents If you use this in production, coordinate with your security/DB teams and apply server-side controls: - Use least-privilege database roles (ideally read-only, schema-limited). - Enforce statement timeouts / max execution time and other resource limits at the role or session level. - Apply query guardrails (e.g., restrict accessible schemas/tables, limit concurrency, and monitor/alert on slow queries). Client-side timeouts do not always guarantee that a running statement is cancelled on the database server. Args: llm: Language model to use for the agent. If agent_type is "tool-calling" then llm is expected to support tool calling. toolkit: SQLDatabaseToolkit for the agent to use. Must provide exactly one of 'toolkit' or 'db'. Specify 'toolkit' if you want to use a different model for the agent and the toolkit. agent_type: One of "tool-calling", "openai-tools", "openai-functions", or "zero-shot-react-description". Defaults to "zero-shot-react-description". "tool-calling" is recommended over the legacy "openai-tools" and "openai-functions" types. callback_manager: DEPRECATED. Pass "callbacks" key into 'agent_executor_kwargs' instead to pass constructor callbacks to AgentExecutor. prefix: Prompt prefix string. Must contain variables "top_k" and "dialect". suffix: Prompt suffix string. Default depends on agent type. format_instructions: Formatting instructions to pass to ZeroShotAgent.create_prompt() when 'agent_type' is "zero-shot-react-description". Otherwise ignored. input_variables: DEPRECATED. top_k: Number of rows to query for by default. max_iterations: Passed to AgentExecutor init. max_execution_time: Passed to AgentExecutor init. early_stopping_method: Passed to AgentExecutor init. verbose: AgentExecutor verbosity. agent_executor_kwargs: Arbitrary additional AgentExecutor args. extra_tools: Additional tools to give to agent on top of the ones that come with SQLDatabaseToolkit. db: SQLDatabase from which to create a SQLDatabaseToolkit. Toolkit is created using 'db' and 'llm'. Must provide exactly one of 'db' or 'toolkit'. prompt: Complete agent prompt. prompt and {prefix, suffix, format_instructions, input_variables} are mutually exclusive. **kwargs: Arbitrary additional Agent args. Returns: An AgentExecutor with the specified agent_type agent. Example: .. code-block:: python from langchain_openai import ChatOpenAI from langchain_community.agent_toolkits import create_sql_agent from langchain_community.utilities import SQLDatabase db = SQLDatabase.from_uri("sqlite:///Chinook.db") llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0) agent_executor = create_sql_agent(llm, db=db, agent_type="tool-calling", verbose=True) """ # noqa: E501 from langchain_classic.agents import ( create_openai_functions_agent, create_openai_tools_agent, create_react_agent, create_tool_calling_agent, ) from langchain_classic.agents.agent import ( AgentExecutor, RunnableAgent, RunnableMultiActionAgent, ) from langchain_classic.agents.agent_types import AgentType if toolkit is None and db is None: raise ValueError( "Must provide exactly one of 'toolkit' or 'db'. Received neither." ) if toolkit and db: raise ValueError( "Must provide exactly one of 'toolkit' or 'db'. Received both." ) toolkit = toolkit or SQLDatabaseToolkit(llm=llm, db=db) # type: ignore[arg-type] agent_type = agent_type or AgentType.ZERO_SHOT_REACT_DESCRIPTION tools = toolkit.get_tools() + list(extra_tools) if prefix is None: prefix = SQL_PREFIX if prompt is None: prefix = prefix.format(dialect=toolkit.dialect, top_k=top_k) else: if "top_k" in prompt.input_variables: prompt = prompt.partial(top_k=str(top_k)) if "dialect" in prompt.input_variables: prompt = prompt.partial(dialect=toolkit.dialect) if any(key in prompt.input_variables for key in ["table_info", "table_names"]): db_context = toolkit.get_context() if "table_info" in prompt.input_variables: prompt = prompt.partial(table_info=db_context["table_info"]) tools = [ tool for tool in tools if not isinstance(tool, InfoSQLDatabaseTool) ] if "table_names" in prompt.input_variables: prompt = prompt.partial(table_names=db_context["table_names"]) tools = [ tool for tool in tools if not isinstance(tool, ListSQLDatabaseTool) ] if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION: if prompt is None: from langchain_classic.agents.mrkl import prompt as react_prompt format_instructions = ( format_instructions or react_prompt.FORMAT_INSTRUCTIONS ) template = "\n\n".join( [ prefix, "{tools}", format_instructions, suffix or SQL_SUFFIX, ] ) prompt = PromptTemplate.from_template(template) agent = RunnableAgent( runnable=create_react_agent(llm, tools, prompt), input_keys_arg=["input"], return_keys_arg=["output"], **kwargs, ) elif agent_type == AgentType.OPENAI_FUNCTIONS: if prompt is None: messages: List = [ SystemMessage(content=cast(str, prefix)), HumanMessagePromptTemplate.from_template("{input}"), AIMessage(content=suffix or SQL_FUNCTIONS_SUFFIX), MessagesPlaceholder(variable_name="agent_scratchpad"), ] prompt = ChatPromptTemplate.from_messages(messages) agent = RunnableAgent( runnable=create_openai_functions_agent(llm, tools, prompt), # type: ignore[arg-type] input_keys_arg=["input"], return_keys_arg=["output"], **kwargs, ) elif agent_type in ("openai-tools", "tool-calling"): if prompt is None: messages = [ SystemMessage(content=cast(str, prefix)), HumanMessagePromptTemplate.from_template("{input}"), AIMessage(content=suffix or SQL_FUNCTIONS_SUFFIX), MessagesPlaceholder(variable_name="agent_scratchpad"), ] prompt = ChatPromptTemplate.from_messages(messages) if agent_type == "openai-tools": runnable = create_openai_tools_agent(llm, tools, prompt) # type: ignore[arg-type] else: runnable = create_tool_calling_agent(llm, tools, prompt) # type: ignore[arg-type] agent = RunnableMultiActionAgent( # type: ignore[assignment] runnable=runnable, input_keys_arg=["input"], return_keys_arg=["output"], **kwargs, ) else: raise ValueError( f"Agent type {agent_type} not supported at the moment. Must be one of " "'tool-calling', 'openai-tools', 'openai-functions', or " "'zero-shot-react-description'." ) return AgentExecutor( name="SQL Agent Executor", agent=agent, tools=tools, callback_manager=callback_manager, verbose=verbose, max_iterations=max_iterations, max_execution_time=max_execution_time, early_stopping_method=early_stopping_method, **(agent_executor_kwargs or {}), )