diff --git a/langchain/agents/initialize.py b/langchain/agents/initialize.py index b0b9a78c805..872400316b1 100644 --- a/langchain/agents/initialize.py +++ b/langchain/agents/initialize.py @@ -14,6 +14,7 @@ def initialize_agent( agent: Optional[str] = None, callback_manager: Optional[BaseCallbackManager] = None, agent_path: Optional[str] = None, + agent_kwargs: Optional[dict] = None, **kwargs: Any, ) -> AgentExecutor: """Load agent given tools and LLM. @@ -50,8 +51,9 @@ def initialize_agent( f"Valid types are: {AGENT_TO_CLASS.keys()}." ) agent_cls = AGENT_TO_CLASS[agent] + agent_kwargs = agent_kwargs or {} agent_obj = agent_cls.from_llm_and_tools( - llm, tools, callback_manager=callback_manager + llm, tools, callback_manager=callback_manager, **agent_kwargs ) elif agent_path is not None: agent_obj = load_agent( diff --git a/tests/unit_tests/agents/test_agent.py b/tests/unit_tests/agents/test_agent.py index b1a70ffb850..68fa1505afe 100644 --- a/tests/unit_tests/agents/test_agent.py +++ b/tests/unit_tests/agents/test_agent.py @@ -196,3 +196,29 @@ def test_agent_tool_return_direct() -> None: output = agent.run("when was langchain made") assert output == "misalignment" + + +def test_agent_with_new_prefix_suffix() -> None: + """Test agent initilization kwargs with new prefix and suffix.""" + fake_llm = FakeListLLM( + responses=["FooBarBaz\nAction: Search\nAction Input: misalignment"] + ) + tools = [ + Tool("Search", lambda x: x, "Useful for searching", return_direct=True), + ] + prefix = "FooBarBaz" + + suffix = "Begin now!\nInput: {input}\nThought: {agent_scratchpad}" + + agent = initialize_agent( + tools=tools, + llm=fake_llm, + agent="zero-shot-react-description", + agent_kwargs={"prefix": prefix, "suffix": suffix}, + ) + + # avoids "BasePromptTemplate" has no attribute "template" error + assert hasattr(agent.agent.llm_chain.prompt, "template") + prompt_str = agent.agent.llm_chain.prompt.template + assert prompt_str.startswith(prefix), "Prompt does not start with prefix" + assert prompt_str.endswith(suffix), "Prompt does not end with suffix"