{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb0735c0",
   "metadata": {},
   "source": [
    "# Few shot examples for chat models\n",
    "\n",
    "This notebook covers how to use few shot examples in chat models. There does not appear to be solid consensus on how best to do few shot prompting, and the optimal prompt compilation will likely vary by model. Because of this, we provide few-shot prompt templates like the [FewShotChatMessagePromptTemplate](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.few_shot.FewShotChatMessagePromptTemplate.html) as a flexible starting point, and you can modify or replace them as you see fit.\n",
    "\n",
    "The goal of few-shot prompt templates are to dynamically select examples based on an input, and then format the examples in a final prompt to provide for the model.\n",
    "\n",
    "\n",
    "**Note:** The following code examples are for chat models. For similar few-shot prompt examples for completion models (LLMs), see the [few-shot prompt templates](few_shot_examples) guide."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d716f2de-cc29-4823-9360-a808c7bfdb86",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Fixed Examples\n",
    "\n",
    "The most basic (and common) few-shot prompting technique is to use a fixed prompt example. This way you can select a chain, evaluate it, and avoid worrying about additional moving parts in production.\n",
    "\n",
    "The basic components of the template are:\n",
    "- `examples`: A list of dictionary examples to include in the final prompt.\n",
    "- `example_prompt`: converts each example into 1 or more messages through its [`format_messages`](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.chat.ChatPromptTemplate.html#langchain.prompts.chat.ChatPromptTemplate.format_messages) method. A common example would be to convert each example into one human message and one AI message response, or a human message followed by a function call message.\n",
    "\n",
    "Below is a simple demonstration. First, import the modules for this example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "91f1ca7f-a748-44c7-a1c6-a89a2d1414ba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.schema import SystemMessage\n",
    "from langchain.prompts import (\n",
    "    FewShotChatMessagePromptTemplate,\n",
    "    HumanMessagePromptTemplate,\n",
    "    AIMessagePromptTemplate,\n",
    "    SystemMessagePromptTemplate,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2844d5ed-c3cc-4bc3-9462-384fc1618b45",
   "metadata": {},
   "source": [
    "Then, define the examples you'd like to include."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0fc5a02a-6249-4e92-95c3-30fff9671e8b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples = [\n",
    "    {\"input\": \"2+2\", \"output\": \"4\"},\n",
    "    {\"input\": \"2+3\", \"output\": \"5\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8710ecc-2aa0-4172-a74c-250f6bc3d9e2",
   "metadata": {},
   "source": [
    "Next, assemble them into the few-shot prompt template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "65e72ad1-9060-47d0-91a1-bc130c8b98ac",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: 2+2\n",
      "AI: 4\n",
      "Human: 2+3\n",
      "AI: 5\n"
     ]
    }
   ],
   "source": [
    "# This is a prompt template used to format each individual example.\n",
    "example_prompt = HumanMessagePromptTemplate.from_template(\n",
    "    \"{input}\"\n",
    ") + AIMessagePromptTemplate.from_template(\"{output}\")\n",
    "few_shot_prompt = FewShotChatMessagePromptTemplate(\n",
    "    example_prompt=example_prompt,\n",
    "    examples=examples,\n",
    ")\n",
    "\n",
    "print(few_shot_prompt.format())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5490bd59-b28f-46a4-bbdf-0191802dd3c5",
   "metadata": {},
   "source": [
    "Finally, assemble your final prompt and use it with a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9f86d6d9-50de-41b6-b6c7-0f9980cc0187",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "final_prompt = (\n",
    "    SystemMessagePromptTemplate.from_template(\"You are wonderous wizard of math.\")\n",
    "    + few_shot_prompt\n",
    "    + HumanMessagePromptTemplate.from_template(\"{input}\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "97d443b1-6fae-4b36-bede-3ff7306288a3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' Triangles do not have a \"square\". A square refers to a shape with 4 equal sides and 4 right angles. Triangles have 3 sides and 3 angles.\\n\\nThe area of a triangle can be calculated using the formula:\\n\\nA = 1/2 * b * h\\n\\nWhere:\\n\\nA is the area \\nb is the base (the length of one of the sides)\\nh is the height (the length from the base to the opposite vertex)\\n\\nSo the area depends on the specific dimensions of the triangle. There is no single \"square of a triangle\". The area can vary greatly between different triangles.', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chat_models import ChatAnthropic\n",
    "\n",
    "chain = final_prompt | ChatAnthropic(temperature=0.0)\n",
    "\n",
    "chain.invoke({\"input\": \"What's the square of a triangle?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70ab7114-f07f-46be-8874-3705a25aba5f",
   "metadata": {},
   "source": [
    "## Dynamic Few-shot Prompting\n",
    "\n",
    "Sometimes you may want to condition which examples are shown based on the input. For this, you can replace the `examples` with an `example_selector`. The other components remain the same as above! To review, the dynamic few-shot prompt template would look like:\n",
    "\n",
    "- `example_selector`: responsible for selecting few-shot examples (and the order in which they are returned) for a given input. These implement the [BaseExampleSelector](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.example_selector.base.BaseExampleSelector.html#langchain.prompts.example_selector.base.BaseExampleSelector) interface. A common example is the vectorstore-backed [SemanticSimilarityExampleSelector](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.example_selector.semantic_similarity.SemanticSimilarityExampleSelector.html#langchain.prompts.example_selector.semantic_similarity.SemanticSimilarityExampleSelector)\n",
    "- `example_prompt`: convert each example into 1 or more messages through its [`format_messages`](https://api.python.langchain.com/en/latest/prompts/langchain.prompts.chat.ChatPromptTemplate.html#langchain.prompts.chat.ChatPromptTemplate.format_messages) method. A common example would be to convert each example into one human message and one AI message response, or a human message followed by a function call message.\n",
    "\n",
    "These once again can be composed with other messages and chat templates to assemble your final prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "6f7b5e86-4ca7-4edd-bf2b-9663030b2393",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.prompts import SemanticSimilarityExampleSelector\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.vectorstores import Chroma"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "303b3f81-8d17-4fa2-81b1-e10bf34dd514",
   "metadata": {},
   "source": [
    "Since we are using a vectorstore to select examples based on semantic similarity, we will want to first populate the store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ad66f06a-66fd-4fcc-8166-5d0e3c801e57",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "examples = [\n",
    "    {\"input\": \"2+2\", \"output\": \"4\"},\n",
    "    {\"input\": \"2+3\", \"output\": \"5\"},\n",
    "    {\"input\": \"2+4\", \"output\": \"6\"},\n",
    "    {\"input\": \"What did the cow say to the moon?\", \"output\": \"nothing at all\"},\n",
    "    {\n",
    "        \"input\": \"Write me a poem about the moon\",\n",
    "        \"output\": \"One for the moon, and one for me, who are we to talk about the moon?\",\n",
    "    },\n",
    "]\n",
    "\n",
    "to_vectorize = [\" \".join(example.values()) for example in examples]\n",
    "embeddings = OpenAIEmbeddings()\n",
    "vectorstore = Chroma.from_texts(to_vectorize, embeddings, metadatas=examples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f7e384a-2031-432b-951c-7ea8cf9262f1",
   "metadata": {},
   "source": [
    "#### Create the `example_selector`\n",
    "\n",
    "With a vectorstore created, you can create the `example_selector`. Here we will isntruct it to only fetch the top 2 examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7790303a-f722-452e-8921-b14bdf20bdff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'input': 'What did the cow say to the moon?', 'output': 'nothing at all'},\n",
       " {'input': '2+4', 'output': '6'}]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_selector = SemanticSimilarityExampleSelector(\n",
    "    vectorstore=vectorstore,\n",
    "    k=2,\n",
    ")\n",
    "\n",
    "# The prompt template will load examples by passing the input do the `select_examples` method\n",
    "example_selector.select_examples({\"input\": \"horse\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc77c40f-3f58-40a2-b757-a2a2ea43f24a",
   "metadata": {},
   "source": [
    "#### Create prompt template\n",
    "\n",
    "Assemble the prompt template, using the `example_selector` created above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "253c255e-41d7-45f6-9d88-c7a0ced4b1bd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.schema import SystemMessage\n",
    "from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate\n",
    "from langchain.prompts.few_shot import FewShotChatMessagePromptTemplate\n",
    "\n",
    "\n",
    "# Define the few-shot prompt.\n",
    "few_shot_prompt = FewShotChatMessagePromptTemplate(\n",
    "    # The input variables select the values to pass to the example_selector\n",
    "    input_variables=[\"input\"],\n",
    "    example_selector=example_selector,\n",
    "    # Define how each example will be formatted.\n",
    "    # In this case, each example will become 2 messages:\n",
    "    # 1 human, and 1 AI\n",
    "    example_prompt=(\n",
    "        HumanMessagePromptTemplate.from_template(\"{input}\")\n",
    "        + AIMessagePromptTemplate.from_template(\"{output}\")\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d960a471-1e1d-4742-ae49-dd0afcdb34d5",
   "metadata": {},
   "source": [
    "Below is an example of how this would be assembled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "860bf682-c469-40e9-b657-27bfe7026099",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: 2+3\n",
      "AI: 5\n",
      "Human: 2+2\n",
      "AI: 4\n"
     ]
    }
   ],
   "source": [
    "print(few_shot_prompt.format(input=\"What's 3+3?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "339cae7d-0eb0-44a6-852f-0267c5ff72b3",
   "metadata": {},
   "source": [
    "Assemble the final prompt template:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e731cb45-f0ea-422c-be37-42af2a6cb2c4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "final_prompt = (\n",
    "    SystemMessagePromptTemplate.from_template(\"You are wonderous wizard of math.\")\n",
    "    + few_shot_prompt\n",
    "    + HumanMessagePromptTemplate.from_template(\"{input}\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e6cc4199-8947-42d7-91f0-375de1e15bd9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: 2+3\n",
      "AI: 5\n",
      "Human: 2+2\n",
      "AI: 4\n"
     ]
    }
   ],
   "source": [
    "print(few_shot_prompt.format(input=\"What's 3+3?\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2408ea69-1880-4ef5-a0fa-ffa8d2026aa9",
   "metadata": {},
   "source": [
    "#### Use with an LLM\n",
    "\n",
    "Now, you can connect your model to the few-shot prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0568cbc6-5354-47f1-ab4d-dfcc616cf583",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' 3 + 3 = 6', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chat_models import ChatAnthropic\n",
    "\n",
    "chain = final_prompt | ChatAnthropic(temperature=0.0)\n",
    "\n",
    "chain.invoke({\"input\": \"What's 3+3?\"})"
   ]
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
