{
 "cells": [
  {
   "attachments": {
    "67b615fe-0c25-4410-9d58-835982547001.png": {
     "image/png": "[encoded data, 160040 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "16dc0e41-80bd-4453-b421-dcf315741bf4",
   "metadata": {},
   "source": [
    "# Code generation with RAG and 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 start with a set of documentation specified by a user\n",
    "2. We use a long context LLM to ingest it and perform RAG to answer a question based upon it\n",
    "3. We will invoke a tool to produce a structured output\n",
    "4. We will perform two unit tests (check imports and code execution) prior returning the solution to the user \n",
    "\n",
    "![Screenshot 2024-05-23 at 2.17.42 PM.png](attachment:67b615fe-0c25-4410-9d58-835982547001.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3900420",
   "metadata": {},
   "outputs": [],
   "source": ["! pip install -U langchain_community langchain-openai langchain-anthropic langchain langgraph bs4"]
  },
  {
   "cell_type": "markdown",
   "id": "38330223-d8c8-4156-82b6-93e63343bc01",
   "metadata": {},
   "source": [
    "## Docs\n",
    "\n",
    "Load [LangChain Expression Language](https://python.langchain.com/v0.2/docs/concepts/#langchain-expression-language-lcel) (LCEL) docs as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c2eb35d1-4990-47dc-a5c4-208bae588a82",
   "metadata": {},
   "outputs": [],
   "source": ["from bs4 import BeautifulSoup as Soup\nfrom langchain_community.document_loaders.recursive_url_loader import RecursiveUrlLoader\n\n# LCEL docs\nurl = \"https://python.langchain.com/v0.2/docs/concepts/#langchain-expression-language-lcel\"\nloader = RecursiveUrlLoader(\n    url=url, max_depth=20, extractor=lambda x: Soup(x, \"html.parser\").text\n)\ndocs = loader.load()\n\n# Sort the list based on the URLs and get the text\nd_sorted = sorted(docs, key=lambda x: x.metadata[\"source\"])\nd_reversed = list(reversed(d_sorted))\nconcatenated_content = \"\\n\\n\\n --- \\n\\n\\n\".join(\n    [doc.page_content for doc in d_reversed]\n)"]
  },