mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-17 23:41:46 +00:00
docs: add csv use case (#16756)
This commit is contained in:
Bagatur
GitHub
@@ -1,26 +1,55 @@
|
||||
from io import IOBase
|
||||
from typing import Any, List, Optional, Union
|
||||
from __future__ import annotations
|
||||
|
||||
from langchain.agents.agent import AgentExecutor
|
||||
from langchain_core.language_models import BaseLanguageModel
|
||||
from io import IOBase
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Union
|
||||
|
||||
from langchain_experimental.agents.agent_toolkits.pandas.base import (
|
||||
create_pandas_dataframe_agent,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain.agents.agent import AgentExecutor
|
||||
from langchain_core.language_models import LanguageModelLike
|
||||
|
||||
|
||||
def create_csv_agent(
|
||||
llm: BaseLanguageModel,
|
||||
llm: LanguageModelLike,
|
||||
path: Union[str, IOBase, List[Union[str, IOBase]]],
|
||||
pandas_kwargs: Optional[dict] = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentExecutor:
|
||||
"""Create csv agent by loading to a dataframe and using pandas agent."""
|
||||
"""Create pandas dataframe agent by loading csv to a dataframe.
|
||||
|
||||
Args:
|
||||
llm: Language model to use for the agent.
|
||||
path: A string path, file-like object or a list of string paths/file-like
|
||||
objects that can be read in as pandas DataFrames with pd.read_csv().
|
||||
pandas_kwargs: Named arguments to pass to pd.read_csv().
|
||||
**kwargs: Additional kwargs to pass to langchain_experimental.agents.agent_toolkits.pandas.base.create_pandas_dataframe_agent().
|
||||
|
||||
Returns:
|
||||
An AgentExecutor with the specified agent_type agent and access to
|
||||
a PythonAstREPLTool with the loaded DataFrame(s) and any user-provided extra_tools.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_experimental.agents import create_csv_agent
|
||||
|
||||
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
|
||||
agent_executor = create_pandas_dataframe_agent(
|
||||
llm,
|
||||
"titanic.csv",
|
||||
agent_type="openai-tools",
|
||||
verbose=True
|
||||
)
|
||||
""" # noqa: E501
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"pandas package not found, please install with `pip install pandas`"
|
||||
"pandas package not found, please install with `pip install pandas`."
|
||||
)
|
||||
|
||||
_kwargs = pandas_kwargs or {}
|
||||
|
||||
@@ -1,16 +1,26 @@
|
||||
"""Agent for working with pandas objects."""
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||
import warnings
|
||||
from typing import Any, Dict, List, Literal, Optional, Sequence, Union
|
||||
|
||||
from langchain.agents.agent import AgentExecutor, BaseSingleActionAgent
|
||||
from langchain.agents import AgentType, create_openai_tools_agent, create_react_agent
|
||||
from langchain.agents.agent import (
|
||||
AgentExecutor,
|
||||
BaseMultiActionAgent,
|
||||
BaseSingleActionAgent,
|
||||
RunnableAgent,
|
||||
RunnableMultiActionAgent,
|
||||
)
|
||||
from langchain.agents.mrkl.base import ZeroShotAgent
|
||||
from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
|
||||
from langchain.agents.types import AgentType
|
||||
from langchain.callbacks.base import BaseCallbackManager
|
||||
from langchain.chains.llm import LLMChain
|
||||
from langchain.schema import BasePromptTemplate
|
||||
from langchain.tools import BaseTool
|
||||
from langchain_core.language_models import BaseLanguageModel
|
||||
from langchain.agents.openai_functions_agent.base import (
|
||||
OpenAIFunctionsAgent,
|
||||
create_openai_functions_agent,
|
||||
)
|
||||
from langchain_core.callbacks import BaseCallbackManager
|
||||
from langchain_core.language_models import LanguageModelLike
|
||||
from langchain_core.messages import SystemMessage
|
||||
from langchain_core.prompts import BasePromptTemplate, ChatPromptTemplate
|
||||
from langchain_core.tools import BaseTool
|
||||
from langchain_core.utils.interactive_env import is_interactive_env
|
||||
|
||||
from langchain_experimental.agents.agent_toolkits.pandas.prompt import (
|
||||
FUNCTIONS_WITH_DF,
|
||||
@@ -28,257 +38,121 @@ from langchain_experimental.tools.python.tool import PythonAstREPLTool
|
||||
|
||||
def _get_multi_prompt(
|
||||
dfs: List[Any],
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
input_variables: Optional[List[str]] = None,
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
extra_tools: Sequence[BaseTool] = (),
|
||||
) -> Tuple[BasePromptTemplate, List[BaseTool]]:
|
||||
num_dfs = len(dfs)
|
||||
tools: Sequence[BaseTool] = (),
|
||||
) -> BasePromptTemplate:
|
||||
if suffix is not None:
|
||||
suffix_to_use = suffix
|
||||
include_dfs_head = True
|
||||
elif include_df_in_prompt:
|
||||
suffix_to_use = SUFFIX_WITH_MULTI_DF
|
||||
include_dfs_head = True
|
||||
else:
|
||||
suffix_to_use = SUFFIX_NO_DF
|
||||
include_dfs_head = False
|
||||
if input_variables is None:
|
||||
input_variables = ["input", "agent_scratchpad", "num_dfs"]
|
||||
if include_dfs_head:
|
||||
input_variables += ["dfs_head"]
|
||||
prefix = prefix if prefix is not None else MULTI_DF_PREFIX
|
||||
|
||||
if prefix is None:
|
||||
prefix = MULTI_DF_PREFIX
|
||||
|
||||
df_locals = {}
|
||||
for i, dataframe in enumerate(dfs):
|
||||
df_locals[f"df{i + 1}"] = dataframe
|
||||
tools = [PythonAstREPLTool(locals=df_locals)] + list(extra_tools)
|
||||
prompt = ZeroShotAgent.create_prompt(
|
||||
tools,
|
||||
prefix=prefix,
|
||||
suffix=suffix_to_use,
|
||||
input_variables=input_variables,
|
||||
)
|
||||
partial_prompt = prompt.partial()
|
||||
if "dfs_head" in input_variables:
|
||||
if "dfs_head" in partial_prompt.input_variables:
|
||||
dfs_head = "\n\n".join([d.head(number_of_head_rows).to_markdown() for d in dfs])
|
||||
partial_prompt = partial_prompt.partial(num_dfs=str(num_dfs), dfs_head=dfs_head)
|
||||
if "num_dfs" in input_variables:
|
||||
partial_prompt = partial_prompt.partial(num_dfs=str(num_dfs))
|
||||
return partial_prompt, tools
|
||||
partial_prompt = partial_prompt.partial(dfs_head=dfs_head)
|
||||
if "num_dfs" in partial_prompt.input_variables:
|
||||
partial_prompt = partial_prompt.partial(num_dfs=str(len(dfs)))
|
||||
return partial_prompt
|
||||
|
||||
|
||||
def _get_single_prompt(
|
||||
df: Any,
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
input_variables: Optional[List[str]] = None,
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
extra_tools: Sequence[BaseTool] = (),
|
||||
) -> Tuple[BasePromptTemplate, List[BaseTool]]:
|
||||
tools: Sequence[BaseTool] = (),
|
||||
) -> BasePromptTemplate:
|
||||
if suffix is not None:
|
||||
suffix_to_use = suffix
|
||||
include_df_head = True
|
||||
elif include_df_in_prompt:
|
||||
suffix_to_use = SUFFIX_WITH_DF
|
||||
include_df_head = True
|
||||
else:
|
||||
suffix_to_use = SUFFIX_NO_DF
|
||||
include_df_head = False
|
||||
|
||||
if input_variables is None:
|
||||
input_variables = ["input", "agent_scratchpad"]
|
||||
if include_df_head:
|
||||
input_variables += ["df_head"]
|
||||
|
||||
if prefix is None:
|
||||
prefix = PREFIX
|
||||
|
||||
tools = [PythonAstREPLTool(locals={"df": df})] + list(extra_tools)
|
||||
prefix = prefix if prefix is not None else PREFIX
|
||||
|
||||
prompt = ZeroShotAgent.create_prompt(
|
||||
tools,
|
||||
prefix=prefix,
|
||||
suffix=suffix_to_use,
|
||||
input_variables=input_variables,
|
||||
)
|
||||
|
||||
partial_prompt = prompt.partial()
|
||||
if "df_head" in input_variables:
|
||||
partial_prompt = partial_prompt.partial(
|
||||
df_head=str(df.head(number_of_head_rows).to_markdown())
|
||||
)
|
||||
return partial_prompt, tools
|
||||
if "df_head" in partial_prompt.input_variables:
|
||||
df_head = str(df.head(number_of_head_rows).to_markdown())
|
||||
partial_prompt = partial_prompt.partial(df_head=df_head)
|
||||
return partial_prompt
|
||||
|
||||
|
||||
def _get_prompt_and_tools(
|
||||
df: Any,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
input_variables: Optional[List[str]] = None,
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
extra_tools: Sequence[BaseTool] = (),
|
||||
) -> Tuple[BasePromptTemplate, List[BaseTool]]:
|
||||
try:
|
||||
import pandas as pd
|
||||
|
||||
pd.set_option("display.max_columns", None)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"pandas package not found, please install with `pip install pandas`"
|
||||
)
|
||||
|
||||
if include_df_in_prompt is not None and suffix is not None:
|
||||
raise ValueError("If suffix is specified, include_df_in_prompt should not be.")
|
||||
|
||||
if isinstance(df, list):
|
||||
for item in df:
|
||||
if not isinstance(item, pd.DataFrame):
|
||||
raise ValueError(f"Expected pandas object, got {type(df)}")
|
||||
return _get_multi_prompt(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
input_variables=input_variables,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
extra_tools=extra_tools,
|
||||
)
|
||||
else:
|
||||
if not isinstance(df, pd.DataFrame):
|
||||
raise ValueError(f"Expected pandas object, got {type(df)}")
|
||||
return _get_single_prompt(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
input_variables=input_variables,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
extra_tools=extra_tools,
|
||||
)
|
||||
def _get_prompt(df: Any, **kwargs: Any) -> BasePromptTemplate:
|
||||
return (
|
||||
_get_multi_prompt(df, **kwargs)
|
||||
if isinstance(df, list)
|
||||
else _get_single_prompt(df, **kwargs)
|
||||
)
|
||||
|
||||
|
||||
def _get_functions_single_prompt(
|
||||
df: Any,
|
||||
*,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
suffix: str = "",
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
) -> Tuple[BasePromptTemplate, List[PythonAstREPLTool]]:
|
||||
if suffix is not None:
|
||||
suffix_to_use = suffix
|
||||
if include_df_in_prompt:
|
||||
suffix_to_use = suffix_to_use.format(
|
||||
df_head=str(df.head(number_of_head_rows).to_markdown())
|
||||
)
|
||||
elif include_df_in_prompt:
|
||||
suffix_to_use = FUNCTIONS_WITH_DF.format(
|
||||
df_head=str(df.head(number_of_head_rows).to_markdown())
|
||||
)
|
||||
else:
|
||||
suffix_to_use = ""
|
||||
|
||||
if prefix is None:
|
||||
prefix = PREFIX_FUNCTIONS
|
||||
|
||||
tools = [PythonAstREPLTool(locals={"df": df})]
|
||||
system_message = SystemMessage(content=prefix + suffix_to_use)
|
||||
) -> ChatPromptTemplate:
|
||||
if include_df_in_prompt:
|
||||
df_head = str(df.head(number_of_head_rows).to_markdown())
|
||||
suffix = (suffix or FUNCTIONS_WITH_DF).format(df_head=df_head)
|
||||
prefix = prefix if prefix is not None else PREFIX_FUNCTIONS
|
||||
system_message = SystemMessage(content=prefix + suffix)
|
||||
prompt = OpenAIFunctionsAgent.create_prompt(system_message=system_message)
|
||||
return prompt, tools
|
||||
return prompt
|
||||
|
||||
|
||||
def _get_functions_multi_prompt(
|
||||
dfs: Any,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
*,
|
||||
prefix: str = "",
|
||||
suffix: str = "",
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
) -> Tuple[BasePromptTemplate, List[PythonAstREPLTool]]:
|
||||
if suffix is not None:
|
||||
suffix_to_use = suffix
|
||||
if include_df_in_prompt:
|
||||
dfs_head = "\n\n".join(
|
||||
[d.head(number_of_head_rows).to_markdown() for d in dfs]
|
||||
)
|
||||
suffix_to_use = suffix_to_use.format(
|
||||
dfs_head=dfs_head,
|
||||
)
|
||||
elif include_df_in_prompt:
|
||||
) -> ChatPromptTemplate:
|
||||
if include_df_in_prompt:
|
||||
dfs_head = "\n\n".join([d.head(number_of_head_rows).to_markdown() for d in dfs])
|
||||
suffix_to_use = FUNCTIONS_WITH_MULTI_DF.format(
|
||||
dfs_head=dfs_head,
|
||||
)
|
||||
else:
|
||||
suffix_to_use = ""
|
||||
|
||||
if prefix is None:
|
||||
prefix = MULTI_DF_PREFIX_FUNCTIONS
|
||||
prefix = prefix.format(num_dfs=str(len(dfs)))
|
||||
|
||||
df_locals = {}
|
||||
for i, dataframe in enumerate(dfs):
|
||||
df_locals[f"df{i + 1}"] = dataframe
|
||||
tools = [PythonAstREPLTool(locals=df_locals)]
|
||||
system_message = SystemMessage(content=prefix + suffix_to_use)
|
||||
suffix = (suffix or FUNCTIONS_WITH_MULTI_DF).format(dfs_head=dfs_head)
|
||||
prefix = (prefix or MULTI_DF_PREFIX_FUNCTIONS).format(num_dfs=str(len(dfs)))
|
||||
system_message = SystemMessage(content=prefix + suffix)
|
||||
prompt = OpenAIFunctionsAgent.create_prompt(system_message=system_message)
|
||||
return prompt, tools
|
||||
return prompt
|
||||
|
||||
|
||||
def _get_functions_prompt_and_tools(
|
||||
df: Any,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
input_variables: Optional[List[str]] = None,
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
) -> Tuple[BasePromptTemplate, List[PythonAstREPLTool]]:
|
||||
try:
|
||||
import pandas as pd
|
||||
|
||||
pd.set_option("display.max_columns", None)
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"pandas package not found, please install with `pip install pandas`"
|
||||
)
|
||||
if input_variables is not None:
|
||||
raise ValueError("`input_variables` is not supported at the moment.")
|
||||
|
||||
if include_df_in_prompt is not None and suffix is not None:
|
||||
raise ValueError("If suffix is specified, include_df_in_prompt should not be.")
|
||||
|
||||
if isinstance(df, list):
|
||||
for item in df:
|
||||
if not isinstance(item, pd.DataFrame):
|
||||
raise ValueError(f"Expected pandas object, got {type(df)}")
|
||||
return _get_functions_multi_prompt(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
)
|
||||
else:
|
||||
if not isinstance(df, pd.DataFrame):
|
||||
raise ValueError(f"Expected pandas object, got {type(df)}")
|
||||
return _get_functions_single_prompt(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
)
|
||||
def _get_functions_prompt(df: Any, **kwargs: Any) -> ChatPromptTemplate:
|
||||
return (
|
||||
_get_functions_multi_prompt(df, **kwargs)
|
||||
if isinstance(df, list)
|
||||
else _get_functions_single_prompt(df, **kwargs)
|
||||
)
|
||||
|
||||
|
||||
def create_pandas_dataframe_agent(
|
||||
llm: BaseLanguageModel,
|
||||
llm: LanguageModelLike,
|
||||
df: Any,
|
||||
agent_type: AgentType = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||||
agent_type: Union[
|
||||
AgentType, Literal["openai-tools"]
|
||||
] = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||||
callback_manager: Optional[BaseCallbackManager] = None,
|
||||
prefix: Optional[str] = None,
|
||||
suffix: Optional[str] = None,
|
||||
@@ -292,54 +166,131 @@ def create_pandas_dataframe_agent(
|
||||
include_df_in_prompt: Optional[bool] = True,
|
||||
number_of_head_rows: int = 5,
|
||||
extra_tools: Sequence[BaseTool] = (),
|
||||
**kwargs: Dict[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> AgentExecutor:
|
||||
"""Construct a pandas agent from an LLM and dataframe."""
|
||||
agent: BaseSingleActionAgent
|
||||
base_tools: Sequence[BaseTool]
|
||||
if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION:
|
||||
prompt, base_tools = _get_prompt_and_tools(
|
||||
"""Construct a Pandas agent from an LLM and dataframe(s).
|
||||
|
||||
Args:
|
||||
llm: Language model to use for the agent.
|
||||
df: Pandas dataframe or list of Pandas dataframes.
|
||||
agent_type: One of "openai-tools", "openai-functions", or
|
||||
"zero-shot-react-description". Defaults to "zero-shot-react-description".
|
||||
"openai-tools" is recommended over "openai-functions".
|
||||
callback_manager: DEPRECATED. Pass "callbacks" key into 'agent_executor_kwargs'
|
||||
instead to pass constructor callbacks to AgentExecutor.
|
||||
prefix: Prompt prefix string.
|
||||
suffix: Prompt suffix string.
|
||||
input_variables: DEPRECATED. Input variables automatically inferred from
|
||||
constructed prompt.
|
||||
verbose: AgentExecutor verbosity.
|
||||
return_intermediate_steps: Passed to AgentExecutor init.
|
||||
max_iterations: Passed to AgentExecutor init.
|
||||
max_execution_time: Passed to AgentExecutor init.
|
||||
early_stopping_method: Passed to AgentExecutor init.
|
||||
agent_executor_kwargs: Arbitrary additional AgentExecutor args.
|
||||
include_df_in_prompt: Whether to include the first number_of_head_rows in the
|
||||
prompt. Must be None if suffix is not None.
|
||||
number_of_head_rows: Number of initial rows to include in prompt if
|
||||
include_df_in_prompt is True.
|
||||
extra_tools: Additional tools to give to agent on top of a PythonAstREPLTool.
|
||||
**kwargs: DEPRECATED. Not used, kept for backwards compatibility.
|
||||
|
||||
Returns:
|
||||
An AgentExecutor with the specified agent_type agent and access to
|
||||
a PythonAstREPLTool with the DataFrame(s) and any user-provided extra_tools.
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_experimental.agents import create_pandas_dataframe_agent
|
||||
import pandas as pd
|
||||
|
||||
df = pd.read_csv("titanic.csv")
|
||||
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
|
||||
agent_executor = create_pandas_dataframe_agent(
|
||||
llm,
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
input_variables=input_variables,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
extra_tools=extra_tools,
|
||||
agent_type="openai-tools",
|
||||
verbose=True
|
||||
)
|
||||
tools = base_tools
|
||||
llm_chain = LLMChain(
|
||||
llm=llm,
|
||||
prompt=prompt,
|
||||
callback_manager=callback_manager,
|
||||
)
|
||||
tool_names = [tool.name for tool in tools]
|
||||
agent = ZeroShotAgent(
|
||||
llm_chain=llm_chain,
|
||||
allowed_tools=tool_names,
|
||||
callback_manager=callback_manager,
|
||||
**kwargs,
|
||||
)
|
||||
elif agent_type == AgentType.OPENAI_FUNCTIONS:
|
||||
_prompt, base_tools = _get_functions_prompt_and_tools(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
input_variables=input_variables,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
)
|
||||
tools = list(base_tools) + list(extra_tools)
|
||||
agent = OpenAIFunctionsAgent(
|
||||
llm=llm,
|
||||
prompt=_prompt,
|
||||
tools=tools,
|
||||
callback_manager=callback_manager,
|
||||
**kwargs,
|
||||
|
||||
""" # noqa: E501
|
||||
try:
|
||||
import pandas as pd
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"pandas package not found, please install with `pip install pandas`"
|
||||
) from e
|
||||
|
||||
if is_interactive_env():
|
||||
pd.set_option("display.max_columns", None)
|
||||
|
||||
for _df in df if isinstance(df, list) else [df]:
|
||||
if not isinstance(_df, pd.DataFrame):
|
||||
raise ValueError(f"Expected pandas DataFrame, got {type(_df)}")
|
||||
|
||||
if input_variables:
|
||||
kwargs = kwargs or {}
|
||||
kwargs["input_variables"] = input_variables
|
||||
if kwargs:
|
||||
warnings.warn(
|
||||
f"Received additional kwargs {kwargs} which are no longer supported."
|
||||
)
|
||||
|
||||
df_locals = {}
|
||||
if isinstance(df, list):
|
||||
for i, dataframe in enumerate(df):
|
||||
df_locals[f"df{i + 1}"] = dataframe
|
||||
else:
|
||||
raise ValueError(f"Agent type {agent_type} not supported at the moment.")
|
||||
return AgentExecutor.from_agent_and_tools(
|
||||
df_locals["df"] = df
|
||||
tools = [PythonAstREPLTool(locals=df_locals)] + list(extra_tools)
|
||||
|
||||
if agent_type == AgentType.ZERO_SHOT_REACT_DESCRIPTION:
|
||||
if include_df_in_prompt is not None and suffix is not None:
|
||||
raise ValueError(
|
||||
"If suffix is specified, include_df_in_prompt should not be."
|
||||
)
|
||||
prompt = _get_prompt(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
tools=tools,
|
||||
)
|
||||
agent: Union[BaseSingleActionAgent, BaseMultiActionAgent] = RunnableAgent(
|
||||
runnable=create_react_agent(llm, tools, prompt), # type: ignore
|
||||
input_keys_arg=["input"],
|
||||
return_keys_arg=["output"],
|
||||
)
|
||||
elif agent_type in (AgentType.OPENAI_FUNCTIONS, "openai-tools"):
|
||||
prompt = _get_functions_prompt(
|
||||
df,
|
||||
prefix=prefix,
|
||||
suffix=suffix,
|
||||
include_df_in_prompt=include_df_in_prompt,
|
||||
number_of_head_rows=number_of_head_rows,
|
||||
)
|
||||
if agent_type == AgentType.OPENAI_FUNCTIONS:
|
||||
agent = RunnableAgent(
|
||||
runnable=create_openai_functions_agent(llm, tools, prompt), # type: ignore
|
||||
input_keys_arg=["input"],
|
||||
return_keys_arg=["output"],
|
||||
)
|
||||
else:
|
||||
agent = RunnableMultiActionAgent(
|
||||
runnable=create_openai_tools_agent(llm, tools, prompt), # type: ignore
|
||||
input_keys_arg=["input"],
|
||||
return_keys_arg=["output"],
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Agent type {agent_type} not supported at the moment. Must be one of "
|
||||
"'openai-tools', 'openai-functions', or 'zero-shot-react-description'."
|
||||
)
|
||||
return AgentExecutor(
|
||||
agent=agent,
|
||||
tools=tools,
|
||||
callback_manager=callback_manager,
|
||||
|
||||
Reference in New Issue
Block a user