Clean up code for Claude3

examples/code_assistant/langgraph_code_assistant.ipynb

@@ -47,7 +47,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "id": "c2eb35d1-4990-47dc-a5c4-208bae588a82",
    "metadata": {},
    "outputs": [],
@@ -86,7 +86,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "id": "3ba3df70-f6b4-4ea5-a210-e10944960bc6",
    "metadata": {},
    "outputs": [],
@@ -94,7 +94,9 @@
     "from langchain_openai import ChatOpenAI\n",
     "from langchain_anthropic import ChatAnthropic\n",
     "from langchain_core.prompts import ChatPromptTemplate\n",
-    "from langchain_core.output_parsers import StrOutputParser\n",
+    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
+    "\n",
+    "### OpenAI\n",
     "\n",
     "# Grader prompt \n",
     "code_gen_prompt = ChatPromptTemplate.from_messages(\n",
@@ -113,20 +115,92 @@
     "    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",
-    "\n",
-    "'''\n",
-    "expt_llm = \"claude3-opus\"\n",
-    "code_gen_llm = ChatAnthropic(temperature=0, model='claude-3-opus-20240229')\n",
-    "expt_llm = \"claude3-haiku\"\n",
-    "code_gen_llm = ChatAnthropic(temperature=0, model=\"claude-3-haiku-20240307\")\n",
-    "'''\n",
+    "    description = \"Schema for code solutions to questions about LCEL.\"\n",
     "\n",
     "expt_llm = \"gpt-4-0125-preview\"\n",
     "llm = ChatOpenAI(temperature=0, model=expt_llm)\n",
     "code_gen_chain = code_gen_prompt | llm.with_structured_output(code)\n",
     "question = \"How do I build a RAG chain in LCEL?\"\n",
-    "solution = code_gen_chain.invoke({\"context\":concatenated_content,\"messages\":[(\"user\",question)]})\n",
-    "solution"
+    "# solution = code_gen_chain_oai.invoke({\"context\":concatenated_content,\"messages\":[(\"user\",question)]})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "cd30b67d-96db-4e51-a540-ae23fcc1f878",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "### Anthropic\n",
+    "\n",
+    "# Important for getting tool use\n",
+    "code_gen_prompt_claude = ChatPromptTemplate.from_messages(\n",
+    "    [(\"system\",\"\"\"<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",
+    "    (\"placeholder\", \"{messages}\"),])\n",
+    "\n",
+    "# LLM\n",
+    "expt_llm = \"claude-3-haiku-20240307\" # claude-3-opus-20240229\n",
+    "llm = ChatAnthropic(\n",
+    "    model=expt_llm,\n",
+    "    default_headers={\"anthropic-beta\": \"tools-2024-04-04\"},\n",
+    ")\n",
+    "\n",
+    "structured_llm_claude = llm.with_structured_output(code, include_raw=True)\n",
+    "\n",
+    "# Check for errors\n",
+    "def 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(code_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",
+    "            f\"You did not use the provided tool! Be sure to invoke the tool to structure the output.\"\n",
+    "        )\n",
+    "    return tool_output\n",
+    "\n",
+    "# Chain with output check\n",
+    "code_chain_claude_raw = code_gen_prompt_claude | structured_llm_claude | check_claude_output\n",
+    "\n",
+    "def insert_errors(inputs):\n",
+    "    \"\"\"Insert errors in the messages\"\"\"\n",
+    "    \n",
+    "    # Get errors\n",
+    "    error = inputs[\"error\"]\n",
+    "    messages = inputs[\"messages\"]\n",
+    "    messages += [\n",
+    "        (\n",
+    "            \"user\",\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",
+    "# This will be run as a fallback chain\n",
+    "fallback_chain = insert_errors | code_chain_claude_raw\n",
+    "N = 3  # Max re-tries\n",
+    "code_gen_chain = code_chain_claude_raw.with_fallbacks(fallbacks=[fallback_chain] * N, exception_key=\"error\")\n",
+    "\n",
+    "# Test\n",
+    "question = \"How do I build a RAG chain in LCEL?\"\n",
+    "# solution = code_gen_chain_claude.invoke({\"context\":concatenated_content,\"messages\":[(\"user\",question)]})"
    ]
   },
   {
@@ -141,7 +215,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "id": "c185f1a2-e943-4bed-b833-4243c9c64092",
    "metadata": {},
    "outputs": [],