Toggle Menu
Prev Up Next

# Source code for langchain_experimental.agents.agent_toolkits.pandas.base

    """Agent for working with pandas objects."""
    import warnings
    from typing import Any, Dict, List, Literal, Optional, Sequence, Union, cast

    from langchain.agents import (
        AgentType,
        create_openai_tools_agent,
        create_react_agent,
        create_tool_calling_agent,
    )
    from langchain.agents.agent import (
        AgentExecutor,
        BaseMultiActionAgent,
        BaseSingleActionAgent,
        RunnableAgent,
        RunnableMultiActionAgent,
    )
    from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS
    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 BaseLanguageModel, LanguageModelLike
    from langchain_core.messages import SystemMessage
    from langchain_core.prompts import (
        BasePromptTemplate,
        ChatPromptTemplate,
        PromptTemplate,
    )
    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,
        FUNCTIONS_WITH_MULTI_DF,
        MULTI_DF_PREFIX,
        MULTI_DF_PREFIX_FUNCTIONS,
        PREFIX,
        PREFIX_FUNCTIONS,
        SUFFIX_NO_DF,
        SUFFIX_WITH_DF,
        SUFFIX_WITH_MULTI_DF,
    )
    from langchain_experimental.tools.python.tool import PythonAstREPLTool

    def _get_multi_prompt(
        dfs: List[Any],
        *,
        prefix: Optional[str] = None,
        suffix: Optional[str] = None,
        include_df_in_prompt: Optional[bool] = True,
        number_of_head_rows: int = 5,
    ) -> BasePromptTemplate:
        if suffix is not None:
            suffix_to_use = suffix
        elif include_df_in_prompt:
            suffix_to_use = SUFFIX_WITH_MULTI_DF
        else:
            suffix_to_use = SUFFIX_NO_DF
        prefix = prefix if prefix is not None else MULTI_DF_PREFIX

        template = "\n\n".join([prefix, "{tools}", FORMAT_INSTRUCTIONS, suffix_to_use])
        prompt = PromptTemplate.from_template(template)
        partial_prompt = prompt.partial()
        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(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,
        include_df_in_prompt: Optional[bool] = True,
        number_of_head_rows: int = 5,
    ) -> BasePromptTemplate:
        if suffix is not None:
            suffix_to_use = suffix
        elif include_df_in_prompt:
            suffix_to_use = SUFFIX_WITH_DF
        else:
            suffix_to_use = SUFFIX_NO_DF
        prefix = prefix if prefix is not None else PREFIX

        template = "\n\n".join([prefix, "{tools}", FORMAT_INSTRUCTIONS, suffix_to_use])
        prompt = PromptTemplate.from_template(template)

        partial_prompt = prompt.partial()
        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(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: str = "",
        include_df_in_prompt: Optional[bool] = True,
        number_of_head_rows: int = 5,
    ) -> 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

    def _get_functions_multi_prompt(
        dfs: Any,
        *,
        prefix: str = "",
        suffix: str = "",
        include_df_in_prompt: Optional[bool] = True,
        number_of_head_rows: int = 5,
    ) -> ChatPromptTemplate:
        if include_df_in_prompt:
            dfs_head = "\n\n".join([d.head(number_of_head_rows).to_markdown() for d in dfs])
            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

    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: LanguageModelLike,
        df: Any,
        agent_type: Union[
            AgentType, Literal["openai-tools", "tool-calling"]
        ] = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        callback_manager: Optional[BaseCallbackManager] = None,
        prefix: Optional[str] = None,
        suffix: Optional[str] = None,
        input_variables: Optional[List[str]] = None,
        verbose: bool = False,
        return_intermediate_steps: bool = False,
        max_iterations: Optional[int] = 15,
        max_execution_time: Optional[float] = None,
        early_stopping_method: str = "force",
        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] = (),
        engine: Literal["pandas", "modin"] = "pandas",
        **kwargs: Any,
    ) -> AgentExecutor:
        """Construct a Pandas agent from an LLM and dataframe(s).

        Args:
            llm: Language model to use for the agent. If agent_type is "tool-calling" then
                llm is expected to support tool calling.
            df: Pandas dataframe or list of Pandas dataframes.
            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.
            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.
            engine: One of "modin" or "pandas". Defaults to "pandas".
            **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,
                    agent_type="tool-calling",
                    verbose=True
                )

        """  # noqa: E501
        try:
            if engine == "modin":
                import modin.pandas as pd
            elif engine == "pandas":
                import pandas as pd
            else:
                raise ValueError(
                    f"Unsupported engine {engine}. It must be one of 'modin' or 'pandas'."
                )
        except ImportError as e:
            raise ImportError(
                f"`{engine}` package not found, please install with `pip install {engine}`"
            ) 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:
            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,
            )
            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", "tool-calling"):
            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:
                runnable = create_openai_functions_agent(
                    cast(BaseLanguageModel, llm), tools, prompt
                )
                agent = RunnableAgent(
                    runnable=runnable,
                    input_keys_arg=["input"],
                    return_keys_arg=["output"],
                )
            else:
                if agent_type == "openai-tools":
                    runnable = create_openai_tools_agent(
                        cast(BaseLanguageModel, llm), tools, prompt
                    )
                else:
                    runnable = create_tool_calling_agent(
                        cast(BaseLanguageModel, llm), tools, prompt
                    )
                agent = RunnableMultiActionAgent(
                    runnable=runnable,
                    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 "
                "'tool-calling', 'openai-tools', 'openai-functions', or "
                "'zero-shot-react-description'."
            )
        return AgentExecutor(
            agent=agent,
            tools=tools,
            callback_manager=callback_manager,
            verbose=verbose,
            return_intermediate_steps=return_intermediate_steps,
            max_iterations=max_iterations,
            max_execution_time=max_execution_time,
            early_stopping_method=early_stopping_method,
            **(agent_executor_kwargs or {}),
        )