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():