{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1d38cbab",
   "metadata": {},
   "source": [
    "This directory is retained purely for archival purposes and is no longer updated. Please see the newly [consolidated LangChain documentation](https://docs.langchain.com/oss/python/langgraph/overview) for the most current information and resources."
   ]
  },
  {
   "attachments": {
    "15d3ac32-cdf3-4800-a30c-f26d828d69c8.png": {
     "image/png": "[encoded data, 174944 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "6101c9d6-b7d1-46af-afab-47d05bfadeae",
   "metadata": {},
   "source": [
    "# Code generation with self-correction\n",
    "\n",
    "AlphaCodium presented an approach for code generation that uses control flow.\n",
    "\n",
    "Main idea: [construct an answer to a coding question iteratively.](https://x.com/karpathy/status/1748043513156272416?s=20). \n",
    "\n",
    "[AlphaCodium](https://github.com/Codium-ai/AlphaCodium) iteravely tests and improves an answer on public and AI-generated tests for a particular question. \n",
    "\n",
    "We will implement some of these ideas from scratch using [LangGraph](https://langchain-ai.github.io/langgraph/):\n",
    "\n",
    "1. We show how to route user questions to different types of documentation\n",
    "2. We we will show how to perform inline unit tests to confirm imports and code execution work\n",
    "3. We will show how to use LangGraph to orchestrate this\n",
    "\n",
    "![Screenshot 2024-05-23 at 2.17.51 PM.png](attachment:15d3ac32-cdf3-4800-a30c-f26d828d69c8.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e501686f-323f-4b87-8f9c-8ba89133078b",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U langchain_community langchain-mistralai langchain langgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ef4fb67-113a-4b88-9f93-7e3a95cee035",
   "metadata": {},
   "source": [
    "### LLM\n",
    "\n",
    "We'll use the Mistral API and `Codestral` instruct model, which support tool use!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "982e4609-86e4-4934-828f-e03d89c20393",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n",
    "mistral_api_key = os.getenv(\"MISTRAL_API_KEY\")  # Ensure this is set"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a20a3eb-6dc5-4a61-9705-9daf68172c0b",
   "metadata": {},
   "source": [
    "### Tracing\n",
    "\n",
    "Optionally, we'll use LangSmith for tracing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b172d2-3a9d-49a8-898c-22ed0cb45c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = \"<your-api-key>\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Mistral-code-gen-testing\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a3dca9-4485-4ae5-aa87-cd1f02bad8b9",
   "metadata": {},
   "source": [
    "## Code Generation\n",
    "\n",
    "Test with structured output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a188c8ca-c053-4e6d-b7af-38a3b6b371c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select LLM\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_mistralai import ChatMistralAI\n",
    "\n",
    "mistral_model = \"mistral-large-latest\"\n",
    "llm = ChatMistralAI(model=mistral_model, temperature=0)\n",
    "\n",
    "# Prompt\n",
    "code_gen_prompt_claude = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"\"\"You are a coding assistant. Ensure any code you provide can be executed with all required imports and variables \\n\n",
    "            defined. Structure your answer: 1) a prefix describing the code solution, 2) the imports, 3) the functioning code block.\n",
    "            \\n Here is the user question:\"\"\",\n",
    "        ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "# Data model\n",
    "class code(BaseModel):\n",
    "    \"\"\"Code output\"\"\"\n",
    "\n",
    "    prefix: str = Field(description=\"Description of the problem and approach\")\n",
    "    imports: str = Field(description=\"Code block import statements\")\n",
    "    code: str = Field(description=\"Code block not including import statements\")\n",
    "    description = \"Schema for code solutions to questions about LCEL.\"\n",
    "\n",
    "\n",
    "# LLM\n",
    "code_gen_chain = llm.with_structured_output(code, include_raw=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9fc0290d-5a04-4514-8664-91f9dbf2da7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Write a function for fibonacci.\"\n",
    "messages = [(\"user\", question)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "973281bd-e74b-4386-98c6-210af5e31982",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "code(prefix='A function to calculate the nth Fibonacci number.', imports='', code='def fibonacci(n):\\n    if n <= 0:\\n        return \"Input should be positive integer\"\\n    elif n == 1:\\n        return 0\\n    elif n == 2:\\n        return 1\\n    else:\\n        a, b = 0, 1\\n        for _ in range(2, n):\\n            a, b = b, a + b\\n        return b', description='Schema for code solutions to questions about LCEL.')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test\n",
    "result = code_gen_chain.invoke(messages)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4235eb2d-f5b3-4cd0-bbb1-a889eb5564d7",
   "metadata": {},
   "source": [
    "## State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "183d77b8-f180-4815-b39f-8ef507ec0534",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, TypedDict\n",
    "\n",
    "from langgraph.graph.message import AnyMessage, add_messages\n",
    "\n",
    "\n",
    "class GraphState(TypedDict):\n",
    "    \"\"\"\n",
    "    Represents the state of our graph.\n",
    "\n",
    "    Attributes:\n",
    "        error : Binary flag for control flow to indicate whether test error was tripped\n",
    "        messages : With user question, error messages, reasoning\n",
    "        generation : Code solution\n",
    "        iterations : Number of tries\n",
    "    \"\"\"\n",
    "\n",
    "    error: str\n",
    "    messages: Annotated[list[AnyMessage], add_messages]\n",
    "    generation: str\n",
    "    iterations: int"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55043d78-c012-4280-bc8b-259f04a29cb4",
   "metadata": {},
   "source": [
    "## Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "14bc89d1-3ca6-4847-a048-1803e0e4600e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "### Parameters\n",
    "max_iterations = 3\n",
    "\n",
    "\n",
    "### Nodes\n",
    "def generate(state: GraphState):\n",
    "    \"\"\"\n",
    "    Generate a code solution\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, generation\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---GENERATING CODE SOLUTION---\")\n",
    "\n",
    "    # State\n",
    "    messages = state[\"messages\"]\n",
    "    iterations = state[\"iterations\"]\n",
    "\n",
    "    # Solution\n",
    "    code_solution = code_gen_chain.invoke(messages)\n",
    "    messages += [\n",
    "        (\n",
    "            \"assistant\",\n",
    "            f\"Here is my attempt to solve the problem: {code_solution.prefix} \\n Imports: {code_solution.imports} \\n Code: {code_solution.code}\",\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    # Increment\n",
    "    iterations = iterations + 1\n",
    "    return {\"generation\": code_solution, \"messages\": messages, \"iterations\": iterations}\n",
    "\n",
    "\n",
    "def code_check(state: GraphState):\n",
    "    \"\"\"\n",
    "    Check code\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, error\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECKING CODE---\")\n",
    "\n",
    "    # State\n",
    "    messages = state[\"messages\"]\n",
    "    code_solution = state[\"generation\"]\n",
    "    iterations = state[\"iterations\"]\n",
    "\n",
    "    # Get solution components\n",
    "    imports = code_solution.imports\n",
    "    code = code_solution.code\n",
    "\n",
    "    # Check imports\n",
    "    try:\n",
    "        exec(imports)\n",
    "    except Exception as e:\n",
    "        print(\"---CODE IMPORT CHECK: FAILED---\")\n",
    "        error_message = [\n",
    "            (\n",
    "                \"user\",\n",
    "                f\"Your solution failed the import test. Here is the error: {e}. Reflect on this error and your prior attempt to solve the problem. (1) State what you think went wrong with the prior solution and (2) try to solve this problem again. Return the FULL SOLUTION. Use the code tool to structure the output with a prefix, imports, and code block:\",\n",
    "            )\n",
    "        ]\n",
    "        messages += error_message\n",
    "        return {\n",
    "            \"generation\": code_solution,\n",
    "            \"messages\": messages,\n",
    "            \"iterations\": iterations,\n",
    "            \"error\": \"yes\",\n",
    "        }\n",
    "\n",
    "    # Check execution\n",
    "    try:\n",
    "        combined_code = f\"{imports}\\n{code}\"\n",
    "        print(f\"CODE TO TEST: {combined_code}\")\n",
    "        # Use a shared scope for exec\n",
    "        global_scope = {}\n",
    "        exec(combined_code, global_scope)\n",
    "    except Exception as e:\n",
    "        print(\"---CODE BLOCK CHECK: FAILED---\")\n",
    "        error_message = [\n",
    "            (\n",
    "                \"user\",\n",
    "                f\"Your solution failed the code execution test: {e}) Reflect on this error and your prior attempt to solve the problem. (1) State what you think went wrong with the prior solution and (2) try to solve this problem again. Return the FULL SOLUTION. Use the code tool to structure the output with a prefix, imports, and code block:\",\n",
    "            )\n",
    "        ]\n",
    "        messages += error_message\n",
    "        return {\n",
    "            \"generation\": code_solution,\n",
    "            \"messages\": messages,\n",
    "            \"iterations\": iterations,\n",
    "            \"error\": \"yes\",\n",
    "        }\n",
    "\n",
    "    # No errors\n",
    "    print(\"---NO CODE TEST FAILURES---\")\n",
    "    return {\n",
    "        \"generation\": code_solution,\n",
    "        \"messages\": messages,\n",
    "        \"iterations\": iterations,\n",
    "        \"error\": \"no\",\n",
    "    }\n",
    "\n",
    "\n",
    "### Conditional edges\n",
    "\n",
    "\n",
    "def decide_to_finish(state: GraphState):\n",
    "    \"\"\"\n",
    "    Determines whether to finish.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Next node to call\n",
    "    \"\"\"\n",
    "    error = state[\"error\"]\n",
    "    iterations = state[\"iterations\"]\n",
    "\n",
    "    if error == \"no\" or iterations == max_iterations:\n",
    "        print(\"---DECISION: FINISH---\")\n",
    "        return \"end\"\n",
    "    else:\n",
    "        print(\"---DECISION: RE-TRY SOLUTION---\")\n",
    "        return \"generate\"\n",
    "\n",
    "\n",
    "### Utilities\n",
    "\n",
    "\n",
    "def _print_event(event: dict, _printed: set, max_length=1500):\n",
    "    current_state = event.get(\"dialog_state\")\n",
    "    if current_state:\n",
    "        print(\"Currently in: \", current_state[-1])\n",
    "    message = event.get(\"messages\")\n",
    "    if message:\n",
    "        if isinstance(message, list):\n",
    "            message = message[-1]\n",
    "        if message.id not in _printed:\n",
    "            msg_repr = message.pretty_repr(html=True)\n",
    "            if len(msg_repr) > max_length:\n",
    "                msg_repr = msg_repr[:max_length] + \" ... (truncated)\"\n",
    "            print(msg_repr)\n",
    "            _printed.add(message.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2dff2209-44c7-4e2c-b607-ba6675f9e45f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import InMemorySaver\n",
    "from langgraph.graph import END, StateGraph, START\n",
    "\n",
    "builder = StateGraph(GraphState)\n",
    "\n",
    "# Define the nodes\n",
    "builder.add_node(\"generate\", generate)  # generation solution\n",
    "builder.add_node(\"check_code\", code_check)  # check code\n",
    "\n",
    "# Build graph\n",
    "builder.add_edge(START, \"generate\")\n",
    "builder.add_edge(\"generate\", \"check_code\")\n",
    "builder.add_conditional_edges(\n",
    "    \"check_code\",\n",
    "    decide_to_finish,\n",
    "    {\n",
    "        \"end\": END,\n",
    "        \"generate\": \"generate\",\n",
    "    },\n",
    ")\n",
    "\n",
    "memory = InMemorySaver()\n",
    "graph = builder.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d4bb21cd-af20-4d4d-89ff-384db034b7c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 11928 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(graph.get_graph(xray=True).draw_mermaid_png()))\n",
    "except Exception:\n",
    "    # This requires some extra dependencies and is optional\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "242aa2f0-2c31-462f-a958-ff9ae0cf7c62",
   "metadata": {},
   "outputs": [],
   "source": [
    "_printed = set()\n",
    "thread_id = str(uuid.uuid4())\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        # Checkpoints are accessed by thread_id\n",
    "        \"thread_id\": thread_id,\n",
    "    }\n",
    "}\n",
    "\n",
    "question = \"Write a Python program that prints 'Hello, World!' to the console.\"\n",
    "events = graph.stream(\n",
    "    {\"messages\": [(\"user\", question)], \"iterations\": 0}, config, stream_mode=\"values\"\n",
    ")\n",
    "for event in events:\n",
    "    _print_event(event, _printed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6924e707-5970-4254-a748-fa75628916f2",
   "metadata": {},
   "source": [
    "`Trace:`\n",
    "\n",
    "https://smith.langchain.com/public/53bcdaab-e3c5-4423-9908-c44595325c38/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "390b2768-f395-4aea-8b0e-9d36212a31ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "_printed = set()\n",
    "thread_id = str(uuid.uuid4())\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        # Checkpoints are accessed by thread_id\n",
    "        \"thread_id\": thread_id,\n",
    "    }\n",
    "}\n",
    "\n",
    "question = \"\"\"Create a Python program that checks if a given string is a palindrome. A palindrome is a word, phrase, number, or other sequence of characters that reads the same forward and backward (ignoring spaces, punctuation, and capitalization).\n",
    "\n",
    "Requirements:\n",
    "The program should define a function is_palindrome(s) that takes a string s as input.\n",
    "The function should return True if the string is a palindrome and False otherwise.\n",
    "Ignore spaces, punctuation, and case differences when checking for palindromes.\n",
    "\n",
    "Give an example of it working on an example input word.\"\"\"\n",
    "\n",
    "events = graph.stream(\n",
    "    {\"messages\": [(\"user\", question)], \"iterations\": 0}, config, stream_mode=\"values\"\n",
    ")\n",
    "for event in events:\n",
    "    _print_event(event, _printed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f96e0137-6df3-4a2a-8711-c9cb7dc66831",
   "metadata": {},
   "source": [
    "Trace:\n",
    "\n",
    "https://smith.langchain.com/public/e749936d-7746-49de-b980-c41b17986e79/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a3f946b-e2f2-44d9-905b-09f36980cf9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "_printed = set()\n",
    "thread_id = str(uuid.uuid4())\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        # Checkpoints are accessed by thread_id\n",
    "        \"thread_id\": thread_id,\n",
    "    }\n",
    "}\n",
    "\n",
    "question = \"\"\"Write a program that prints the numbers from 1 to 100. \n",
    "But for multiples of three, print \"Fizz\" instead of the number, and for the multiples of five, print \"Buzz\". \n",
    "For numbers which are multiples of both three and five, print \"FizzBuzz\".\"\"\"\n",
    "\n",
    "events = graph.stream(\n",
    "    {\"messages\": [(\"user\", question)], \"iterations\": 0}, config, stream_mode=\"values\"\n",
    ")\n",
    "for event in events:\n",
    "    _print_event(event, _printed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f3ef03b-0a07-49f5-9cbf-15e2503d020e",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/f5c19708-7592-4512-9f00-9696ab34a9eb/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bb883df-540b-46ab-9415-fe27db68456f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "_printed = set()\n",
    "thread_id = str(uuid.uuid4())\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        # Checkpoints are accessed by thread_id\n",
    "        \"thread_id\": thread_id,\n",
    "    }\n",
    "}\n",
    "\n",
    "question = \"\"\"I want to vectorize a function\n",
    "\n",
    "        frame = np.zeros((out_h, out_w, 3), dtype=np.uint8)\n",
    "        for i, val1 in enumerate(rows):\n",
    "            for j, val2 in enumerate(cols):\n",
    "                for j, val3 in enumerate(ch):\n",
    "                    # Assuming you want to store the pair as tuples in the matrix\n",
    "                    frame[i, j, k] = image[val1, val2, val3]\n",
    "\n",
    "        out.write(np.array(frame))\n",
    "\n",
    "with a simple numpy function that does something like this what is it called. Show me a test case with this working.\"\"\"\n",
    "\n",
    "events = graph.stream(\n",
    "    {\"messages\": [(\"user\", question)], \"iterations\": 0}, config, stream_mode=\"values\"\n",
    ")\n",
    "for event in events:\n",
    "    _print_event(event, _printed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "750a3292-1e0e-49cf-8b28-bef179afe6a2",
   "metadata": {},
   "source": [
    "Trace w/ good example of self-correction:\n",
    "\n",
    "https://smith.langchain.com/public/b54778a0-d267-4f09-bc28-71761201c522/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee05da1f-c272-405d-8a7b-552cfc3106e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "_printed = set()\n",
    "thread_id = str(uuid.uuid4())\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        # Checkpoints are accessed by thread_id\n",
    "        \"thread_id\": thread_id,\n",
    "    }\n",
    "}\n",
    "\n",
    "question = \"\"\"Create a Python program that allows two players to play a game of Tic-Tac-Toe. The game should be played on a 3x3 grid. The program should:\n",
    "\n",
    "- Allow players to take turns to input their moves.\n",
    "- Check for invalid moves (e.g., placing a marker on an already occupied space).\n",
    "- Determine and announce the winner or if the game ends in a draw.\n",
    "\n",
    "Requirements:\n",
    "- Use a 2D list to represent the Tic-Tac-Toe board.\n",
    "- Use functions to modularize the code.\n",
    "- Validate player input.\n",
    "- Check for win conditions and draw conditions after each move.\"\"\"\n",
    "\n",
    "events = graph.stream(\n",
    "    {\"messages\": [(\"user\", question)], \"iterations\": 0}, config, stream_mode=\"values\"\n",
    ")\n",
    "for event in events:\n",
    "    _print_event(event, _printed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d900cd6-2df9-467d-8e74-803527269008",
   "metadata": {},
   "source": [
    "Trace w/ good example of failure to correct:\n",
    "\n",
    "https://smith.langchain.com/public/871ae736-2f77-44d4-b0da-a600d8f5377d/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "814fc2a4-8e5b-4faa-8f52-3977226bd09a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
