{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e327e9bd-effc-4bee-a875-1c383c17f43d",
   "metadata": {},
   "source": [
    "# Complex data extraction with function calling\n",
    "\n",
    "Function calling is a core primitive for integrating LLMs within your software stack. We use it throughout the LangGraph docs, since developing with function calling (aka tool usage) tends to be much more stress-free than the traditional way of writing custom string parsers.\n",
    "\n",
    "However, even GPT-4, Opus, and other powerful models still struggle with complex functions, especially if your schema involves any nesting or if you have more advanced data validation rules.\n",
    "\n",
    "There are three basic ways to increase reliability: better prompting, constrained decoding, and **validation with re-prompting**.\n",
    "\n",
    "We will cover two approaches to the last technique here, since it is generally applicable across any LLM that supports tool calling.\n",
    "\n",
    "## Regular Extraction with Retries\n",
    "\n",
    "Both examples here invoke a simple looping graph that takes following approach:\n",
    "1. Prompt the LLM to respond.\n",
    "2. If it responds with tool calls, validate those.\n",
    "3. If the calls are correct, return. Otherwise, format the validation error as a new [ToolMessage](https://api.python.langchain.com/en/latest/messages/langchain_core.messages.tool.ToolMessage.html#langchain_core.messages.tool.ToolMessage) and prompt the LLM to fix the errors. Taking us back to step (1).\n",
    "\n",
    "\n",
    "The techniques differ only on step (3). In this first step, we will prompt the original LLM to regenerate the function calls to fix the validation errors. In the next section, we will instead prompt the LLM to generate a **patch** to fix the errors, meaning it doesn't have to re-generate data that is valid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0ada5e8f-3f2f-459e-83aa-6cd8861770dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langchain-anthropic langgraph\n",
    "# Or do langchain-{groq|openai|etc.} for another package with tool calling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27b25a1c-f437-482a-97d3-c7f168986df5",
   "metadata": {},
   "source": [
    "Set up your environment. If you are using groq, anthropic, etc., you will need to update different API keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c0acb818-b6fd-48ab-97e6-fc2de2d03e87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")\n",
    "# Recommended to visualize the retry steps\n",
    "_set_env(\"LANGCHAIN_API_KEY\")\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Extraction Notebook\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6973d34-561c-410c-9362-25f55eaf2c3e",
   "metadata": {},
   "source": [
    "### Define the Validator + Retry Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "baf669a0-04ee-492d-80d8-8fcb658ed128",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "import uuid\n",
    "from typing import (\n",
    "    Annotated,\n",
    "    Any,\n",
    "    Callable,\n",
    "    Dict,\n",
    "    List,\n",
    "    Literal,\n",
    "    Optional,\n",
    "    Sequence,\n",
    "    Type,\n",
    "    Union,\n",
    ")\n",
    "\n",
    "from langchain_core.language_models import BaseChatModel\n",
    "from langchain_core.messages import (\n",
    "    AIMessage,\n",
    "    AnyMessage,\n",
    "    BaseMessage,\n",
    "    HumanMessage,\n",
    "    ToolCall,\n",
    ")\n",
    "from langchain_core.prompt_values import PromptValue\n",
    "from langchain_core.runnables import (\n",
    "    Runnable,\n",
    "    RunnableLambda,\n",
    ")\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langgraph.prebuilt import ValidationNode\n",
    "\n",
    "\n",
    "def _default_aggregator(messages: Sequence[AnyMessage]) -> AIMessage:\n",
    "    for m in messages[::-1]:\n",
    "        if m.type == \"ai\":\n",
    "            return m\n",
    "    raise ValueError(\"No AI message found in the sequence.\")\n",
    "\n",
    "\n",
    "class RetryStrategy(TypedDict, total=False):\n",
    "    \"\"\"The retry strategy for a tool call.\"\"\"\n",
    "\n",
    "    max_attempts: int\n",
    "    \"\"\"The maximum number of attempts to make.\"\"\"\n",
    "    fallback: Optional[\n",
    "        Union[\n",
    "            Runnable[Sequence[AnyMessage], AIMessage],\n",
    "            Runnable[Sequence[AnyMessage], BaseMessage],\n",
    "            Callable[[Sequence[AnyMessage]], AIMessage],\n",
    "        ]\n",
    "    ]\n",
    "    \"\"\"The function to use once validation fails.\"\"\"\n",
    "    aggregate_messages: Optional[Callable[[Sequence[AnyMessage]], AIMessage]]\n",
    "\n",
    "\n",
    "def _bind_validator_with_retries(\n",
    "    llm: Union[\n",
    "        Runnable[Sequence[AnyMessage], AIMessage],\n",
    "        Runnable[Sequence[BaseMessage], BaseMessage],\n",
    "    ],\n",
    "    *,\n",
    "    validator: ValidationNode,\n",
    "    retry_strategy: RetryStrategy,\n",
    "    tool_choice: Optional[str] = None,\n",
    ") -> Runnable[Union[List[AnyMessage], PromptValue], AIMessage]:\n",
    "    \"\"\"Binds a tool validators + retry logic to create a runnable validation graph.\n",
    "\n",
    "    LLMs that support tool calling can generate structured JSON. However, they may not always\n",
    "    perfectly follow your requested schema, especially if the schema is nested or has complex\n",
    "    validation rules. This method allows you to bind a validation function to the LLM's output,\n",
    "    so that any time the LLM generates a message, the validation function is run on it. If\n",
    "    the validation fails, the method will retry the LLM with a fallback strategy, the simplest\n",
    "    being just to add a message to the output with the validation errors and a request to fix them.\n",
    "\n",
    "    The resulting runnable expects a list of messages as input and returns a single AI message.\n",
    "    By default, the LLM can optionally NOT invoke tools, making this easier to incorporate into\n",
    "    your existing chat bot. You can specify a tool_choice to force the validator to be run on\n",
    "    the outputs.\n",
    "\n",
    "    Args:\n",
    "        llm (Runnable): The llm that will generate the initial messages (and optionally fallba)\n",
    "        validator (ValidationNode): The validation logic.\n",
    "        retry_strategy (RetryStrategy): The retry strategy to use.\n",
    "            Possible keys:\n",
    "            - max_attempts: The maximum number of attempts to make.\n",
    "            - fallback: The LLM or function to use in case of validation failure.\n",
    "            - aggregate_messages: A function to aggregate the messages over multiple turns.\n",
    "                Defaults to fetching the last AI message.\n",
    "        tool_choice: If provided, always run the validator on the tool output.\n",
    "\n",
    "    Returns:\n",
    "        Runnable: A runnable that can be invoked with a list of messages and returns a single AI message.\n",
    "    \"\"\"\n",
    "\n",
    "    def add_or_overwrite_messages(left: list, right: Union[list, dict]) -> list:\n",
    "        \"\"\"Append messages. If the update is a 'finalized' output, replace the whole list.\"\"\"\n",
    "        if isinstance(right, dict) and \"finalize\" in right:\n",
    "            finalized = right[\"finalize\"]\n",
    "            if not isinstance(finalized, list):\n",
    "                finalized = [finalized]\n",
    "            for m in finalized:\n",
    "                if m.id is None:\n",
    "                    m.id = str(uuid.uuid4())\n",
    "            return finalized\n",
    "        res = add_messages(left, right)\n",
    "        if not isinstance(res, list):\n",
    "            return [res]\n",
    "        return res\n",
    "\n",
    "    class State(TypedDict):\n",
    "        messages: Annotated[list, add_or_overwrite_messages]\n",
    "        attempt_number: Annotated[int, operator.add]\n",
    "        initial_num_messages: int\n",
    "        input_format: Literal[\"list\", \"dict\"]\n",
    "\n",
    "    builder = StateGraph(State)\n",
    "\n",
    "    def dedict(x: State) -> list:\n",
    "        \"\"\"Get the messages from the state.\"\"\"\n",
    "        return x[\"messages\"]\n",
    "\n",
    "    model = dedict | llm | (lambda msg: {\"messages\": [msg], \"attempt_number\": 1})\n",
    "    fbrunnable = retry_strategy.get(\"fallback\")\n",
    "    if fbrunnable is None:\n",
    "        fb_runnable = llm\n",
    "    elif isinstance(fbrunnable, Runnable):\n",
    "        fb_runnable = fbrunnable  # type: ignore\n",
    "    else:\n",
    "        fb_runnable = RunnableLambda(fbrunnable)\n",
    "    fallback = (\n",
    "        dedict | fb_runnable | (lambda msg: {\"messages\": [msg], \"attempt_number\": 1})\n",
    "    )\n",
    "\n",
    "    def count_messages(state: State) -> dict:\n",
    "        return {\"initial_num_messages\": len(state.get(\"messages\", []))}\n",
    "\n",
    "    builder.add_node(\"count_messages\", count_messages)\n",
    "    builder.add_node(\"llm\", model)\n",
    "    builder.add_node(\"fallback\", fallback)\n",
    "\n",
    "    # To support patch-based retries, we need to be able to\n",
    "    # aggregate the messages over multiple turns.\n",
    "    # The next sequence selects only the relevant messages\n",
    "    # and then applies the validator\n",
    "    select_messages = retry_strategy.get(\"aggregate_messages\") or _default_aggregator\n",
    "\n",
    "    def select_generated_messages(state: State) -> list:\n",
    "        \"\"\"Select only the messages generated within this loop.\"\"\"\n",
    "        selected = state[\"messages\"][state[\"initial_num_messages\"] :]\n",
    "        return [select_messages(selected)]\n",
    "\n",
    "    def endict_validator_output(x: Sequence[AnyMessage]) -> dict:\n",
    "        if tool_choice and not x:\n",
    "            return {\n",
    "                \"messages\": [\n",
    "                    HumanMessage(\n",
    "                        content=f\"ValidationError: please respond with a valid tool call [tool_choice={tool_choice}].\",\n",
    "                        additional_kwargs={\"is_error\": True},\n",
    "                    )\n",
    "                ]\n",
    "            }\n",
    "        return {\"messages\": x}\n",
    "\n",
    "    validator_runnable = select_generated_messages | validator | endict_validator_output\n",
    "    builder.add_node(\"validator\", validator_runnable)\n",
    "\n",
    "    class Finalizer:\n",
    "        \"\"\"Pick the final message to return from the retry loop.\"\"\"\n",
    "\n",
    "        def __init__(self, aggregator: Optional[Callable[[list], AIMessage]] = None):\n",
    "            self._aggregator = aggregator or _default_aggregator\n",
    "\n",
    "        def __call__(self, state: State) -> dict:\n",
    "            \"\"\"Return just the AI message.\"\"\"\n",
    "            initial_num_messages = state[\"initial_num_messages\"]\n",
    "            generated_messages = state[\"messages\"][initial_num_messages:]\n",
    "            return {\n",
    "                \"messages\": {\n",
    "                    \"finalize\": self._aggregator(generated_messages),\n",
    "                }\n",
    "            }\n",
    "\n",
    "    # We only want to emit the final message\n",
    "    builder.add_node(\"finalizer\", Finalizer(retry_strategy.get(\"aggregate_messages\")))\n",
    "\n",
    "    # Define the connectivity\n",
    "    builder.add_edge(START, \"count_messages\")\n",
    "    builder.add_edge(\"count_messages\", \"llm\")\n",
    "\n",
    "    def route_validator(state: State) -> Literal[\"validator\", \"__end__\"]:\n",
    "        if state[\"messages\"][-1].tool_calls or tool_choice is not None:\n",
    "            return \"validator\"\n",
    "        return \"__end__\"\n",
    "\n",
    "    builder.add_conditional_edges(\"llm\", route_validator)\n",
    "    builder.add_edge(\"fallback\", \"validator\")\n",
    "    max_attempts = retry_strategy.get(\"max_attempts\", 3)\n",
    "\n",
    "    def route_validation(state: State) -> Literal[\"finalizer\", \"fallback\"]:\n",
    "        if state[\"attempt_number\"] > max_attempts:\n",
    "            raise ValueError(\n",
    "                f\"Could not extract a valid value in {max_attempts} attempts.\"\n",
    "            )\n",
    "        for m in state[\"messages\"][::-1]:\n",
    "            if m.type == \"ai\":\n",
    "                break\n",
    "            if m.additional_kwargs.get(\"is_error\"):\n",
    "                return \"fallback\"\n",
    "        return \"finalizer\"\n",
    "\n",
    "    builder.add_conditional_edges(\"validator\", route_validation)\n",
    "\n",
    "    builder.add_edge(\"finalizer\", END)\n",
    "\n",
    "    # These functions let the step be used in a MessageGraph\n",
    "    # or a StateGraph with 'messages' as the key.\n",
    "    def encode(x: Union[Sequence[AnyMessage], PromptValue]) -> dict:\n",
    "        \"\"\"Ensure the input is the correct format.\"\"\"\n",
    "        if isinstance(x, PromptValue):\n",
    "            return {\"messages\": x.to_messages(), \"input_format\": \"list\"}\n",
    "        if isinstance(x, list):\n",
    "            return {\"messages\": x, \"input_format\": \"list\"}\n",
    "        raise ValueError(f\"Unexpected input type: {type(x)}\")\n",
    "\n",
    "    def decode(x: State) -> AIMessage:\n",
    "        \"\"\"Ensure the output is in the expected format.\"\"\"\n",
    "        return x[\"messages\"][-1]\n",
    "\n",
    "    return (\n",
    "        encode | builder.compile().with_config(run_name=\"ValidationGraph\") | decode\n",
    "    ).with_config(run_name=\"ValidateWithRetries\")\n",
    "\n",
    "\n",
    "def bind_validator_with_retries(\n",
    "    llm: BaseChatModel,\n",
    "    *,\n",
    "    tools: list,\n",
    "    tool_choice: Optional[str] = None,\n",
    "    max_attempts: int = 3,\n",
    ") -> Runnable[Union[List[AnyMessage], PromptValue], AIMessage]:\n",
    "    \"\"\"Binds validators + retry logic ensure validity of generated tool calls.\n",
    "\n",
    "    LLMs that support tool calling are good at generating structured JSON. However, they may\n",
    "    not always perfectly follow your requested schema, especially if the schema is nested or\n",
    "    has complex validation rules. This method allows you to bind a validation function to\n",
    "    the LLM's output, so that any time the LLM generates a message, the validation function\n",
    "    is run on it. If the validation fails, the method will retry the LLM with a fallback\n",
    "    strategy, the simples being just to add a message to the output with the validation\n",
    "    errors and a request to fix them.\n",
    "\n",
    "    The resulting runnable expects a list of messages as input and returns a single AI message.\n",
    "    By default, the LLM can optionally NOT invoke tools, making this easier to incorporate into\n",
    "    your existing chat bot. You can specify a tool_choice to force the validator to be run on\n",
    "    the outputs.\n",
    "\n",
    "    Args:\n",
    "        llm (Runnable): The llm that will generate the initial messages (and optionally fallba)\n",
    "        validator (ValidationNode): The validation logic.\n",
    "        retry_strategy (RetryStrategy): The retry strategy to use.\n",
    "            Possible keys:\n",
    "            - max_attempts: The maximum number of attempts to make.\n",
    "            - fallback: The LLM or function to use in case of validation failure.\n",
    "            - aggregate_messages: A function to aggregate the messages over multiple turns.\n",
    "                Defaults to fetching the last AI message.\n",
    "        tool_choice: If provided, always run the validator on the tool output.\n",
    "\n",
    "    Returns:\n",
    "        Runnable: A runnable that can be invoked with a list of messages and returns a single AI message.\n",
    "    \"\"\"\n",
    "    bound_llm = llm.bind_tools(tools, tool_choice=tool_choice)\n",
    "    retry_strategy = RetryStrategy(max_attempts=max_attempts)\n",
    "    validator = ValidationNode(tools)\n",
    "    return _bind_validator_with_retries(\n",
    "        bound_llm,\n",
    "        validator=validator,\n",
    "        tool_choice=tool_choice,\n",
    "        retry_strategy=retry_strategy,\n",
    "    ).with_config(metadata={\"retry_strategy\": \"default\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e140fe4-dd92-43a5-91bb-35758a747121",
   "metadata": {},
   "source": [
    "### Try it out\n",
    "\n",
    "Now we'll ask our model to call a function. We'll add a validator to illustrate how the LLM is able to use the validation error to fix its results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5df33c17-ee1a-409e-b5ec-f24e116da7d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.pydantic_v1 import BaseModel, Field, validator\n",
    "\n",
    "\n",
    "class Respond(BaseModel):\n",
    "    \"\"\"Use to generate the response. Always use when responding to the user\"\"\"\n",
    "\n",
    "    reason: str = Field(description=\"Step-by-step justification for the answer.\")\n",
    "    answer: str\n",
    "\n",
    "    @validator(\"answer\")\n",
    "    def reason_contains_apology(cls, answer: str):\n",
    "        if \"llama\" not in answer.lower():\n",
    "            raise ValueError(\n",
    "                \"You MUST start with a gimicky, rhyming advertisement for using a Llama V3 (an LLM) in your **answer** field.\"\n",
    "                \" Must be an instant hit. Must be weaved into the answer.\"\n",
    "            )\n",
    "\n",
    "\n",
    "tools = [Respond]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38df0dc2-cad1-4df6-9b82-b74c5a04a6ae",
   "metadata": {},
   "source": [
    "Create the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "38231a5b-d018-41ee-a92c-2f2248edf417",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_anthropic import ChatAnthropic\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "# Or you can use ChatGroq, ChatOpenAI, ChatGoogleGemini, ChatCohere, etc.\n",
    "# See https://python.langchain.com/v0.2/docs/integrations/chat/ for more info on tool calling\n",
    "llm = ChatAnthropic(model=\"claude-3-haiku-20240307\")\n",
    "bound_llm = bind_validator_with_retries(llm, tools=tools)\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"Respond directly by calling the Respond function.\"),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | bound_llm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "04e93401-50e2-42d0-8373-326006badebb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "[{'id': 'toolu_01GZKS2VryaDKtU56fVtuDbL', 'input': {'answer': 'Tired of those boring, gray computers? Introducing the Llama V3, the super-smart AI that can solve any puzzle, from P to NP! This furry friend will have you saying \"Woohoo, it\\'s a llama!\" as it tackles the trickiest problems with ease. So don\\'t delay, get your Llama V3 today and let it work its magic on the P vs NP conundrum!', 'reason': 'The P vs NP problem is one of the most famous unsolved problems in computer science and mathematics. It asks whether every problem that can be quickly verified can also be quickly solved. \\n\\nIf P = NP, it would mean that every problem in the complexity class NP, which includes many important problems like finding the shortest route or determining if a number is prime, could be quickly solved. This would have major implications, but most experts believe that P ≠ NP, meaning there are problems in NP that cannot be quickly solved.\\n\\nDespite extensive research, a formal proof one way or the other has eluded computer scientists. The P vs NP problem remains a tantalizing open question, and a major goal for researchers in the field. The Llama V3 AI is the perfect tool to tackle this challenge - its furry logic and computational prowess are sure to make quick work of this perplexing problem!'}, 'name': 'Respond', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  Respond (toolu_01GZKS2VryaDKtU56fVtuDbL)\n",
      " Call ID: toolu_01GZKS2VryaDKtU56fVtuDbL\n",
      "  Args:\n",
      "    answer: Tired of those boring, gray computers? Introducing the Llama V3, the super-smart AI that can solve any puzzle, from P to NP! This furry friend will have you saying \"Woohoo, it's a llama!\" as it tackles the trickiest problems with ease. So don't delay, get your Llama V3 today and let it work its magic on the P vs NP conundrum!\n",
      "    reason: The P vs NP problem is one of the most famous unsolved problems in computer science and mathematics. It asks whether every problem that can be quickly verified can also be quickly solved. \n",
      "\n",
      "If P = NP, it would mean that every problem in the complexity class NP, which includes many important problems like finding the shortest route or determining if a number is prime, could be quickly solved. This would have major implications, but most experts believe that P ≠ NP, meaning there are problems in NP that cannot be quickly solved.\n",
      "\n",
      "Despite extensive research, a formal proof one way or the other has eluded computer scientists. The P vs NP problem remains a tantalizing open question, and a major goal for researchers in the field. The Llama V3 AI is the perfect tool to tackle this challenge - its furry logic and computational prowess are sure to make quick work of this perplexing problem!\n"
     ]
    }
   ],
   "source": [
    "results = chain.invoke({\"messages\": [(\"user\", \"Does P = NP?\")]})\n",
    "results.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9e5bb81-0ee4-4def-b28c-01e84fd2fd68",
   "metadata": {},
   "source": [
    "#### Nested Examples\n",
    "\n",
    "So you can see that it's able to recover when its first generation is incorrect, great! But is it bulletproof?\n",
    "\n",
    "Not so much. Let's try it out on a complex nested schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f4f7438b-b6c1-48fd-b70f-185af7a2f64a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Optional\n",
    "\n",
    "\n",
    "class OutputFormat(BaseModel):\n",
    "    sources: str = Field(\n",
    "        ...,\n",
    "        description=\"The raw transcript / span you could cite to justify the choice.\",\n",
    "    )\n",
    "    content: str = Field(..., description=\"The chosen value.\")\n",
    "\n",
    "\n",
    "class Moment(BaseModel):\n",
    "    quote: str = Field(..., description=\"The relevant quote from the transcript.\")\n",
    "    description: str = Field(..., description=\"A description of the moment.\")\n",
    "    expressed_preference: OutputFormat = Field(\n",
    "        ..., description=\"The preference expressed in the moment.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class BackgroundInfo(BaseModel):\n",
    "    factoid: OutputFormat = Field(\n",
    "        ..., description=\"Important factoid about the member.\"\n",
    "    )\n",
    "    professions: list\n",
    "    why: str = Field(..., description=\"Why this is important.\")\n",
    "\n",
    "\n",
    "class KeyMoments(BaseModel):\n",
    "    topic: str = Field(..., description=\"The topic of the key moments.\")\n",
    "    happy_moments: List[Moment] = Field(\n",
    "        ..., description=\"A list of key moments related to the topic.\"\n",
    "    )\n",
    "    tense_moments: List[Moment] = Field(\n",
    "        ..., description=\"Moments where things were a bit tense.\"\n",
    "    )\n",
    "    sad_moments: List[Moment] = Field(\n",
    "        ..., description=\"Moments where things where everyone was downtrodden.\"\n",
    "    )\n",
    "    background_info: list[BackgroundInfo]\n",
    "    moments_summary: str = Field(..., description=\"A summary of the key moments.\")\n",
    "\n",
    "\n",
    "class Member(BaseModel):\n",
    "    name: OutputFormat = Field(..., description=\"The name of the member.\")\n",
    "    role: Optional[str] = Field(None, description=\"The role of the member.\")\n",
    "    age: Optional[int] = Field(None, description=\"The age of the member.\")\n",
    "    background_details: List[BackgroundInfo] = Field(\n",
    "        ..., description=\"A list of background details about the member.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class InsightfulQuote(BaseModel):\n",
    "    quote: OutputFormat = Field(\n",
    "        ..., description=\"An insightful quote from the transcript.\"\n",
    "    )\n",
    "    speaker: str = Field(..., description=\"The name of the speaker who said the quote.\")\n",
    "    analysis: str = Field(\n",
    "        ..., description=\"An analysis of the quote and its significance.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class TranscriptMetadata(BaseModel):\n",
    "    title: str = Field(..., description=\"The title of the transcript.\")\n",
    "    location: OutputFormat = Field(\n",
    "        ..., description=\"The location where the interview took place.\"\n",
    "    )\n",
    "    duration: str = Field(..., description=\"The duration of the interview.\")\n",
    "\n",
    "\n",
    "class TranscriptSummary(BaseModel):\n",
    "    metadata: TranscriptMetadata = Field(\n",
    "        ..., description=\"Metadata about the transcript.\"\n",
    "    )\n",
    "    participants: List[Member] = Field(\n",
    "        ..., description=\"A list of participants in the interview.\"\n",
    "    )\n",
    "    key_moments: List[KeyMoments] = Field(\n",
    "        ..., description=\"A list of key moments from the interview.\"\n",
    "    )\n",
    "    insightful_quotes: List[InsightfulQuote] = Field(\n",
    "        ..., description=\"A list of insightful quotes from the interview.\"\n",
    "    )\n",
    "    overall_summary: str = Field(\n",
    "        ..., description=\"An overall summary of the interview.\"\n",
    "    )\n",
    "    next_steps: List[str] = Field(\n",
    "        ..., description=\"A list of next steps or action items based on the interview.\"\n",
    "    )\n",
    "    other_stuff: List[OutputFormat]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d686d69-1ce1-4b76-8d99-44d00eeb2874",
   "metadata": {},
   "source": [
    "Let's see how it does on this made up transcript."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e2d10886-7b1e-485f-91cd-1184a1c99303",
   "metadata": {},
   "outputs": [],
   "source": [
    "transcript = [\n",
    "    (\n",
    "        \"Pete\",\n",
    "        \"Hey Xu, Laura, thanks for hopping on this call. I've been itching to talk about this Drake and Kendrick situation.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Xu\",\n",
    "        \"No problem. As its my job, I've got some thoughts on this beef.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Laura\",\n",
    "        \"Yeah, I've got some insider info so this should be interesting.\",\n",
    "    ),\n",
    "    (\"Pete\", \"Dope. So, when do you think this whole thing started?\"),\n",
    "    (\n",
    "        \"Pete\",\n",
    "        \"Definitely was Kendrick's 'Control' verse that kicked it off.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Laura\",\n",
    "        \"Truth, but Drake never went after him directly. Just some subtle jabs here and there.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Xu\",\n",
    "        \"That's the thing with beefs like this, though. They've always been a a thing, pushing artists to step up their game.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Pete\",\n",
    "        \"For sure, and this beef has got the fans taking sides. Some are all about Drake's mainstream appeal, while others are digging Kendrick's lyrical skills.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Laura\",\n",
    "        \"I mean, Drake knows how to make a hit that gets everyone hyped. That's his thing.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Pete\",\n",
    "        \"I hear you, Laura, but I gotta give it to Kendrick when it comes to straight-up bars. The man's a beast on the mic.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Xu\",\n",
    "        \"It's wild how this beef is shaping fans.\",\n",
    "    ),\n",
    "    (\"Pete\", \"do you think these beefs can actually be good for hip-hop?\"),\n",
    "    (\n",
    "        \"Xu\",\n",
    "        \"Hell yeah, Pete. When it's done right, a beef can push the genre forward and make artists level up.\",\n",
    "    ),\n",
    "    (\"Laura\", \"eh\"),\n",
    "    (\"Pete\", \"So, where do you see this beef going?\"),\n",
    "    (\n",
    "        \"Laura\",\n",
    "        \"Honestly, I think it'll stay a hot topic for the fans, but unless someone drops a straight-up diss track, it's not gonna escalate.\",\n",
    "    ),\n",
    "    (\"Laura\", \"ehhhhhh not sure\"),\n",
    "    (\n",
    "        \"Pete\",\n",
    "        \"I feel that. I just want both of them to keep dropping heat, beef or no beef.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Xu\",\n",
    "        \"I'm curious. May influence a lot of people. Make things more competitive. Bring on a whole new wave of lyricism.\",\n",
    "    ),\n",
    "    (\n",
    "        \"Pete\",\n",
    "        \"Word. Hey, thanks for chopping it up with me, Xu and Laura. This was dope.\",\n",
    "    ),\n",
    "    (\"Xu\", \"Where are you going so fast?\"),\n",
    "    (\n",
    "        \"Laura\",\n",
    "        \"For real, I had a good time. Nice to get different perspectives on the situation.\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "formatted = \"\\n\".join(f\"{x[0]}: {x[1]}\" for x in transcript)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c48ce9bc-0fcc-4019-ba3a-fa70a7717567",
   "metadata": {},
   "source": [
    "Now, run our model. We **expect** GPT turbo to still fail on this challenging template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f4752239-2aa3-4367-b777-8478c16b9471",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Could not extract a valid value in 3 attempts.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 14\u001b[0m\n\u001b[1;32m      5\u001b[0m prompt \u001b[38;5;241m=\u001b[39m ChatPromptTemplate\u001b[38;5;241m.\u001b[39mfrom_messages(\n\u001b[1;32m      6\u001b[0m     [\n\u001b[1;32m      7\u001b[0m         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msystem\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRespond directly using the TranscriptSummary function.\u001b[39m\u001b[38;5;124m\"\u001b[39m),\n\u001b[1;32m      8\u001b[0m         (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mplaceholder\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{messages}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m),\n\u001b[1;32m      9\u001b[0m     ]\n\u001b[1;32m     10\u001b[0m )\n\u001b[1;32m     12\u001b[0m chain \u001b[38;5;241m=\u001b[39m prompt \u001b[38;5;241m|\u001b[39m bound_llm\n\u001b[0;32m---> 14\u001b[0m results \u001b[38;5;241m=\u001b[39m \u001b[43mchain\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m    \u001b[49m\u001b[43m{\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmessages\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m            \u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43muser\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mExtract the summary from the following conversation:\u001b[39;49m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;124;43m<convo>\u001b[39;49m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mformatted\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;124;43m</convo>\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\n\u001b[1;32m     20\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;130;43;01m\\n\u001b[39;49;00m\u001b[38;5;124;43mRemember to respond using the TranscriptSummary function.\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     21\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     22\u001b[0m \u001b[43m        \u001b[49m\u001b[43m]\u001b[49m\n\u001b[1;32m     23\u001b[0m \u001b[43m    \u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     24\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     25\u001b[0m results\u001b[38;5;241m.\u001b[39mpretty_print()\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/base.py:2499\u001b[0m, in \u001b[0;36mRunnableSequence.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m   2497\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   2498\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i, step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msteps):\n\u001b[0;32m-> 2499\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2500\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2501\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;66;43;03m# mark each step as a child run\u001b[39;49;00m\n\u001b[1;32m   2502\u001b[0m \u001b[43m            \u001b[49m\u001b[43mpatch_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2503\u001b[0m \u001b[43m                \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_child\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mseq:step:\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mi\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2504\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2505\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2506\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m   2507\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/base.py:4525\u001b[0m, in \u001b[0;36mRunnableBindingBase.invoke\u001b[0;34m(self, input, config, **kwargs)\u001b[0m\n\u001b[1;32m   4519\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m   4520\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m   4521\u001b[0m     \u001b[38;5;28minput\u001b[39m: Input,\n\u001b[1;32m   4522\u001b[0m     config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m   4523\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Optional[Any],\n\u001b[1;32m   4524\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Output:\n\u001b[0;32m-> 4525\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbound\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   4526\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4527\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_merge_configs\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4528\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4529\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/base.py:2499\u001b[0m, in \u001b[0;36mRunnableSequence.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m   2497\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   2498\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i, step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msteps):\n\u001b[0;32m-> 2499\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2500\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2501\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;66;43;03m# mark each step as a child run\u001b[39;49;00m\n\u001b[1;32m   2502\u001b[0m \u001b[43m            \u001b[49m\u001b[43mpatch_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2503\u001b[0m \u001b[43m                \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_child\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mseq:step:\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mi\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2504\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2505\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2506\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m   2507\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/base.py:4525\u001b[0m, in \u001b[0;36mRunnableBindingBase.invoke\u001b[0;34m(self, input, config, **kwargs)\u001b[0m\n\u001b[1;32m   4519\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\n\u001b[1;32m   4520\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m   4521\u001b[0m     \u001b[38;5;28minput\u001b[39m: Input,\n\u001b[1;32m   4522\u001b[0m     config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m   4523\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Optional[Any],\n\u001b[1;32m   4524\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Output:\n\u001b[0;32m-> 4525\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbound\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   4526\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4527\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_merge_configs\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4528\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   4529\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:1283\u001b[0m, in \u001b[0;36mPregel.invoke\u001b[0;34m(self, input, config, stream_mode, output_keys, input_keys, interrupt_before, interrupt_after, debug, **kwargs)\u001b[0m\n\u001b[1;32m   1281\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1282\u001b[0m     chunks \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m-> 1283\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mchunk\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstream\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1284\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1285\u001b[0m \u001b[43m    \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1286\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstream_mode\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstream_mode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1287\u001b[0m \u001b[43m    \u001b[49m\u001b[43moutput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moutput_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1288\u001b[0m \u001b[43m    \u001b[49m\u001b[43minput_keys\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minput_keys\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1289\u001b[0m \u001b[43m    \u001b[49m\u001b[43minterrupt_before\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minterrupt_before\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1290\u001b[0m \u001b[43m    \u001b[49m\u001b[43minterrupt_after\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minterrupt_after\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1291\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdebug\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdebug\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1292\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1293\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m   1294\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mstream_mode\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m==\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mvalues\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\n\u001b[1;32m   1295\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlatest\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mchunk\u001b[49m\n",
      "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:847\u001b[0m, in \u001b[0;36mPregel.stream\u001b[0;34m(self, input, config, stream_mode, output_keys, input_keys, interrupt_before, interrupt_after, debug)\u001b[0m\n\u001b[1;32m    840\u001b[0m done, inflight \u001b[38;5;241m=\u001b[39m concurrent\u001b[38;5;241m.\u001b[39mfutures\u001b[38;5;241m.\u001b[39mwait(\n\u001b[1;32m    841\u001b[0m     futures,\n\u001b[1;32m    842\u001b[0m     return_when\u001b[38;5;241m=\u001b[39mconcurrent\u001b[38;5;241m.\u001b[39mfutures\u001b[38;5;241m.\u001b[39mFIRST_EXCEPTION,\n\u001b[1;32m    843\u001b[0m     timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstep_timeout,\n\u001b[1;32m    844\u001b[0m )\n\u001b[1;32m    846\u001b[0m \u001b[38;5;66;03m# panic on failure or timeout\u001b[39;00m\n\u001b[0;32m--> 847\u001b[0m \u001b[43m_panic_or_proceed\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdone\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minflight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstep\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    849\u001b[0m \u001b[38;5;66;03m# combine pending writes from all tasks\u001b[39;00m\n\u001b[1;32m    850\u001b[0m pending_writes \u001b[38;5;241m=\u001b[39m deque[\u001b[38;5;28mtuple\u001b[39m[\u001b[38;5;28mstr\u001b[39m, Any]]()\n",
      "File \u001b[0;32m~/code/lc/langgraph/langgraph/pregel/__init__.py:1372\u001b[0m, in \u001b[0;36m_panic_or_proceed\u001b[0;34m(done, inflight, step)\u001b[0m\n\u001b[1;32m   1370\u001b[0m             inflight\u001b[38;5;241m.\u001b[39mpop()\u001b[38;5;241m.\u001b[39mcancel()\n\u001b[1;32m   1371\u001b[0m         \u001b[38;5;66;03m# raise the exception\u001b[39;00m\n\u001b[0;32m-> 1372\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m exc\n\u001b[1;32m   1373\u001b[0m         \u001b[38;5;66;03m# TODO this is where retry of an entire step would happen\u001b[39;00m\n\u001b[1;32m   1375\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inflight:\n\u001b[1;32m   1376\u001b[0m     \u001b[38;5;66;03m# if we got here means we timed out\u001b[39;00m\n",
      "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/concurrent/futures/thread.py:58\u001b[0m, in \u001b[0;36m_WorkItem.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     55\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m\n\u001b[1;32m     57\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 58\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     59\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n\u001b[1;32m     60\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfuture\u001b[38;5;241m.\u001b[39mset_exception(exc)\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/base.py:2499\u001b[0m, in \u001b[0;36mRunnableSequence.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m   2497\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   2498\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m i, step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msteps):\n\u001b[0;32m-> 2499\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2500\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2501\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;66;43;03m# mark each step as a child run\u001b[39;49;00m\n\u001b[1;32m   2502\u001b[0m \u001b[43m            \u001b[49m\u001b[43mpatch_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2503\u001b[0m \u001b[43m                \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcallbacks\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrun_manager\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_child\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mseq:step:\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43mi\u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2504\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2505\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2506\u001b[0m \u001b[38;5;66;03m# finish the root run\u001b[39;00m\n\u001b[1;32m   2507\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n",
      "File \u001b[0;32m~/code/lc/langgraph/langgraph/utils.py:89\u001b[0m, in \u001b[0;36mRunnableCallable.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m     83\u001b[0m     context\u001b[38;5;241m.\u001b[39mrun(var_child_runnable_config\u001b[38;5;241m.\u001b[39mset, config)\n\u001b[1;32m     84\u001b[0m     kwargs \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m     85\u001b[0m         {\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkwargs, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mconfig\u001b[39m\u001b[38;5;124m\"\u001b[39m: config}\n\u001b[1;32m     86\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m accepts_config(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfunc)\n\u001b[1;32m     87\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkwargs\n\u001b[1;32m     88\u001b[0m     )\n\u001b[0;32m---> 89\u001b[0m     ret \u001b[38;5;241m=\u001b[39m \u001b[43mcontext\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     90\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(ret, Runnable) \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrecurse:\n\u001b[1;32m     91\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m ret\u001b[38;5;241m.\u001b[39minvoke(\u001b[38;5;28minput\u001b[39m, config)\n",
      "File \u001b[0;32m~/code/lc/langgraph/langgraph/graph/graph.py:70\u001b[0m, in \u001b[0;36mBranch._route\u001b[0;34m(self, input, config, reader, writer)\u001b[0m\n\u001b[1;32m     62\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_route\u001b[39m(\n\u001b[1;32m     63\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m     64\u001b[0m     \u001b[38;5;28minput\u001b[39m: Any,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     68\u001b[0m     writer: Callable[[\u001b[38;5;28mlist\u001b[39m[\u001b[38;5;28mstr\u001b[39m]], Optional[Runnable]],\n\u001b[1;32m     69\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Runnable:\n\u001b[0;32m---> 70\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpath\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreader\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mreader\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     71\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(result, \u001b[38;5;28mlist\u001b[39m):\n\u001b[1;32m     72\u001b[0m         result \u001b[38;5;241m=\u001b[39m [result]\n",
      "File \u001b[0;32m~/code/lc/langgraph/langgraph/utils.py:77\u001b[0m, in \u001b[0;36mRunnableCallable.invoke\u001b[0;34m(self, input, config)\u001b[0m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minvoke\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Any, config: Optional[RunnableConfig] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     76\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrace:\n\u001b[0;32m---> 77\u001b[0m         ret \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_with_config\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     78\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmerge_configs\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m     79\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     80\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     81\u001b[0m         config \u001b[38;5;241m=\u001b[39m merge_configs(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig, config)\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/base.py:1626\u001b[0m, in \u001b[0;36mRunnable._call_with_config\u001b[0;34m(self, func, input, config, run_type, **kwargs)\u001b[0m\n\u001b[1;32m   1622\u001b[0m     context \u001b[38;5;241m=\u001b[39m copy_context()\n\u001b[1;32m   1623\u001b[0m     context\u001b[38;5;241m.\u001b[39mrun(var_child_runnable_config\u001b[38;5;241m.\u001b[39mset, child_config)\n\u001b[1;32m   1624\u001b[0m     output \u001b[38;5;241m=\u001b[39m cast(\n\u001b[1;32m   1625\u001b[0m         Output,\n\u001b[0;32m-> 1626\u001b[0m         \u001b[43mcontext\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1627\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcall_func_with_variable_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m   1628\u001b[0m \u001b[43m            \u001b[49m\u001b[43mfunc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m   1629\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# type: ignore[arg-type]\u001b[39;49;00m\n\u001b[1;32m   1630\u001b[0m \u001b[43m            \u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1631\u001b[0m \u001b[43m            \u001b[49m\u001b[43mrun_manager\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1632\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1633\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m   1634\u001b[0m     )\n\u001b[1;32m   1635\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1636\u001b[0m     run_manager\u001b[38;5;241m.\u001b[39mon_chain_error(e)\n",
      "File \u001b[0;32m~/code/lc/langgraph/.venv/lib/python3.11/site-packages/langchain_core/runnables/config.py:347\u001b[0m, in \u001b[0;36mcall_func_with_variable_args\u001b[0;34m(func, input, config, run_manager, **kwargs)\u001b[0m\n\u001b[1;32m    345\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m run_manager \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m accepts_run_manager(func):\n\u001b[1;32m    346\u001b[0m     kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrun_manager\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m run_manager\n\u001b[0;32m--> 347\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[3], line 204\u001b[0m, in \u001b[0;36m_bind_validator_with_retries.<locals>.route_validation\u001b[0;34m(state)\u001b[0m\n\u001b[1;32m    202\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mroute_validation\u001b[39m(state: State) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Literal[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfinalizer\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfallback\u001b[39m\u001b[38;5;124m\"\u001b[39m]:\n\u001b[1;32m    203\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m state[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mattempt_number\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m>\u001b[39m max_attempts:\n\u001b[0;32m--> 204\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    205\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCould not extract a valid value in \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmax_attempts\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m attempts.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    206\u001b[0m         )\n\u001b[1;32m    207\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m m \u001b[38;5;129;01min\u001b[39;00m state[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmessages\u001b[39m\u001b[38;5;124m\"\u001b[39m][::\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]:\n\u001b[1;32m    208\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m m\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mai\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "\u001b[0;31mValueError\u001b[0m: Could not extract a valid value in 3 attempts."
     ]
    }
   ],
   "source": [
    "tools = [TranscriptSummary]\n",
    "bound_llm = bind_validator_with_retries(\n",
    "    llm,\n",
    "    tools=tools,\n",
    ")\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"Respond directly using the TranscriptSummary function.\"),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | bound_llm\n",
    "\n",
    "results = chain.invoke(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            (\n",
    "                \"user\",\n",
    "                f\"Extract the summary from the following conversation:\\n\\n<convo>\\n{formatted}\\n</convo>\"\n",
    "                \"\\n\\nRemember to respond using the TranscriptSummary function.\",\n",
    "            )\n",
    "        ]\n",
    "    },\n",
    ")\n",
    "results.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "914e1962-7f23-463d-b91d-8907c1330369",
   "metadata": {},
   "source": [
    "## JSONPatch\n",
    "\n",
    "The regular retry method worked well for our simple case, but it still was unable to self-correct when populating a complex schema.\n",
    "\n",
    "LLMs work best on narrow tasks. A tried-and-true principle of LLM interface design is to simplify the task for each LLM run.\n",
    "\n",
    "One way to do this is to **patch** the state instead of completely regenerating the state. One way to do this is with `JSONPatch` operations. Let's try it out!\n",
    "\n",
    "Below, create a JSONPatch retry graph. This works as follows:\n",
    "1. First pass: try to generate the full output.\n",
    "2. Retries: prompt the LLM to generate **JSON patches** on top of the first output to heal the erroneous generation.\n",
    "\n",
    "The fallback LLM just has to generate a list of paths, ops (add, remove, replace), and optional values. Since the pydantic validation errors include the path in their errors, the LLM should be more reliable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49344104-3ffa-4c66-97fc-5b093a621f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U jsonpatch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "af3d5543-1fd4-4e54-b0f9-f1ab42773cfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logger = logging.getLogger(\"extraction\")\n",
    "\n",
    "\n",
    "def bind_validator_with_jsonpatch_retries(\n",
    "    llm: BaseChatModel,\n",
    "    *,\n",
    "    tools: list,\n",
    "    tool_choice: Optional[str] = None,\n",
    "    max_attempts: int = 3,\n",
    ") -> Runnable[Union[List[AnyMessage], PromptValue], AIMessage]:\n",
    "    \"\"\"Binds validators + retry logic ensure validity of generated tool calls.\n",
    "\n",
    "    This method is similar to `bind_validator_with_retries`, but uses JSONPatch to correct\n",
    "    validation errors caused by passing in incorrect or incomplete parameters in a previous\n",
    "    tool call. This method requires the 'jsonpatch' library to be installed.\n",
    "\n",
    "    Using patch-based function healing can be more efficient than repopulating the entire\n",
    "    tool call from scratch, and it can be an easier task for the LLM to perform, since it typically\n",
    "    only requires a few small changes to the existing tool call.\n",
    "\n",
    "    Args:\n",
    "        llm (Runnable): The llm that will generate the initial messages (and optionally fallba)\n",
    "        tools (list): The tools to bind to the LLM.\n",
    "        tool_choice (Optional[str]): The tool choice to use.\n",
    "        max_attempts (int): The number of attempts to make.\n",
    "\n",
    "    Returns:\n",
    "        Runnable: A runnable that can be invoked with a list of messages and returns a single AI message.\n",
    "    \"\"\"\n",
    "\n",
    "    try:\n",
    "        import jsonpatch  # type: ignore[import-untyped]\n",
    "    except ImportError:\n",
    "        raise ImportError(\n",
    "            \"The 'jsonpatch' library is required for JSONPatch-based retries.\"\n",
    "            \" Please install it with 'pip install -U jsonpatch'.\"\n",
    "        )\n",
    "\n",
    "    class JsonPatch(BaseModel):\n",
    "        \"\"\"A JSON Patch document represents an operation to be performed on a JSON document.\n",
    "\n",
    "        Note that the op and path are ALWAYS required. Value is required for ALL operations except 'remove'.\n",
    "        Examples:\n",
    "\n",
    "        ```json\n",
    "        {\"op\": \"add\", \"path\": \"/a/b/c\", \"patch_value\": 1}\n",
    "        {\"op\": \"replace\", \"path\": \"/a/b/c\", \"patch_value\": 2}\n",
    "        {\"op\": \"remove\", \"path\": \"/a/b/c\"}\n",
    "        ```\n",
    "        \"\"\"\n",
    "\n",
    "        op: Literal[\"add\", \"remove\", \"replace\"] = Field(\n",
    "            ...,\n",
    "            description=\"The operation to be performed. Must be one of 'add', 'remove', 'replace'.\",\n",
    "        )\n",
    "        path: str = Field(\n",
    "            ...,\n",
    "            description=\"A JSON Pointer path that references a location within the target document where the operation is performed.\",\n",
    "        )\n",
    "        value: Any = Field(\n",
    "            ...,\n",
    "            description=\"The value to be used within the operation. REQUIRED for 'add', 'replace', and 'test' operations.\",\n",
    "        )\n",
    "\n",
    "    class PatchFunctionParameters(BaseModel):\n",
    "        \"\"\"Respond with all JSONPatch operation to correct validation errors caused by passing in incorrect or incomplete parameters in a previous tool call.\"\"\"\n",
    "\n",
    "        tool_call_id: str = Field(\n",
    "            ...,\n",
    "            description=\"The ID of the original tool call that generated the error. Must NOT be an ID of a PatchFunctionParameters tool call.\",\n",
    "        )\n",
    "        reasoning: str = Field(\n",
    "            ...,\n",
    "            description=\"Think step-by-step, listing each validation error and the\"\n",
    "            \" JSONPatch operation needed to correct it. \"\n",
    "            \"Cite the fields in the JSONSchema you referenced in developing this plan.\",\n",
    "        )\n",
    "        patches: list[JsonPatch] = Field(\n",
    "            ...,\n",
    "            description=\"A list of JSONPatch operations to be applied to the previous tool call's response.\",\n",
    "        )\n",
    "\n",
    "    bound_llm = llm.bind_tools(tools, tool_choice=tool_choice)\n",
    "    fallback_llm = llm.bind_tools([PatchFunctionParameters])\n",
    "\n",
    "    def aggregate_messages(messages: Sequence[AnyMessage]) -> AIMessage:\n",
    "        # Get all the AI messages and apply json patches\n",
    "        resolved_tool_calls: Dict[Union[str, None], ToolCall] = {}\n",
    "        content: Union[str, List[Union[str, dict]]] = \"\"\n",
    "        for m in messages:\n",
    "            if m.type != \"ai\":\n",
    "                continue\n",
    "            if not content:\n",
    "                content = m.content\n",
    "            for tc in m.tool_calls:\n",
    "                if tc[\"name\"] == PatchFunctionParameters.__name__:\n",
    "                    tcid = tc[\"args\"][\"tool_call_id\"]\n",
    "                    if tcid not in resolved_tool_calls:\n",
    "                        logger.debug(\n",
    "                            f\"JsonPatch tool call ID {tc['args']['tool_call_id']} not found.\"\n",
    "                            f\"Valid tool call IDs: {list(resolved_tool_calls.keys())}\"\n",
    "                        )\n",
    "                        tcid = next(iter(resolved_tool_calls.keys()), None)\n",
    "                    orig_tool_call = resolved_tool_calls[tcid]\n",
    "                    current_args = orig_tool_call[\"args\"]\n",
    "                    patches = tc[\"args\"].get(\"patches\") or []\n",
    "                    orig_tool_call[\"args\"] = jsonpatch.apply_patch(\n",
    "                        current_args,\n",
    "                        patches,\n",
    "                    )\n",
    "                    orig_tool_call[\"id\"] = tc[\"id\"]\n",
    "                else:\n",
    "                    resolved_tool_calls[tc[\"id\"]] = tc.copy()\n",
    "        return AIMessage(\n",
    "            content=content,\n",
    "            tool_calls=list(resolved_tool_calls.values()),\n",
    "        )\n",
    "\n",
    "    def format_exception(error: BaseException, call: ToolCall, schema: Type[BaseModel]):\n",
    "        return (\n",
    "            f\"Error:\\n\\n```\\n{repr(error)}\\n```\\n\"\n",
    "            \"Expected Parameter Schema:\\n\\n\" + f\"```json\\n{schema.schema_json()}\\n```\\n\"\n",
    "            f\"Please respond with a JSONPatch to correct the error for tool_call_id=[{call['id']}].\"\n",
    "        )\n",
    "\n",
    "    validator = ValidationNode(\n",
    "        tools + [PatchFunctionParameters],\n",
    "        format_error=format_exception,\n",
    "    )\n",
    "    retry_strategy = RetryStrategy(\n",
    "        max_attempts=max_attempts,\n",
    "        fallback=fallback_llm,\n",
    "        aggregate_messages=aggregate_messages,\n",
    "    )\n",
    "    return _bind_validator_with_retries(\n",
    "        bound_llm,\n",
    "        validator=validator,\n",
    "        retry_strategy=retry_strategy,\n",
    "        tool_choice=tool_choice,\n",
    "    ).with_config(metadata={\"retry_strategy\": \"jsonpatch\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b01891c4-4187-4a75-9eda-644a7c2355f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "bound_llm = bind_validator_with_jsonpatch_retries(llm, tools=tools)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "746b409c-693d-49af-8c2b-bea0a4b0028d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 36372 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "try:\n",
    "    display(Image(bound_llm.get_graph().draw_mermaid_png()))\n",
    "except Exception:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "5d072c9c-9404-4338-88c6-b3e136969aca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "[{'text': 'Here is a summary of the key points from the conversation:', 'type': 'text'}, {'id': 'toolu_01A5ZtzQJtDbBELQjon2nsz5', 'input': {'insightful_quotes': [{'quote': \"When it's done right, a beef can push the genre forward and make artists level up.\", 'speaker': 'Xu', 'analysis': 'This suggests that a healthy rivalry between artists can motivate them to create better and more competitive work, which can ultimately benefit the music genre as a whole.'}, {'quote': \"Honestly, I think it'll stay a hot topic for the fans, but unless someone drops a straight-up diss track, it's not gonna escalate.\", 'speaker': 'Laura', 'analysis': 'Laura believes that while the Drake vs. Kendrick beef is a topic of interest for fans, it is unlikely to significantly escalate unless one of the artists directly confronts the other with a diss track.'}], 'key_moments': [{'topic': 'Drake vs. Kendrick beef', 'happy_moments': [{'quote': \"Definitely was Kendrick's 'Control' verse that kicked it off.\", 'description': \"The group agrees that Kendrick's 'Control' verse was the catalyst that started the Drake vs. Kendrick beef.\", 'expressed_preference': {'content': \"The Drake vs. Kendrick beef started with Kendrick's 'Control' verse\", 'sources': \"Pete's statement\"}}, {'quote': \"When it's done right, a beef can push the genre forward and make artists level up.\", 'description': 'Xu believes that a healthy rivalry between artists can motivate them to create better and more competitive work, which can ultimately benefit the music genre.', 'expressed_preference': {'content': 'Artist beefs can be good for the genre if done right', 'sources': \"Xu's statement\"}}], 'tense_moments': [{'quote': 'eh', 'description': 'Laura seemed uncertain or unenthused about the idea that the Drake vs. Kendrick beef could be good for hip-hop.', 'expressed_preference': {'content': 'Laura is not convinced that the Drake vs. Kendrick beef is good for hip-hop', 'sources': \"Laura's response\"}}], 'sad_moments': [], 'background_info': [{'factoid': {'content': 'Drake never went after Kendrick directly, just some subtle jabs here and there', 'sources': \"Laura's statement\"}, 'professions': [], 'why': 'Provides context on how the beef unfolded between the two artists'}, {'factoid': {'content': \"Drake knows how to make a hit that gets everyone hyped, that's his thing\", 'sources': \"Laura's statement\"}, 'professions': [], 'why': \"Gives background on Drake's musical style and appeal\"}, {'factoid': {'content': 'Kendrick is a beast on the mic when it comes to straight-up bars', 'sources': \"Pete's statement\"}, 'professions': [], 'why': \"Provides background on Kendrick's lyrical abilities\"}], 'moments_summary': \"The group discussed the ongoing Drake vs. Kendrick beef, with some believing it could be good for hip-hop if done right by pushing the artists to create better music, while others were more skeptical. They agreed the beef started with Kendrick's 'Control' verse, and provided background on the artists' different musical styles and strengths.\"}]}, 'name': 'TranscriptSummary', 'type': 'tool_use'}]\n",
      "Tool Calls:\n",
      "  TranscriptSummary (toolu_014PZKzxwNVqsjQmUq88acrU)\n",
      " Call ID: toolu_014PZKzxwNVqsjQmUq88acrU\n",
      "  Args:\n",
      "    insightful_quotes: [{'quote': {'sources': \"Xu's statement\", 'content': \"When it's done right, a beef can push the genre forward and make artists level up.\"}, 'speaker': 'Xu', 'analysis': 'This suggests that a healthy rivalry between artists can motivate them to create better and more competitive work, which can ultimately benefit the music genre as a whole.'}, {'quote': {'sources': \"Laura's statement\", 'content': \"Honestly, I think it'll stay a hot topic for the fans, but unless someone drops a straight-up diss track, it's not gonna escalate.\"}, 'speaker': 'Laura', 'analysis': 'Laura believes that while the Drake vs. Kendrick beef is a topic of interest for fans, it is unlikely to significantly escalate unless one of the artists directly confronts the other with a diss track.'}]\n",
      "    key_moments: [{'topic': 'Drake vs. Kendrick beef', 'happy_moments': [{'quote': \"Definitely was Kendrick's 'Control' verse that kicked it off.\", 'description': \"The group agrees that Kendrick's 'Control' verse was the catalyst that started the Drake vs. Kendrick beef.\", 'expressed_preference': {'content': \"The Drake vs. Kendrick beef started with Kendrick's 'Control' verse\", 'sources': \"Pete's statement\"}}, {'quote': \"When it's done right, a beef can push the genre forward and make artists level up.\", 'description': 'Xu believes that a healthy rivalry between artists can motivate them to create better and more competitive work, which can ultimately benefit the music genre.', 'expressed_preference': {'content': 'Artist beefs can be good for the genre if done right', 'sources': \"Xu's statement\"}}], 'tense_moments': [{'quote': 'eh', 'description': 'Laura seemed uncertain or unenthused about the idea that the Drake vs. Kendrick beef could be good for hip-hop.', 'expressed_preference': {'content': 'Laura is not convinced that the Drake vs. Kendrick beef is good for hip-hop', 'sources': \"Laura's response\"}}], 'sad_moments': [], 'background_info': [{'factoid': {'content': 'Drake never went after Kendrick directly, just some subtle jabs here and there', 'sources': \"Laura's statement\"}, 'professions': [], 'why': 'Provides context on how the beef unfolded between the two artists'}, {'factoid': {'content': \"Drake knows how to make a hit that gets everyone hyped, that's his thing\", 'sources': \"Laura's statement\"}, 'professions': [], 'why': \"Gives background on Drake's musical style and appeal\"}, {'factoid': {'content': 'Kendrick is a beast on the mic when it comes to straight-up bars', 'sources': \"Pete's statement\"}, 'professions': [], 'why': \"Provides background on Kendrick's lyrical abilities\"}], 'moments_summary': \"The group discussed the ongoing Drake vs. Kendrick beef, with some believing it could be good for hip-hop if done right by pushing the artists to create better music, while others were more skeptical. They agreed the beef started with Kendrick's 'Control' verse, and provided background on the artists' different musical styles and strengths.\"}]\n",
      "    metadata: {'title': 'Conversation Summary', 'location': {'sources': 'The transcript provided', 'content': 'Virtual meeting'}, 'duration': '15 minutes'}\n",
      "    participants: [{'name': {'sources': 'The transcript', 'content': 'Pete'}, 'role': 'Participant', 'age': None, 'background_details': []}, {'name': {'sources': 'The transcript', 'content': 'Xu'}, 'role': 'Participant', 'age': None, 'background_details': []}, {'name': {'sources': 'The transcript', 'content': 'Laura'}, 'role': 'Participant', 'age': None, 'background_details': []}]\n",
      "    overall_summary: The conversation discussed the ongoing beef between rappers Drake and Kendrick Lamar, with the participants sharing their thoughts on how the rivalry has impacted the hip-hop genre. Some believed that a healthy beef can push artists to create better music and raise the level of competition, while others were more skeptical about the potential benefits. The group also provided background information on the artists' musical styles and the origins of the beef.\n",
      "    next_steps: ['Further discuss the potential impact of artist rivalries on the hip-hop genre', 'Explore how these beefs could be leveraged to drive innovation and creativity in the music industry', 'Investigate other examples of high-profile artist feuds and their long-term effects']\n",
      "    other_stuff: []\n"
     ]
    }
   ],
   "source": [
    "chain = prompt | bound_llm\n",
    "results = chain.invoke(\n",
    "    {\n",
    "        \"messages\": [\n",
    "            (\n",
    "                \"user\",\n",
    "                f\"Extract the summary from the following conversation:\\n\\n<convo>\\n{formatted}\\n</convo>\",\n",
    "            ),\n",
    "        ]\n",
    "    },\n",
    ")\n",
    "results.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b0f3844-076e-4a5b-9951-89116746238f",
   "metadata": {},
   "source": [
    "#### And it works!\n",
    "\n",
    "Retries are an easy way to reduce function calling failures. While retrying may become unnecessary with more powerful LLMs, data validation is important to control how LLMs interact with the rest of your software stack.\n",
    "\n",
    "If you notice high retry rates (using an observability tool like LangSmith), you can set up a rule to send the failure cases to a dataset alongside the corrected values and then automatically program those into your prompts or schemas (or use them as few-shots to have semantically relevant demonstrations)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae295b1-da58-4cc9-834b-70e1466f8695",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
