From 3b975c6ebe3be15e62eb4caef26a6f36e37571bf Mon Sep 17 00:00:00 2001 From: Massimiliano Pronesti Date: Fri, 8 Mar 2024 04:23:07 +0100 Subject: [PATCH] experimental[minor]: add support for modin in pandas agent (#18749) Added support for Intel's [modin](https://github.com/modin-project/modin) in `create_pandas_dataframe_agent`. --- .../agents/agent_toolkits/pandas/base.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/experimental/langchain_experimental/agents/agent_toolkits/pandas/base.py b/libs/experimental/langchain_experimental/agents/agent_toolkits/pandas/base.py index 6fa67dc60f..92fc997374 100644 --- a/libs/experimental/langchain_experimental/agents/agent_toolkits/pandas/base.py +++ b/libs/experimental/langchain_experimental/agents/agent_toolkits/pandas/base.py @@ -162,6 +162,7 @@ def create_pandas_dataframe_agent( 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). @@ -189,6 +190,7 @@ def create_pandas_dataframe_agent( 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: @@ -213,10 +215,17 @@ def create_pandas_dataframe_agent( """ # noqa: E501 try: - import pandas as pd + 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( - "pandas package not found, please install with `pip install pandas`" + f"`{engine}` package not found, please install with `pip install {engine}`" ) from e if is_interactive_env():