Support extra tools argument for pandas agent toolkit (#11040)

**Description** 

We support adding new tools in some toolkits already like the [SQLAgent
toolkit](https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/agents/agent_toolkits/sql/base.py#L27).

Related
[SO](https://stackoverflow.com/questions/76583163/are-langchain-toolkits-able-to-be-modified-can-we-add-tools-to-a-pandas-datafra)
thread
This replicates the same functionality here, so users can add custom
bespoke tools.
pull/10843/head^2
tanujtiwari-at 12 months ago committed by GitHub
parent c4471d1877
commit a79f595543
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -1,5 +1,5 @@
"""Agent for working with pandas objects.""" """Agent for working with pandas objects."""
from typing import Any, Dict, List, Optional, Tuple from typing import Any, Dict, List, Optional, Sequence, Tuple
from langchain.agents.agent import AgentExecutor, BaseSingleActionAgent from langchain.agents.agent import AgentExecutor, BaseSingleActionAgent
from langchain.agents.agent_toolkits.pandas.prompt import ( from langchain.agents.agent_toolkits.pandas.prompt import (
@ -21,6 +21,7 @@ from langchain.chains.llm import LLMChain
from langchain.schema import BasePromptTemplate from langchain.schema import BasePromptTemplate
from langchain.schema.language_model import BaseLanguageModel from langchain.schema.language_model import BaseLanguageModel
from langchain.schema.messages import SystemMessage from langchain.schema.messages import SystemMessage
from langchain.tools import BaseTool
from langchain.tools.python.tool import PythonAstREPLTool from langchain.tools.python.tool import PythonAstREPLTool
@ -280,12 +281,13 @@ def create_pandas_dataframe_agent(
agent_executor_kwargs: Optional[Dict[str, Any]] = None, agent_executor_kwargs: Optional[Dict[str, Any]] = None,
include_df_in_prompt: Optional[bool] = True, include_df_in_prompt: Optional[bool] = True,
number_of_head_rows: int = 5, number_of_head_rows: int = 5,
extra_tools: Sequence[BaseTool] = (),
**kwargs: Dict[str, Any], **kwargs: Dict[str, Any],
) -> AgentExecutor: ) -> AgentExecutor:
"""Construct a pandas agent from an LLM and dataframe.""" """Construct a pandas agent from an LLM and dataframe."""
agent: BaseSingleActionAgent agent: BaseSingleActionAgent
if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION: if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION:
prompt, tools = _get_prompt_and_tools( prompt, base_tools = _get_prompt_and_tools(
df, df,
prefix=prefix, prefix=prefix,
suffix=suffix, suffix=suffix,
@ -293,6 +295,7 @@ def create_pandas_dataframe_agent(
include_df_in_prompt=include_df_in_prompt, include_df_in_prompt=include_df_in_prompt,
number_of_head_rows=number_of_head_rows, number_of_head_rows=number_of_head_rows,
) )
tools = base_tools + list(extra_tools)
llm_chain = LLMChain( llm_chain = LLMChain(
llm=llm, llm=llm,
prompt=prompt, prompt=prompt,
@ -306,7 +309,7 @@ def create_pandas_dataframe_agent(
**kwargs, **kwargs,
) )
elif agent_type == AgentType.OPENAI_FUNCTIONS: elif agent_type == AgentType.OPENAI_FUNCTIONS:
_prompt, tools = _get_functions_prompt_and_tools( _prompt, base_tools = _get_functions_prompt_and_tools(
df, df,
prefix=prefix, prefix=prefix,
suffix=suffix, suffix=suffix,
@ -314,6 +317,7 @@ def create_pandas_dataframe_agent(
include_df_in_prompt=include_df_in_prompt, include_df_in_prompt=include_df_in_prompt,
number_of_head_rows=number_of_head_rows, number_of_head_rows=number_of_head_rows,
) )
tools = base_tools + list(extra_tools)
agent = OpenAIFunctionsAgent( agent = OpenAIFunctionsAgent(
llm=llm, llm=llm,
prompt=_prompt, prompt=_prompt,

Loading…
Cancel
Save