From c5cc09d4e343a9596d1835f2f66e25d6c043af6c Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Mon, 1 May 2023 20:28:43 -0700 Subject: [PATCH] Harrison/agent exec kwargs (#3917) Co-authored-by: Zach Schillaci <40636930+zachschillaci27@users.noreply.github.com> --- langchain/agents/agent_toolkits/json/base.py | 11 ++++++--- .../agents/agent_toolkits/openapi/base.py | 9 ++++--- .../agents/agent_toolkits/openapi/planner.py | 15 ++++++++++-- .../agents/agent_toolkits/pandas/base.py | 8 ++++--- .../agents/agent_toolkits/powerbi/base.py | 6 ++--- .../agent_toolkits/powerbi/chat_base.py | 11 +++++---- .../agents/agent_toolkits/python/base.py | 13 +++++++--- langchain/agents/agent_toolkits/sql/base.py | 7 ++++-- .../agents/agent_toolkits/vectorstore/base.py | 24 +++++++++++++++---- 9 files changed, 76 insertions(+), 28 deletions(-) diff --git a/langchain/agents/agent_toolkits/json/base.py b/langchain/agents/agent_toolkits/json/base.py index 48d3fd9cc0..e46e994f88 100644 --- a/langchain/agents/agent_toolkits/json/base.py +++ b/langchain/agents/agent_toolkits/json/base.py @@ -1,5 +1,5 @@ """Json agent.""" -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from langchain.agents.agent import AgentExecutor from langchain.agents.agent_toolkits.json.prompt import JSON_PREFIX, JSON_SUFFIX @@ -20,7 +20,8 @@ def create_json_agent( format_instructions: str = FORMAT_INSTRUCTIONS, input_variables: Optional[List[str]] = None, verbose: bool = False, - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a json agent from an LLM and tools.""" tools = toolkit.get_tools() @@ -39,5 +40,9 @@ def create_json_agent( tool_names = [tool.name for tool in tools] agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) return AgentExecutor.from_agent_and_tools( - agent=agent, tools=toolkit.get_tools(), verbose=verbose + agent=agent, + tools=tools, + callback_manager=callback_manager, + verbose=verbose, + **(agent_executor_kwargs or {}), ) diff --git a/langchain/agents/agent_toolkits/openapi/base.py b/langchain/agents/agent_toolkits/openapi/base.py index fcaf3a13d1..7e18c70b64 100644 --- a/langchain/agents/agent_toolkits/openapi/base.py +++ b/langchain/agents/agent_toolkits/openapi/base.py @@ -1,5 +1,5 @@ """OpenAPI spec agent.""" -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from langchain.agents.agent import AgentExecutor from langchain.agents.agent_toolkits.openapi.prompt import ( @@ -27,7 +27,8 @@ def create_openapi_agent( early_stopping_method: str = "force", verbose: bool = False, return_intermediate_steps: bool = False, - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a json agent from an LLM and tools.""" tools = toolkit.get_tools() @@ -47,10 +48,12 @@ def create_openapi_agent( agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) return AgentExecutor.from_agent_and_tools( agent=agent, - tools=toolkit.get_tools(), + 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 {}), ) diff --git a/langchain/agents/agent_toolkits/openapi/planner.py b/langchain/agents/agent_toolkits/openapi/planner.py index a6d9fd098f..fcfacc1436 100644 --- a/langchain/agents/agent_toolkits/openapi/planner.py +++ b/langchain/agents/agent_toolkits/openapi/planner.py @@ -2,7 +2,7 @@ import json import re from functools import partial -from typing import Callable, List, Optional +from typing import Any, Callable, Dict, List, Optional import yaml from pydantic import Field @@ -29,6 +29,7 @@ from langchain.agents.agent_toolkits.openapi.spec import ReducedOpenAPISpec from langchain.agents.mrkl.base import ZeroShotAgent from langchain.agents.tools import Tool from langchain.base_language import BaseLanguageModel +from langchain.callbacks.base import BaseCallbackManager from langchain.chains.llm import LLMChain from langchain.llms.openai import OpenAI from langchain.memory import ReadOnlySharedMemory @@ -261,7 +262,10 @@ def create_openapi_agent( requests_wrapper: RequestsWrapper, llm: BaseLanguageModel, shared_memory: Optional[ReadOnlySharedMemory] = None, + callback_manager: Optional[BaseCallbackManager] = None, verbose: bool = True, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Instantiate API planner and controller for a given spec. @@ -288,5 +292,12 @@ def create_openapi_agent( agent = ZeroShotAgent( llm_chain=LLMChain(llm=llm, prompt=prompt, memory=shared_memory), allowed_tools=[tool.name for tool in tools], + **kwargs, + ) + return AgentExecutor.from_agent_and_tools( + agent=agent, + tools=tools, + callback_manager=callback_manager, + verbose=verbose, + **(agent_executor_kwargs or {}), ) - return AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=verbose) diff --git a/langchain/agents/agent_toolkits/pandas/base.py b/langchain/agents/agent_toolkits/pandas/base.py index e6a4326e6d..fdeb840fbd 100644 --- a/langchain/agents/agent_toolkits/pandas/base.py +++ b/langchain/agents/agent_toolkits/pandas/base.py @@ -1,5 +1,5 @@ """Agent for working with pandas objects.""" -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from langchain.agents.agent import AgentExecutor from langchain.agents.agent_toolkits.pandas.prompt import PREFIX, SUFFIX @@ -22,7 +22,8 @@ def create_pandas_dataframe_agent( max_iterations: Optional[int] = 15, max_execution_time: Optional[float] = None, early_stopping_method: str = "force", - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a pandas agent from an LLM and dataframe.""" import pandas as pd @@ -51,10 +52,11 @@ def create_pandas_dataframe_agent( return AgentExecutor.from_agent_and_tools( 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, - callback_manager=callback_manager, + **(agent_executor_kwargs or {}), ) diff --git a/langchain/agents/agent_toolkits/powerbi/base.py b/langchain/agents/agent_toolkits/powerbi/base.py index 43fd515f38..578b58e0be 100644 --- a/langchain/agents/agent_toolkits/powerbi/base.py +++ b/langchain/agents/agent_toolkits/powerbi/base.py @@ -27,7 +27,7 @@ def create_pbi_agent( input_variables: Optional[List[str]] = None, top_k: int = 10, verbose: bool = False, - agent_kwargs: Optional[Dict[str, Any]] = None, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a pbi agent from an LLM and tools.""" @@ -51,12 +51,12 @@ def create_pbi_agent( verbose=verbose, ), allowed_tools=[tool.name for tool in tools], - **(agent_kwargs or {}), + **kwargs, ) return AgentExecutor.from_agent_and_tools( agent=agent, tools=tools, callback_manager=callback_manager, verbose=verbose, - **kwargs, + **(agent_executor_kwargs or {}), ) diff --git a/langchain/agents/agent_toolkits/powerbi/chat_base.py b/langchain/agents/agent_toolkits/powerbi/chat_base.py index da1fbac103..70956ca403 100644 --- a/langchain/agents/agent_toolkits/powerbi/chat_base.py +++ b/langchain/agents/agent_toolkits/powerbi/chat_base.py @@ -2,6 +2,7 @@ from typing import Any, Dict, List, Optional from langchain.agents import AgentExecutor +from langchain.agents.agent import AgentOutputParser from langchain.agents.agent_toolkits.powerbi.prompt import ( POWERBI_CHAT_PREFIX, POWERBI_CHAT_SUFFIX, @@ -20,6 +21,7 @@ def create_pbi_chat_agent( toolkit: Optional[PowerBIToolkit], powerbi: Optional[PowerBIDataset] = None, callback_manager: Optional[BaseCallbackManager] = None, + output_parser: Optional[AgentOutputParser] = None, prefix: str = POWERBI_CHAT_PREFIX, suffix: str = POWERBI_CHAT_SUFFIX, examples: Optional[str] = None, @@ -27,7 +29,7 @@ def create_pbi_chat_agent( memory: Optional[BaseChatMemory] = None, top_k: int = 10, verbose: bool = False, - agent_kwargs: Optional[Dict[str, Any]] = None, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a pbi agent from an Chat LLM and tools. @@ -43,11 +45,12 @@ def create_pbi_chat_agent( llm=llm, tools=tools, system_message=prefix.format(top_k=top_k), - user_message=suffix, + human_message=suffix, input_variables=input_variables, callback_manager=callback_manager, + output_parser=output_parser, verbose=verbose, - **(agent_kwargs or {}), + **kwargs, ) return AgentExecutor.from_agent_and_tools( agent=agent, @@ -56,5 +59,5 @@ def create_pbi_chat_agent( memory=memory or ConversationBufferMemory(memory_key="chat_history", return_messages=True), verbose=verbose, - **kwargs, + **(agent_executor_kwargs or {}), ) diff --git a/langchain/agents/agent_toolkits/python/base.py b/langchain/agents/agent_toolkits/python/base.py index 644dd786da..a4663898bb 100644 --- a/langchain/agents/agent_toolkits/python/base.py +++ b/langchain/agents/agent_toolkits/python/base.py @@ -1,6 +1,6 @@ """Python agent.""" -from typing import Any, Optional +from typing import Any, Dict, Optional from langchain.agents.agent import AgentExecutor from langchain.agents.agent_toolkits.python.prompt import PREFIX @@ -17,7 +17,8 @@ def create_python_agent( callback_manager: Optional[BaseCallbackManager] = None, verbose: bool = False, prefix: str = PREFIX, - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a python agent from an LLM and tool.""" tools = [tool] @@ -29,4 +30,10 @@ def create_python_agent( ) tool_names = [tool.name for tool in tools] agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) - return AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=verbose) + return AgentExecutor.from_agent_and_tools( + agent=agent, + tools=tools, + callback_manager=callback_manager, + verbose=verbose, + **(agent_executor_kwargs or {}), + ) diff --git a/langchain/agents/agent_toolkits/sql/base.py b/langchain/agents/agent_toolkits/sql/base.py index 5f45fabebd..697c8dbd80 100644 --- a/langchain/agents/agent_toolkits/sql/base.py +++ b/langchain/agents/agent_toolkits/sql/base.py @@ -1,5 +1,5 @@ """SQL agent.""" -from typing import Any, List, Optional +from typing import Any, Dict, List, Optional from langchain.agents.agent import AgentExecutor from langchain.agents.agent_toolkits.sql.prompt import SQL_PREFIX, SQL_SUFFIX @@ -24,7 +24,8 @@ def create_sql_agent( max_execution_time: Optional[float] = None, early_stopping_method: str = "force", verbose: bool = False, - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a sql agent from an LLM and tools.""" tools = toolkit.get_tools() @@ -46,8 +47,10 @@ def create_sql_agent( return AgentExecutor.from_agent_and_tools( agent=agent, tools=tools, + callback_manager=callback_manager, verbose=verbose, max_iterations=max_iterations, max_execution_time=max_execution_time, early_stopping_method=early_stopping_method, + **(agent_executor_kwargs or {}), ) diff --git a/langchain/agents/agent_toolkits/vectorstore/base.py b/langchain/agents/agent_toolkits/vectorstore/base.py index 3b31fab8b5..524976259a 100644 --- a/langchain/agents/agent_toolkits/vectorstore/base.py +++ b/langchain/agents/agent_toolkits/vectorstore/base.py @@ -1,5 +1,5 @@ """VectorStore agent.""" -from typing import Any, Optional +from typing import Any, Dict, Optional from langchain.agents.agent import AgentExecutor from langchain.agents.agent_toolkits.vectorstore.prompt import PREFIX, ROUTER_PREFIX @@ -19,7 +19,8 @@ def create_vectorstore_agent( callback_manager: Optional[BaseCallbackManager] = None, prefix: str = PREFIX, verbose: bool = False, - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a vectorstore agent from an LLM and tools.""" tools = toolkit.get_tools() @@ -31,7 +32,13 @@ def create_vectorstore_agent( ) tool_names = [tool.name for tool in tools] agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) - return AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=verbose) + return AgentExecutor.from_agent_and_tools( + agent=agent, + tools=tools, + callback_manager=callback_manager, + verbose=verbose, + **(agent_executor_kwargs or {}), + ) def create_vectorstore_router_agent( @@ -40,7 +47,8 @@ def create_vectorstore_router_agent( callback_manager: Optional[BaseCallbackManager] = None, prefix: str = ROUTER_PREFIX, verbose: bool = False, - **kwargs: Any, + agent_executor_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Dict[str, Any], ) -> AgentExecutor: """Construct a vectorstore router agent from an LLM and tools.""" tools = toolkit.get_tools() @@ -52,4 +60,10 @@ def create_vectorstore_router_agent( ) tool_names = [tool.name for tool in tools] agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs) - return AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=verbose) + return AgentExecutor.from_agent_and_tools( + agent=agent, + tools=tools, + callback_manager=callback_manager, + verbose=verbose, + **(agent_executor_kwargs or {}), + )