Skip to content

Commit 54f9e42

Browse files
jxnlJason Liu
andauthored
Pass kwargs from initialize_agent into agent classmethod (#799)
# Problem I noticed that in order to change the prefix of the prompt in the `zero-shot-react-description` agent we had to dig around to subset strings deep into the agent's attributes. It requires the user to inspect a long chain of attributes and classes. `initialize_agent -> AgentExecutor -> Agent -> LLMChain -> Prompt from Agent.create_prompt` ``` python agent = initialize_agent( tools=tools, llm=fake_llm, agent="zero-shot-react-description" ) prompt_str = agent.agent.llm_chain.prompt.template new_prompt_str = change_prefix(prompt_str) agent.agent.llm_chain.prompt.template = new_prompt_str ``` # Implemented Solution `initialize_agent` accepts `**kwargs` but passes it to `AgentExecutor` but not `ZeroShotAgent`, by simply giving the kwargs to the agent class methods we can support changing the prefix and suffix for one agent while allowing future agents to take advantage of `initialize_agent`. ``` agent = initialize_agent( tools=tools, llm=fake_llm, agent="zero-shot-react-description", agent_kwargs={"prefix": prefix, "suffix": suffix} ) ``` To be fair, this was before finding docs around custom agents here: https://langchain.readthedocs.io/en/latest/modules/agents/examples/custom_agent.html?highlight=custom%20#custom-llmchain but i find that my use case just needed to change the prefix a little. # Changes * Pass kwargs to Agent class method * Added a test to check suffix and prefix --------- Co-authored-by: Jason Liu <[email protected]>
1 parent c331009 commit 54f9e42

File tree

2 files changed

+29
-1
lines changed

2 files changed

+29
-1
lines changed

langchain/agents/initialize.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ def initialize_agent(
1414
agent: Optional[str] = None,
1515
callback_manager: Optional[BaseCallbackManager] = None,
1616
agent_path: Optional[str] = None,
17+
agent_kwargs: Optional[dict] = None,
1718
**kwargs: Any,
1819
) -> AgentExecutor:
1920
"""Load agent given tools and LLM.
@@ -50,8 +51,9 @@ def initialize_agent(
5051
f"Valid types are: {AGENT_TO_CLASS.keys()}."
5152
)
5253
agent_cls = AGENT_TO_CLASS[agent]
54+
agent_kwargs = agent_kwargs or {}
5355
agent_obj = agent_cls.from_llm_and_tools(
54-
llm, tools, callback_manager=callback_manager
56+
llm, tools, callback_manager=callback_manager, **agent_kwargs
5557
)
5658
elif agent_path is not None:
5759
agent_obj = load_agent(

tests/unit_tests/agents/test_agent.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,3 +196,29 @@ def test_agent_tool_return_direct() -> None:
196196

197197
output = agent.run("when was langchain made")
198198
assert output == "misalignment"
199+
200+
201+
def test_agent_with_new_prefix_suffix() -> None:
202+
"""Test agent initilization kwargs with new prefix and suffix."""
203+
fake_llm = FakeListLLM(
204+
responses=["FooBarBaz\nAction: Search\nAction Input: misalignment"]
205+
)
206+
tools = [
207+
Tool("Search", lambda x: x, "Useful for searching", return_direct=True),
208+
]
209+
prefix = "FooBarBaz"
210+
211+
suffix = "Begin now!\nInput: {input}\nThought: {agent_scratchpad}"
212+
213+
agent = initialize_agent(
214+
tools=tools,
215+
llm=fake_llm,
216+
agent="zero-shot-react-description",
217+
agent_kwargs={"prefix": prefix, "suffix": suffix},
218+
)
219+
220+
# avoids "BasePromptTemplate" has no attribute "template" error
221+
assert hasattr(agent.agent.llm_chain.prompt, "template")
222+
prompt_str = agent.agent.llm_chain.prompt.template
223+
assert prompt_str.startswith(prefix), "Prompt does not start with prefix"
224+
assert prompt_str.endswith(suffix), "Prompt does not end with suffix"

0 commit comments

Comments
 (0)