Adding custom tools to SQL Agent (#10198)

Changes in:
- `create_sql_agent` function so that user can easily add custom tools
as complement for the toolkit.
- updating **sql use case** notebook to showcase 2 examples of extra
tools.

Motivation for these changes is having the possibility of including
domain expert knowledge to the agent, which improves accuracy and
reduces time/tokens.

---------

Co-authored-by: Manuel Soria <manuel.soria@greyscaleai.com>
Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
Manuel Soria
2023-09-04 19:28:28 -03:00
committed by GitHub
parent 5dbae94e04
commit dde1992fdd
2 changed files with 394 additions and 3 deletions

View File

@@ -1,5 +1,5 @@
"""SQL agent."""
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Optional, Sequence
from langchain.agents.agent import AgentExecutor, BaseSingleActionAgent
from langchain.agents.agent_toolkits.sql.prompt import (
@@ -21,6 +21,7 @@ from langchain.prompts.chat import (
)
from langchain.schema.language_model import BaseLanguageModel
from langchain.schema.messages import AIMessage, SystemMessage
from langchain.tools import BaseTool
def create_sql_agent(
@@ -38,10 +39,11 @@ def create_sql_agent(
early_stopping_method: str = "force",
verbose: bool = False,
agent_executor_kwargs: Optional[Dict[str, Any]] = None,
extra_tools: Sequence[BaseTool] = (),
**kwargs: Dict[str, Any],
) -> AgentExecutor:
"""Construct an SQL agent from an LLM and tools."""
tools = toolkit.get_tools()
tools = toolkit.get_tools() + list(extra_tools)
prefix = prefix.format(dialect=toolkit.dialect, top_k=top_k)
agent: BaseSingleActionAgent