{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "69ad077f-af4e-49b8-b549-fb3112299aa8",
   "metadata": {},
   "source": [
    "# Competitive Programming\n",
    "\n",
    "In this tutorial, you will build a computing olympiad agent that leverages three complementary techniques to boost performance: **reflection**, **retrieval**, and **human-in-the-loop** collaboration. These techniques and data are all adapted from the paper \"Can Language Models Solve Olympiad Programming?\" by Quan Shi, Michael Tang, Karthik Narasimhan, and Shunyu Yao. You can check out their paper at the following link:\n",
    "\n",
    "[![arXiv](http://img.shields.io/badge/cs.CL-arXiv%3A2404.10952v1-B31B1B.svg)](https://arxiv.org/abs/2404.10952v1)\n",
    "\n",
    "You will construct an agentic graph capable of answering programming questions of increasing difficulty.\n",
    "\n",
    "1. **Reflection**: In part 1, you will create a zero-shot tool calling agent and prompt it to reflect on the test case results to correct its initial errors. This is similar to the agent the paper reported as having a pass rate of 12.38 on the USACO benchmark.\n",
    "2. **Retrieval**: In Part 2, you will implement an initial retrieval step as \"episodic memory\" for the agent that retrieves high-quality few-shot examples from our corpora of programming problems to help solve the **bronze** level question. This agent is similar to the one the paper benchmarked at 20.2.\n",
    "3. **Human-in-the-loop**: In part 3, you will use `interrupt_after` to let the user copilot the agent to a better answer. The benchmark performance then is constrained only by the competitiveness of the human it is paired with.\n",
    "\n",
    "Your final agent graph will be structured like the diagram below:\n",
    "\n",
    "![diagram](./img/diagram.png)\n",
    "\n",
    "Parts 1 and 2 are analogous to the systems benchmarked in the paper as having a pass rate of 12.38 and 20.2 respectively.\n",
    "\n",
    "![Benchmark system results](./img/benchmark.png)\n",
    "\n",
    "\n",
    "While LLMs are not yet capable of autonomously solving all these problems, we can design the system that far surpasses the capabilities of a basic ReAct agent at answering these questions. \n",
    "\n",
    "Before diving in, let's set up our machine. This will involve installing dependencies, fetching the dataset, and defining a utility function.\n",
    "\n",
    "## Setup\n",
    "\n",
    "For this tutorial, we will need to install some dependencies, fetch the Olympiad dataset, and define a utility function to help run the candidate solutions to see if they pass the test cases.\n",
    "\n",
    "First, install the requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c686827a-8078-4fd4-af7a-638ca1362796",
   "metadata": {},
   "outputs": [],
   "source": ["%%capture --no-stderr\n%pip install -U langgraph langsmith langchain_anthropic datasets langchain langchainhub"]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e2e542bb-a99e-44d3-8ebb-6a952dcbf2bf",
   "metadata": {},
   "outputs": [],
   "source": ["import getpass\nimport os\n\n\ndef _get_env(var: str):\n    if not os.environ.get(var):\n        os.environ[var] = getpass.getpass(f\"{var}: \")\n\n\n_get_env(\"ANTHROPIC_API_KEY\")\n# Recommended\n_get_env(\"LANGCHAIN_API_KEY\")\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\""]
  },
  {
   "cell_type": "markdown",
   "id": "c86ebbeb-c070-45d4-99ef-de33b53d447d",
   "metadata": {},
   "source": [
    "#### Data\n",
    "\n",
    "Fetch the USACO benchmark data using the util below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f7a0c7bd-512d-4e5b-ab43-1bc3b8c97fd4",
   "metadata": {},
   "outputs": [],
   "source": ["import os\nimport zipfile\n\nimport datasets\nimport requests\n\nusaco_url = \"https://storage.googleapis.com/benchmarks-artifacts/usaco/usaco_sampled_with_tests.zip\"\nzip_path = \"usaco.zip\"\nextract_path = \"usaco_datasets\"\n\nresponse = requests.get(usaco_url)\nwith open(zip_path, \"wb\") as file:\n    file.write(response.content)\n\nwith zipfile.ZipFile(zip_path, \"r\") as zip_ref:\n    zip_ref.extractall(extract_path)\n\nos.remove(zip_path)\n\nds = datasets.load_from_disk(os.path.join(extract_path, \"usaco_v3_sampled_with_tests\"))"]
  },
  {
   "cell_type": "markdown",
   "id": "2e4ad035-2c1a-4e31-93b4-58793d219bc9",
   "metadata": {},
   "source": [
    "#### Test Evaluation Utils\n",
    "\n",
    "We also need a way to evaluate our generated code. We will use this unsafe code execution program to run the generated code against our test cases.\n",
    "**Note:** The code below runs arbitrary code on your local machine! Proceed with caution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "54f9d037-121e-412f-857a-3e0ccc73892e",
   "metadata": {},
   "outputs": [],
   "source": ["import multiprocessing\nimport queue\nimport subprocess\nimport sys\nimport time\nimport traceback\n\nmultiprocessing.set_start_method(\"fork\", force=True)\n# WARNING\n# This program exists to execute untrusted model-generated code. Although\n# it is highly unlikely that model-generated code will do something overtly\n# malicious in response to this test suite, model-generated code may act\n# destructively due to a lack of model capability or alignment.\n# Users are strongly encouraged to sandbox this evaluation suite so that it\n# does not perform destructive actions on their host or network.\n# Proceed at your own risk:\n\n\ndef exec_program(q, program, input_data, expected_output, timeout):\n    try:\n        start_time = time.time()\n        process = subprocess.Popen(\n            [sys.executable, \"-c\", program],\n            stdin=subprocess.PIPE,\n            stdout=subprocess.PIPE,\n            stderr=subprocess.PIPE,\n            text=True,\n        )\n        stdout, stderr = process.communicate(input=input_data, timeout=timeout)\n        if time.time() - start_time > timeout:\n            raise TimeoutError(\"Execution timed out.\")\n        if process.returncode != 0:\n            q.put(f\"failed: {stderr}\")\n        else:\n            if stdout.strip() == expected_output.strip():\n                q.put(\"passed\")\n            else:\n                q.put(f\"wrong answer. Expected '{expected_output}', got '{stdout}'\")\n    except subprocess.TimeoutExpired:\n        process.kill()\n        q.put(\"timed out\")\n    except Exception:\n        q.put(f\"failed: {traceback.format_exc()}\")\n\n\ndef check_correctness(\n    program: str, input_data: str, expected_output: str, timeout: float\n) -> str:\n    q = multiprocessing.Queue()\n    process = multiprocessing.Process(\n        target=exec_program, args=(q, program, input_data, expected_output, timeout)\n    )\n    process.start()\n    process.join(timeout=timeout + 1)\n    if process.is_alive():\n        process.terminate()\n        process.join()\n        result = \"timed out\"\n    else:\n        try:\n            result = q.get_nowait()\n        except queue.Empty:\n            result = \"no result returned\"\n    return result"]
  },
  {
   "cell_type": "markdown",
   "id": "1e799866-6334-4c3a-8d03-7b7b1cf730ab",
   "metadata": {},
   "source": [
    "Let's check an example program and output to see how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "411cfc6a-d430-4642-8f48-2d6335430dd9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example 1:  passed\n",
      "Example 2:  wrong answer. Expected 'hi there', got 'goodbye\n",
      "'\n"
     ]
    }
   ],
   "source": ["program_code = \"print('hello, world!')\"\ninput_data = \"\"\nexpected_output = \"hello, world!\"\ntimeout = 2\n\ntest_result = check_correctness(program_code, input_data, expected_output, timeout)\nprint(\"Example 1: \", test_result)\ntest_result = check_correctness(\"print('goodbye')\", input_data, \"hi there\", timeout)\nprint(\"Example 2: \", test_result)"]
  },
  {
   "cell_type": "markdown",
   "id": "1d5da52f-f804-47b9-989a-25a1b2f461f6",
   "metadata": {},
   "source": [
    "## Part 1: Zero-Shot with Reflection\n",
    "\n",
    "In our first section, we will build a simple zero-shot tool-calling agent to try to solve these problems. We will incorporate a simple form of [reflection](https://www.youtube.com/watch?v=v5ymBTXNqtk) directly in the agent's tool calling schema by adding a \"reasoning\" field. Furthermore, Claude was trained to \"reason\" with freeform text prior to invoking any tools. Together, this should induce reflective \"chain-of-thought\" prompting.\n",
    "\n",
    "_Note: this diverges somewhat from the paper's implementation, which uses an explicit reflection step with a variation of the [Reflexion](../reflexion/reflexion.ipynb) prompt._\n",
    "\n",
    "By the end of this section, we will have built a reflective zero-shot programming agent that looks like the section marked \"Part 1\" in the system diagram below:\n",
    "\n",
    "![Part 1 diagram](./img/diagram-part-1.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00f91dac-d13b-4221-be1b-9254ca849c8d",
   "metadata": {},
   "source": [
    "### State\n",
    "\n",
    "LangGraph's main primitive is the `StateGraph`, which you use to define an agent as a controllable state machine.  The graph has `node`'s (python functions) that perform the work, and `edge`s that define how to route between the nodes.\n",
    "The `State` defines the interface between each node and carries all the information your agent needs.\n",
    "\n",
    "Below, define a `State` for our programming olympiad agent. The `messages` will track the sequence of submissions (and test case feedback) as chat history. The `status` field will flip from `in_progress` to `success` if the submission passes all test cases.\n",
    "The other fields (test_cases, runtime_limit) are used by the `evaluation` node to test the agent's submissions. These values are not seen by the agent itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f43d68d9-10be-4544-879a-88a33db18bea",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import Annotated\n\nfrom typing_extensions import TypedDict\n\nfrom langgraph.graph.message import AnyMessage, add_messages\n\n\nclass TestCase(TypedDict):\n    inputs: str\n    outputs: str\n\n\nclass State(TypedDict):\n    # Append-only chat memory so the agent can try to recover from initial mistakes.\n    messages: Annotated[list[AnyMessage], add_messages]\n    # From the dataset. These are used for testing.\n    test_cases: list[TestCase]\n    runtime_limit: int\n    status: str"]
  },
  {
   "cell_type": "markdown",
   "id": "64921a60-411a-4a9b-aaf0-8476d02d8a3a",
   "metadata": {},
   "source": [
    "Now, convert the dataset into inputs our graph will accept."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6d56776f-993b-4ca7-89ef-21dec01dc9d3",
   "metadata": {},
   "outputs": [],
   "source": ["input_states = [\n    {\n        \"messages\": [(\"user\", row[\"description\"])],\n        \"test_cases\": row[\"test_cases\"],\n        \"runtime_limit\": row[\"runtime_limit\"],\n        \"status\": \"in_progress\",\n        \"problem_level\": row[\"problem_level\"],\n    }\n    for row in ds\n]"]
  },
  {
   "cell_type": "markdown",
   "id": "9e7883ee-b6b3-4d89-b5a5-2e139fe361c9",
   "metadata": {},
   "source": [
    "#### Node 1: Solver\n",
    "\n",
    "Create a `solver` node that prompts an LLM \"agent\" to use a [writePython tool](https://python.langchain.com/v0.2/docs/integrations/chat/anthropic/#integration-details) to generate the submitted code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7b9e7742-16a3-4ad2-bc63-5f9cd4fd734b",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.language_models import BaseChatModel\nfrom langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\n\n\nclass writePython(BaseModel):\n    \"\"\"Write python code that resolves the problem.\"\"\"\n\n    reasoning: str = Field(..., description=\"Conceptual solution.\")\n    pseudocode: str = Field(..., description=\"Detailed English pseudocode.\")\n    code: str = Field(..., description=\"Valid Python 3 solution to the problem\")\n\n\nclass Solver:\n    def __init__(self, llm: BaseChatModel, prompt: ChatPromptTemplate):\n        self.runnable = prompt | llm.bind_tools([writePython])\n\n    def __call__(self, state: State) -> dict:\n        # Our agent only can see the \"messages\" and will ignore the test info\n        return {\"messages\": [self.runnable.invoke({\"messages\": state[\"messages\"]})]}"]
  },
  {
   "cell_type": "markdown",
   "id": "c22eaa6d-36ae-4526-9b80-49d615fc055c",
   "metadata": {},
   "source": [
    "Now, create the solver below. We'll use Claude Opus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6cc472f1-b9b3-4f81-a797-c64704bb07d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***********************************Prompt***********************************\n",
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "You are a world-class competitive programmer.\n",
      "Please reply with a Python 3 solution to the problem below. \n",
      "First, reason through the problem and conceptualize a solution.\n",
      "Then write detailed pseudocode to uncover any potential logical errors or omissions.\n",
      "Finally output the working Python code for your solution, ensuring to fix any errors uncovered while writing pseudocode.\n",
      "\n",
      "No outside libraries are allowed.\u001b[33;1m\u001b[1;3m{examples}\u001b[0m\n",
      "\n",
      "=============================\u001b[1m Messages Placeholder \u001b[0m=============================\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{messages}\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/wfh/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: The function `bind_tools` is in beta. It is actively being worked on, so the API may change.\n",
      "  warn_beta(\n"
     ]
    }
   ],
   "source": ["from langchain import hub\nfrom langchain_anthropic import ChatAnthropic\n\n# For this section, we are testing zero-shot performance and won't have\n# any examples. Partial them out to pre-fill the template.\nprompt = hub.pull(\"wfh/usaco-draft-solver\").partial(examples=\"\")\nprint(\"*\" * 35 + \"Prompt\" + \"*\" * 35)\nprompt.pretty_print()\n\n# Use Haiku if you want to save $$ while (almost) never correctly answering the question\n# llm = ChatAnthropic(model=\"claude-3-haiku-20240307\")\nllm = ChatAnthropic(model=\"claude-3-opus-20240229\")\n\nsolver = Solver(llm, prompt)"]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5d9560ba-900a-43d1-ad38-f132fd660337",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "********************************** Example **********************************\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "[{'text': \"<thinking>\\nTo address this problem, we need to use the writePython function, which requires the following parameters:\\n- reasoning: a conceptual solution to the problem\\n- pseudocode: detailed pseudocode for the solution\\n- code: working Python code implementing the solution\\n\\nThe key aspects to address in the solution are:\\n1. We have an infinite stream, so we can't store all elements. Need an online algorithm.\\n2. Need to ensure each element has an equal probability of being in the final sample.\\n\\nI believe I have enough information to provide values for all the required parameters.\\n</thinking>\", 'type': 'text'}, {'id': 'toolu_01UqpLYyueky5GtYMidS9oLF', 'input': {'reasoning': 'To get a perfectly random sample of size k from an infinite stream:\\n\\n1. Store the first k elements in an array (reservoir). \\n2. For each ith element after the kth element (i > k):\\n   - Generate a random integer j between 0 and i (inclusive)\\n   - If j < k, replace the jth element of the reservoir with the ith element\\n3. At the end, the reservoir contains the random sample.\\n\\nThis works because for any element, when we process the nth element, the probability that it is in the reservoir is:\\n- k/n when n <= k (first k elements always selected)\\n- k/n * k/(n-1) * k/(n-2) * ... * k/(k+1) = k/n when n > k\\n\\nSo any element has k/n probability of being in final reservoir, giving a perfectly random sample.', 'pseudocode': '```\\nfunction selectKItems(stream, k):\\n    reservoir = [0..k-1]  # store first k elements\\n\\n    i = k\\n    while stream has next item:\\n        item = stream.next()\\n        j = random(0, i)  # generate random index between 0 and i\\n        if j < k:\\n            reservoir[j] = item  # replace element at random index with new item\\n        i += 1\\n\\n    return reservoir\\n```', 'code': 'import random\\n\\ndef reservoir_sampling(stream, k):\\n    reservoir = []\\n    \\n    # Store first k elements in reservoir\\n    for i in range(k):\\n        reservoir.append(next(stream))\\n\\n    i = k\\n    for item in stream:\\n        # Generate random index between 0 and i\\n        j = random.randint(0, i) \\n        \\n        # Replace element at random index with new item\\n        if j < k:\\n            reservoir[j] = item\\n        i += 1\\n\\n    return reservoir'}, 'name': 'writePython', 'type': 'tool_use'}]\n"
     ]
    }
   ],
   "source": ["print(\"*\" * 34 + \" Example \" + \"*\" * 34)\nresult = solver(\n    {\n        \"messages\": [\n            (\n                \"user\",\n                \"How do I get a perfectly random sample from an infinite stream\",\n            )\n        ]\n    }\n)\nresult[\"messages\"][0].pretty_print()\n# Could expand to include (1)\n# 1. Restate the problem in plain English\n# 2. Closely following the explanation, restate and explain the solution in plain English\n# 3. Write a pseudocode solution\n# 4. Output the final Python solution with your solution steps in comments."]
  },
  {
   "cell_type": "markdown",
   "id": "58c5fa7a-a2eb-4954-be34-4194a58f00d3",
   "metadata": {},
   "source": [
    "#### Node 2: Evaluate\n",
    "\n",
    "Now define the \"`evaluate`\" node. This node takes the `solver`'s submitted code and executes it against the `test_cases` in our `State`.\n",
    "This uses the unsafe `check_correctness` utility we defined in the setup above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1785015b-24f8-415f-b950-e229b5137887",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.messages import AIMessage, HumanMessage, ToolMessage\n\n\n# This is the node we will add to the graph.\n# Most tool-calling APIs require that the `ToolMessage` contain the ID\n# of the\ndef format_tool_message(response: str, ai_message: AIMessage):\n    return ToolMessage(\n        content=response + \"\\nMake all fixes using the writePython tool.\",\n        tool_call_id=ai_message.tool_calls[0][\"id\"],\n    )\n\n\ndef evaluate(state: State):\n    test_cases = state[\"test_cases\"]\n    ai_message: AIMessage = state[\"messages\"][-1]\n    if not ai_message.tool_calls:\n        return {\n            \"messages\": [\n                HumanMessage(\n                    content=\"No code submitted. Please try again using the correct python code.\"\n                )\n            ]\n        }\n    try:\n        code = ai_message.tool_calls[0][\"args\"][\"code\"]\n    except Exception as e:\n        return {\"messages\": [format_tool_message(repr(e), ai_message)]}\n    num_test_cases = len(test_cases)\n    succeeded = 0\n    test_results = []\n    # TODO: Multiprocess\n    for test_case in test_cases:\n        input_data = test_case[\"inputs\"]\n        expected_output = test_case[\"outputs\"]\n        test_result = check_correctness(code, input_data, expected_output, timeout)\n        test_results.append(test_result)\n        if test_result == \"passed\":\n            succeeded += 1\n    pass_rate = succeeded / num_test_cases if num_test_cases else \"N/A\"\n    if pass_rate == 1:\n        return {\"status\": \"success\"}\n\n    responses = \"\\n\".join(\n        [f\"<test id={i}>\\n{r}\\n</test>\" for i, r in enumerate(test_results)]\n    )\n    response = f\"Incorrect submission. Please respond with updated code.\\nPass rate: {succeeded}/{num_test_cases}\\nResults:\\n{responses}\"\n    formatted_message = format_tool_message(response, ai_message)\n    return {\"messages\": [formatted_message]}"]
  },
  {
   "cell_type": "markdown",
   "id": "76f2a46d-c3e2-49d4-b44b-cadb6fd4e5a0",
   "metadata": {},
   "source": [
    "#### Create Graph\n",
    "\n",
    "Now, put it all together! Once you've defined each node, defining the connectivity / state transitions is fairly easy.\n",
    "\n",
    "Our Zero-shot graph defines a loop. If we visualize the data flow, we want the logic to:\n",
    "1. First go to the `solver`, which attempts a first solution.\n",
    "2. Next go to the `evaluate` node, which tests the solution.\n",
    "3. If the solution passes, end, otherwise, return to the `solver` to try again.\n",
    "\n",
    "In LangGraph, we use `conditional_edges` to define state transitions that contain conditional logic.\n",
    "Below, define the graph, adding a `control_edge` to handle step (3) above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "caf1560e-1517-4229-8a43-186816da6a3a",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.graph import END, StateGraph, START\n\nbuilder = StateGraph(State)\nbuilder.add_node(\"solver\", solver)\nbuilder.add_edge(START, \"solver\")\nbuilder.add_node(\"evaluate\", evaluate)\nbuilder.add_edge(\"solver\", \"evaluate\")\n\n\ndef control_edge(state: State):\n    if state.get(\"status\") == \"success\":\n        return END\n    return \"solver\"\n\n\nbuilder.add_conditional_edges(\"evaluate\", control_edge, {END: END, \"solver\": \"solver\"})\ngraph = builder.compile()"]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7275a2c3-1818-4d14-a7a5-97bc56243a9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 16544 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(graph.get_graph().draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
  },
  {
   "cell_type": "markdown",
   "id": "fc4b1bb4-262f-41c3-822d-e015daf0744a",
   "metadata": {},
   "source": [
    "Now that we've created our graph, let's see the type of question it will have to solve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "308d6e18-69ad-4e8d-9c5f-06111b0806ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Farmer John has $N$ ($1 \\leq N \\leq 2 \\cdot 10^5$) farms, numbered from $1$ to\n",
      "$N$. It is known that FJ closes farm $i$ at time $c_i$. Bessie wakes up at time\n",
      "$S$, and wants to maximize the productivity of her day by visiting as many farms\n",
      "as possible before they close. She plans to visit farm $i$ on time $t_i + S$.\n",
      "Bessie must arrive at a farm strictly before Farmer John closes it to actually visit it.\n",
      "\n",
      "Bessie has $Q$ $(1 \\leq Q \\leq 2 \\cdot 10^5)$ queries. For each query, she gives\n",
      "you two integers $S$ and $V$. For each query, output whether Bessie can visit at\n",
      "least $V$ farms if she wakes up at time $S$.\n",
      "\n",
      "INPUT FORMAT (input arrives from the terminal / stdin):\n",
      "The first line consists of $N$ and $Q$.\n",
      "\n",
      "The second line consists of $c_1, c_2, c_3 \\dots c_N$ ($1 \\leq c_i \\leq 10^6$).\n",
      "\n",
      "The third line consists of $t_1, t_2, t_3 \\dots t_N$ ($1 \\leq t_i \\leq 10^6$).\n",
      "\n",
      "The next $Q$ lines each consist of two integers $V$ ($1 \\leq V \\leq N$) and $S$\n",
      "($1 \\leq S \\leq 10^6$).\n",
      "\n",
      "OUTPUT FORMAT (print output to the terminal / stdout):\n",
      "For each of the $Q$ queries, output YES or NO on a new line.\n",
      "\n",
      "SAMPLE INPUT:\n",
      "5 5\n",
      "3 5 7 9 12\n",
      "4 2 3 3 8\n",
      "1 5\n",
      "1 6\n",
      "3 3\n",
      "4 2\n",
      "5 1\n",
      "SAMPLE OUTPUT: \n",
      "YES\n",
      "NO\n",
      "YES\n",
      "YES\n",
      "NO\n",
      "\n",
      "For the first query, Bessie will visit the farms at time $t = [9, 7, 8, 8, 13]$,\n",
      "so she will only get to visit farm $4$ on time before FJ closes the farm.\n",
      "\n",
      "For the second query, Bessie will not be able to visit any of the farms on time.\n",
      "\n",
      "For the third query, Bessie will visit farms $3, 4, 5$ on time.\n",
      "\n",
      "For the fourth and fifth queries, Bessie will be able to visit all but the first\n",
      "farm on time.\n",
      "\n",
      "SCORING:\n",
      "Inputs 2-4: $N,Q\\le 10^3$Inputs 5-9: $c_i, t_i \\le 20$Inputs 10-17: No additional constraints.\n",
      "\n",
      "\n",
      "Problem credits: Chongtian Ma\n",
      "\n"
     ]
    }
   ],
   "source": ["input_state = input_states[0].copy()\n# We will reduce the test cases to speed this notebook up\ninput_state[\"test_cases\"] = input_state[\"test_cases\"][:3]\nprint(input_state[\"messages\"][0][1])"]
  },
  {
   "cell_type": "markdown",
   "id": "3c444be0-4c7d-49cc-9aaa-2b99dabe3922",
   "metadata": {},
   "source": [
    "Pretty difficult! Let's run our simple \"zero-shot\" agent below to see how it fares. **It most likely will not be able to solve this question** (unless you are using a more powerful model than what I had available at the time of writing this tutorial (2024/04/20).\n",
    "We will trace the trajectory to LangSmith to review the series of submissions. To reduce the packet size, we will use \"`hide_inputs`\" and filter out the test_cases. All this is optional but useful for development. \n",
    "\n",
    "**Note:** We _expect_ a **GraphRecursionError** here from it not being able to answer it correctly in the allocated number of steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "2ebe0da0-8f27-4805-829c-54bc1dc29ba4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Assistant: [{'text': '<thinking>\\nThe key steps to solve this\n",
      "Assistant: KeyError('code')\\nMake all fixes using the writePy\n",
      "Assistant: [{'id': 'toolu_01KimhKt8aqQjGZJmrHVnAtE', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01CMZTqAd7BZQ2nSgtk9djRW', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01Kbaq9gX4BnHvps6TMfVGHL', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01MiSnpiGK5Yy4Cpp6GGbjmT', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01GWuvJezXLMVurUBG84odDP', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01W8DGmhcpFVctySmx58scf9', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_018bhYtCKDK6S4MHiAxUZCrb', 'input':\n",
      "Assistant: KeyError('code')\\nMake all fixes using the writePy\n",
      "Assistant: [{'id': 'toolu_01LCwaCjX9uZBV3jt9eAkmAa', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01WqJvdE2WDeTZXoKp2V7PWb', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_01DGevkunt9zWx7SVDCHdBuv', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'id': 'toolu_013comYKVxNSzTM4ZbH3L3FP', 'input':\n",
      "Assistant: Incorrect submission. Please respond with updated \n"
     ]
    },
    {
     "ename": "GraphRecursionError",
     "evalue": "Recursion limit of 25 reachedwithout hitting a stop condition. You can increase the limit by setting the `recursion_limit` config key.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mGraphRecursionError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[25], line 17\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m tracing_v2_enabled(client\u001b[38;5;241m=\u001b[39mclient):\n\u001b[1;32m     16\u001b[0m     events \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39mstream(input_state)\n\u001b[0;32m---> 17\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevent\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevents\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevent\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalues\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m            \u001b[49m\u001b[43mmessages\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\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\n",
      "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langgraph/pregel/__init__.py:645\u001b[0m, in \u001b[0;36mPregel.stream\u001b[0;34m(self, input, config, stream_mode, output_keys, input_keys, interrupt_before_nodes, interrupt_after_nodes, debug)\u001b[0m\n\u001b[1;32m    643\u001b[0m         \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m    644\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m step \u001b[38;5;241m==\u001b[39m config[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrecursion_limit\u001b[39m\u001b[38;5;124m\"\u001b[39m]:\n\u001b[0;32m--> 645\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m GraphRecursionError(\n\u001b[1;32m    646\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRecursion limit of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mconfig[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mrecursion_limit\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m reached\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    647\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwithout hitting a stop condition. You can increase the \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    648\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlimit by setting the `recursion_limit` config key.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    649\u001b[0m     )\n\u001b[1;32m    651\u001b[0m \u001b[38;5;66;03m# before execution, check if we should interrupt\u001b[39;00m\n\u001b[1;32m    652\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _should_interrupt(\n\u001b[1;32m    653\u001b[0m     checkpoint,\n\u001b[1;32m    654\u001b[0m     interrupt_before_nodes,\n\u001b[1;32m    655\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstream_channels_list,\n\u001b[1;32m    656\u001b[0m     next_tasks,\n\u001b[1;32m    657\u001b[0m ):\n",
      "\u001b[0;31mGraphRecursionError\u001b[0m: Recursion limit of 25 reachedwithout hitting a stop condition. You can increase the limit by setting the `recursion_limit` config key."
     ]
    }
   ],
   "source": ["from langchain_core.tracers.context import tracing_v2_enabled\nfrom langsmith import Client\n\n\n# We don't need to include all the test cases in our traces.\ndef _hide_test_cases(inputs):\n    copied = inputs.copy()\n    # These are tens of MB in size. No need to send them up\n    copied[\"test_cases\"] = \"...\"\n    return copied\n\n\nclient = Client(hide_inputs=_hide_test_cases, hide_outputs=_hide_test_cases)\nwith tracing_v2_enabled(client=client):\n    events = graph.stream(input_state)\n    for event in events:\n        for value in event.values():\n            messages = value.get(\"messages\")\n            if messages:\n                if isinstance(messages, list):\n                    messages = value[\"messages\"][-1]\n                print(\n                    \"Assistant:\",\n                    str(messages.content).replace(\"\\n\", \"\\\\n\")[:50],\n                )"]
  },
  {
   "cell_type": "markdown",
   "id": "2c496d9a-a95b-4cab-86fa-daafc7ddd065",
   "metadata": {},
   "source": [
    "It wasn't able to solve it in time **but that's OK**! If it were easy, this paper would be a lot shorter :)\n",
    "\n",
    "You can view the [agent's full LangSmith trace](https://smith.langchain.com/public/61c84ad0-51db-40f1-b50d-6983d9481ca1/r) at the provided link.\n",
    "\n",
    "In the next section we will add an improvement the paper terms \"episodic memory\", which in this case is really few-shot retrieval."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "090a1639-09c5-4682-813f-71100f236eb3",
   "metadata": {},
   "source": [
    "## Part 2: Few-shot Retrieval\n",
    "\n",
    "Even with reflective tool calling, our baseline agent from part 1 struggled with this difficult task. One way to \"teach\" an LLM how to better perform a task is through demonstrations, also known as \"few-shot examples.\"\n",
    "\n",
    "What the authors of the USACO paper call \"episodic memory\" **is really just few-shot prompting over similar examples.**\n",
    "\n",
    "Each examples in this case is a different problems + solution within the dataset. The term \"episodic memory\" makes sense if you pretend your agent has already \"solved\" these problems and is recalling its solutions to them.\n",
    "\n",
    "This section adds the \"Episodic Memory\" components from \"Part 2\" in the diagram below.\n",
    "\n",
    "![Part 2 diagram](./img/diagram-part-2.png)\n",
    "\n",
    "Note that this memory step is performed **one time**,  **before** the logic of our zero-shot loop from part 1. The steps are as follows:\n",
    "\n",
    "1. Prompt the LLM to generate a candidate solution.\n",
    "2. Use the text of the candidate solution to retrieve the N most similar (problem, solution) pairs.\n",
    "3. Format this result in the Zero-shot agent's prompt.\n",
    "\n",
    "Below, let's implement our episodic memory as a retriever. We will follow the paper's retriever selection and use [BM25](https://en.wikipedia.org/wiki/Okapi_BM25)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d612dd8d-31af-426c-944b-203acd55ace0",
   "metadata": {},
   "outputs": [],
   "source": ["%%capture --no-stderr\n%pip install --upgrade --quiet  rank_bm25"]
  },
  {
   "cell_type": "markdown",
   "id": "32e0c485-8eba-41e6-ae29-64bcea98d19e",
   "metadata": {},
   "source": [
    "#### State\n",
    "\n",
    "The state is mostly recycled from part 1. Add additional \"candidate\" and \"examples\" fields to store the information for the memory steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "16937fef-58b9-4ab2-bbfc-5237aad235ec",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import Annotated\n\nfrom typing_extensions import TypedDict\n\nfrom langgraph.graph.message import AnyMessage, add_messages\n\n\nclass TestCase(TypedDict):\n    inputs: str\n    outputs: str\n\n\nclass State(TypedDict):\n    # NEW! Candidate for retrieval + formatted fetched examples as \"memory\"\n    candidate: AIMessage\n    examples: str\n    # Repeated from Part 1\n    messages: Annotated[list[AnyMessage], add_messages]\n    test_cases: list[TestCase]\n    runtime_limit: int\n    status: str"]
  },
  {
   "cell_type": "markdown",
   "id": "61fb9104-e73f-47db-9af5-84a375a8d323",
   "metadata": {},
   "source": [
    "#### Nodes 1 and 3: Draft & Solver\n",
    "\n",
    "Let's create our \"agent\". We will modify the `Solver` from Part 1 to reuse it for  for the agent node and for the candidate program generation node (\"draft\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "25f947a7-15bb-4119-a47e-b5c33ca0a249",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain import hub\nfrom langchain_anthropic import ChatAnthropic\n\n\nclass Solver:\n    def __init__(self, llm: BaseChatModel, prompt: ChatPromptTemplate):\n        self.runnable = prompt | llm.bind_tools([writePython])\n\n    def __call__(self, state: State) -> dict:\n        # Our agent only can see the \"messages\" and will ignore the test info\n        inputs = {\"messages\": state[\"messages\"]}\n        has_examples = bool(state.get(\"examples\"))\n        output_key = \"candidate\"  # Used in the draft node\n        if has_examples:\n            output_key = \"messages\"\n            # Used in the solve node\n            inputs[\"examples\"] = state[\"examples\"]\n        response = self.runnable.invoke(inputs)\n        if not response.content:\n            return {\n                output_key: AIMessage(\n                    content=\"I'll need to think about this step by step.\"\n                )\n            }\n        return {output_key: response}\n\n\nprompt = hub.pull(\"wfh/usaco-draft-solver\")\nllm = ChatAnthropic(model=\"claude-3-opus-20240229\")\n\ndraft_solver = Solver(llm, prompt.partial(examples=\"\"))\nsolver = Solver(llm, prompt)"]
  },
  {
   "cell_type": "markdown",
   "id": "273a526b-3b52-4630-a58e-0317a0034609",
   "metadata": {},
   "source": [
    "#### Node 2: Retrieve\n",
    "\n",
    "The retrieve node takes a candidate solution (made by the 'solver' node), uses _this_ to search for similar examples, then formats those in the message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e5e0aa40-79a4-4071-9ad2-9aa2f36599ce",
   "metadata": {},
   "outputs": [],
   "source": ["# We will test our agent on index 0 (the same as above).\n# Later, we will test on index 2 (the first 'silver difficulty' question)\ntest_indices = [0, 2]\ntrain_ds = [row for i, row in enumerate(ds) if i not in test_indices]\ntest_ds = [row for i, row in enumerate(ds) if i in test_indices]"]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "96a1ff96-7556-4959-9f54-1ade3bd1c01a",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_community.retrievers import BM25Retriever\n\n\ndef format_example(row):\n    question = row[\"description\"]\n    answer = row[\"solution\"]\n    return f\"\"\"<problem>\n{question}\n</problem>\n<solution>\n{answer}\n</solution>\"\"\"\n\n\n# Skip our 'test examples' to avoid cheating\n# This is \"simulating\" having seen other in-context examples\nretriever = BM25Retriever.from_texts([format_example(row) for row in train_ds])"]
  },
  {
   "cell_type": "markdown",
   "id": "6ad20d3a-291c-41a5-a343-bd14284f76f2",
   "metadata": {},
   "source": [
    "Now define the node. Any node can optionally accept a second `config` positional argument. This contains `configurable` params you can adjust when invoking the graph. For instance, we can\n",
    "adjust the top `k` examples to retrieve for our agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "af42962d-c06e-4b6e-96df-72ad48f17617",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.runnables import RunnableConfig\n\n\ndef retrieve_examples(state: State, config: RunnableConfig):\n    top_k = config[\"configurable\"].get(\"k\") or 2\n    ai_message: AIMessage = state[\"candidate\"]\n    if not ai_message.tool_calls:\n        # We err here. To make more robust, you could loop back\n        raise ValueError(\"Draft agent did not produce a valid code block\")\n    code = ai_message.tool_calls[0][\"args\"][\"code\"]\n    examples_str = \"\\n\".join(\n        [doc.page_content for doc in retriever.invoke(code)[:top_k]]\n    )\n    examples_str = f\"\"\"\nYou previously solved the following problems in this competition:\n<Examples>\n{examples_str}\n<Examples>\nApproach this new question with similar sophistication.\"\"\"\n    return {\"examples\": examples_str}"]
  },
  {
   "cell_type": "markdown",
   "id": "df5cbe72-fca0-4974-a769-284267d3df91",
   "metadata": {},
   "source": [
    "#### Graph\n",
    "\n",
    "Now let's put it all together. The graph is slightly more complicated than in part 1, since we have to add the initial \"draft\" and \"retrieve\" nodes to our agent loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e6e73e85-1232-4848-beba-3139ac7d0a64",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.checkpoint.memory import MemorySaver\nfrom langgraph.graph import END, StateGraph, START\n\nbuilder = StateGraph(State)\nbuilder.add_node(\"draft\", draft_solver)\nbuilder.add_edge(START, \"draft\")\nbuilder.add_node(\"retrieve\", retrieve_examples)\nbuilder.add_node(\"solve\", solver)\nbuilder.add_node(\"evaluate\", evaluate)\n# Add connectivity\nbuilder.add_edge(\"draft\", \"retrieve\")\nbuilder.add_edge(\"retrieve\", \"solve\")\nbuilder.add_edge(\"solve\", \"evaluate\")\n\n\ndef control_edge(state: State):\n    if state.get(\"status\") == \"success\":\n        return END\n    return \"solve\"\n\n\nbuilder.add_conditional_edges(\"evaluate\", control_edge, {END: END, \"solve\": \"solve\"})\n\n\ncheckpointer = MemorySaver()\ngraph = builder.compile(checkpointer=checkpointer)"]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "57acf78d-5e68-46dd-adae-2259e5c5d3f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 20864 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(graph.get_graph().draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
  },
  {
   "cell_type": "markdown",
   "id": "3bb633e0-0d02-4050-96a8-18265594385b",
   "metadata": {},
   "source": [
    "Let's try again on this problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "dc6f0455-a6e9-46a0-8f8b-03de00fb1f8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'text': \"<thinking>\\nThis problem essentially asks to find the number of farms Bessie can visit before they close at each query. The key insights are:\\n\\n1. Bessie's arrival time at each farm is S +\n",
      "Retrieved examples:\n",
      "\n",
      " \n",
      "You previously solved the following problems in this competition:\n",
      "<Examples>\n",
      "<problem>\n",
      "\n",
      "Farmer John...\n",
      "Assistant: [{'text': \"<thinking>\\nThe key information given i\n"
     ]
    }
   ],
   "source": ["config = {\"configurable\": {\"thread_id\": \"question-recall\", \"k\": 3}}\nwith tracing_v2_enabled(client=client):\n    events = graph.stream(input_state, config)\n    for event in events:\n        for value in event.values():\n            messages = value.get(\"messages\")\n            if messages:\n                if isinstance(messages, list):\n                    messages = value[\"messages\"][-1]\n                print(\n                    \"Assistant:\",\n                    str(messages.content).replace(\"\\n\", \"\\\\n\")[:50],\n                )\n            elif value.get(\"examples\"):\n                print(\"Retrieved examples:\\n\\n\", value[\"examples\"][:100] + \"...\")\n            elif value.get(\"candidate\"):\n                print(str(value[\"candidate\"].content)[:200])"]
  },
  {
   "cell_type": "markdown",
   "id": "1389ffb0-839e-4646-be68-dbc79279c8ae",
   "metadata": {},
   "source": [
    "**No recursion error!** You can view the [full LangSmith trace](https://smith.langchain.com/public/1f1c4db3-b53c-49bf-a287-a2b51c081156/r/31f90ddd-8ae9-4b23-a2b5-b0c0d67c5cc3) of the graph's execution at the provided link to confirm the results. You can also check the graph state to confirm that it passed all test cases successfully:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d55a1d95-9fa4-45b2-9f5e-7ca02fb8415b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'success'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["checkpoint = graph.get_state(config)\ncheckpoint.values[\"status\"]"]
  },
  {
   "cell_type": "markdown",
   "id": "50348c66-4d38-44d0-bb20-ae1a14517f78",
   "metadata": {},
   "source": [
    "**Congrats!** You added \"episodic memory\" to your agent to fetch few-shot examples and solve this bronze level programming olympiad question!\n",
    "\n",
    "Our agent is still limited, however. Let's test it out on a more challenging 🪙🏆silver✨ level question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1ef0b06f-c448-49a5-8f1f-f7041a5d6b87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'silver'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["silver_row = test_ds[1]\nsilver_row[\"problem_level\"]"]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "2ae55ddc-e629-4be7-badd-756608eec50e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'text': \"<thinking>\\nThe relevant tool for this problem is writePython. It requires the following parameters:\\n- reasoning: To solve this problem, we need to simulate the cruise by following the seq\n",
      "Retrieved examples:\n",
      "\n",
      " \n",
      "You previously solved the following problems in this competition:\n",
      "<Examples>\n",
      "<problem>\n",
      "\n",
      "Farmer John...\n",
      "Assistant: [{'text': \"<thinking>\\nTo solve this problem, we n\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nAfter reviewing the failed \n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nAfter reviewing the latest \n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nOops, looks like I made a s\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nHmm, some of the test cases\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': '<thinking>\\nOops, looks like I accident\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nLooks like the code is now \n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': '<thinking>\\nOops, looks like I accident\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nHmm, the optimization to si\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nOops, I did it again - acci\n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Assistant: [{'text': \"<thinking>\\nHmm, the latest code is sti\n",
      "Assistant: Incorrect submission. Please respond with updated \n"
     ]
    },
    {
     "ename": "GraphRecursionError",
     "evalue": "Recursion limit of 25 reachedwithout hitting a stop condition. You can increase the limit by setting the `recursion_limit` config key.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mGraphRecursionError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[37], line 12\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m tracing_v2_enabled(client\u001b[38;5;241m=\u001b[39mclient):\n\u001b[1;32m     11\u001b[0m     events \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39mstream(silver_input, config)\n\u001b[0;32m---> 12\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevent\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevents\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mevent\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalues\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m            \u001b[49m\u001b[43mmessages\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\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\n",
      "File \u001b[0;32m~/.pyenv/versions/3.11.2/lib/python3.11/site-packages/langgraph/pregel/__init__.py:645\u001b[0m, in \u001b[0;36mPregel.stream\u001b[0;34m(self, input, config, stream_mode, output_keys, input_keys, interrupt_before_nodes, interrupt_after_nodes, debug)\u001b[0m\n\u001b[1;32m    643\u001b[0m         \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m    644\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m step \u001b[38;5;241m==\u001b[39m config[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mrecursion_limit\u001b[39m\u001b[38;5;124m\"\u001b[39m]:\n\u001b[0;32m--> 645\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m GraphRecursionError(\n\u001b[1;32m    646\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mRecursion limit of \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mconfig[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mrecursion_limit\u001b[39m\u001b[38;5;124m'\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m reached\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    647\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwithout hitting a stop condition. You can increase the \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    648\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlimit by setting the `recursion_limit` config key.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    649\u001b[0m     )\n\u001b[1;32m    651\u001b[0m \u001b[38;5;66;03m# before execution, check if we should interrupt\u001b[39;00m\n\u001b[1;32m    652\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _should_interrupt(\n\u001b[1;32m    653\u001b[0m     checkpoint,\n\u001b[1;32m    654\u001b[0m     interrupt_before_nodes,\n\u001b[1;32m    655\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstream_channels_list,\n\u001b[1;32m    656\u001b[0m     next_tasks,\n\u001b[1;32m    657\u001b[0m ):\n",
      "\u001b[0;31mGraphRecursionError\u001b[0m: Recursion limit of 25 reachedwithout hitting a stop condition. You can increase the limit by setting the `recursion_limit` config key."
     ]
    }
   ],
   "source": ["silver_input = {\n    \"messages\": [(\"user\", silver_row[\"description\"])],\n    \"test_cases\": silver_row[\"test_cases\"],\n    \"runtime_limit\": silver_row[\"runtime_limit\"],\n    \"status\": \"in_progress\",\n}\n\n\nconfig = {\"configurable\": {\"thread_id\": \"silver-question-1\", \"k\": 2}}\nwith tracing_v2_enabled(client=client):\n    events = graph.stream(silver_input, config)\n    for event in events:\n        for value in event.values():\n            messages = value.get(\"messages\")\n            if messages:\n                if isinstance(messages, list):\n                    messages = value[\"messages\"][-1]\n                print(\n                    \"Assistant:\",\n                    str(messages.content).replace(\"\\n\", \"\\\\n\")[:50],\n                )\n            elif value.get(\"examples\"):\n                print(\"Retrieved examples:\\n\\n\", value[\"examples\"][:100] + \"...\")\n            elif value.get(\"candidate\"):\n                print(str(value[\"candidate\"].content)[:200])"]
  },
  {
   "cell_type": "markdown",
   "id": "7b4c9217-38de-4f48-9070-2092fd9ecc15",
   "metadata": {},
   "source": [
    "**Still too hard!** AGI not achieved yet. To investigate our agent's trajectory in detail, check out the [full LangSmith trace](https://smith.langchain.com/public/13018b44-0c4f-4f1a-9e6d-dea1f3fd4705/r).\n",
    "\n",
    "Our agent isn't good enough to be autonomous. The great thing about LangGraph is you don't have to decide between \"autonomous agent\" and \"simple DAG\": you can inject control and user-interfaces wherever it can usefully benefit your application."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59177162-3fb8-4307-8e7f-ea67e785d4cf",
   "metadata": {},
   "source": [
    "## Part 3: Human-in-the-loop\n",
    "\n",
    "Our retrieval-enhanced agent was able to solve the `bronze`-level question but still failed for those with the more challenging **silver** difficulty. \n",
    "\n",
    "Recall that the paper presented 3 complementary techniques that improved performance:\n",
    "\n",
    "1. Reflection: explicitly prompting the LLM to \"reflect\" on its mistakes can help it\n",
    "2. Few-shot prompting: retrieving relevant, high-quality examples as \"memory\"\n",
    "3. **Human-in-the-loop collaboration:**  without giving the correct answer, the human is allowed to help the agent reflect on its approach and point it in a better direction.\n",
    "\n",
    "\n",
    "In this section, we will add the \"human\" node (marked as \"part 3\" in the diagram below), completing our agent graph:\n",
    "\n",
    "![Diagram](./img/diagram.png)\n",
    "\n",
    "From an ML perspective, this is a bit of a [clever hans](https://en.wikipedia.org/wiki/Clever_Hans), but from the application designer's perspective, where the primary goal is to achieve a higher combined success rate, letting the human interject with thoughts and insights is only natural. \n",
    "\n",
    "In either case, adding a human check to a LangGraph instance requires no extra lines of code. Let's do so by instructing the graph to `interrupt_after` the \"`evaluate`\" node to give the user a chance to modify the trajectory.\n",
    "\n",
    "Start assembling your graph below. The following section is identical to our application in part 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "3c6456ba-363c-4133-8631-6dabb042b6ce",
   "metadata": {},
   "outputs": [],
   "source": ["# This is all the same as before\nfrom langgraph.checkpoint.memory import MemorySaver\nfrom langgraph.graph import END, StateGraph, START\n\nbuilder = StateGraph(State)\nprompt = hub.pull(\"wfh/usaco-draft-solver\")\nllm = ChatAnthropic(model=\"claude-3-opus-20240229\", max_tokens_to_sample=4000)\n\ndraft_solver = Solver(llm, prompt.partial(examples=\"\"))\nbuilder.add_node(\"draft\", draft_solver)\nbuilder.add_edge(START, \"draft\")\nbuilder.add_node(\"retrieve\", retrieve_examples)\nsolver = Solver(llm, prompt)\nbuilder.add_node(\"solve\", solver)\nbuilder.add_node(\"evaluate\", evaluate)\nbuilder.add_edge(\"draft\", \"retrieve\")\nbuilder.add_edge(\"retrieve\", \"solve\")\nbuilder.add_edge(\"solve\", \"evaluate\")\n\n\ndef control_edge(state: State):\n    if state.get(\"status\") == \"success\":\n        return END\n    return \"solve\"\n\n\nbuilder.add_conditional_edges(\"evaluate\", control_edge, {END: END, \"solve\": \"solve\"})\ncheckpointer = MemorySaver()"]
  },
  {
   "cell_type": "markdown",
   "id": "3d837103-ede4-4515-a0cd-0b3319382013",
   "metadata": {},
   "source": [
    "Now finish by compiling the graph. Set`interrupt_after=[\"evaluate\"]` to instruct the agent to wait for human input before continuing execution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "461c13ba-01cc-44e1-b837-6a64d03069d9",
   "metadata": {},
   "outputs": [],
   "source": ["graph = builder.compile(\n    checkpointer=checkpointer,\n    # New: this tells the graph to break any time it goes to the \"human\" node\n    interrupt_after=[\"evaluate\"],\n)"]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "040b2100-5b2e-40c9-b6af-a1b680e16ee3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 20864 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(graph.get_graph().draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
  },
  {
   "cell_type": "markdown",
   "id": "372462bf-c74e-4c18-bbee-c21c8132b2e5",
   "metadata": {},
   "source": [
    "As you can see in the graph above, the structure is the same as Part 2, except that we've inserted a \"`human`\" breakpoint between the \"`evaluate`\" and \"`solve`\" nodes.\n",
    "\n",
    "Let's try this question again!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5f9ad0d0-cdaf-4ba2-9527-b5b9739e7b67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'text': \"<thinking>\\nTo solve this problem, we need to:\\n1. Read in the input data - number of ports N, length of direction sequence M, number of repetitions K, the port connections, and the directi\n",
      "Retrieved examples:\n",
      "\n",
      " \n",
      "You previously solved the following problems in this competition:\n",
      "<Examples>\n",
      "<problem>\n",
      "Farmer John ...\n",
      "Assistant: [{'text': '<thinking>\\nTo determine where Bessie e\n",
      "Assistant: Incorrect submission. Please respond with updated \n"
     ]
    }
   ],
   "source": ["config = {\"configurable\": {\"thread_id\": \"silver-hl-1\", \"k\": 2}}\nwith tracing_v2_enabled(client=client):\n    events = graph.stream(silver_input, config)\n    for event in events:\n        for value in event.values():\n            messages = value.get(\"messages\")\n            if messages:\n                if isinstance(messages, list):\n                    messages = value[\"messages\"][-1]\n                print(\n                    \"Assistant:\",\n                    str(messages.content).replace(\"\\n\", \"\\\\n\")[:50],\n                )\n            elif value.get(\"examples\"):\n                print(\"Retrieved examples:\\n\\n\", value[\"examples\"][:100] + \"...\")\n            elif value.get(\"candidate\"):\n                print(str(value[\"candidate\"].content)[:200])"]
  },
  {
   "cell_type": "markdown",
   "id": "947e3089-6ae5-40e4-b993-7ef7e8630a12",
   "metadata": {},
   "source": [
    "**⏰Time to weigh in⏰:** our model failed in its first attempt, so we have the opportunity to give it some advice.\n",
    "\n",
    "Recall the original question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "4fcdf8c9-6a5a-4463-90ae-eec89a57cd05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem 3: Luxury River Cruise [Josh Alman and Nathan Pinsker, 2013]\n",
      "\n",
      "Farmer John is taking Bessie and the cows on a cruise! They are sailing on a \n",
      "network of rivers with N ports (1 <= N <= 1,000) labeled 1..N, and Bessie \n",
      "starts at port 1. Each port has exactly two rivers leading out of it which \n",
      "lead directly to other ports, and rivers can only be sailed one way.\n",
      "\n",
      "At each port, the tour guides choose either the \"left\" river or the \"right\" \n",
      "river to sail down next, but they keep repeating the same choices over and \n",
      "over. More specifically, the tour guides have chosen a short sequence of M \n",
      "directions (1 <= M <= 500), each either \"left\" or \"right\", and have\n",
      "repeated it K times (1 <= K <= 1,000,000,000). Bessie thinks she is going\n",
      "in circles -- help her figure out where she ends up!\n",
      "\n",
      "PROBLEM NAME: cruise\n",
      "\n",
      "INPUT FORMAT:\n",
      "\n",
      "* Line 1: Three space-separated integers N, M, and K.\n",
      "\n",
      "* Lines 2..N+1: Line i+1 has two space-separated integers,\n",
      "        representing the number of the ports that port i's left and\n",
      "        right rivers lead to, respectively.\n",
      "\n",
      "* Line N+2: M space-separated characters, either 'L' or 'R'. 'L'\n",
      "        represents a choice of  'left' and 'R' represents a choice of\n",
      "        'right'.\n",
      "\n",
      "SAMPLE INPUT:\n",
      "\n",
      "4 3 3\n",
      "2 4\n",
      "3 1\n",
      "4 2\n",
      "1 3\n",
      "L L R\n",
      "\n",
      "INPUT DETAILS:\n",
      "\n",
      "The port numbers are arranged clockwise in a circle, with 'L' being a \n",
      "clockwise rotation and 'R' being a counterclockwise rotation. The sequence \n",
      "taken is LLRLLRLLR.\n",
      "\n",
      "OUTPUT FORMAT:\n",
      "\n",
      "* Line 1: A single integer giving the number of the port where\n",
      "        Bessie's cruise ends.\n",
      "\n",
      "SAMPLE OUTPUT:\n",
      "\n",
      "4\n",
      "\n",
      "OUTPUT DETAILS:\n",
      "\n",
      "After the first iteration of the sequence of directions, Bessie is at port\n",
      "2 (1 -> 2 -> 3 -> 2); after the second, she is at port 3 (2 -> 3 -> 4 ->\n",
      "3), and at the end she is at port 4 (3 -> 4 -> 1 -> 4).\n",
      "\n"
     ]
    }
   ],
   "source": ["snapshot = graph.get_state(config)\nprint(snapshot.values[\"messages\"][0].content)"]
  },
  {
   "cell_type": "markdown",
   "id": "232ed165-e2cf-495f-a4bb-a0f3f87a719c",
   "metadata": {},
   "source": [
    "And then review the agent's current submission:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "2049980a-0a0c-4135-98c7-d4de1af4757b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<thinking>\n",
      "To determine where Bessie ends up, we need to:\n",
      "1. Simulate the cruise by following the sequence of left/right directions\n",
      "2. Repeat this sequence K times to find the final destination port\n",
      "\n",
      "The problem provides:\n",
      "- The number of ports N\n",
      "- The connections between ports (left and right rivers for each port)\n",
      "- The sequence of M directions (L or R) to follow\n",
      "- The number of times K to repeat the sequence\n",
      "\n",
      "With this information, we have everything needed to simulate the cruise and find the ending port. The key steps will be:\n",
      "1. Read in the input data to initialize the river connections and direction sequence \n",
      "2. Iterate K times:\n",
      "   - For each direction in the M-length sequence:\n",
      "     - Move to the next port based on the current port and direction \n",
      "3. Output the final port number after K iterations\n",
      "\n",
      "The solution will require loops to repeat the sequence K times and follow the M directions. Since K can be up to 1 billion, simulating all K iterations directly would be too slow. Instead, we can find a pattern in how the port changes after each M-length sequence, and then \"fast-forward\" by calculating which port we reach after K repetitions of the pattern.\n",
      "</thinking>\n",
      "\n",
      "\n",
      "Code:\n",
      "\n",
      "\n",
      "N, M, K = map(int, input().split())\n",
      "\n",
      "ports = []\n",
      "for _ in range(N):\n",
      "  left, right = map(int, input().split())\n",
      "  ports.append((left, right))\n",
      "\n",
      "directions = input().split()\n",
      "\n",
      "cur = 1\n",
      "pattern = []\n",
      "seen = set() \n",
      "steps = 0\n",
      "\n",
      "while cur not in seen:\n",
      "  seen.add(cur)\n",
      "  for d in directions:\n",
      "    steps += 1\n",
      "    if d == 'L': \n",
      "      cur = ports[cur-1][0]\n",
      "    else:\n",
      "      cur = ports[cur-1][1]\n",
      "  pattern.append((cur, steps))\n",
      "\n",
      "K %= steps\n",
      "for port, step in pattern:\n",
      "  if step > K:\n",
      "    cur = port\n",
      "    break\n",
      "  K -= step\n",
      "  \n",
      "print(cur)\n"
     ]
    }
   ],
   "source": ["snapshot = graph.get_state(config)\nprint(snapshot.values[\"messages\"][-2].content[0][\"text\"])\nprint(\"\\n\\nCode:\\n\\n\")\nprint(snapshot.values[\"messages\"][-2].tool_calls[0][\"args\"][\"code\"])"]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "69fd7efa-acbc-4450-927e-fb2b760066b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Incorrect submission. Please respond with updated code.\n",
      "Pass rate: 4/10\n",
      "Results:\n",
      "<test id=0>\n",
      "wrong answer. Expected '4\n",
      "', got '3\n",
      "'\n",
      "</test>\n",
      "<test id=1>\n",
      "wrong answer. Expected '50\n",
      "', got '2\n",
      "'\n",
      "</test>\n",
      "<t\n"
     ]
    }
   ],
   "source": ["print(snapshot.values[\"messages\"][-1].content[:200])"]
  },
  {
   "cell_type": "markdown",
   "id": "0afb3b34-8c67-4ceb-95ff-4ac65053c8a4",
   "metadata": {},
   "source": [
    "The agent failed. It's on the right track but clearly doesn't handle all the edge cases.\n",
    "\n",
    "The agent needs to remember that simulation should include the cycle + whatever steps led up to the example. It could use the \"tortoise and hare\" algo for cycle detection, use the simulated path and break if and when a repeat is detected, and then \n",
    "\n",
    "Let's let the agent know this by **updating the graph state**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "b10fcbc9-6dd4-41ad-98f7-1cf1685035e6",
   "metadata": {},
   "outputs": [],
   "source": ["updated_config = graph.update_state(\n    config,\n    values={\n        \"messages\": [\n            (\n                \"user\",\n                \"\"\"Consider breaking down the algorithm into separate parts: reading inputs, detecting cycles using the tortoise and hare algorithm, and determining Bessie's final position by skipping ahead K steps.\n\nRead the inputs into three arrays:\n- Two arrays L and R for the ports (adjust for 0-based indexing)\n- A third array S for the direction sequence\n\nOptimize by multiplying K by M before the main loop to convert the number of repetitions into the total number of steps.\n\nUse the tortoise and hare algorithm to detect the cycle:\n- Define a helper function get_next(v) that returns the next position and direction index\n- Initialize two pointers s0 and s1 to (0, 0)\n- In each iteration:\n  - Move s0 by 1 step and s1 by 2 steps using get_next()\n  - If s0 equals s1, decrement K by 1 and break out of the loop\n  - Otherwise, decrement K by 1\n- After the loop, if K is not 0, there is a cycle\n\nTo find the cycle length:\n- Initialize a counter variable rho to 1\n- Move s0 by 1 step using get_next()\n- Enter a loop:\n  - Move s0 by 1 step using get_next()\n  - Increment rho\n  - If s0 equals s1, break out of the loop\n\nSkip ahead by reducing K modulo rho.\n\nSimulate the remaining steps:\n- While K > 0, move s0 to the next position using get_next() and decrement K\n\nPrint the final position (converted to 1-based indexing).\n\nPay close attention to the initialization and movement of pointers during cycle detection and length calculation. Ensure that the logic is correct and handles all cases accurately.\"\"\",\n            )\n        ]\n    },\n)"]
  },
  {
   "cell_type": "markdown",
   "id": "e0542e9b-0992-407e-967f-c43bf1a75cc6",
   "metadata": {},
   "source": [
    "Now the graph's state contains our new message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "818e93f2-2204-4704-a832-f5c102d330f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HumanMessage(content=\"Consider breaking down the algorithm into separate parts: reading inputs, detecting cycles using the tortoise and hare algorithm, and determining Bessie's final position by skipping ahead K steps.\\n\\nRead the inputs into three arrays:\\n- Two arrays L and R for the ports (adjust for 0-based indexing)\\n- A third array S for the direction sequence\\n\\nOptimize by multiplying K by M before the main loop to convert the number of repetitions into the total number of steps.\\n\\nUse the tortoise and hare algorithm to detect the cycle:\\n- Define a helper function get_next(v) that returns the next position and direction index\\n- Initialize two pointers s0 and s1 to (0, 0)\\n- In each iteration:\\n  - Move s0 by 1 step and s1 by 2 steps using get_next()\\n  - If s0 equals s1, decrement K by 1 and break out of the loop\\n  - Otherwise, decrement K by 1\\n- After the loop, if K is not 0, there is a cycle\\n\\nTo find the cycle length:\\n- Initialize a counter variable rho to 1\\n- Move s0 by 1 step using get_next()\\n- Enter a loop:\\n  - Move s0 by 1 step using get_next()\\n  - Increment rho\\n  - If s0 equals s1, break out of the loop\\n\\nSkip ahead by reducing K modulo rho.\\n\\nSimulate the remaining steps:\\n- While K > 0, move s0 to the next position using get_next() and decrement K\\n\\nPrint the final position (converted to 1-based indexing).\\n\\nPay close attention to the initialization and movement of pointers during cycle detection and length calculation. Ensure that the logic is correct and handles all cases accurately.\", id='98888982-a469-4c5a-ab65-743d2f2608dc')"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["graph.get_state(config).values[\"messages\"][-1]"]
  },
  {
   "cell_type": "markdown",
   "id": "4eb67198-c84f-458b-8baf-783d7246dddc",
   "metadata": {},
   "source": [
    "Let's let the agent try again. Call `stream` with `None` to just use the inputs loaded from the memory. We will skip our human review for the next few attempats\n",
    "to see if it can correct itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "d5d76d8f-e49b-46bf-a762-a6c6978ee96c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Assistant: [{'text': '<thinking>\\nThank you for the detailed \n",
      "Assistant: Incorrect submission. Please respond with updated \n",
      "Continuing...\n"
     ]
    }
   ],
   "source": ["num_trials = 1\nwith tracing_v2_enabled(client=client):\n    for _ in range(num_trials):\n        events = graph.stream(None, updated_config)\n        for event in events:\n            for value in event.values():\n                messages = value.get(\"messages\")\n                if messages:\n                    if isinstance(messages, list):\n                        messages = value[\"messages\"][-1]\n                    print(\n                        \"Assistant:\",\n                        str(messages.content).replace(\"\\n\", \"\\\\n\")[:50],\n                    )\n                elif value.get(\"examples\"):\n                    print(\"Retrieved examples:\\n\\n\", value[\"examples\"][:100] + \"...\")\n                elif value.get(\"candidate\"):\n                    print(str(value[\"candidate\"].content)[:200])\n        if graph.get_state(config).values[\"status\"] == \"success\":\n            break\n        print(\"Continuing...\")"]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "20ee7535-1bc8-4105-87c4-0e7a89a011ff",
   "metadata": {},
   "outputs": [],
   "source": ["most_recent_state = list(graph.get_state_history(config))[0]"]
  },
  {
   "cell_type": "markdown",
   "id": "f9a492fd-03fc-4de5-8ba5-3ccdaeb1791c",
   "metadata": {},
   "source": [
    "OK so the agent tried again. Check out the [LangSmith trace](https://smith.langchain.com/public/707be522-9eaf-4b6a-994e-1742f421a433/r/add3d8e7-85b1-40cf-bbd3-e78c50f835e8) from this step to see its update."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1f691881-f56e-4e2c-8b6a-5febb7ccf3ae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'text': '<thinking>\\nThank you for the detailed algorithm breakdown! Let me go through each step to make sure I understand and have the necessary information to implement the solution.\\n\\nReading inputs:\\n- Read N, M, K and store in separate variables\\n- Create arrays L and R to store the left and right port connections (adjust for 0-based indexing)\\n- Create array S to store the M-length direction sequence \\n- Multiply K by M upfront to get the total number of steps\\n\\nDetecting cycles with tortoise and hare:\\n- Define get_next(v) to return the next position and direction index\\n  - It will use the current position and direction to look up the next port in L/R\\n- Initialize two pointers s0 and s1 to (0, 0) \\n- Loop until s0 equals s1 or all K steps are taken:\\n  - Move s0 by 1 step and s1 by 2 steps using get_next()\\n  - Decrement K\\n- After the loop, check if K is 0 to determine if a cycle was found\\n\\nFinding cycle length:\\n- If a cycle was found, initialize rho to 1\\n- Move s0 by 1 step \\n- Loop until s0 equals s1 again:\\n  - Move s0 by 1 step and increment rho\\n- rho will equal the cycle length\\n\\nSkipping ahead:\\n- Reduce K by taking it modulo rho\\n\\nSimulating remaining steps:\\n- While K is greater than 0:\\n  - Move s0 using get_next()\\n  - Decrement K\\n- s0 will hold the final position\\n\\nPrinting result:\\n- Add 1 to the final position to convert back to 1-based indexing before printing\\n\\nThe key aspects are:\\n- Handling the input format and 0-based indexing \\n- Defining get_next() to handle moving to the next port based on direction\\n- Correctly implementing the tortoise and hare cycle detection\\n- Finding the cycle length after detection\\n- Skipping ahead with modulo and simulating any remaining steps\\n- Adjusting the output back to 1-based indexing\\n\\nI believe I have all the necessary pieces to implement this solution now. Let me code it up using the writePython tool.\\n</thinking>', 'type': 'text'}, {'id': 'toolu_01EDrYeHJU7GxApRb1QfMA1b', 'input': {'reasoning': \"Here's the problem-solving approach:\\n\\n1. Read in the input data:\\n   - N ports, M-length direction sequence, K repetitions\\n   - L and R arrays for left/right port connections\\n   - S array for direction sequence\\n   - Multiply K by M to get total steps\\n\\n2. Define get_next(v) helper function:\\n   - Takes current position and direction index\\n   - Returns next position and incremented direction index\\n   - Looks up next port in L/R arrays based on current direction\\n\\n3. Detect cycle using tortoise and hare algorithm:\\n   - Initialize s0 and s1 pointers to (0, 0)\\n   - Loop until match or all steps taken:\\n     - Move s0 by 1 step, s1 by 2 steps\\n     - Decrement K\\n   - Check if K is 0 after loop\\n\\n4. If cycle found, find cycle length:\\n   - Initialize rho to 1\\n   - Move s0 by 1 step\\n   - Loop until s0 equals s1 again:\\n     - Move s0 and increment rho\\n   - rho is the cycle length\\n\\n5. Skip ahead by K % rho steps\\n\\n6. Simulate remaining steps:\\n   - While K > 0:\\n     - Move s0 with get_next()\\n     - Decrement K\\n   \\n7. Print final position (+1 for 1-based indexing)\\n\\nKey points:\\n- Multiplying K*M avoids nested loop\\n- get_next() handles port transitions \\n- Tortoise and hare finds cycles\\n- Modulo skips ahead in cycle\\n- Adjust 0-based indexing for input/output\", 'pseudocode': \"1. Read input:\\n   N, M, K = read_ints()\\n   L = [0] * N\\n   R = [0] * N\\n   for i in 0..N-1:\\n     L[i], R[i] = read_ints()\\n   S = read_direction_sequence()\\n   K *= M\\n\\n2. Define get_next(v):\\n   def get_next(pos, dir_idx):\\n     if S[dir_idx] == 'L':\\n       next_pos = L[pos]\\n     else:\\n       next_pos = R[pos]\\n     next_dir_idx = (dir_idx + 1) % M\\n     return (next_pos, next_dir_idx)\\n\\n3. Find cycle:\\n   s0 = (0, 0)\\n   s1 = (0, 0)  \\n   while K:\\n     s0 = get_next(s0[0], s0[1])\\n     s1 = get_next(s1[0], get_next(s1[0], s1[1])[1])\\n     K -= 1\\n     if s0 == s1: break\\n   if K != 0: no cycle, print s0[0] + 1\\n\\n4. Find cycle length:\\n   rho = 1\\n   s0 = get_next(s0[0], s0[1])\\n   while s0 != s1:\\n     s0 = get_next(s0[0], s0[1]) \\n     rho += 1\\n\\n5. Skip steps:\\n   K %= rho\\n\\n6. Remaining steps:  \\n   while K:\\n     s0 = get_next(s0[0], s0[1])\\n     K -= 1\\n     \\n7. Print result:\\n   print(s0[0] + 1)\", 'code': \"def read_ints():\\n  return map(int, input().split())\\n\\nN, M, K = read_ints()\\n\\nL = [0] * N\\nR = [0] * N\\nfor i in range(N):\\n  L[i], R[i] = read_ints()\\n  L[i] -= 1\\n  R[i] -= 1\\n\\nS = input().split()\\n\\nK *= M\\n\\ndef get_next(pos, dir_idx):\\n  if S[dir_idx] == 'L':\\n    next_pos = L[pos] \\n  else:\\n    next_pos = R[pos]\\n  next_dir_idx = (dir_idx + 1) % M\\n  return (next_pos, next_dir_idx)\\n\\ns0 = (0, 0)  \\ns1 = (0, 0)\\n\\nwhile K:\\n  if s0 == s1: break\\n  \\n  s0 = get_next(s0[0], s0[1])\\n  s1 = get_next(s1[0], get_next(s1[0], s1[1])[1])\\n  \\n  K -= 1\\n  \\nif K:\\n  rho = 1\\n  s0 = get_next(s0[0], s0[1])\\n  while s0 != s1:\\n    s0 = get_next(s0[0], s0[1])\\n    rho += 1\\n  \\n  K %= rho\\n  \\nwhile K:  \\n  s0 = get_next(s0[0], s0[1])\\n  K -= 1\\n  \\nprint(s0[0] + 1)\"}, 'name': 'writePython', 'type': 'tool_use'}]\n",
      "\n",
      "\n",
      "Code:\n",
      "\n",
      "\n",
      "def read_ints():\n",
      "  return map(int, input().split())\n",
      "\n",
      "N, M, K = read_ints()\n",
      "\n",
      "L = [0] * N\n",
      "R = [0] * N\n",
      "for i in range(N):\n",
      "  L[i], R[i] = read_ints()\n",
      "  L[i] -= 1\n",
      "  R[i] -= 1\n",
      "\n",
      "S = input().split()\n",
      "\n",
      "K *= M\n",
      "\n",
      "def get_next(pos, dir_idx):\n",
      "  if S[dir_idx] == 'L':\n",
      "    next_pos = L[pos] \n",
      "  else:\n",
      "    next_pos = R[pos]\n",
      "  next_dir_idx = (dir_idx + 1) % M\n",
      "  return (next_pos, next_dir_idx)\n",
      "\n",
      "s0 = (0, 0)  \n",
      "s1 = (0, 0)\n",
      "\n",
      "while K:\n",
      "  if s0 == s1: break\n",
      "  \n",
      "  s0 = get_next(s0[0], s0[1])\n",
      "  s1 = get_next(s1[0], get_next(s1[0], s1[1])[1])\n",
      "  \n",
      "  K -= 1\n",
      "  \n",
      "if K:\n",
      "  rho = 1\n",
      "  s0 = get_next(s0[0], s0[1])\n",
      "  while s0 != s1:\n",
      "    s0 = get_next(s0[0], s0[1])\n",
      "    rho += 1\n",
      "  \n",
      "  K %= rho\n",
      "  \n",
      "while K:  \n",
      "  s0 = get_next(s0[0], s0[1])\n",
      "  K -= 1\n",
      "  \n",
      "print(s0[0] + 1)\n"
     ]
    }
   ],
   "source": ["snapshot = graph.get_state(most_recent_state.config)\nai_message = snapshot.values[\"messages\"][-2]\nif ai_message.content:\n    print(ai_message.content)\nprint(\"\\n\\nCode:\\n\\n\")\nprint(ai_message.tool_calls[0][\"args\"][\"code\"] if ai_message.tool_calls else \"N/A\")"]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "b6fdebe7-557e-4751-97f3-15901c95600a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Incorrect submission. Please respond with updated code.\n",
      "Pass rate: 3/10\n",
      "Results:\n",
      "<test id=0>\n",
      "passed\n",
      "</test>\n",
      "<test id=1>\n",
      "timed out\n",
      "</test>\n",
      "<test id=2>\n",
      "timed out\n",
      "</test>\n",
      "<test id=3>\n",
      "timed out\n",
      "</test>\n",
      "<t\n"
     ]
    }
   ],
   "source": ["print(snapshot.values[\"messages\"][-1].content[:200])"]
  },
  {
   "cell_type": "markdown",
   "id": "ff2f9068-3c42-408c-a414-5914f923b78f",
   "metadata": {},
   "source": [
    "Still getting most test cases wrong.\n",
    "\n",
    "Let's provide more feedback."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "6eb46517-cdd1-4716-8a9e-df72cdd9ba67",
   "metadata": {},
   "outputs": [],
   "source": ["updated_config = graph.update_state(\n    updated_config,\n    values={\n        \"messages\": [\n            (\n                \"user\",\n                \"\"\"That's better, but you're still getting some errors. Let's double check some things:\n                       \n1. When calculating the cycle length, make sure the initialization and movement of the pointers is correct. Double-check the logic there and see if you can spot any discrepancies.\n2. Check the condition for whether there's a cycle after the main loop to ensure it covers all cases, like if  K becomes 0 in the last iteration.\n\nThink step by step through youur implementation and update using the writePython tool.\"\"\",\n            )\n        ]\n    },\n)"]
  },
  {
   "cell_type": "markdown",
   "id": "cd222b55-9bc5-4db0-8d28-58f366826dd4",
   "metadata": {},
   "source": [
    "Now that we've provided this feedback, let's give the agent a few attempts at solving it before we weigh in again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "ab89905b-ab75-4820-bde8-20ae01410aa7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Assistant: [{'text': \"<thinking>\\nThe algorithm looks mostly \n"
     ]
    }
   ],
   "source": ["num_trials = 2\nwith tracing_v2_enabled(client=client):\n    for _ in range(num_trials):\n        events = graph.stream(None, updated_config)\n        for event in events:\n            for value in event.values():\n                messages = value.get(\"messages\")\n                if messages:\n                    if isinstance(messages, list):\n                        messages = value[\"messages\"][-1]\n                    print(\n                        \"Assistant:\",\n                        str(messages.content).replace(\"\\n\", \"\\\\n\")[:50],\n                    )\n                elif value.get(\"examples\"):\n                    print(\"Retrieved examples:\\n\\n\", value[\"examples\"][:100] + \"...\")\n                elif value.get(\"candidate\"):\n                    print(str(value[\"candidate\"].content)[:200])\n        if graph.get_state(config).values[\"status\"] == \"success\":\n            break\n        print(\"Continuing...\")"]
  },
  {
   "cell_type": "markdown",
   "id": "28446ca4-0063-4ac2-bd87-bb9319fef84c",
   "metadata": {},
   "source": [
    "You can review [a LangSmith trace (link)](https://smith.langchain.com/public/d383e743-f8f1-4206-9dce-47627f152612/r/3f89582f-9107-461a-a34e-608d52641eeb) of the agent's response to your feedback at the provided link."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "28afec75-d661-4b88-b295-edb886231693",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "success\n"
     ]
    }
   ],
   "source": ["snapshot = graph.get_state(config)\nprint(snapshot.values[\"status\"])"]
  },
  {
   "cell_type": "markdown",
   "id": "72327753-b8b6-4cbc-b5cc-b3cd90b7dd2e",
   "metadata": {},
   "source": [
    "**Success!** - the LLM really wouldn't have been able to come to the correct answer without detailed human involvement."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1e2d6e8-3daf-4cbe-acb9-387e3058e602",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Congrats on making it to the end! In this tutorial, you implemented an agent in LangGraph capable of solving challenging programming problems. You did so by leveraging a few common techniques to improve performance, including:\n",
    "\n",
    "1. **Reflection**: while we didn't implement an explicit reflection step, our prompt and tool invocation was designed to encourage critique of previous outputs. You added this in Part 1.\n",
    "2. **Retrieval**: the \"episodic memory\" of the agent retrieves high-quality few-shot examples from our corpora of programming problems to help solve the **bronze** level question. In Part 2, you implemented a retrieval memory as an initial step.\n",
    "3. **Human-in-the-loop**: LLM-powered agents are still too weak to answer all these questions autonomously, but at times, they can get most of the way there and land on the right answer with human feedback. In Part 3, you used `interrupt_after` on the `evaluate` node and then included your feedback by using `update_state` on the graph.\n",
    "\n",
    "\n",
    "LLMs are not capable of solving all these problems autonomously, but through better prompting and clever engineering, you can create a system that is able to more reliably arrive at the proper solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71b4ba7-96ba-4643-b2de-eb2acdf4daba",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
