diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index 453e06438..4632e3488 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -1,9 +1,9 @@ import json -from typing import Annotated, Optional, Sequence, TypedDict, Union +from typing import Annotated, Callable, Optional, Sequence, TypedDict, Union from langchain_core.language_models import LanguageModelLike from langchain_core.messages import BaseMessage, FunctionMessage, SystemMessage -from langchain_core.runnables import RunnableLambda +from langchain_core.runnables import Runnable, RunnableLambda from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_function @@ -137,7 +137,7 @@ def create_function_calling_executor( def create_tool_calling_executor( model: LanguageModelLike, tools: Union[ToolExecutor, Sequence[BaseTool]], - system_message: Optional[Union[str, SystemMessage]] = None, + messages_modifier: Optional[Union[SystemMessage, str, Callable, Runnable]] = None, checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, @@ -148,8 +148,13 @@ def create_tool_calling_executor( Args: model (LanguageModelLike): The chat model that supports OpenAI tool calling. tools (Union[ToolExecutor, Sequence[BaseTool]]): A list of tools or a ToolExecutor instance. - system_message: (Optional[Union[str, SystemMessage]]): An optional system message to pass in - to the model. Is appended at the start of the messages. + messages_modifier: (Optional[Union[SystemMessage, str, Callable, Runnable]]): An optional + messages modifier. This applies to messages BEFORE they are passed into the LLM. + 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. checkpointer (Optional[BaseCheckpointSaver]): An optional checkpoint saver object. interrupt_before (Optional[Sequence[str]]): An optional list of node names to interrupt before. interrupt_after (Optional[Sequence[str]]): An optional list of node names to interrupt after. @@ -192,22 +197,31 @@ def create_tool_calling_executor( else: return "continue" + # Add the message modifier, if exists + if messages_modifier is None: + model_runnable = model + elif isinstance(messages_modifier, str): + _system_message: BaseMessage = SystemMessage(content=messages_modifier) + model_runnable = (lambda messages: [_system_message] + messages) | model + elif isinstance(messages_modifier, SystemMessage): + model_runnable = (lambda messages: [messages_modifier] + messages) | model + elif isinstance(messages_modifier, (Callable, Runnable)): + model_runnable = messages_modifier | model + else: + raise ValueError( + f"Got unexpected type for `messages_modifier`: {type(messages_modifier)}" + ) + # Define the function that calls the model def call_model(state: AgentState): messages = state["messages"] - if system_message is not None: - if isinstance(system_message, str): - _system_message: BaseMessage = SystemMessage(content=system_message) - else: - _system_message = system_message - messages = [_system_message] + list(messages) - response = model.invoke(messages) + response = model_runnable.invoke(messages) # We return a list, because this will get added to the existing list return {"messages": [response]} async def acall_model(state: AgentState): messages = state["messages"] - response = await model.ainvoke(messages) + response = await model_runnable.ainvoke(messages) # We return a list, because this will get added to the existing list return {"messages": [response]} diff --git a/tests/test_prebuilt.py b/tests/test_prebuilt.py new file mode 100644 index 000000000..57a1d1159 --- /dev/null +++ b/tests/test_prebuilt.py @@ -0,0 +1,97 @@ +from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union + +from langchain_core.callbacks import ( + CallbackManagerForLLMRun, +) +from langchain_core.language_models import ( + BaseChatModel, + LanguageModelInput, +) +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_core.pydantic_v1 import BaseModel +from langchain_core.runnables import Runnable, RunnableLambda +from langchain_core.tools import BaseTool + +from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor + + +class FakeToolCallingModel(BaseChatModel): + def _generate( + self, + messages: List[BaseMessage], + stop: Optional[List[str]] = None, + run_manager: Optional[CallbackManagerForLLMRun] = None, + **kwargs: Any, + ) -> ChatResult: + """Top Level call""" + messages_string = "-".join([m.content for m in messages]) + message = AIMessage(content=messages_string, id="0") + return ChatResult(generations=[ChatGeneration(message=message)]) + + @property + def _llm_type(self) -> str: + return "fake-tool-call-model" + + def bind_tools( + self, + tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]], + **kwargs: Any, + ) -> Runnable[LanguageModelInput, BaseMessage]: + if len(tools) > 0: + raise ValueError("Not supported yet!") + return self + + +def test_no_modifier(): + model = FakeToolCallingModel() + agent = create_tool_calling_executor(model, []) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} + assert response == expected_response + + +def test_system_message_modifier(): + model = FakeToolCallingModel() + messages_modifier = SystemMessage(content="Foo") + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Foo-hi?", id="0")]} + assert response == expected_response + + +def test_system_message_string_modifier(): + model = FakeToolCallingModel() + messages_modifier = "Foo" + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Foo-hi?", id="0")]} + assert response == expected_response + + +def test_callable_modifier(): + model = FakeToolCallingModel() + + def messages_modifier(messages): + return [HumanMessage(content="Bar")] + + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Bar", id="0")]} + assert response == expected_response + + +def test_runnable_modifier(): + model = FakeToolCallingModel() + + messages_modifier = RunnableLambda(lambda x: [HumanMessage(content="Baz")]) + + agent = create_tool_calling_executor(model, [], messages_modifier=messages_modifier) + inputs = [HumanMessage("hi?")] + response = agent.invoke({"messages": inputs}) + expected_response = {"messages": inputs + [AIMessage(content="Baz", id="0")]} + assert response == expected_response