mirror of
https://github.com/hwchase17/langchain.git
synced 2025-07-13 16:36:06 +00:00
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.
This commit is contained in:
parent
c4471d1877
commit
a79f595543
@ -1,5 +1,5 @@
|
||||
"""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_toolkits.pandas.prompt import (
|
||||
@ -21,6 +21,7 @@ from langchain.chains.llm import LLMChain
|
||||
from langchain.schema import BasePromptTemplate
|
||||
from langchain.schema.language_model import BaseLanguageModel
|
||||
from langchain.schema.messages import SystemMessage
|
||||
from langchain.tools import BaseTool
|
||||
from langchain.tools.python.tool import PythonAstREPLTool
|
||||
|
||||
|
||||
@ -280,12 +281,13 @@ def create_pandas_dataframe_agent(
|
||||
agent_executor_kwargs: Optional[Dict[str, Any]] = None,
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
extra_tools: Sequence[BaseTool] = (),
|
||||
**kwargs: Dict[str, Any],
|
||||
) -> AgentExecutor:
|
||||
"""Construct a pandas agent from an LLM and dataframe."""
|
||||
agent: BaseSingleActionAgent
|
||||
if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION:
|
||||
prompt, tools = _get_prompt_and_tools(
|
||||
prompt, base_tools = _get_prompt_and_tools(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
@ -293,6 +295,7 @@ def create_pandas_dataframe_agent(
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
)
|
||||
tools = base_tools + list(extra_tools)
|
||||
llm_chain = LLMChain(
|
||||
llm=llm,
|
||||
prompt=prompt,
|
||||
@ -306,7 +309,7 @@ def create_pandas_dataframe_agent(
|
||||
**kwargs,
|
||||
)
|
||||
elif agent_type == AgentType.OPENAI_FUNCTIONS:
|
||||
_prompt, tools = _get_functions_prompt_and_tools(
|
||||
_prompt, base_tools = _get_functions_prompt_and_tools(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
@ -314,6 +317,7 @@ def create_pandas_dataframe_agent(
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
)
|
||||
tools = base_tools + list(extra_tools)
|
||||
agent = OpenAIFunctionsAgent(
|
||||
llm=llm,
|
||||
prompt=_prompt,
|
||||
|
Loading…
Reference in New Issue
Block a user