From abfca72c0bf83c541b83d541d294098733986220 Mon Sep 17 00:00:00 2001 From: Zach Jones Date: Wed, 12 Apr 2023 20:09:42 -0400 Subject: [PATCH] Add max_execution_time to openapi, pandas, and sql creators (#2779) In #2399 we added the ability to set `max_execution_time` when creating an AgentExecutor. This PR adds the `max_execution_time` argument to the built-in pandas, sql, and openapi agents. Co-authored-by: Zachary Jones --- langchain/agents/agent_toolkits/openapi/base.py | 6 ++++++ langchain/agents/agent_toolkits/pandas/base.py | 2 ++ langchain/agents/agent_toolkits/sql/base.py | 2 ++ 3 files changed, 10 insertions(+) diff --git a/langchain/agents/agent_toolkits/openapi/base.py b/langchain/agents/agent_toolkits/openapi/base.py index 32232881..fcaf3a13 100644 --- a/langchain/agents/agent_toolkits/openapi/base.py +++ b/langchain/agents/agent_toolkits/openapi/base.py @@ -22,6 +22,9 @@ def create_openapi_agent( suffix: str = OPENAPI_SUFFIX, format_instructions: str = FORMAT_INSTRUCTIONS, input_variables: Optional[List[str]] = None, + max_iterations: Optional[int] = 15, + max_execution_time: Optional[float] = None, + early_stopping_method: str = "force", verbose: bool = False, return_intermediate_steps: bool = False, **kwargs: Any, @@ -47,4 +50,7 @@ def create_openapi_agent( tools=toolkit.get_tools(), verbose=verbose, return_intermediate_steps=return_intermediate_steps, + max_iterations=max_iterations, + max_execution_time=max_execution_time, + early_stopping_method=early_stopping_method, ) diff --git a/langchain/agents/agent_toolkits/pandas/base.py b/langchain/agents/agent_toolkits/pandas/base.py index 40d0c329..eac19e9f 100644 --- a/langchain/agents/agent_toolkits/pandas/base.py +++ b/langchain/agents/agent_toolkits/pandas/base.py @@ -20,6 +20,7 @@ def create_pandas_dataframe_agent( verbose: bool = False, return_intermediate_steps: bool = False, max_iterations: Optional[int] = 15, + max_execution_time: Optional[float] = None, early_stopping_method: str = "force", **kwargs: Any, ) -> AgentExecutor: @@ -48,5 +49,6 @@ def create_pandas_dataframe_agent( verbose=verbose, return_intermediate_steps=return_intermediate_steps, max_iterations=max_iterations, + max_execution_time=max_execution_time, early_stopping_method=early_stopping_method, ) diff --git a/langchain/agents/agent_toolkits/sql/base.py b/langchain/agents/agent_toolkits/sql/base.py index 6a24a430..5f45fabe 100644 --- a/langchain/agents/agent_toolkits/sql/base.py +++ b/langchain/agents/agent_toolkits/sql/base.py @@ -21,6 +21,7 @@ def create_sql_agent( input_variables: Optional[List[str]] = None, top_k: int = 10, max_iterations: Optional[int] = 15, + max_execution_time: Optional[float] = None, early_stopping_method: str = "force", verbose: bool = False, **kwargs: Any, @@ -47,5 +48,6 @@ def create_sql_agent( tools=tools, verbose=verbose, max_iterations=max_iterations, + max_execution_time=max_execution_time, early_stopping_method=early_stopping_method, )