1-langchain-structured-output-examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using LangChain to get structured outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "from langchain_ollama import ChatOllama\n",
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "import streamlit as st"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "claude_api_key = \"<API KEY>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_model = ChatAnthropic(model=\"claude-3-haiku-20240307\", api_key=claude_api_key)\n",
    "# llm_model = ChatOllama(model=\"llama3.2\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 1: Structured output using the tool-calling API under the hood\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define a Pydantic model and the output will be returned as a Pydantic object with validation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Because it wanted to be the purr-cussionist!', punchline='Why did the cat join a band?', rating=8)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Optional\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "class Joke(BaseModel):\n",
    "    \"\"\"Joke to tell user.\"\"\"\n",
    "\n",
    "    setup: str = Field(description=\"The setup of the joke\")\n",
    "    punchline: str = Field(description=\"The punchline to the joke\")\n",
    "    rating: int = Field(description=\"How funny the joke is, from 1 to 10\")\n",
    "\n",
    "\n",
    "structured_llm = llm_model.with_structured_output(Joke)\n",
    "structured_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining the schema using a TypedDict parses the JSON output into a Python dict not a Pydantic object so there's no schema validation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'punchline': 'Why did the cat join a band?',\n",
       " 'rating': 8,\n",
       " 'setup': 'Because it wanted to be the purr-cussionist!'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing_extensions import Annotated, TypedDict\n",
    "\n",
    "\n",
    "class JokeTD(TypedDict):\n",
    "    \"\"\"Joke to tell user.\"\"\"\n",
    "\n",
    "    setup: Annotated[str, ..., \"The setup of the joke\"]\n",
    "    punchline: Annotated[str, ..., \"The punchline of the joke\"]\n",
    "    rating: Annotated[Optional[int], ..., \"How funny the joke is, from 1 to 10\"]\n",
    "\n",
    "\n",
    "structured_llm = llm_model.with_structured_output(JokeTD)\n",
    "structured_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or just extract the JSON Schema object\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'punchline': 'Why did the cat join a band?',\n",
       " 'rating': 8,\n",
       " 'setup': 'Because it wanted to be the purr-cussionist!'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "structured_llm = llm_model.with_structured_output(Joke.model_json_schema())\n",
    "structured_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try a more complicated structure with nested types\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ArticleResponse(title='The History of Texas in America', context='Texas has a rich and diverse history that spans thousands of years, from the earliest Native American inhabitants to its current status as the second-largest state in the US.', historical_timeline=['The Caddo and Comanche tribes inhabited the region for centuries before European exploration', 'In 1528, Álvar Núñez Cabeza de Vaca became the first European to visit Texas', 'In 1690, Spanish explorer Francisco Vásquez de Coronado arrived in Texas', 'Texas declared its independence from Mexico in 1836 and became the Republic of Texas', 'The US annexed Texas in 1845 and it became the 28th state in 1845'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class ArticleResponse(BaseModel):\n",
    "    \"\"\"A clear and concise answer to the users question.\"\"\"\n",
    "\n",
    "    title: str = Field(description=\"Title of the article\")\n",
    "    context: str = Field(\n",
    "        description=\"Provide a brief historical context to answer the question.\"\n",
    "    )\n",
    "    historical_timeline: list[str] = Field(\n",
    "        description=\"Provide a list of historical events relevant to the question\"\n",
    "    )\n",
    "\n",
    "\n",
    "structured_llm = llm_model.with_structured_output(ArticleResponse)\n",
    "structured_llm.invoke(\"Tell me the history of the state of Texas in America\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By specifying `include_raw=True` we get back the full data not just the parsed Pydantic object. This is useful if there are errors. Also we can clearly see the structures output is piggy-backing off the tool calling itnerface.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "structured_llm = llm_model.with_structured_output(ArticleResponse, include_raw=True)\n",
    "results = structured_llm.invoke(\"Tell me the history of the state of Texas in America\")\n",
    "\n",
    "# Get data from tool call argyments\n",
    "raw_output = results[\"raw\"].tool_calls[0][\"args\"]\n",
    "\n",
    "try:\n",
    "    print(ArticleResponse(**raw_output))\n",
    "except Exception as e:\n",
    "    print(f\"{type(e).__name__}: {str(e)}\")\n",
    "    print(f\"\\nRaw output:\\n{raw_output}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "structured_llm = llm_model.with_structured_output(ArticleResponse, include_raw=True)\n",
    "results = structured_llm.invoke(\"Tell me the history of the state of Texas in America\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can directly create the JSON schema object from the Pydantic object and we get the raw dict output without Pydantic validation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'context': 'Texas has a rich and diverse history that spans thousands of years, from the earliest Native American inhabitants to its current status as the second-largest state in the US.',\n",
       " 'historical_timeline': ['The Caddo and Comanche tribes inhabited the region for centuries before European exploration',\n",
       "  'In 1528, Álvar Núñez Cabeza de Vaca became the first European to visit Texas',\n",
       "  'In 1690, Spanish explorer Francisco Vásquez de Coronado arrived in Texas',\n",
       "  'Texas declared its independence from Mexico in 1836 and became the Republic of Texas',\n",
       "  'The US annexed Texas in 1845 and it became a state in 1845'],\n",
       " 'title': 'The History of Texas in America'}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "structured_llm_js = llm_model.with_structured_output(\n",
    "    ArticleResponse.model_json_schema()\n",
    ")\n",
    "structured_llm_js.invoke(\"Tell me the history of the state of Texas in America\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Under the hood: How Pydantic models are converted to JSONSchema\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The JSON schema representation is quite straightforward\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'description': 'Joke to tell user.',\n",
       " 'properties': {'setup': {'description': 'The setup of the joke',\n",
       "   'title': 'Setup',\n",
       "   'type': 'string'},\n",
       "  'punchline': {'description': 'The punchline to the joke',\n",
       "   'title': 'Punchline',\n",
       "   'type': 'string'},\n",
       "  'rating': {'description': 'How funny the joke is, from 1 to 10',\n",
       "   'title': 'Rating',\n",
       "   'type': 'integer'}},\n",
       " 'required': ['setup', 'punchline', 'rating'],\n",
       " 'title': 'Joke',\n",
       " 'type': 'object'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Joke.model_json_schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the same schema is contained in the format instructions, expect for 'title' and 'type'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The output should be formatted as a JSON instance that conforms to the JSON schema below.\n",
      "\n",
      "As an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}\n",
      "the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.\n",
      "\n",
      "Here is the output schema:\n",
      "```\n",
      "{\"description\": \"Joke to tell user.\", \"properties\": {\"setup\": {\"description\": \"The setup of the joke\", \"title\": \"Setup\", \"type\": \"string\"}, \"punchline\": {\"description\": \"The punchline to the joke\", \"title\": \"Punchline\", \"type\": \"string\"}, \"rating\": {\"description\": \"How funny the joke is, from 1 to 10\", \"title\": \"Rating\", \"type\": \"integer\"}}, \"required\": [\"setup\", \"punchline\", \"rating\"]}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.output_parsers import PydanticOutputParser\n",
    "\n",
    "output_parser = PydanticOutputParser(pydantic_object=Joke)\n",
    "print(output_parser.get_format_instructions())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method 2: JSON formating instructions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the PydanticOutputParser allows us to specify JSON outputs for other models that don't support tool calling.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why did the cat join a band?', punchline='Because it wanted to be the purr-cussionist!', rating=8)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser\n",
    "from langchain_core.exceptions import OutputParserException\n",
    "\n",
    "parser = PydanticOutputParser(pydantic_object=Joke)\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"Answer the user query. Wrap the output in `json` tags\\n{format_instructions}\",\n",
    "        ),\n",
    "        (\"human\", \"{query}\"),\n",
    "    ]\n",
    ").partial(format_instructions=parser.get_format_instructions())\n",
    "\n",
    "chain_llm = prompt | llm_model | parser\n",
    "chain_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"description\": \"Why did the cat join a band?\", \"properties\": {\"setup\": {\"description\": \"The setup of the joke\", \"title\": \"Setup\", \"type\": \"string\"}, \"punchline\": {\"description\": \"The punchline to the joke\", \"title\": \"Punchline\", \"type\": \"string\"}, \"rating\": {\"description\": \"How funny the joke is, from 1 to 10\", \"title\": \"Rating\", \"type\": \"integer\"}}, \"required\": [\"setup\", \"punchline\", \"rating\"]}\n",
      "\n",
      "{\n",
      "  \"setup\": \"Because it wanted to be the purr-cussionist.\",\n",
      "  \"punchline\": \"It was a cat-astrophe!\",\n",
      "  \"rating\": 8\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "prompt_user_format = ChatPromptTemplate.from_template(\n",
    "    \"{input} \\n{format_instructions}\"\n",
    ").partial(format_instructions=parser.get_format_instructions())\n",
    "\n",
    "structured_llm = prompt_user_format | llm_model | StrOutputParser()\n",
    "print(structured_llm.invoke(\"Tell me a joke about cats\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Structure output with Pydantic validation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "setup='Why did the cat join a band?' punchline='Because it wanted to be the purr-cussionist!' rating=8\n"
     ]
    }
   ],
   "source": [
    "chain = prompt_user_format | llm_model.with_structured_output(schema=Joke)\n",
    "\n",
    "try:\n",
    "    output = chain_llm.invoke(\"Tell me a joke about cats\")\n",
    "    print(output)\n",
    "except Exception as e:\n",
    "    print(f\"{type(e).__name__}: {str(e)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Structured output without Pydantic validation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'dict'>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'punchline': 'Because it wanted to be the purr-cussionist.',\n",
       " 'rating': 8,\n",
       " 'setup': 'Why did the cat join a band?'}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain_llm_novalid = prompt_user_format | llm_model.with_structured_output(\n",
    "    schema=Joke.model_json_schema()\n",
    ")\n",
    "output = chain_llm_novalid.invoke(\"Tell me a joke about cats\")\n",
    "print(type(output))\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Structured output using output parsers\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using a system prompt seems much less reliable than just inserting the format instructions into a user prompt. Why is this?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'punchline': 'Because he was caught horsing around!', 'rating': 8, 'setup': 'Why did the donkey get kicked out of the movie theater?'}\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.output_parsers import PydanticOutputParser, JsonOutputParser\n",
    "from langchain.output_parsers.fix import OutputFixingParser\n",
    "\n",
    "parser = JsonOutputParser(pydantic_object=Joke)\n",
    "prompt = prompt_user_format.partial(\n",
    "    format_instructions=parser.get_format_instructions()\n",
    ")\n",
    "\n",
    "structured_llm = prompt | llm_model | parser\n",
    "\n",
    "try:\n",
    "    output = chain_llm_novalid.invoke(\"Tell me a joke about donkeys\")\n",
    "    print(output)\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"{type(e).__name__}: {str(e)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'punchline': 'Because he was an ear-resistible dancer!', 'rating': 8, 'setup': 'Why did the aardvark go to the party?'}\n"
     ]
    }
   ],
   "source": [
    "parser = JsonOutputParser(pydantic_object=Joke)\n",
    "prompt = prompt_user_format.partial(\n",
    "    format_instructions=parser.get_format_instructions()\n",
    ")\n",
    "\n",
    "structured_llm = prompt | llm_model | parser\n",
    "try:\n",
    "    output = chain_llm_novalid.invoke(\"Tell me a joke about aardvarks\")\n",
    "    print(output)\n",
    "except Exception as e:\n",
    "    print(f\"{type(e).__name__}: {str(e)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fixing the output with `OutputFixingParser`, it could be better to use another model with lower temperature instead of the original model?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'punchline': 'Because it was a shell of a good time!', 'rating': 8, 'setup': 'Why did the weevil go to the party?'}\n"
     ]
    }
   ],
   "source": [
    "parser = PydanticOutputParser(pydantic_object=Joke)\n",
    "prompt = prompt_user_format.partial(\n",
    "    format_instructions=parser.get_format_instructions()\n",
    ")\n",
    "llm_model_fix = ChatOllama(model=\"llama3.2\", temperature=0)\n",
    "\n",
    "parser_fix = OutputFixingParser.from_llm(parser=parser, llm=llm_model_fix)\n",
    "\n",
    "try:\n",
    "    structured_llm = prompt | llm_model | parser_fix\n",
    "    output = chain_llm_novalid.invoke(\"Tell me a joke about weevils\")\n",
    "    print(output)\n",
    "except Exception as e:\n",
    "    print(f\"{type(e).__name__}: {str(e)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Which models support what?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_models = {\n",
    "    # \"Anthropic_Haiku\": ChatAnthropic(model=\"claude-3-haiku-20240307\", api_key=claude_api_key),\n",
    "    \"Ollama_llama32\": ChatOllama(model=\"llama3.2\", temperature=0),\n",
    "    \"Ollama_llama32_json\": ChatOllama(model=\"llama3.2\", format=\"json\", temperature=0),\n",
    "    \"Ollama_gemma2\": ChatOllama(model=\"gemma2\", temperature=0),\n",
    "    \"Ollama_gemma2_json\": ChatOllama(model=\"gemma2\", format=\"json\", temperature=0),\n",
    "    \"Ollama_phi3\": ChatOllama(model=\"phi3\", temperature=0),\n",
    "    \"Ollama_phi3_json\": ChatOllama(model=\"phi3\", format=\"json\", temperature=0),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: ChatOllama(model='llama3.2', temperature=0.0)\n",
      "  Tool use support\n",
      "Model: ChatOllama(model='llama3.2', temperature=0.0, format='json')\n",
      "  Tool use support\n",
      "Model: ChatOllama(model='gemma2', temperature=0.0)\n",
      "registry.ollama.ai/library/gemma2:latest does not support tools (status code: 400)\n",
      "  No tool use\n",
      "Model: ChatOllama(model='gemma2', temperature=0.0, format='json')\n",
      "registry.ollama.ai/library/gemma2:latest does not support tools (status code: 400)\n",
      "  No tool use\n",
      "Model: ChatOllama(model='phi3', temperature=0.0)\n",
      "registry.ollama.ai/library/phi3:latest does not support tools (status code: 400)\n",
      "  No tool use\n",
      "Model: ChatOllama(model='phi3', temperature=0.0, format='json')\n",
      "registry.ollama.ai/library/phi3:latest does not support tools (status code: 400)\n",
      "  No tool use\n"
     ]
    }
   ],
   "source": [
    "for llm_model in llm_models.values():\n",
    "    print(f\"Model: {llm_model.__repr__()}\")\n",
    "    test_structured_llm = llm_model.with_structured_output(JokeTD)\n",
    "\n",
    "    try:\n",
    "        output = test_structured_llm.invoke(\"Tell me a joke about cats\")\n",
    "        print(\"  Tool use support\")\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\"  No tool use\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}