From 554e7639946507a03f507ab071652ca33a017757 Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Mon, 27 Jan 2025 14:03:01 -0500 Subject: [PATCH] langgraph: use 'prompt' param for model input preprocessing in create_react_agent (#3173) --- .../langgraph/prebuilt/chat_agent_executor.py | 152 ++++++++---------- libs/langgraph/tests/test_prebuilt.py | 151 ++++++++--------- 2 files changed, 146 insertions(+), 157 deletions(-) diff --git a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py index 005ae59d3..17b9652bb 100644 --- a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py @@ -1,5 +1,7 @@ +import functools import inspect from typing import ( + Any, Callable, Literal, Optional, @@ -21,7 +23,6 @@ from langchain_core.tools import BaseTool from pydantic import BaseModel from typing_extensions import Annotated, TypedDict -from langgraph._api.deprecation import deprecated_parameter from langgraph.errors import ErrorCode, create_error_message from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph @@ -35,6 +36,7 @@ from langgraph.utils.runnable import RunnableCallable StructuredResponse = Union[dict, BaseModel] StructuredResponseSchema = Union[dict, type[BaseModel]] +F = TypeVar("F", bound=Callable[..., Any]) # We create the AgentState that we will pass around @@ -56,7 +58,7 @@ class AgentState(TypedDict): StateSchema = TypeVar("StateSchema", bound=AgentState) StateSchemaType = Type[StateSchema] -STATE_MODIFIER_RUNNABLE_NAME = "StateModifier" +PROMPT_RUNNABLE_NAME = "Prompt" MessagesModifier = Union[ SystemMessage, @@ -65,7 +67,7 @@ MessagesModifier = Union[ Runnable[Sequence[BaseMessage], Sequence[BaseMessage]], ] -StateModifier = Union[ +Prompt = Union[ SystemMessage, str, Callable[[StateSchema], Sequence[BaseMessage]], @@ -73,81 +75,84 @@ StateModifier = Union[ ] -def _get_state_modifier_runnable( - state_modifier: Optional[StateModifier], store: Optional[BaseStore] = None -) -> Runnable: - state_modifier_runnable: Runnable - if state_modifier is None: - state_modifier_runnable = RunnableCallable( - lambda state: state["messages"], name=STATE_MODIFIER_RUNNABLE_NAME +def _get_prompt_runnable(prompt: Optional[Prompt]) -> Runnable: + prompt_runnable: Runnable + if prompt is None: + prompt_runnable = RunnableCallable( + lambda state: state["messages"], name=PROMPT_RUNNABLE_NAME ) - elif isinstance(state_modifier, str): - _system_message: BaseMessage = SystemMessage(content=state_modifier) - state_modifier_runnable = RunnableCallable( + elif isinstance(prompt, str): + _system_message: BaseMessage = SystemMessage(content=prompt) + prompt_runnable = RunnableCallable( lambda state: [_system_message] + state["messages"], - name=STATE_MODIFIER_RUNNABLE_NAME, + name=PROMPT_RUNNABLE_NAME, ) - elif isinstance(state_modifier, SystemMessage): - state_modifier_runnable = RunnableCallable( - lambda state: [state_modifier] + state["messages"], - name=STATE_MODIFIER_RUNNABLE_NAME, + elif isinstance(prompt, SystemMessage): + prompt_runnable = RunnableCallable( + lambda state: [prompt] + state["messages"], + name=PROMPT_RUNNABLE_NAME, ) - elif inspect.iscoroutinefunction(state_modifier): - state_modifier_runnable = RunnableCallable( + elif inspect.iscoroutinefunction(prompt): + prompt_runnable = RunnableCallable( None, - state_modifier, - name=STATE_MODIFIER_RUNNABLE_NAME, + prompt, + name=PROMPT_RUNNABLE_NAME, ) - elif callable(state_modifier): - state_modifier_runnable = RunnableCallable( - state_modifier, - name=STATE_MODIFIER_RUNNABLE_NAME, + elif callable(prompt): + prompt_runnable = RunnableCallable( + prompt, + name=PROMPT_RUNNABLE_NAME, ) - elif isinstance(state_modifier, Runnable): - state_modifier_runnable = state_modifier + elif isinstance(prompt, Runnable): + prompt_runnable = prompt else: - raise ValueError( - f"Got unexpected type for `state_modifier`: {type(state_modifier)}" - ) + raise ValueError(f"Got unexpected type for `prompt`: {type(prompt)}") - return state_modifier_runnable + return prompt_runnable -def _convert_messages_modifier_to_state_modifier( +def _convert_messages_modifier_to_prompt( messages_modifier: MessagesModifier, -) -> StateModifier: - state_modifier: StateModifier +) -> Prompt: + prompt: Prompt if isinstance(messages_modifier, (str, SystemMessage)): return messages_modifier elif callable(messages_modifier): - def state_modifier(state: AgentState) -> Sequence[BaseMessage]: + def prompt(state: AgentState) -> Sequence[BaseMessage]: return messages_modifier(state["messages"]) - return state_modifier + return prompt elif isinstance(messages_modifier, Runnable): - state_modifier = (lambda state: state["messages"]) | messages_modifier - return state_modifier + prompt = (lambda state: state["messages"]) | messages_modifier + return prompt raise ValueError( f"Got unexpected type for `messages_modifier`: {type(messages_modifier)}" ) -def _get_model_preprocessing_runnable( - state_modifier: Optional[StateModifier], - messages_modifier: Optional[MessagesModifier], - store: Optional[BaseStore], -) -> Runnable: - # Add the state or message modifier, if exists - if state_modifier is not None and messages_modifier is not None: - raise ValueError( - "Expected value for either state_modifier or messages_modifier, got values for both" - ) +def _convert_modifier_to_prompt(func: F) -> F: + """Decorator that converts state_modifier/messages_modifier kwargs to prompt kwarg.""" - if state_modifier is None and messages_modifier is not None: - state_modifier = _convert_messages_modifier_to_state_modifier(messages_modifier) + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + prompt = kwargs.get("prompt") + state_modifier = kwargs.pop("state_modifier", None) + messages_modifier = kwargs.pop("messages_modifier", None) + if sum(p is not None for p in (prompt, state_modifier, messages_modifier)) > 1: + raise ValueError( + "Expected only one of prompt, state_modifier, or messages_modifier, got multiple values" + ) - return _get_state_modifier_runnable(state_modifier, store) + if state_modifier is not None: + prompt = state_modifier + elif messages_modifier is not None: + prompt = _convert_messages_modifier_to_prompt(messages_modifier) + + kwargs["prompt"] = prompt + return func(*args, **kwargs) + + return cast(F, wrapper) def _should_bind_tools(model: LanguageModelLike, tools: Sequence[BaseTool]) -> bool: @@ -228,14 +233,13 @@ def _validate_chat_history( raise ValueError(error_message) -@deprecated_parameter("messages_modifier", "0.1.9", "state_modifier", removal="0.3.0") +@_convert_modifier_to_prompt def create_react_agent( model: Union[str, LanguageModelLike], tools: Union[ToolExecutor, Sequence[BaseTool], ToolNode], *, state_schema: Optional[StateSchemaType] = None, - messages_modifier: Optional[MessagesModifier] = None, - state_modifier: Optional[StateModifier] = None, + prompt: Optional[Prompt] = None, response_format: Optional[ Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]] ] = None, @@ -254,26 +258,15 @@ def create_react_agent( state_schema: An optional state schema that defines graph state. Must have `messages` and `is_last_step` keys. Defaults to `AgentState` that defines those two keys. - messages_modifier: An optional - messages modifier. This applies to messages BEFORE they are passed into the LLM. + prompt: An optional prompt for the LLM. Can take a few different forms: - Can take a few different forms: - - - SystemMessage: this is added to the beginning of the list of messages. - - str: This is converted to a SystemMessage and added to the beginning of the list of messages. - - Callable: This function should take in a list of messages and the output is then passed to the language model. - - Runnable: This runnable should take in a list of messages and the output is then passed to the language model. - !!! Warning - `messages_modifier` parameter is deprecated as of version 0.1.9 and will be removed in 0.2.0 - state_modifier: An optional - state modifier. This takes full graph state BEFORE the LLM is called and prepares the input to LLM. - - Can take a few different forms: - - - SystemMessage: this is added to the beginning of the list of messages in state["messages"]. - str: This is converted to a SystemMessage and added to the beginning of the list of messages in state["messages"]. + - SystemMessage: this is added to the beginning of the list of messages in state["messages"]. - Callable: This function should take in full graph state and the output is then passed to the language model. - Runnable: This runnable should take in full graph state and the output is then passed to the language model. + + !!! Note + Prior to `v0.2.68`, the prompt was set using `state_modifier` / `messages_modifier` parameters. response_format: An optional schema for the final agent output. If provided, output will be formatted to match the given schema and returned in the 'structured_response' state key. @@ -389,7 +382,7 @@ def create_react_agent( ```pycon >>> system_prompt = "You are a helpful bot named Fred." - >>> graph = create_react_agent(model, tools, state_modifier=system_prompt) + >>> graph = create_react_agent(model, tools, prompt=system_prompt) >>> inputs = {"messages": [("user", "What's your name? And what's the weather in SF?")]} >>> for s in graph.stream(inputs, stream_mode="values"): ... message = s["messages"][-1] @@ -421,11 +414,8 @@ def create_react_agent( ... ("placeholder", "{messages}"), ... ("user", "Remember, always be polite!"), ... ]) - >>> def format_for_model(state: AgentState): - ... # You can do more complex modifications here - ... return prompt.invoke({"messages": state["messages"]}) >>> - >>> graph = create_react_agent(model, tools, state_modifier=format_for_model) + >>> graph = create_react_agent(model, tools, prompt=prompt) >>> inputs = {"messages": [("user", "What's your name? And what's the weather in SF?")]} >>> for s in graph.stream(inputs, stream_mode="values"): ... message = s["messages"][-1] @@ -453,7 +443,7 @@ def create_react_agent( ... messages: Annotated[list[BaseMessage], add_messages] ... is_last_step: IsLastStep >>> - >>> graph = create_react_agent(model, tools, state_schema=CustomState, state_modifier=prompt) + >>> graph = create_react_agent(model, tools, state_schema=CustomState, prompt=prompt) >>> inputs = {"messages": [("user", "What's today's date? And what's the weather in SF?")], "today": "July 16, 2004"} >>> for s in graph.stream(inputs, stream_mode="values"): ... message = s["messages"][-1] @@ -539,7 +529,7 @@ def create_react_agent( >>> from langgraph.checkpoint.memory import MemorySaver >>> from langgraph.store.memory import InMemoryStore >>> store = InMemoryStore() - >>> graph = create_react_agent(model, [save_memory], state_modifier=prepare_model_inputs, store=store, checkpointer=MemorySaver()) + >>> graph = create_react_agent(model, [save_memory], prompt=prepare_model_inputs, store=store, checkpointer=MemorySaver()) >>> config = {"configurable": {"thread_id": "thread-1", "user_id": "1"}} >>> inputs = {"messages": [("user", "Hey I'm Will, how's it going?")]} @@ -619,11 +609,7 @@ def create_react_agent( if _should_bind_tools(model, tool_classes) and tool_calling_enabled: model = cast(BaseChatModel, model).bind_tools(tool_classes) - # we're passing store here for validation - preprocessor = _get_model_preprocessing_runnable( - state_modifier, messages_modifier, store - ) - model_runnable = preprocessor | model + model_runnable = _get_prompt_runnable(prompt) | model # If any of the tools are configured to return_directly after running, # our graph needs to check if these were called diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 04ce3ee69..6a5327ca6 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -147,7 +147,7 @@ class FakeToolCallingModel(BaseChatModel): @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> None: +def test_no_prompt(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer: BaseCheckpointSaver = request.getfixturevalue( "checkpointer_" + checkpointer_name ) @@ -181,7 +181,7 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_no_modifier_async(checkpointer_name: str) -> None: +async def test_no_prompt_async(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: model = FakeToolCallingModel() @@ -217,16 +217,17 @@ def test_passing_two_modifiers(): with pytest.raises(ValueError): create_react_agent(model, [], messages_modifier="Foo", state_modifier="Bar") + with pytest.raises(ValueError): + create_react_agent(model, [], messages_modifier="Foo", prompt="Bar") -def test_system_message_modifier(): - messages_modifier = SystemMessage(content="Foo") - agent_1 = create_react_agent( - FakeToolCallingModel(), [], messages_modifier=messages_modifier - ) - agent_2 = create_react_agent( - FakeToolCallingModel(), [], state_modifier=messages_modifier - ) - for agent in [agent_1, agent_2]: + +def test_system_message_prompt(): + prompt = SystemMessage(content="Foo") + for agent in ( + create_react_agent(FakeToolCallingModel(), [], prompt=prompt), + create_react_agent(FakeToolCallingModel(), [], messages_modifier=prompt), + create_react_agent(FakeToolCallingModel(), [], state_modifier=prompt), + ): inputs = [HumanMessage("hi?")] response = agent.invoke({"messages": inputs}) expected_response = { @@ -235,15 +236,13 @@ def test_system_message_modifier(): assert response == expected_response -def test_system_message_string_modifier(): - messages_modifier = "Foo" - agent_1 = create_react_agent( - FakeToolCallingModel(), [], messages_modifier=messages_modifier - ) - agent_2 = create_react_agent( - FakeToolCallingModel(), [], state_modifier=messages_modifier - ) - for agent in [agent_1, agent_2]: +def test_string_prompt(): + prompt = "Foo" + for agent in ( + create_react_agent(FakeToolCallingModel(), [], prompt=prompt), + create_react_agent(FakeToolCallingModel(), [], messages_modifier=prompt), + create_react_agent(FakeToolCallingModel(), [], state_modifier=prompt), + ): inputs = [HumanMessage("hi?")] response = agent.invoke({"messages": inputs}) expected_response = { @@ -252,63 +251,71 @@ def test_system_message_string_modifier(): assert response == expected_response -def test_callable_messages_modifier(): - model = FakeToolCallingModel() - +def test_callable_prompt(): def messages_modifier(messages): modified_message = f"Bar {messages[-1].content}" return [HumanMessage(content=modified_message)] - agent = create_react_agent(model, [], messages_modifier=messages_modifier) - inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}) - expected_response = {"messages": inputs + [AIMessage(content="Bar hi?", id="0")]} - assert response == expected_response - - -def test_callable_state_modifier(): - model = FakeToolCallingModel() - - def state_modifier(state): + def prompt(state): modified_message = f"Bar {state['messages'][-1].content}" return [HumanMessage(content=modified_message)] - agent = create_react_agent(model, [], state_modifier=state_modifier) - inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}) - expected_response = {"messages": inputs + [AIMessage(content="Bar hi?", id="0")]} - assert response == expected_response + for agent in ( + create_react_agent(FakeToolCallingModel(), [], prompt=prompt), + create_react_agent(FakeToolCallingModel(), [], state_modifier=prompt), + create_react_agent( + FakeToolCallingModel(), [], messages_modifier=messages_modifier + ), + ): + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = { + "messages": inputs + [AIMessage(content="Bar hi?", id="0")] + } + assert response == expected_response -def test_runnable_messages_modifier(): - model = FakeToolCallingModel() +async def test_callable_prompt_async(): + async def prompt(state): + modified_message = f"Bar {state['messages'][-1].content}" + return [HumanMessage(content=modified_message)] + for agent in ( + create_react_agent(FakeToolCallingModel(), [], prompt=prompt), + create_react_agent(FakeToolCallingModel(), [], state_modifier=prompt), + ): + inputs = [HumanMessage("hi?")] + response = await agent.ainvoke({"messages": inputs}) + expected_response = { + "messages": inputs + [AIMessage(content="Bar hi?", id="0")] + } + assert response == expected_response + + +def test_runnable_prompt(): messages_modifier = RunnableLambda( lambda messages: [HumanMessage(content=f"Baz {messages[-1].content}")] ) - - agent = create_react_agent(model, [], messages_modifier=messages_modifier) - inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}) - expected_response = {"messages": inputs + [AIMessage(content="Baz hi?", id="0")]} - assert response == expected_response - - -def test_runnable_state_modifier(): - model = FakeToolCallingModel() - - state_modifier = RunnableLambda( + prompt = RunnableLambda( lambda state: [HumanMessage(content=f"Baz {state['messages'][-1].content}")] ) - agent = create_react_agent(model, [], state_modifier=state_modifier) - inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}) - expected_response = {"messages": inputs + [AIMessage(content="Baz hi?", id="0")]} - assert response == expected_response + for agent in ( + create_react_agent(FakeToolCallingModel(), [], prompt=prompt), + create_react_agent(FakeToolCallingModel(), [], state_modifier=prompt), + create_react_agent( + FakeToolCallingModel(), [], messages_modifier=messages_modifier + ), + ): + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = { + "messages": inputs + [AIMessage(content="Baz hi?", id="0")] + } + assert response == expected_response -def test_state_modifier_with_store(): +def test_prompt_with_store(): def add(a: int, b: int): """Adds a and b""" return a + b @@ -317,20 +324,18 @@ def test_state_modifier_with_store(): in_memory_store.put(("memories", "1"), "user_name", {"data": "User name is Alice"}) in_memory_store.put(("memories", "2"), "user_name", {"data": "User name is Bob"}) - def modify(state, config, *, store): + def prompt(state, config, *, store): user_id = config["configurable"]["user_id"] system_str = store.get(("memories", user_id), "user_name").value["data"] return [SystemMessage(system_str)] + state["messages"] - def modify_no_store(state, config): + def prompt_no_store(state, config): return SystemMessage("foo") + state["messages"] model = FakeToolCallingModel() # test state modifier that uses store works - agent = create_react_agent( - model, [add], state_modifier=modify, store=in_memory_store - ) + agent = create_react_agent(model, [add], prompt=prompt, store=in_memory_store) response = agent.invoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "1"}} ) @@ -338,7 +343,7 @@ def test_state_modifier_with_store(): # test state modifier that doesn't use store works agent = create_react_agent( - model, [add], state_modifier=modify_no_store, store=in_memory_store + model, [add], prompt=prompt_no_store, store=in_memory_store ) response = agent.invoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "2"}} @@ -346,7 +351,7 @@ def test_state_modifier_with_store(): assert response["messages"][-1].content == "foo-hi" -async def test_state_modifier_with_store_async(): +async def test_prompt_with_store_async(): async def add(a: int, b: int): """Adds a and b""" return a + b @@ -359,22 +364,20 @@ async def test_state_modifier_with_store_async(): ("memories", "2"), "user_name", {"data": "User name is Bob"} ) - async def modify(state, config, *, store): + async def prompt(state, config, *, store): user_id = config["configurable"]["user_id"] system_str = (await store.aget(("memories", user_id), "user_name")).value[ "data" ] return [SystemMessage(system_str)] + state["messages"] - async def modify_no_store(state, config): + async def prompt_no_store(state, config): return SystemMessage("foo") + state["messages"] model = FakeToolCallingModel() # test state modifier that uses store works - agent = create_react_agent( - model, [add], state_modifier=modify, store=in_memory_store - ) + agent = create_react_agent(model, [add], prompt=prompt, store=in_memory_store) response = await agent.ainvoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "1"}} ) @@ -382,7 +385,7 @@ async def test_state_modifier_with_store_async(): # test state modifier that doesn't use store works agent = create_react_agent( - model, [add], state_modifier=modify_no_store, store=in_memory_store + model, [add], prompt=prompt_no_store, store=in_memory_store ) response = await agent.ainvoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "2"}} @@ -1663,7 +1666,7 @@ def test_react_agent_update_state(): } ) - def state_modifier(state: State): + def prompt(state: State): user_name = state.get("user_name") if user_name is None: return state["messages"] @@ -1678,7 +1681,7 @@ def test_react_agent_update_state(): model, [get_user_name], state_schema=State, - state_modifier=state_modifier, + prompt=prompt, checkpointer=checkpointer, ) config = {"configurable": {"thread_id": "1"}}