
  {
   "cell_type": "markdown",
   "id": "662d4ff4-1709-412f-bfed-5eb2b8d3d3dc",
   "metadata": {},
   "source": [
    "## LLMs\n",
    "\n",
    "### Code solution\n",
    "\n",
    "Try OpenAI and [Claude3](https://docs.anthropic.com/en/docs/about-claude/models) with function calling.\n",
    "\n",
    "Create `code_gen_chain` w/ either OpenAI or Claude and test here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3ba3df70-f6b4-4ea5-a210-e10944960bc6",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\nfrom langchain_openai import ChatOpenAI\n\n### OpenAI\n\n# Grader prompt\ncode_gen_prompt = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"You are a coding assistant with expertise in LCEL, LangChain expression language. \\n \n    Here is a full set of LCEL documentation:  \\n ------- \\n  {context} \\n ------- \\n Answer the user \n    question based on the above provided documentation. Ensure any code you provide can be executed \\n \n    with all required imports and variables defined. Structure your answer with a description of the code solution. \\n\n    Then list the imports. And finally list the functioning code block. Here is the user question:\"\"\",\n        ),\n        (\"placeholder\", \"{messages}\"),\n    ]\n)\n\n\n# Data model\nclass 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\nexpt_llm = \"gpt-4-0125-preview\"\nllm = ChatOpenAI(temperature=0, model=expt_llm)\ncode_gen_chain = code_gen_prompt | llm.with_structured_output(code)\nquestion = \"How do I build a RAG chain in LCEL?\"\n# solution = code_gen_chain_oai.invoke({\"context\":concatenated_content,\"messages\":[(\"user\",question)]})"]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cd30b67d-96db-4e51-a540-ae23fcc1f878",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_anthropic import ChatAnthropic\nfrom langchain_core.prompts import ChatPromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\n\n### Anthropic\n\n# Prompt to enforce tool use\ncode_gen_prompt_claude = ChatPromptTemplate.from_messages(\n    [\n        (\n            \"system\",\n            \"\"\"<instructions> You are a coding assistant with expertise in LCEL, LangChain expression language. \\n \n    Here is the LCEL documentation:  \\n ------- \\n  {context} \\n ------- \\n Answer the user  question based on the \\n \n    above provided documentation. 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    Invoke the code tool to structure the output correctly. </instructions> \\n Here is the user question:\"\"\",\n        ),\n        (\"placeholder\", \"{messages}\"),\n    ]\n)\n\n\n# Data model\nclass 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# expt_llm = \"claude-3-haiku-20240307\"\nexpt_llm = \"claude-3-opus-20240229\"\nllm = ChatAnthropic(\n    model=expt_llm,\n    default_headers={\"anthropic-beta\": \"tools-2024-04-04\"},\n)\n\nstructured_llm_claude = llm.with_structured_output(code, include_raw=True)\n\n\n# Optional: Check for errors in case tool use is flaky\ndef check_claude_output(tool_output):\n    \"\"\"Check for parse error or failure to call the tool\"\"\"\n\n    # Error with parsing\n    if tool_output[\"parsing_error\"]:\n        # Report back output and parsing errors\n        print(\"Parsing error!\")\n        raw_output = str(tool_output[\"raw\"].content)\n        error = tool_output[\"parsing_error\"]\n        raise ValueError(\n            f\"Error parsing your output! Be sure to invoke the tool. Output: {raw_output}. \\n Parse error: {error}\"\n        )\n\n    # Tool was not invoked\n    elif not tool_output[\"parsed\"]:\n        print(\"Failed to invoke tool!\")\n        raise ValueError(\n            \"You did not use the provided tool! Be sure to invoke the tool to structure the output.\"\n        )\n    return tool_output\n\n\n# Chain with output check\ncode_chain_claude_raw = (\n    code_gen_prompt_claude | structured_llm_claude | check_claude_output\n)\n\n\ndef insert_errors(inputs):\n    \"\"\"Insert errors for tool parsing in the messages\"\"\"\n\n    # Get errors\n    error = inputs[\"error\"]\n    messages = inputs[\"messages\"]\n    messages += [\n        (\n            \"assistant\",\n            f\"Retry. You are required to fix the parsing errors: {error} \\n\\n You must invoke the provided tool.\",\n        )\n    ]\n    return {\n        \"messages\": messages,\n        \"context\": inputs[\"context\"],\n    }\n\n\n# This will be run as a fallback chain\nfallback_chain = insert_errors | code_chain_claude_raw\nN = 3  # Max re-tries\ncode_gen_chain_re_try = code_chain_claude_raw.with_fallbacks(\n    fallbacks=[fallback_chain] * N, exception_key=\"error\"\n)\n\n\ndef parse_output(solution):\n    \"\"\"When we add 'include_raw=True' to structured output,\n    it will return a dict w 'raw', 'parsed', 'parsing_error'.\"\"\"\n\n    return solution[\"parsed\"]\n\n\n# Optional: With re-try to correct for failure to invoke tool\ncode_gen_chain = code_gen_chain_re_try | parse_output\n\n# No re-try\ncode_gen_chain = code_gen_prompt_claude | structured_llm_claude | parse_output"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f14750f-dddc-485b-ba29-5392cdf4ba43",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": ["# Test\nquestion = \"How do I build a RAG chain in LCEL?\"\nsolution = code_gen_chain.invoke(\n    {\"context\": concatenated_content, \"messages\": [(\"user\", question)]}\n)\nsolution"]
  },
  {
   "cell_type": "markdown",
   "id": "131f2055-2f64-4d19-a3d1-2d3cb8b42894",
   "metadata": {},
   "source": [
    "## State \n",
    "\n",
    "Our state is a dict that will contain keys (errors, question, code generation) relevant to code generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c185f1a2-e943-4bed-b833-4243c9c64092",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import List, TypedDict\n\n\nclass 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: List\n    generation: str\n    iterations: int"]
  },