{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1b79ff35-50a3-40cd-86d9-703f1f8cd2c5",
   "metadata": {},
   "source": [
    "# How to get a RAG application to add citations\n",
    "\n",
    "This guide reviews methods to get a model to cite which parts of the source documents it referenced in generating its response.\n",
    "\n",
    "We will cover five methods:\n",
    "\n",
    "1. Using tool-calling to cite document IDs;\n",
    "2. Using tool-calling to cite documents IDs and provide text snippets;\n",
    "3. Direct prompting;\n",
    "4. Retrieval post-processing (i.e., compressing the retrieved context to make it more relevant);\n",
    "5. Generation post-processing (i.e., issuing a second LLM call to annotate a generated answer with citations).\n",
    "\n",
    "We generally suggest using the first item of the list that works for your use-case. That is, if your model supports tool-calling, try methods 1 or 2; otherwise, or if those fail, advance down the list.\n",
    "\n",
    "Let's first create a simple RAG chain. To start we'll just retrieve from Wikipedia using the [WikipediaRetriever](https://api.python.langchain.com/en/latest/retrievers/langchain_community.retrievers.wikipedia.WikipediaRetriever.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a70c423-f61f-4230-b70a-d3605b31afab",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First we'll need to install some dependencies and set environment vars for the models we'll be using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f1d26ded-e8d5-4f80-86b9-26d464869175",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU langchain langchain-openai langchain-anthropic langchain-community wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8732a85a-dd1a-483c-8da7-a81251276aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "os.environ[\"ANTHROPIC_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "# Uncomment if you want to log to LangSmith\n",
    "# os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\n",
    "# os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30a4401f-7feb-4bd9-9409-77c3859c4292",
   "metadata": {},
   "source": [
    "Let's first select a LLM:\n",
    "\n",
    "```{=mdx}\n",
    "import ChatModelTabs from \"@theme/ChatModelTabs\";\n",
    "\n",
    "<ChatModelTabs customVarName=\"llm\" />\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dd00165d-0b32-466d-8f75-ec26326a9e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | output: false\n",
    "# | echo: false\n",
    "\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4e17c3f6-8ce6-4767-b615-50a57c84c7b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m System Message \u001b[0m================================\n",
      "\n",
      "You're a helpful AI assistant. Given a user question and some Wikipedia article snippets, answer the user question. If none of the articles answer the question, just say you don't know.\n",
      "\n",
      "Here are the Wikipedia articles: \u001b[33;1m\u001b[1;3m{context}\u001b[0m\n",
      "\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "\u001b[33;1m\u001b[1;3m{input}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.retrievers import WikipediaRetriever\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "system_prompt = (\n",
    "    \"You're a helpful AI assistant. Given a user question \"\n",
    "    \"and some Wikipedia article snippets, answer the user \"\n",
    "    \"question. If none of the articles answer the question, \"\n",
    "    \"just say you don't know.\"\n",
    "    \"\\n\\nHere are the Wikipedia articles: \"\n",
    "    \"{context}\"\n",
    ")\n",
    "\n",
    "retriever = WikipediaRetriever(top_k_results=6, doc_content_chars_max=2000)\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system_prompt),\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "prompt.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c89e2045-9244-43e6-bf3f-59af22658529",
   "metadata": {},
   "source": [
    "Now that we've got a model, retriver and prompt, let's chain them all together. We'll need to add some logic for formatting our retrieved Documents to a string that can be passed to our prompt. Following the how-to guide on [adding citations](/docs/how_to/qa_citations) to a RAG application, we'll make it so our chain returns both the answer and the retrieved Documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4cd55e1c-a6b7-44b7-9dde-5f42abe714ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "\n",
    "def format_docs(docs: List[Document]):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "\n",
    "rag_chain_from_docs = (\n",
    "    RunnablePassthrough.assign(context=(lambda x: format_docs(x[\"context\"])))\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "retrieve_docs = (lambda x: x[\"input\"]) | retriever\n",
    "\n",
    "chain = RunnablePassthrough.assign(context=retrieve_docs).assign(\n",
    "    answer=rag_chain_from_docs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "42b28717-d34c-42de-b923-155ac60529a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = chain.invoke({\"input\": \"How fast are cheetahs?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8b20cf8e-dccd-45d1-aef0-25f1ad1aca6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['input', 'context', 'answer'])\n"
     ]
    }
   ],
   "source": [
    "print(result.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ae5ed9a7-c72a-480d-80c6-0a6bd38b9941",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a' metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}\n"
     ]
    }
   ],
   "source": [
    "print(result[\"context\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "31f20897-0a7a-44e8-aeac-75d54f6e3789",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cheetahs are capable of running at speeds of 93 to 104 km/h (58 to 65 mph). They have evolved specialized adaptations for speed, including a light build, long thin legs, and a long tail.\n"
     ]
    }
   ],
   "source": [
    "print(result[\"answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f1f9a49-8f3f-44dd-98df-0218b5fb93a6",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/0472c5d1-49dc-4c1c-8100-61910067d7ed/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7619ba1-33bd-48bf-8637-be409c94037f",
   "metadata": {},
   "source": [
    "## Function-calling\n",
    "\n",
    "If your LLM of choice implements a [tool-calling](/docs/concepts#functiontool-calling) feature, you can use it to make the model specify which of the provided documents it's referencing when generating its answer. LangChain tool-calling models implement a `.with_structured_output` method which will force generation adhering to a desired schema (see for example [here](https://api.python.langchain.com/en/latest/chat_models/langchain_openai.chat_models.base.ChatOpenAI.html#langchain_openai.chat_models.base.ChatOpenAI.with_structured_output)).\n",
    "\n",
    "### Cite documents\n",
    "\n",
    "To cite documents using an identifier, we format the identifiers into the prompt, then use `.with_structured_output` to coerce the LLM to reference these identifiers in its output.\n",
    "\n",
    "First we define a schema for the output. The `.with_structured_output` supports multiple formats, including JSON schema and Pydantic. Here we will use Pydantic:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0af2c3a1-870c-428e-95da-0c2fd04d5616",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class CitedAnswer(BaseModel):\n",
    "    \"\"\"Answer the user question based only on the given sources, and cite the sources used.\"\"\"\n",
    "\n",
    "    answer: str = Field(\n",
    "        ...,\n",
    "        description=\"The answer to the user question, which is based only on the given sources.\",\n",
    "    )\n",
    "    citations: List[int] = Field(\n",
    "        ...,\n",
    "        description=\"The integer IDs of the SPECIFIC sources which justify the answer.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68b95186-faf5-46f1-8715-ebbc38207d5d",
   "metadata": {},
   "source": [
    "Let's see what the model output is like when we pass in our functions and a user input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2e2b7a87-3642-4ed8-9445-684daa93b0d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CitedAnswer(answer='Brian\\'s height is 5\\'11\".', citations=[1, 3])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "structured_llm = llm.with_structured_output(CitedAnswer)\n",
    "\n",
    "example_q = \"\"\"What Brian's height?\n",
    "\n",
    "Source: 1\n",
    "Information: Suzy is 6'2\"\n",
    "\n",
    "Source: 2\n",
    "Information: Jeremiah is blonde\n",
    "\n",
    "Source: 3\n",
    "Information: Brian is 3 inches shorter than Suzy\"\"\"\n",
    "result = structured_llm.invoke(example_q)\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b847b53-987e-4d3a-9621-77e613d49cfd",
   "metadata": {},
   "source": [
    "Or as a dict:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3ee49bbd-567f-41cc-8798-d5aad0fe1cea",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'answer': 'Brian\\'s height is 5\\'11\".', 'citations': [1, 3]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result.dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb8bbbb5-2afc-401f-a140-648c3d2c4522",
   "metadata": {},
   "source": [
    "Now we structure the source identifiers into the prompt to replicate with our chain. We will make three changes:\n",
    "\n",
    "1. Update the prompt to include source identifiers;\n",
    "2. Use the `structured_llm` (i.e., `llm.with_structured_output(CitedAnswer));\n",
    "3. Remove the `StrOutputParser`, to retain the Pydantic object in the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3cb835f3-3cf5-4144-bf6b-24558b9faf31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_docs_with_id(docs: List[Document]) -> str:\n",
    "    formatted = [\n",
    "        f\"Source ID: {i}\\nArticle Title: {doc.metadata['title']}\\nArticle Snippet: {doc.page_content}\"\n",
    "        for i, doc in enumerate(docs)\n",
    "    ]\n",
    "    return \"\\n\\n\" + \"\\n\\n\".join(formatted)\n",
    "\n",
    "\n",
    "rag_chain_from_docs = (\n",
    "    RunnablePassthrough.assign(context=(lambda x: format_docs_with_id(x[\"context\"])))\n",
    "    | prompt\n",
    "    | structured_llm\n",
    ")\n",
    "\n",
    "retrieve_docs = (lambda x: x[\"input\"]) | retriever\n",
    "\n",
    "chain = RunnablePassthrough.assign(context=retrieve_docs).assign(\n",
    "    answer=rag_chain_from_docs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3e259b2f-5147-4c3c-9c26-b4eb8143e5f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = chain.invoke({\"input\": \"How fast are cheetahs?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2d8d2a01-608d-479f-85f1-eb8d14b11bc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "answer='Cheetahs can run at speeds of 93 to 104 km/h (58 to 65 mph). They are known as the fastest land animals.' citations=[0]\n"
     ]
    }
   ],
   "source": [
    "print(result[\"answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da8341f5-a48a-4c07-8445-a313e20c36a2",
   "metadata": {},
   "source": [
    "We can inspect the document at index 0, which the model cited:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "02d19f2b-2e15-492f-b44b-577990d15a86",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "page_content='The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned a' metadata={'title': 'Cheetah', 'summary': 'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.', 'source': 'https://en.wikipedia.org/wiki/Cheetah'}\n"
     ]
    }
   ],
   "source": [
    "print(result[\"context\"][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94f2898a-ef4d-423a-b002-910fef7a65c9",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/aff39dc7-3e09-4d64-8083-87026d975534/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdbd1407-8a5b-4c35-aa2b-9d26424edb93",
   "metadata": {},
   "source": [
    "### Cite snippets\n",
    "\n",
    "To return text spans (perhaps in addition to source identifiers), we can use the same approach. The only change will be to build a more complex output schema, here using Pydantic, that includes a \"quote\" alongside a source identifier.\n",
    "\n",
    "*Aside: Note that if we break up our documents so that we have many documents with only a sentence or two instead of a few long documents, citing documents becomes roughly equivalent to citing snippets, and may be easier for the model because the model just needs to return an identifier for each snippet instead of the actual text. Probably worth trying both approaches and evaluating.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fbf708aa-e8ac-4dea-bb57-82229597e2e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Citation(BaseModel):\n",
    "    source_id: int = Field(\n",
    "        ...,\n",
    "        description=\"The integer ID of a SPECIFIC source which justifies the answer.\",\n",
    "    )\n",
    "    quote: str = Field(\n",
    "        ...,\n",
    "        description=\"The VERBATIM quote from the specified source that justifies the answer.\",\n",
    "    )\n",
    "\n",
    "\n",
    "class QuotedAnswer(BaseModel):\n",
    "    \"\"\"Answer the user question based only on the given sources, and cite the sources used.\"\"\"\n",
    "\n",
    "    answer: str = Field(\n",
    "        ...,\n",
    "        description=\"The answer to the user question, which is based only on the given sources.\",\n",
    "    )\n",
    "    citations: List[Citation] = Field(\n",
    "        ..., description=\"Citations from the given sources that justify the answer.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "beabab7b-7b6b-4eef-b874-e92d1ed8707c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_chain_from_docs = (\n",
    "    RunnablePassthrough.assign(context=(lambda x: format_docs_with_id(x[\"context\"])))\n",
    "    | prompt\n",
    "    | llm.with_structured_output(QuotedAnswer)\n",
    ")\n",
    "\n",
    "retrieve_docs = (lambda x: x[\"input\"]) | retriever\n",
    "\n",
    "chain = RunnablePassthrough.assign(context=retrieve_docs).assign(\n",
    "    answer=rag_chain_from_docs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9709ee6d-416f-4bd3-89c6-23667b9f3cca",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = chain.invoke({\"input\": \"How fast are cheetahs?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42ba8c6-4214-49f5-b920-f0e028f301c2",
   "metadata": {},
   "source": [
    "Here we see that the model has extracted a relevant snippet of text from source 0:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "56b01963-8680-4782-9c3f-384c197f0c2d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuotedAnswer(answer='Cheetahs can run at speeds of 93 to 104 km/h (58 to 65 mph).', citations=[Citation(source_id=0, quote='The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.')])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"answer\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28676cf1-4a2e-44d2-8b2f-36303a12a371",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/0f638cc9-8409-4a53-9010-86ac28144129/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb2d90a4-0370-4598-9f4b-e8e9a554346e",
   "metadata": {},
   "source": [
    "## Direct prompting\n",
    "\n",
    "Many models don't support function-calling. We can achieve similar results with direct prompting. Let's try instructing a model to generate structured XML for its output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4e95bd8a-2f15-4e20-a1d9-225974b8d598",
   "metadata": {},
   "outputs": [],
   "source": [
    "xml_system = \"\"\"You're a helpful AI assistant. Given a user question and some Wikipedia article snippets, \\\n",
    "answer the user question and provide citations. If none of the articles answer the question, just say you don't know.\n",
    "\n",
    "Remember, you must return both an answer and citations. A citation consists of a VERBATIM quote that \\\n",
    "justifies the answer and the ID of the quote article. Return a citation for every quote across all articles \\\n",
    "that justify the answer. Use the following format for your final output:\n",
    "\n",
    "<cited_answer>\n",
    "    <answer></answer>\n",
    "    <citations>\n",
    "        <citation><source_id></source_id><quote></quote></citation>\n",
    "        <citation><source_id></source_id><quote></quote></citation>\n",
    "        ...\n",
    "    </citations>\n",
    "</cited_answer>\n",
    "\n",
    "Here are the Wikipedia articles:{context}\"\"\"\n",
    "xml_prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", xml_system), (\"human\", \"{input}\")]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d3bd0f7-e249-4bc6-bd46-6fb74ebf0118",
   "metadata": {},
   "source": [
    "We now make similar small updates to our chain:\n",
    "\n",
    "1. We update the formatting function to wrap the retrieved context in XML tags;\n",
    "2. We do not use `.with_structured_output` (e.g., because it does not exist for a model);\n",
    "3. We use [XMLOutputParser](https://api.python.langchain.com/en/latest/output_parsers/langchain_core.output_parsers.xml.XMLOutputParser.html) in place of `StrOutputParser` to parse the answer into a dict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5861ca8c-63b7-4918-bdc6-fe4e53fe03ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import XMLOutputParser\n",
    "\n",
    "\n",
    "def format_docs_xml(docs: List[Document]) -> str:\n",
    "    formatted = []\n",
    "    for i, doc in enumerate(docs):\n",
    "        doc_str = f\"\"\"\\\n",
    "    <source id=\\\"{i}\\\">\n",
    "        <title>{doc.metadata['title']}</title>\n",
    "        <article_snippet>{doc.page_content}</article_snippet>\n",
    "    </source>\"\"\"\n",
    "        formatted.append(doc_str)\n",
    "    return \"\\n\\n<sources>\" + \"\\n\".join(formatted) + \"</sources>\"\n",
    "\n",
    "\n",
    "rag_chain_from_docs = (\n",
    "    RunnablePassthrough.assign(context=(lambda x: format_docs_xml(x[\"context\"])))\n",
    "    | xml_prompt\n",
    "    | llm\n",
    "    | XMLOutputParser()\n",
    ")\n",
    "\n",
    "retrieve_docs = (lambda x: x[\"input\"]) | retriever\n",
    "\n",
    "chain = RunnablePassthrough.assign(context=retrieve_docs).assign(\n",
    "    answer=rag_chain_from_docs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f1edb401-6027-4112-82ec-25736e8ebabd",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = chain.invoke({\"input\": \"How fast are cheetahs?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5264571-48c2-492d-a750-640f9fff3e71",
   "metadata": {},
   "source": [
    "Note that citations are again structured into the answer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a2b4bdc9-92dd-434c-b61c-11ec44c92905",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'cited_answer': [{'answer': 'Cheetahs are capable of running at 93 to 104 km/h (58 to 65 mph).'},\n",
       "  {'citations': [{'citation': [{'source_id': '0'},\n",
       "      {'quote': 'The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.'}]}]}]}"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"answer\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "940db8d5-8f43-44dd-9738-04fc7464baac",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/a3636c70-39c6-4c8f-bc83-1c7a174c237e/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d4180b0-5d29-4bfa-85be-2a6161a872c4",
   "metadata": {},
   "source": [
    "## Retrieval post-processing\n",
    "\n",
    "Another approach is to post-process our retrieved documents to compress the content, so that the source content is already minimal enough that we don't need the model to cite specific sources or spans. For example, we could break up each document into a sentence or two, embed those and keep only the most relevant ones. LangChain has some built-in components for this. Here we'll use a [RecursiveCharacterTextSplitter](https://api.python.langchain.com/en/latest/text_splitter/langchain_text_splitters.RecursiveCharacterTextSplitter.html#langchain_text_splitters.RecursiveCharacterTextSplitter), which creates chunks of a sepacified size by splitting on separator substrings, and an [EmbeddingsFilter](https://api.python.langchain.com/en/latest/retrievers/langchain.retrievers.document_compressors.embeddings_filter.EmbeddingsFilter.html#langchain.retrievers.document_compressors.embeddings_filter.EmbeddingsFilter), which keeps only the texts with the most relevant embeddings.\n",
    "\n",
    "This approach effectively swaps our original retriever with an updated one that compresses the documents. To start, we build the retriever:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9b14f817-4454-47b2-9eb0-2b8783a8c252",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail\n",
      "\n",
      "\n",
      "\n",
      "The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in)\n",
      "\n",
      "\n",
      "\n",
      "2 mph), or 171 body lengths per second. The cheetah, the fastest land mammal, scores at only 16 body lengths per second, while Anna's hummingbird has the highest known length-specific velocity attained by any vertebrate\n",
      "\n",
      "\n",
      "\n",
      "It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson's gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year\n",
      "\n",
      "\n",
      "\n",
      "The cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran\n",
      "\n",
      "\n",
      "\n",
      "The cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk\n",
      "\n",
      "\n",
      "\n",
      "The Southeast African cheetah (Acinonyx jubatus jubatus) is the nominate cheetah subspecies native to East and Southern Africa. The Southern African cheetah lives mainly in the lowland areas and deserts of the Kalahari, the savannahs of Okavango Delta, and the grasslands of the Transvaal region in South Africa. In Namibia, cheetahs are mostly found in farmlands\n",
      "\n",
      "\n",
      "\n",
      "Subpopulations have been called \"South African cheetah\" and \"Namibian cheetah.\"\n",
      "\n",
      "\n",
      "\n",
      "In India, four cheetahs of the subspecies are living in Kuno National Park in Madhya Pradesh after having been introduced there\n",
      "\n",
      "\n",
      "\n",
      "Acinonyx jubatus velox proposed in 1913 by Edmund Heller on basis of a cheetah that was shot by Kermit Roosevelt in June 1909 in the Kenyan highlands.\n",
      "Acinonyx rex proposed in 1927 by Reginald Innes Pocock on basis of a specimen from the Umvukwe Range in Rhodesia.\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from langchain.retrievers.document_compressors import EmbeddingsFilter\n",
    "from langchain_core.runnables import RunnableParallel\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=400,\n",
    "    chunk_overlap=0,\n",
    "    separators=[\"\\n\\n\", \"\\n\", \".\", \" \"],\n",
    "    keep_separator=False,\n",
    ")\n",
    "compressor = EmbeddingsFilter(embeddings=OpenAIEmbeddings(), k=10)\n",
    "\n",
    "\n",
    "def split_and_filter(input) -> List[Document]:\n",
    "    docs = input[\"docs\"]\n",
    "    question = input[\"question\"]\n",
    "    split_docs = splitter.split_documents(docs)\n",
    "    stateful_docs = compressor.compress_documents(split_docs, question)\n",
    "    return [stateful_doc for stateful_doc in stateful_docs]\n",
    "\n",
    "\n",
    "new_retriever = (\n",
    "    RunnableParallel(question=RunnablePassthrough(), docs=retriever) | split_and_filter\n",
    ")\n",
    "docs = new_retriever.invoke(\"How fast are cheetahs?\")\n",
    "for doc in docs:\n",
    "    print(doc.page_content)\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "984bc1e1-76fb-4d84-baa9-5fa5abca9da4",
   "metadata": {},
   "source": [
    "Next, we assemble it into our chain as before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "fa2adb01-5d8f-484c-8216-bae35717db0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_chain_from_docs = (\n",
    "    RunnablePassthrough.assign(context=(lambda x: format_docs(x[\"context\"])))\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "chain = RunnablePassthrough.assign(\n",
    "    context=(lambda x: x[\"input\"]) | new_retriever\n",
    ").assign(answer=rag_chain_from_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "1a5b72f8-135b-4604-8777-59f2ef682323",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cheetahs are capable of running at speeds between 93 to 104 km/h (58 to 65 mph), making them the fastest land animals.\n"
     ]
    }
   ],
   "source": [
    "result = chain.invoke({\"input\": \"How fast are cheetahs?\"})\n",
    "\n",
    "print(result[\"answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9ac43ab-db4f-458a-9b5a-fd3e116229bd",
   "metadata": {},
   "source": [
    "Note that the document content is now compressed, although the document objects retain the original content in a \"summary\" key in their metadata. These summaries are not passed to the model; only the condensed content is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "80625506-8764-4adf-a467-33f465d0f51f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"context\"][0].page_content  # passed to model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "672c5691-5d54-4271-9d97-93571eebda91",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The cheetah (Acinonyx jubatus) is a large cat and the fastest land animal. It has a tawny to creamy white or pale buff fur that is marked with evenly spaced, solid black spots. The head is small and rounded, with a short snout and black tear-like facial streaks. It reaches 67–94 cm (26–37 in) at the shoulder, and the head-and-body length is between 1.1 and 1.5 m (3 ft 7 in and 4 ft 11 in). Adults weigh between 21 and 72 kg (46 and 159 lb). The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.\\nThe cheetah was first described in the late 18th century. Four subspecies are recognised today that are native to Africa and central Iran. An African subspecies was introduced to India in 2022. It is now distributed mainly in small, fragmented populations in northwestern, eastern and southern Africa and central Iran. It lives in a variety of habitats such as savannahs in the Serengeti, arid mountain ranges in the Sahara, and hilly desert terrain.\\nThe cheetah lives in three main social groups: females and their cubs, male \"coalitions\", and solitary males. While females lead a nomadic life searching for prey in large home ranges, males are more sedentary and instead establish much smaller territories in areas with plentiful prey and access to females. The cheetah is active during the day, with peaks during dawn and dusk. It feeds on small- to medium-sized prey, mostly weighing under 40 kg (88 lb), and prefers medium-sized ungulates such as impala, springbok and Thomson\\'s gazelles. The cheetah typically stalks its prey within 60–100 m (200–330 ft) before charging towards it, trips it during the chase and bites its throat to suffocate it to death. It breeds throughout the year. After a gestation of nearly three months, females give birth to a litter of three or four cubs. Cheetah cubs are highly vulnerable to predation by other large carnivores. They are weaned at around four months and are independent by around 20 months of age.\\nThe cheetah is threatened by habitat loss, conflict with humans, poaching and high susceptibility to diseases. In 2016, the global cheetah population was estimated at 7,100 individuals in the wild; it is listed as Vulnerable on the IUCN Red List. It has been widely depicted in art, literature, advertising, and animation. It was tamed in ancient Egypt and trained for hunting ungulates in the Arabian Peninsula and India. It has been kept in zoos since the early 19th century.'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"context\"][0].metadata[\"summary\"]  # original document"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ab8fd6-f6b4-4ba5-b022-f10cca983490",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/a61304fa-e5a5-4c64-a268-b0aef1130d53/r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "445722dc-2ecb-45a4-9d4d-c172d0a2fa7d",
   "metadata": {},
   "source": [
    "## Generation post-processing\n",
    "\n",
    "Another approach is to post-process our model generation. In this example we'll first generate just an answer, and then we'll ask the model to annotate it's own answer with citations. The downside of this approach is of course that it is slower and more expensive, because two model calls need to be made.\n",
    "\n",
    "Let's apply this to our initial chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "daff5cb9-7639-4d30-b6e7-d795736a2b58",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Citation(BaseModel):\n",
    "    source_id: int = Field(\n",
    "        ...,\n",
    "        description=\"The integer ID of a SPECIFIC source which justifies the answer.\",\n",
    "    )\n",
    "    quote: str = Field(\n",
    "        ...,\n",
    "        description=\"The VERBATIM quote from the specified source that justifies the answer.\",\n",
    "    )\n",
    "\n",
    "\n",
    "class AnnotatedAnswer(BaseModel):\n",
    "    \"\"\"Annotate the answer to the user question with quote citations that justify the answer.\"\"\"\n",
    "\n",
    "    citations: List[Citation] = Field(\n",
    "        ..., description=\"Citations from the given sources that justify the answer.\"\n",
    "    )\n",
    "\n",
    "\n",
    "structured_llm = llm.with_structured_output(AnnotatedAnswer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "6f505eb9-db02-4c49-add3-1e469844d7ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import MessagesPlaceholder\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system_prompt),\n",
    "        (\"human\", \"{question}\"),\n",
    "        MessagesPlaceholder(\"chat_history\", optional=True),\n",
    "    ]\n",
    ")\n",
    "answer = prompt | llm\n",
    "annotation_chain = prompt | structured_llm\n",
    "\n",
    "chain = (\n",
    "    RunnableParallel(\n",
    "        question=RunnablePassthrough(), docs=(lambda x: x[\"input\"]) | retriever\n",
    "    )\n",
    "    .assign(context=format)\n",
    "    .assign(ai_message=answer)\n",
    "    .assign(\n",
    "        chat_history=(lambda x: [x[\"ai_message\"]]),\n",
    "        answer=(lambda x: x[\"ai_message\"].content),\n",
    "    )\n",
    "    .assign(annotations=annotation_chain)\n",
    "    .pick([\"answer\", \"docs\", \"annotations\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "eb11c422-09b3-4d5a-87eb-3bad2e73cf6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = chain.invoke({\"input\": \"How fast are cheetahs?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "5b8bbc02-f753-4abc-87ec-211aac3dc3d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cheetahs are capable of running at speeds between 93 to 104 km/h (58 to 65 mph). Their specialized adaptations for speed, such as a light build, long thin legs, and a long tail, allow them to be the fastest land animals.\n"
     ]
    }
   ],
   "source": [
    "print(result[\"answer\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c7882b76-db21-40ee-bb31-ff438880adf6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnotatedAnswer(citations=[Citation(source_id=0, quote='The cheetah is capable of running at 93 to 104 km/h (58 to 65 mph); it has evolved specialized adaptations for speed, including a light build, long thin legs and a long tail.')])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result[\"annotations\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "803c6155-48af-40db-b4b0-1ecc5328e99b",
   "metadata": {},
   "source": [
    "LangSmith trace: https://smith.langchain.com/public/bf5e8856-193b-4ff2-af8d-c0f4fbd1d9cb/r"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
