From 961ddd49ed498df7ffaa6f6d688f7214b883b34f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 1 Apr 2024 17:40:55 -0700 Subject: [PATCH] Add test and notebook for using pydantic base model as state object --- examples/state-model.ipynb | 446 +++++++++++++++++++++++++++++++++++++ langgraph/graph/state.py | 5 + langgraph/pregel/read.py | 21 +- tests/test_pregel.py | 134 +++++++++++ 4 files changed, 603 insertions(+), 3 deletions(-) create mode 100644 examples/state-model.ipynb diff --git a/examples/state-model.ipynb b/examples/state-model.ipynb new file mode 100644 index 000000000..931d3c30a --- /dev/null +++ b/examples/state-model.ipynb @@ -0,0 +1,446 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "51466c8d-8ce4-4b3d-be4e-18fdbeda5f53", + "metadata": {}, + "source": [ + "# Pydantic Base Model as State\n", + "\n", + "In this example we will build a chat executor which uses a pydantic base model as the state object. This means all nodes receive an instance of the model as their first arg, and validation is run before each node executes." + ] + }, + { + "cell_type": "markdown", + "id": "7cbd446a-808f-4394-be92-d45ab818953c", + "metadata": {}, + "source": [ + "## Setup\n", + "\n", + "First we need to install the packages required" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "af4ce0ba-7596-4e5f-8bf8-0b0bd6e62833", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.2\u001b[0m\n", + "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n" + ] + } + ], + "source": [ + "!pip install --quiet -U langchain langchain_openai tavily-python" + ] + }, + { + "cell_type": "markdown", + "id": "0abe11f4-62ed-4dc4-8875-3db21e260d1d", + "metadata": {}, + "source": [ + "Next, we need to set API keys for OpenAI (the LLM we will use) and Tavily (the search tool we will use)" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c903a1cf-2977-4e2d-ad7d-8b3946821d89", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OpenAI API Key: ········\n", + "Tavily API Key: ········\n" + ] + } + ], + "source": [ + "import os\n", + "import getpass\n", + "\n", + "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")\n", + "os.environ[\"TAVILY_API_KEY\"] = getpass.getpass(\"Tavily API Key:\")" + ] + }, + { + "cell_type": "markdown", + "id": "f0ed46a8-effe-4596-b0e1-a6a29ee16f5c", + "metadata": {}, + "source": [ + "Optionally, we can set API key for [LangSmith tracing](https://smith.langchain.com/), which will give us best-in-class observability." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "95e25aec-7c9f-4a63-b143-225d0e9a79c3", + "metadata": {}, + "outputs": [], + "source": [ + "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n", + "os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass(\"LangSmith API Key:\")" + ] + }, + { + "cell_type": "markdown", + "id": "21ac643b-cb06-4724-a80c-2862ba4773f1", + "metadata": {}, + "source": [ + "## Set up the tools\n", + "\n", + "We will first define the tools we want to use.\n", + "For this simple example, we will use a built-in search tool via Tavily.\n", + "However, it is really easy to create your own tools - see documentation [here](https://python.langchain.com/docs/modules/agents/tools/custom_tools) on how to do that.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "d7ef57dd-5d6e-4ad3-9377-a92201c1310e", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_community.tools.tavily_search import TavilySearchResults\n", + "\n", + "tools = [TavilySearchResults(max_results=1)]" + ] + }, + { + "cell_type": "markdown", + "id": "01885785-b71a-44d1-b1d6-7b5b14d53b58", + "metadata": {}, + "source": [ + "We can now wrap these tools in a simple ToolExecutor.\n", + "This is a real simple class that takes in a ToolInvocation and calls that tool, returning the output.\n", + "A ToolInvocation is any class with `tool` and `tool_input` attribute.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5cf3331e-ccb3-41c8-aeb9-a840a94d41e7", + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.prebuilt import ToolExecutor\n", + "\n", + "tool_executor = ToolExecutor(tools)" + ] + }, + { + "cell_type": "markdown", + "id": "5497ed70-fce3-47f1-9cad-46f912bad6a5", + "metadata": {}, + "source": [ + "## Set up the model\n", + "\n", + "Now we need to load the chat model we want to use.\n", + "Importantly, this should satisfy two criteria:\n", + "\n", + "1. It should work with messages. We will represent all agent state in the form of messages, so it needs to be able to work well with them.\n", + "2. It should work with OpenAI function calling. This means it should either be an OpenAI model or a model that exposes a similar interface.\n", + "\n", + "Note: these model requirements are not requirements for using LangGraph - they are just requirements for this one example.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "892b54b9-75f0-4804-9ed0-88b5e5532989", + "metadata": {}, + "outputs": [], + "source": [ + "from langchain_openai import ChatOpenAI\n", + "\n", + "# We will set streaming=True so that we can stream tokens\n", + "# See the streaming section for more information on this.\n", + "model = ChatOpenAI(temperature=0, streaming=True)" + ] + }, + { + "cell_type": "markdown", + "id": "a77995c0-bae2-4cee-a036-8688a90f05b9", + "metadata": {}, + "source": [ + "\n", + "After we've done this, we should make sure the model knows that it has these tools available to call.\n", + "We can do this by converting the LangChain tools into the format for OpenAI function calling, and then bind them to the model class.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "cd3cbae5-d92c-4559-a4aa-44721b80d107", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The function `format_tool_to_openai_function` was deprecated in LangChain 0.1.16 and will be removed in 0.2.0. Use langchain_core.utils.function_calling.convert_to_openai_function() instead.\n", + " warn_deprecated(\n" + ] + } + ], + "source": [ + "from langchain.tools.render import format_tool_to_openai_function\n", + "\n", + "functions = [format_tool_to_openai_function(t) for t in tools]\n", + "model = model.bind_functions(functions)" + ] + }, + { + "cell_type": "markdown", + "id": "8e8b9211-93d0-4ad5-aa7a-9c09099c53ff", + "metadata": {}, + "source": [ + "## Define the agent state\n", + "\n", + "The main type of graph in `langgraph` is the `StateGraph`.\n", + "This graph is parameterized by a state object that it passes around to each node.\n", + "Each node then returns operations to update that state.\n", + "These operations can either SET specific attributes on the state (e.g. overwrite the existing values) or ADD to the existing attribute.\n", + "Whether to set or add is denoted by annotating the state object you construct the graph with.\n", + "\n", + "For this example, the state we will track will just be a list of messages.\n", + "We want each node to just add messages to that list.\n", + "Therefore, we will use a `pydantic.BaseModel` with one key (`messages`) and annotate it so that the `messages` attribute is always added to.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "ea793afa-2eab-4901-910d-6eed90cd6564", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Annotated, Sequence\n", + "import operator\n", + "from langchain_core.messages import BaseMessage\n", + "from langchain_core.pydantic_v1 import BaseModel\n", + "\n", + "class AgentState(BaseModel):\n", + " messages: Annotated[Sequence[BaseMessage], operator.add]" + ] + }, + { + "cell_type": "markdown", + "id": "e03c5094-9297-4d19-a04e-3eedc75cefb4", + "metadata": {}, + "source": [ + "## Define the nodes\n", + "\n", + "We now need to define a few different nodes in our graph.\n", + "In `langgraph`, a node can be either a function or a [runnable](https://python.langchain.com/docs/expression_language/).\n", + "There are two main nodes we need for this:\n", + "\n", + "1. The agent: responsible for deciding what (if any) actions to take.\n", + "2. A function to invoke tools: if the agent decides to take an action, this node will then execute that action.\n", + "\n", + "We will also need to define some edges.\n", + "Some of these edges may be conditional.\n", + "The reason they are conditional is that based on the output of a node, one of several paths may be taken.\n", + "The path that is taken is not known until that node is run (the LLM decides).\n", + "\n", + "1. Conditional Edge: after the agent is called, we should either:\n", + " a. If the agent said to take an action, then the function to invoke tools should be called\n", + " b. If the agent said that it was finished, then it should finish\n", + "2. Normal Edge: after the tools are invoked, it should always go back to the agent to decide what to do next\n", + "\n", + "Let's define the nodes, as well as a function to decide how what conditional edge to take.\n", + "\n", + "**MODIFICATION**\n", + "\n", + "We define each node to receive the AgentState base model as its first argument." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "3b541bb9-900c-40d0-964d-7b5dfee30667", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "from langgraph.prebuilt import ToolInvocation\n", + "from langchain_core.messages import FunctionMessage\n", + "\n", + "\n", + "# Define the function that determines whether to continue or not\n", + "def should_continue(state):\n", + " messages = state.messages\n", + " last_message = messages[-1]\n", + " # If there is no function call, then we finish\n", + " if \"function_call\" not in last_message.additional_kwargs:\n", + " return \"end\"\n", + " # Otherwise if there is, we continue\n", + " else:\n", + " return \"continue\"\n", + "\n", + "\n", + "# Define the function that calls the model\n", + "def call_model(state):\n", + " messages = state.messages\n", + " response = model.invoke(messages)\n", + " # We return a list, because this will get added to the existing list\n", + " return {\"messages\": [response]}\n", + "\n", + "\n", + "# Define the function to execute tools\n", + "def call_tool(state):\n", + " messages = state.messages\n", + " # Based on the continue condition\n", + " # we know the last message involves a function call\n", + " last_message = messages[-1]\n", + " # We construct an ToolInvocation from the function_call\n", + " action = ToolInvocation(\n", + " tool=last_message.additional_kwargs[\"function_call\"][\"name\"],\n", + " tool_input=json.loads(\n", + " last_message.additional_kwargs[\"function_call\"][\"arguments\"]\n", + " ),\n", + " )\n", + " # We call the tool_executor and get back a response\n", + " response = tool_executor.invoke(action)\n", + " # We use the response to create a FunctionMessage\n", + " function_message = FunctionMessage(content=str(response), name=action.tool)\n", + " # We return a list, because this will get added to the existing list\n", + " return {\"messages\": [function_message]}" + ] + }, + { + "cell_type": "markdown", + "id": "ffd6e892-946c-4899-8cc0-7c9291c1f73b", + "metadata": {}, + "source": [ + "## Define the graph\n", + "\n", + "We can now put it all together and define the graph!" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "813ae66c-3b58-4283-a02a-36da72a2ab90", + "metadata": {}, + "outputs": [], + "source": [ + "from langgraph.graph import StateGraph, END\n", + "\n", + "# Define a new graph\n", + "workflow = StateGraph(AgentState)\n", + "\n", + "# Define the two nodes we will cycle between\n", + "workflow.add_node(\"agent\", call_model)\n", + "workflow.add_node(\"action\", call_tool)\n", + "\n", + "# Set the entrypoint as `agent`\n", + "# This means that this node is the first one called\n", + "workflow.set_entry_point(\"agent\")\n", + "\n", + "# We now add a conditional edge\n", + "workflow.add_conditional_edges(\n", + " # First, we define the start node. We use `agent`.\n", + " # This means these are the edges taken after the `agent` node is called.\n", + " \"agent\",\n", + " # Next, we pass in the function that will determine which node is called next.\n", + " should_continue,\n", + " # Finally we pass in a mapping.\n", + " # The keys are strings, and the values are other nodes.\n", + " # END is a special node marking that the graph should finish.\n", + " # What will happen is we will call `should_continue`, and then the output of that\n", + " # will be matched against the keys in this mapping.\n", + " # Based on which one it matches, that node will then be called.\n", + " {\n", + " # If `tools`, then we call the tool node.\n", + " \"continue\": \"action\",\n", + " # Otherwise we finish.\n", + " \"end\": END,\n", + " },\n", + ")\n", + "\n", + "# We now add a normal edge from `tools` to `agent`.\n", + "# This means that after `tools` is called, `agent` node is called next.\n", + "workflow.add_edge(\"action\", \"agent\")\n", + "\n", + "# Finally, we compile it!\n", + "# This compiles it into a LangChain Runnable,\n", + "# meaning you can use it as you would any other runnable\n", + "app = workflow.compile()" + ] + }, + { + "cell_type": "markdown", + "id": "547c3931-3dae-4281-ad4e-4b51305594d4", + "metadata": {}, + "source": [ + "## Use it!\n", + "\n", + "We can now use it!\n", + "This now exposes the [same interface](https://python.langchain.com/docs/expression_language/) as all other LangChain runnables." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "8edb04b9-40b6-46f1-a7a8-4b2d8aba7752", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'function_call': {'arguments': '{\"query\":\"weather in San Francisco\"}', 'name': 'tavily_search_results_json'}})]}}\n", + "{'action': {'messages': [FunctionMessage(content=\"[{'url': 'https://forecast.weather.gov/zipcity.php?inputstring=San francisco,CA', 'content': 'NOAA National Weather Service National Weather Service. Toggle navigation. HOME; FORECAST . Local; Graphical; Aviation; Marine; Rivers and Lakes; Hurricanes; Severe Weather; Fire Weather; ... San Francisco CA 37.77°N 122.41°W (Elev. 131 ft) Last Update: 1:27 am PDT Apr 1, 2024. Forecast Valid: 8am PDT Apr 1, 2024-6pm PDT Apr 7, 2024 .'}]\", name='tavily_search_results_json')]}}\n", + "{'agent': {'messages': [AIMessage(content='You can check the weather in San Francisco by visiting the [NOAA National Weather Service website](https://forecast.weather.gov/zipcity.php?inputstring=San%20francisco,CA).')]}}\n" + ] + } + ], + "source": [ + "from langchain_core.messages import HumanMessage\n", + "\n", + "inputs = {\"messages\": [HumanMessage(content=\"what is the weather in sf\")]}\n", + "for chunk in app.stream(inputs):\n", + " print(chunk)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "296c7456-da05-4326-95dc-47d6b312da9d", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 91cd52e9d..fe92c6a99 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -168,6 +168,11 @@ class CompiledStateGraph(CompiledGraph): ChannelRead( state_keys[0] if state_keys == ["__root__"] else state_keys, fresh=True, + mapper=( + None + if state_keys == ["__root__"] + else partial(_coerce_state, self.graph.schema) + ), ), ], ).pipe(node) diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index b7ce6e4af..578d22073 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -26,6 +26,8 @@ class ChannelRead(RunnableLambda): fresh: bool = False + mapper: Optional[Callable[[Any], Any]] = None + @property def config_specs(self) -> list[ConfigurableFieldSpec]: return [ @@ -38,9 +40,16 @@ class ChannelRead(RunnableLambda): ), ] - def __init__(self, channel: Union[str, list[str]], fresh: bool = False) -> None: + def __init__( + self, + channel: Union[str, list[str]], + *, + fresh: bool = False, + mapper: Optional[Callable[[Any], Any]] = None, + ) -> None: super().__init__(func=self._read, afunc=self._aread) self.fresh = fresh + self.mapper = mapper self.channel = channel self.name = f"ChannelRead<{channel}>" @@ -52,7 +61,10 @@ class ChannelRead(RunnableLambda): f"Runnable {self} is not configured with a read function" "Make sure to call in the context of a Pregel process" ) - return read(self.channel, self.fresh) + if self.mapper: + return self.mapper(read(self.channel, self.fresh)) + else: + return read(self.channel, self.fresh) async def _aread(self, _: Any, config: RunnableConfig) -> Any: try: @@ -62,7 +74,10 @@ class ChannelRead(RunnableLambda): f"Runnable {self} is not configured with a read function" "Make sure to call in the context of a Pregel process" ) - return read(self.channel, self.fresh) + if self.mapper: + return self.mapper(read(self.channel, self.fresh)) + else: + return read(self.channel, self.fresh) default_bound: RunnablePassthrough = RunnablePassthrough() diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 90197f1dd..560423936 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2964,6 +2964,140 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] +@pytest.mark.parametrize( + "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] +) +def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( + checkpoint_at: CheckpointAt, +) -> None: + from langchain_core.pydantic_v1 import BaseModel + + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(BaseModel): + query: str + answer: Optional[str] = None + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data.query}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data.query}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + return {"answer": ",".join(data.docs)} + + def decider(data: State) -> str: + assert isinstance(data, State) + return "retriever_two" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges( + "rewrite_query", decider, {"retriever_two": "retriever_two"} + ) + workflow.add_edge(["retriever_one", "retriever_two"], "qa") + workflow.set_finish_point("qa") + + app = workflow.compile() + + assert app.get_graph().draw_ascii() == ( + """ +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + ** ** + ** ** + ** ** ++--------------+ +-----------------------+ +| analyzer_one | | rewrite_query_decider | ++--------------+ +-----------------------+ + * * + * * + * * ++---------------+ +---------------+ +| retriever_one | | retriever_two | ++---------------+ +---------------+ + ** ** + ** ** + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ """ + ) + + assert app.invoke({"query": "what is weather in sf"}) == { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } + + assert [*app.stream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(at=checkpoint_at), + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + ] + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + @pytest.mark.parametrize( "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP] )