
  {
   "cell_type": "markdown",
   "id": "64454465-26a3-40de-ad85-bcf59a2c3086",
   "metadata": {},
   "source": [
    "## Graph \n",
    "\n",
    "Our graph lays out the logical flow shown in the figure above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b70e8301-63ae-4f7e-ad8f-c9a052fe3566",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_core.pydantic_v1 import BaseModel, Field\n\n### Parameter\n\n# Max tries\nmax_iterations = 3\n# Reflect\n# flag = 'reflect'\nflag = \"do not reflect\"\n\n### Nodes\n\n\ndef 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    error = state[\"error\"]\n\n    # We have been routed back to generation with an error\n    if error == \"yes\":\n        messages += [\n            (\n                \"user\",\n                \"Now, try again. Invoke the code tool to structure the output with a prefix, imports, and code block:\",\n            )\n        ]\n\n    # Solution\n    code_solution = code_gen_chain.invoke(\n        {\"context\": concatenated_content, \"messages\": messages}\n    )\n    messages += [\n        (\n            \"assistant\",\n            f\"{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\ndef 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 = [(\"user\", f\"Your solution failed the import test: {e}\")]\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        exec(imports + \"\\n\" + code)\n    except Exception as e:\n        print(\"---CODE BLOCK CHECK: FAILED---\")\n        error_message = [(\"user\", f\"Your solution failed the code execution test: {e}\")]\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\ndef reflect(state: GraphState):\n    \"\"\"\n    Reflect on errors\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    code_solution = state[\"generation\"]\n\n    # Prompt reflection\n\n    # Add reflection\n    reflections = code_gen_chain.invoke(\n        {\"context\": concatenated_content, \"messages\": messages}\n    )\n    messages += [(\"assistant\", f\"Here are reflections on the error: {reflections}\")]\n    return {\"generation\": code_solution, \"messages\": messages, \"iterations\": iterations}\n\n\n### Edges\n\n\ndef 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        if flag == \"reflect\":\n            return \"reflect\"\n        else:\n            return \"generate\""]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f66b4e00-4731-42c8-bc38-72dd0ff7c92c",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.graph import END, StateGraph, START\n\nworkflow = StateGraph(GraphState)\n\n# Define the nodes\nworkflow.add_node(\"generate\", generate)  # generation solution\nworkflow.add_node(\"check_code\", code_check)  # check code\nworkflow.add_node(\"reflect\", reflect)  # reflect\n\n# Build graph\nworkflow.add_edge(START, \"generate\")\nworkflow.add_edge(\"generate\", \"check_code\")\nworkflow.add_conditional_edges(\n    \"check_code\",\n    decide_to_finish,\n    {\n        \"end\": END,\n        \"reflect\": \"reflect\",\n        \"generate\": \"generate\",\n    },\n)\nworkflow.add_edge(\"reflect\", \"generate\")\napp = workflow.compile()"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bcaafe4-ddcf-4fab-8620-2d9b6c508f98",
   "metadata": {},
   "outputs": [],
   "source": ["question = \"How can I directly pass a string to a runnable and use it to construct the input needed for my prompt?\"\napp.invoke({\"messages\": [(\"user\", question)], \"iterations\": 0})"]
  },
  {
   "cell_type": "markdown",
   "id": "744f48a5-9ad3-4342-899f-7dd4266a9a15",
   "metadata": {},
   "source": [
    "## Eval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89852874-b538-4c8d-a4c3-1d68302db492",
   "metadata": {},
   "source": [
    "[Here](https://smith.langchain.com/public/326674a6-62bd-462d-88ae-eea49d503f9d/d) is a public dataset of LCEL questions. \n",
    "\n",
    "I saved this as `test-LCEL-code-gen`.\n",
    "\n",
    "You can also find the csv [here](https://github.com/langchain-ai/lcel-teacher/blob/main/eval/eval.csv)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "678e8954-56b5-4cc6-be26-f7f2a060b242",
   "metadata": {},
   "outputs": [],
   "source": ["import langsmith\n\nclient = langsmith.Client()"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef7cf662-7a6f-4dee-965c-6309d4045feb",
   "metadata": {},
   "outputs": [],
   "source": ["# Clone the dataset to your tenant to use it\npublic_dataset = (\n    \"https://smith.langchain.com/public/326674a6-62bd-462d-88ae-eea49d503f9d/d\"\n)\nclient.clone_public_dataset(public_dataset)"]
  },