324 lines
12 KiB
Python
324 lines
12 KiB
Python
|
"""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,
|
||
|
|
)
|
||
|
|
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
|
||
|
|
from langchain_community.tools.sql_database.tool import (
|
||
|
|
InfoSQLDatabaseTool,
|
||
|
|
ListSQLDatabaseTool,
|
||
|
|
)
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from langchain.agents.agent import AgentExecutor
|
||
|
|
from langchain.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.
|
||
|
|
|
||
|
|
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.agents import (
|
||
|
|
create_openai_functions_agent,
|
||
|
|
create_openai_tools_agent,
|
||
|
|
create_react_agent,
|
||
|
|
create_tool_calling_agent,
|
||
|
|
)
|
||
|
|
from langchain.agents.agent import (
|
||
|
|
AgentExecutor,
|
||
|
|
RunnableAgent,
|
||
|
|
RunnableMultiActionAgent,
|
||
|
|
)
|
||
|
|
from langchain.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 prompt is None:
|
||
|
|
prefix = prefix or SQL_PREFIX
|
||
|
|
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.agents.mrkl import prompt as react_prompt
|
||
|
|
|
||
|
|
format_instructions = (
|
||
|
|
format_instructions or react_prompt.FORMAT_INSTRUCTIONS
|
||
|
|
)
|
||
|
|
template = "\n\n".join(
|
||
|
|
[
|
||
|
|
react_prompt.PREFIX,
|
||
|
|
"{tools}",
|
||
|
|
format_instructions,
|
||
|
|
react_prompt.SUFFIX,
|
||
|
|
]
|
||
|
|
)
|
||
|
|
prompt = PromptTemplate.from_template(template)
|
||
|
|
agent = RunnableAgent(
|
||
|
|
runnable=create_react_agent(llm, tools, prompt, output_parser=ReActSingleInputOutputParserWithOutMarkDown()),
|
||
|
|
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
|
||
|
|
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
|
||
|
|
else:
|
||
|
|
runnable = create_tool_calling_agent(llm, tools, prompt) # type: ignore
|
||
|
|
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 {}),
|
||
|
|
)
|
||
|
|
|
||
|
|
import re
|
||
|
|
|
||
|
|
from langchain_core.agents import AgentAction, AgentFinish
|
||
|
|
from langchain_core.exceptions import OutputParserException
|
||
|
|
|
||
|
|
from langchain.agents.agent import AgentOutputParser
|
||
|
|
from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
|
||
|
|
|
||
|
|
FINAL_ANSWER_ACTION = "Final Answer:"
|
||
|
|
MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE = (
|
||
|
|
"Invalid Format: Missing 'Action:' after 'Thought:"
|
||
|
|
)
|
||
|
|
MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE = (
|
||
|
|
"Invalid Format: Missing 'Action Input:' after 'Action:'"
|
||
|
|
)
|
||
|
|
FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE = (
|
||
|
|
"Parsing LLM output produced both a final answer and a parse-able action:"
|
||
|
|
)
|
||
|
|
class ReActSingleInputOutputParserWithOutMarkDown(AgentOutputParser):
|
||
|
|
def get_format_instructions(self) -> str:
|
||
|
|
return FORMAT_INSTRUCTIONS
|
||
|
|
|
||
|
|
def parse(self, text: str) -> Union[AgentAction, AgentFinish]:
|
||
|
|
includes_answer = FINAL_ANSWER_ACTION in text
|
||
|
|
regex = (
|
||
|
|
r"Action\s*\d*\s*:[\s]*(.*?)[\s]*Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)"
|
||
|
|
)
|
||
|
|
print("text: {}".format(text))
|
||
|
|
action_match = re.search(regex, text, re.DOTALL)
|
||
|
|
if action_match:
|
||
|
|
if includes_answer:
|
||
|
|
raise OutputParserException(
|
||
|
|
f"{FINAL_ANSWER_AND_PARSABLE_ACTION_ERROR_MESSAGE}: {text}"
|
||
|
|
)
|
||
|
|
action = action_match.group(1).strip()
|
||
|
|
action_input = action_match.group(2)
|
||
|
|
tool_input = action_input.strip(" ")
|
||
|
|
tool_input = tool_input.strip('"')
|
||
|
|
|
||
|
|
# Remove markdown code block markers if present
|
||
|
|
tool_input = re.sub(r'(^```\s*sql\s*|^\s*```$)', '', tool_input, flags=re.MULTILINE).strip()
|
||
|
|
tool_input = re.sub(r'(^`\s*|`\s*$)', '', tool_input).strip()
|
||
|
|
|
||
|
|
return AgentAction(action, tool_input, text)
|
||
|
|
|
||
|
|
elif includes_answer:
|
||
|
|
return AgentFinish(
|
||
|
|
{"output": text.split(FINAL_ANSWER_ACTION)[-1].strip()}, text
|
||
|
|
)
|
||
|
|
|
||
|
|
if not re.search(r"Action\s*\d*\s*:[\s]*(.*?)", text, re.DOTALL):
|
||
|
|
raise OutputParserException(
|
||
|
|
f"Could not parse LLM output: `{text}`",
|
||
|
|
observation=MISSING_ACTION_AFTER_THOUGHT_ERROR_MESSAGE,
|
||
|
|
llm_output=text,
|
||
|
|
send_to_llm=True,
|
||
|
|
)
|
||
|
|
elif not re.search(
|
||
|
|
r"[\s]*Action\s*\d*\s*Input\s*\d*\s*:[\s]*(.*)", text, re.DOTALL
|
||
|
|
):
|
||
|
|
raise OutputParserException(
|
||
|
|
f"Could not parse LLM output: `{text}`",
|
||
|
|
observation=MISSING_ACTION_INPUT_AFTER_ACTION_ERROR_MESSAGE,
|
||
|
|
llm_output=text,
|
||
|
|
send_to_llm=True,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
raise OutputParserException(f"Could not parse LLM output: `{text}`")
|
||
|
|
|
||
|
|
@property
|
||
|
|
def _type(self) -> str:
|
||
|
|
return "react-single-input-without-markdown"
|
||
|
|
|
||
|
|
def remove_markdown_code_block(text: str):
|
||
|
|
pattern = re.compile(r"```sql\s*(.*?)\s*```", re.DOTALL)
|
||
|
|
|
||
|
|
# 查找匹配的内容
|
||
|
|
match = pattern.search(text)
|
||
|
|
if match:
|
||
|
|
# 提取 SQL 语句
|
||
|
|
sql_query = match.group(1)
|
||
|
|
return sql_query.strip()
|
||
|
|
else:
|
||
|
|
return "No SQL code block found."
|