Harrison/agent exec kwargs (#3917)

Co-authored-by: Zach Schillaci <40636930+zachschillaci27@users.noreply.github.com>

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 {}),
     )

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 {}),
     )

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)

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 {}),
     )

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 {}),
     )

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 {}),
     )

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 {}),
+    )

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 {}),
     )

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 {}),
+    )