{
 "cells": [
  {
   "cell_type": "raw",
   "id": "27598444",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "---\n",
    "sidebar_position: 3\n",
    "keywords: [structured output, json, information extraction, with_structured_output]\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e3f0f72",
   "metadata": {},
   "source": [
    "# How to return structured data from a model\n",
    "\n",
    ":::info Prerequisites\n",
    "\n",
    "This guide assumes familiarity with the following concepts:\n",
    "- [Chat models](/docs/concepts/#chat-models)\n",
    "- [Function/tool calling](/docs/concepts/#functiontool-calling)\n",
    ":::\n",
    "\n",
    "It is often useful to have a model return output that matches a specific schema. One common use-case is extracting data from text to insert into a database or use with some other downstream system. This guide covers a few strategies for getting structured outputs from a model.\n",
    "\n",
    "## The `.with_structured_output()` method\n",
    "\n",
    "<span data-heading-keywords=\"with_structured_output\"></span>\n",
    "\n",
    ":::info Supported models\n",
    "\n",
    "You can find a [list of models that support this method here](/docs/integrations/chat/).\n",
    "\n",
    ":::\n",
    "\n",
    "This is the easiest and most reliable way to get structured outputs. `with_structured_output()` is implemented for models that provide native APIs for structuring outputs, like tool/function calling or JSON mode, and makes use of these capabilities under the hood.\n",
    "\n",
    "This method takes a schema as input which specifies the names, types, and descriptions of the desired output attributes. The method returns a model-like Runnable, except that instead of outputting strings or Messages it outputs objects corresponding to the given schema. The schema can be specified as a [JSON Schema](https://json-schema.org/) or a Pydantic class. If JSON Schema is used then a dictionary will be returned by the Runnable, and if a Pydantic class is used then Pydantic objects will be returned.\n",
    "\n",
    "As an example, let's get a model to generate a joke and separate the setup from the punchline:\n",
    "\n",
    "```{=mdx}\n",
    "import ChatModelTabs from \"@theme/ChatModelTabs\";\n",
    "\n",
    "<ChatModelTabs\n",
    "  customVarName=\"llm\"\n",
    "/>\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6d55008f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | output: false\n",
    "# | echo: false\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model=\"gpt-4-0125-preview\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a808a401-be1f-49f9-ad13-58dd68f7db5f",
   "metadata": {},
   "source": [
    "If we want the model to return a Pydantic object, we just need to pass in the desired Pydantic class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "070bf702",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why was the cat sitting on the computer?', punchline='To keep an eye on the mouse!', rating=None)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Optional\n",
    "\n",
    "from langchain_core.pydantic_v1 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: Optional[int] = Field(description=\"How funny the joke is, from 1 to 10\")\n",
    "\n",
    "\n",
    "structured_llm = llm.with_structured_output(Joke)\n",
    "\n",
    "structured_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00890a47-3cdf-4805-b8f1-6d110f0633d3",
   "metadata": {},
   "source": [
    ":::tip\n",
    "Beyond just the structure of the Pydantic class, the name of the Pydantic class, the docstring, and the names and provided descriptions of parameters are very important. Most of the time `with_structured_output` is using a model's function/tool calling API, and you can effectively think of all of this information as being added to the model prompt.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deddb6d3",
   "metadata": {},
   "source": [
    "We can also pass in a [JSON Schema](https://json-schema.org/) dict if you prefer not to use Pydantic. In this case, the response is also a dict:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6700994a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'setup': 'Why was the cat sitting on the computer?',\n",
       " 'punchline': 'Because it wanted to keep an eye on the mouse!',\n",
       " 'rating': 8}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "json_schema = {\n",
    "    \"title\": \"joke\",\n",
    "    \"description\": \"Joke to tell user.\",\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"setup\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": \"The setup of the joke\",\n",
    "        },\n",
    "        \"punchline\": {\n",
    "            \"type\": \"string\",\n",
    "            \"description\": \"The punchline to the joke\",\n",
    "        },\n",
    "        \"rating\": {\n",
    "            \"type\": \"integer\",\n",
    "            \"description\": \"How funny the joke is, from 1 to 10\",\n",
    "        },\n",
    "    },\n",
    "    \"required\": [\"setup\", \"punchline\"],\n",
    "}\n",
    "structured_llm = llm.with_structured_output(json_schema)\n",
    "\n",
    "structured_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3da57988",
   "metadata": {},
   "source": [
    "### Choosing between multiple schemas\n",
    "\n",
    "The simplest way to let the model choose from multiple schemas is to create a parent Pydantic class that has a Union-typed attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9194bcf2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Response(output=Joke(setup='Why was the cat sitting on the computer?', punchline='To keep an eye on the mouse!', rating=8))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Union\n",
    "\n",
    "\n",
    "class ConversationalResponse(BaseModel):\n",
    "    \"\"\"Respond in a conversational manner. Be kind and helpful.\"\"\"\n",
    "\n",
    "    response: str = Field(description=\"A conversational response to the user's query\")\n",
    "\n",
    "\n",
    "class Response(BaseModel):\n",
    "    output: Union[Joke, ConversationalResponse]\n",
    "\n",
    "\n",
    "structured_llm = llm.with_structured_output(Response)\n",
    "\n",
    "structured_llm.invoke(\"Tell me a joke about cats\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "84d86132",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Response(output=ConversationalResponse(response=\"I'm just a digital assistant, so I don't have feelings, but I'm here and ready to help you. How can I assist you today?\"))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "structured_llm.invoke(\"How are you today?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e28c14d3",
   "metadata": {},
   "source": [
    "Alternatively, you can use tool calling directly to allow the model to choose between options, if your [chosen model supports it](/docs/integrations/chat/). This involves a bit more parsing and setup but in some instances leads to better performance because you don't have to use nested schemas. See [this how-to guide](/docs/how_to/tool_calling) for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a40f703-7fd2-4fe0-ab2a-fa2d711ba009",
   "metadata": {},
   "source": [
    "### Streaming\n",
    "\n",
    "We can stream outputs from our structured model when the output type is a dict (i.e., when the schema is specified as a JSON Schema dict). \n",
    "\n",
    ":::info\n",
    "\n",
    "Note that what's yielded is already aggregated chunks, not deltas.\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "aff89877-28a3-472f-a1aa-eff893fe7736",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{}\n",
      "{'setup': ''}\n",
      "{'setup': 'Why'}\n",
      "{'setup': 'Why was'}\n",
      "{'setup': 'Why was the'}\n",
      "{'setup': 'Why was the cat'}\n",
      "{'setup': 'Why was the cat sitting'}\n",
      "{'setup': 'Why was the cat sitting on'}\n",
      "{'setup': 'Why was the cat sitting on the'}\n",
      "{'setup': 'Why was the cat sitting on the computer'}\n",
      "{'setup': 'Why was the cat sitting on the computer?'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': ''}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on the'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on the mouse'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on the mouse!'}\n",
      "{'setup': 'Why was the cat sitting on the computer?', 'punchline': 'Because it wanted to keep an eye on the mouse!', 'rating': 8}\n"
     ]
    }
   ],
   "source": [
    "structured_llm = llm.with_structured_output(json_schema)\n",
    "\n",
    "for chunk in structured_llm.stream(\"Tell me a joke about cats\"):\n",
    "    print(chunk)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a526cdf-e736-451b-96be-22e8986d3863",
   "metadata": {},
   "source": [
    "### Few-shot prompting\n",
    "\n",
    "For more complex schemas it's very useful to add few-shot examples to the prompt. This can be done in a few ways.\n",
    "\n",
    "The simplest and most universal way is to add examples to a system message in the prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "283ba784-2072-47ee-9b2c-1119e3c69e8e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'setup': 'Woodpecker',\n",
       " 'punchline': \"Woodpecker goes 'knock knock', but don't worry, they never expect you to answer the door!\",\n",
       " 'rating': 8}"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "system = \"\"\"You are a hilarious comedian. Your specialty is knock-knock jokes. \\\n",
    "Return a joke which has the setup (the response to \"Who's there?\") and the final punchline (the response to \"<setup> who?\").\n",
    "\n",
    "Here are some examples of jokes:\n",
    "\n",
    "example_user: Tell me a joke about planes\n",
    "example_assistant: {{\"setup\": \"Why don't planes ever get tired?\", \"punchline\": \"Because they have rest wings!\", \"rating\": 2}}\n",
    "\n",
    "example_user: Tell me another joke about planes\n",
    "example_assistant: {{\"setup\": \"Cargo\", \"punchline\": \"Cargo 'vroom vroom', but planes go 'zoom zoom'!\", \"rating\": 10}}\n",
    "\n",
    "example_user: Now about caterpillars\n",
    "example_assistant: {{\"setup\": \"Caterpillar\", \"punchline\": \"Caterpillar really slow, but watch me turn into a butterfly and steal the show!\", \"rating\": 5}}\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages([(\"system\", system), (\"human\", \"{input}\")])\n",
    "\n",
    "few_shot_structured_llm = prompt | structured_llm\n",
    "few_shot_structured_llm.invoke(\"what's something funny about woodpeckers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c12b389-153d-44d1-af34-37e5b926d3db",
   "metadata": {},
   "source": [
    "When the underlying method for structuring outputs is tool calling, we can pass in our examples as explicit tool calls. You can check if the model you're using makes use of tool calling in its API reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "d7381cb0-b2c3-4302-a319-ed72d0b9e43f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'setup': 'Crocodile',\n",
       " 'punchline': \"Crocodile 'see you later', but in a while, it becomes an alligator!\",\n",
       " 'rating': 7}"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import AIMessage, HumanMessage, ToolMessage\n",
    "\n",
    "examples = [\n",
    "    HumanMessage(\"Tell me a joke about planes\", name=\"example_user\"),\n",
    "    AIMessage(\n",
    "        \"\",\n",
    "        name=\"example_assistant\",\n",
    "        tool_calls=[\n",
    "            {\n",
    "                \"name\": \"joke\",\n",
    "                \"args\": {\n",
    "                    \"setup\": \"Why don't planes ever get tired?\",\n",
    "                    \"punchline\": \"Because they have rest wings!\",\n",
    "                    \"rating\": 2,\n",
    "                },\n",
    "                \"id\": \"1\",\n",
    "            }\n",
    "        ],\n",
    "    ),\n",
    "    # Most tool-calling models expect a ToolMessage(s) to follow an AIMessage with tool calls.\n",
    "    ToolMessage(\"\", tool_call_id=\"1\"),\n",
    "    # Some models also expect an AIMessage to follow any ToolMessages,\n",
    "    # so you may need to add an AIMessage here.\n",
    "    HumanMessage(\"Tell me another joke about planes\", name=\"example_user\"),\n",
    "    AIMessage(\n",
    "        \"\",\n",
    "        name=\"example_assistant\",\n",
    "        tool_calls=[\n",
    "            {\n",
    "                \"name\": \"joke\",\n",
    "                \"args\": {\n",
    "                    \"setup\": \"Cargo\",\n",
    "                    \"punchline\": \"Cargo 'vroom vroom', but planes go 'zoom zoom'!\",\n",
    "                    \"rating\": 10,\n",
    "                },\n",
    "                \"id\": \"2\",\n",
    "            }\n",
    "        ],\n",
    "    ),\n",
    "    ToolMessage(\"\", tool_call_id=\"2\"),\n",
    "    HumanMessage(\"Now about caterpillars\", name=\"example_user\"),\n",
    "    AIMessage(\n",
    "        \"\",\n",
    "        tool_calls=[\n",
    "            {\n",
    "                \"name\": \"joke\",\n",
    "                \"args\": {\n",
    "                    \"setup\": \"Caterpillar\",\n",
    "                    \"punchline\": \"Caterpillar really slow, but watch me turn into a butterfly and steal the show!\",\n",
    "                    \"rating\": 5,\n",
    "                },\n",
    "                \"id\": \"3\",\n",
    "            }\n",
    "        ],\n",
    "    ),\n",
    "    ToolMessage(\"\", tool_call_id=\"3\"),\n",
    "]\n",
    "system = \"\"\"You are a hilarious comedian. Your specialty is knock-knock jokes. \\\n",
    "Return a joke which has the setup (the response to \"Who's there?\") \\\n",
    "and the final punchline (the response to \"<setup> who?\").\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", system), (\"placeholder\", \"{examples}\"), (\"human\", \"{input}\")]\n",
    ")\n",
    "few_shot_structured_llm = prompt | structured_llm\n",
    "few_shot_structured_llm.invoke({\"input\": \"crocodiles\", \"examples\": examples})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "498d893b-ceaa-47ff-a9d8-4faa60702715",
   "metadata": {},
   "source": [
    "For more on few shot prompting when using tool calling, see [here](/docs/how_to/function_calling/#Few-shot-prompting)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39d7a555",
   "metadata": {},
   "source": [
    "### (Advanced) Specifying the method for structuring outputs\n",
    "\n",
    "For models that support more than one means of structuring outputs (i.e., they support both tool calling and JSON mode), you can specify which method to use with the `method=` argument.\n",
    "\n",
    ":::info JSON mode\n",
    "\n",
    "If using JSON mode you'll have to still specify the desired schema in the model prompt. The schema you pass to `with_structured_output` will only be used for parsing the model outputs, it will not be passed to the model the way it is with tool calling.\n",
    "\n",
    "To see if the model you're using supports JSON mode, check its entry in the [API reference](https://api.python.langchain.com/en/latest/langchain_api_reference.html).\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "df0370e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why was the cat sitting on the computer?', punchline='Because it wanted to keep an eye on the mouse!', rating=None)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "structured_llm = llm.with_structured_output(Joke, method=\"json_mode\")\n",
    "\n",
    "structured_llm.invoke(\n",
    "    \"Tell me a joke about cats, respond in JSON with `setup` and `punchline` keys\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e92a98a",
   "metadata": {},
   "source": [
    "## Prompting and parsing model directly\n",
    "\n",
    "Not all models support `.with_structured_output()`, since not all models have tool calling or JSON mode support. For such models you'll need to directly prompt the model to use a specific format, and use an output parser to extract the structured response from the raw model output.\n",
    "\n",
    "### Using `PydanticOutputParser`\n",
    "\n",
    "The following example uses the built-in [`PydanticOutputParser`](https://api.python.langchain.com/en/latest/output_parsers/langchain_core.output_parsers.pydantic.PydanticOutputParser.html) to parse the output of a chat model prompted to match the given Pydantic schema. Note that we are adding `format_instructions` directly to the prompt from a method on the parser:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "6e514455",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_core.output_parsers import PydanticOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class Person(BaseModel):\n",
    "    \"\"\"Information about a person.\"\"\"\n",
    "\n",
    "    name: str = Field(..., description=\"The name of the person\")\n",
    "    height_in_meters: float = Field(\n",
    "        ..., description=\"The height of the person expressed in meters.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class People(BaseModel):\n",
    "    \"\"\"Identifying information about all people in a text.\"\"\"\n",
    "\n",
    "    people: List[Person]\n",
    "\n",
    "\n",
    "# Set up a parser\n",
    "parser = PydanticOutputParser(pydantic_object=People)\n",
    "\n",
    "# Prompt\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())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082fa166",
   "metadata": {},
   "source": [
    "Let’s take a look at what information is sent to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "3d73d33d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System: Answer the user query. Wrap the output in `json` tags\n",
      "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\": \"Identifying information about all people in a text.\", \"properties\": {\"people\": {\"title\": \"People\", \"type\": \"array\", \"items\": {\"$ref\": \"#/definitions/Person\"}}}, \"required\": [\"people\"], \"definitions\": {\"Person\": {\"title\": \"Person\", \"description\": \"Information about a person.\", \"type\": \"object\", \"properties\": {\"name\": {\"title\": \"Name\", \"description\": \"The name of the person\", \"type\": \"string\"}, \"height_in_meters\": {\"title\": \"Height In Meters\", \"description\": \"The height of the person expressed in meters.\", \"type\": \"number\"}}, \"required\": [\"name\", \"height_in_meters\"]}}}\n",
      "```\n",
      "Human: Anna is 23 years old and she is 6 feet tall\n"
     ]
    }
   ],
   "source": [
    "query = \"Anna is 23 years old and she is 6 feet tall\"\n",
    "\n",
    "print(prompt.invoke(query).to_string())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "081956b9",
   "metadata": {},
   "source": [
    "And now let's invoke it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8d6b3d17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "People(people=[Person(name='Anna', height_in_meters=1.8288)])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = prompt | llm | parser\n",
    "\n",
    "chain.invoke({\"query\": query})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6732dd87",
   "metadata": {},
   "source": [
    "For a deeper dive into using output parsers with prompting techniques for structured output, see [this guide](/docs/how_to/output_parser_structured).\n",
    "\n",
    "### Custom Parsing\n",
    "\n",
    "You can also create a custom prompt and parser with [LangChain Expression Language (LCEL)](/docs/concepts/#langchain-expression-language), using a plain function to parse the output from the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e8d37e15",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import re\n",
    "from typing import List\n",
    "\n",
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class Person(BaseModel):\n",
    "    \"\"\"Information about a person.\"\"\"\n",
    "\n",
    "    name: str = Field(..., description=\"The name of the person\")\n",
    "    height_in_meters: float = Field(\n",
    "        ..., description=\"The height of the person expressed in meters.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class People(BaseModel):\n",
    "    \"\"\"Identifying information about all people in a text.\"\"\"\n",
    "\n",
    "    people: List[Person]\n",
    "\n",
    "\n",
    "# Prompt\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"Answer the user query. Output your answer as JSON that  \"\n",
    "            \"matches the given schema: ```json\\n{schema}\\n```. \"\n",
    "            \"Make sure to wrap the answer in ```json and ``` tags\",\n",
    "        ),\n",
    "        (\"human\", \"{query}\"),\n",
    "    ]\n",
    ").partial(schema=People.schema())\n",
    "\n",
    "\n",
    "# Custom parser\n",
    "def extract_json(message: AIMessage) -> List[dict]:\n",
    "    \"\"\"Extracts JSON content from a string where JSON is embedded between ```json and ``` tags.\n",
    "\n",
    "    Parameters:\n",
    "        text (str): The text containing the JSON content.\n",
    "\n",
    "    Returns:\n",
    "        list: A list of extracted JSON strings.\n",
    "    \"\"\"\n",
    "    text = message.content\n",
    "    # Define the regular expression pattern to match JSON blocks\n",
    "    pattern = r\"```json(.*?)```\"\n",
    "\n",
    "    # Find all non-overlapping matches of the pattern in the string\n",
    "    matches = re.findall(pattern, text, re.DOTALL)\n",
    "\n",
    "    # Return the list of matched JSON strings, stripping any leading or trailing whitespace\n",
    "    try:\n",
    "        return [json.loads(match.strip()) for match in matches]\n",
    "    except Exception:\n",
    "        raise ValueError(f\"Failed to parse: {message}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f1bc8f7",
   "metadata": {},
   "source": [
    "Here is the prompt sent to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c8a30d0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System: Answer the user query. Output your answer as JSON that  matches the given schema: ```json\n",
      "{'title': 'People', 'description': 'Identifying information about all people in a text.', 'type': 'object', 'properties': {'people': {'title': 'People', 'type': 'array', 'items': {'$ref': '#/definitions/Person'}}}, 'required': ['people'], 'definitions': {'Person': {'title': 'Person', 'description': 'Information about a person.', 'type': 'object', 'properties': {'name': {'title': 'Name', 'description': 'The name of the person', 'type': 'string'}, 'height_in_meters': {'title': 'Height In Meters', 'description': 'The height of the person expressed in meters.', 'type': 'number'}}, 'required': ['name', 'height_in_meters']}}}\n",
      "```. Make sure to wrap the answer in ```json and ``` tags\n",
      "Human: Anna is 23 years old and she is 6 feet tall\n"
     ]
    }
   ],
   "source": [
    "query = \"Anna is 23 years old and she is 6 feet tall\"\n",
    "\n",
    "print(prompt.format_prompt(query=query).to_string())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec018893",
   "metadata": {},
   "source": [
    "And here's what it looks like when we invoke it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e1e7baf6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'people': [{'name': 'Anna', 'height_in_meters': 1.8288}]}]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = prompt | llm | extract_json\n",
    "\n",
    "chain.invoke({\"query\": query})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poetry-venv-2",
   "language": "python",
   "name": "poetry-venv-2"
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}