
  {
   "cell_type": "markdown",
   "id": "9d171396-022b-47ec-a741-c782aff9fdae",
   "metadata": {},
   "source": [
    "Custom evals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "455a34ea-52cb-4ae5-9f4a-7e4a08cd0c09",
   "metadata": {},
   "outputs": [],
   "source": ["from langsmith.schemas import Example, Run\n\n\ndef check_import(run: Run, example: Example) -> dict:\n    imports = run.outputs.get(\"imports\")\n    try:\n        exec(imports)\n        return {\"key\": \"import_check\", \"score\": 1}\n    except Exception:\n        return {\"key\": \"import_check\", \"score\": 0}\n\n\ndef check_execution(run: Run, example: Example) -> dict:\n    imports = run.outputs.get(\"imports\")\n    code = run.outputs.get(\"code\")\n    try:\n        exec(imports + \"\\n\" + code)\n        return {\"key\": \"code_execution_check\", \"score\": 1}\n    except Exception:\n        return {\"key\": \"code_execution_check\", \"score\": 0}"]
  },
  {
   "cell_type": "markdown",
   "id": "c90bf261-0d94-4779-bbde-c76adeefe3d7",
   "metadata": {},
   "source": [
    "Compare LangGraph to Context Stuffing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c8fa6bcb-b245-4422-b79a-582cd8a7d7ea",
   "metadata": {},
   "outputs": [],
   "source": ["def predict_base_case(example: dict):\n    \"\"\"Context stuffing\"\"\"\n    solution = code_gen_chain.invoke(\n        {\"context\": concatenated_content, \"messages\": [(\"user\", example[\"question\"])]}\n    )\n    solution_structured = code_gen_chain.invoke([(\"code\", solution)])\n    return {\"imports\": solution_structured.imports, \"code\": solution_structured.code}\n\n\ndef predict_langgraph(example: dict):\n    \"\"\"LangGraph\"\"\"\n    graph = app.invoke({\"messages\": [(\"user\", example[\"question\"])], \"iterations\": 0})\n    solution = graph[\"generation\"]\n    return {\"imports\": solution.imports, \"code\": solution.code}"]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d9c57468-97f6-47d6-a5e9-c09b53bfdd83",
   "metadata": {},
   "outputs": [],
   "source": ["from langsmith.evaluation import evaluate\n\n# Evaluator\ncode_evalulator = [check_import, check_execution]\n\n# Dataset\ndataset_name = \"test-LCEL-code-gen\""]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dacccf0-d73f-4017-aaf0-9806ffe5bd2c",
   "metadata": {},
   "outputs": [],
   "source": ["# Run base case\nexperiment_results_ = evaluate(\n    predict_base_case,\n    data=dataset_name,\n    evaluators=code_evalulator,\n    experiment_prefix=f\"test-without-langgraph-{expt_llm}\",\n    max_concurrency=2,\n    metadata={\n        \"llm\": expt_llm,\n    },\n)"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d90f9e-9dad-410c-a709-093d275029ae",
   "metadata": {},
   "outputs": [],
   "source": ["# Run with langgraph\nexperiment_results = evaluate(\n    predict_langgraph,\n    data=dataset_name,\n    evaluators=code_evalulator,\n    experiment_prefix=f\"test-with-langgraph-{expt_llm}-{flag}\",\n    max_concurrency=2,\n    metadata={\n        \"llm\": expt_llm,\n        \"feedback\": flag,\n    },\n)"]
  },
  {
   "cell_type": "markdown",
   "id": "d69da747-b4ea-455d-9314-60c3d9d30549",
   "metadata": {},
   "source": [
    "`Results:`\n",
    "\n",
    "* `LangGraph outperforms base case`: adding re-try loop improve performance\n",
    "* `Reflection did not help`: reflection prior to re-try regression vs just passing errors directly back to the LLM\n",
    "* `GPT-4 outperforms Claude3`: Claude3 had 3 and 1 run fail due to tool-use error for Opus and Haiku, respectively\n",
    "\n",
    "https://smith.langchain.com/public/78a3d858-c811-4e46-91cb-0f10ef56260b/d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a42333c3-c098-4576-ae2a-0258de64ece2",
   "metadata": {},
   "outputs": [],
   "source": [""]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
