{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "861a4d1f",
   "metadata": {},
   "source": [
    "# LengthBased ExampleSelector\n",
    "\n",
    "This ExampleSelector selects which examples to use based on length. This is useful when you are worried about constructing a prompt that will go over the length of the context window. For longer inputs, it will select fewer examples to include, while for shorter inputs it will select more.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c469c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.prompts import FewShotPromptTemplate\n",
    "from langchain.prompts.example_selector import LengthBasedExampleSelector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0ec6d950",
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are a lot of examples of a pretend task of creating antonyms.\n",
    "examples = [\n",
    "    {\"input\": \"happy\", \"output\": \"sad\"},\n",
    "    {\"input\": \"tall\", \"output\": \"short\"},\n",
    "    {\"input\": \"energetic\", \"output\": \"lethargic\"},\n",
    "    {\"input\": \"sunny\", \"output\": \"gloomy\"},\n",
    "    {\"input\": \"windy\", \"output\": \"calm\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "207e55f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_prompt = PromptTemplate(\n",
    "    input_variables=[\"input\", \"output\"],\n",
    "    template=\"Input: {input}\\nOutput: {output}\",\n",
    ")\n",
    "example_selector = LengthBasedExampleSelector(\n",
    "    # These are the examples it has available to choose from.\n",
    "    examples=examples, \n",
    "    # This is the PromptTemplate being used to format the examples.\n",
    "    example_prompt=example_prompt, \n",
    "    # This is the maximum length that the formatted examples should be.\n",
    "    # Length is measured by the get_text_length function below.\n",
    "    max_length=25,\n",
    "    # This is the function used to get the length of a string, which is used\n",
    "    # to determine which examples to include. It is commented out because\n",
    "    # it is provided as a default value if none is specified.\n",
    "    # get_text_length: Callable[[str], int] = lambda x: len(re.split(\"\\n| \", x))\n",
    ")\n",
    "dynamic_prompt = FewShotPromptTemplate(\n",
    "    # We provide an ExampleSelector instead of examples.\n",
    "    example_selector=example_selector,\n",
    "    example_prompt=example_prompt,\n",
    "    prefix=\"Give the antonym of every input\",\n",
    "    suffix=\"Input: {adjective}\\nOutput:\", \n",
    "    input_variables=[\"adjective\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d00b4385",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Give the antonym of every input\n",
      "\n",
      "Input: happy\n",
      "Output: sad\n",
      "\n",
      "Input: tall\n",
      "Output: short\n",
      "\n",
      "Input: energetic\n",
      "Output: lethargic\n",
      "\n",
      "Input: sunny\n",
      "Output: gloomy\n",
      "\n",
      "Input: windy\n",
      "Output: calm\n",
      "\n",
      "Input: big\n",
      "Output:\n"
     ]
    }
   ],
   "source": [
    "# An example with small input, so it selects all examples.\n",
    "print(dynamic_prompt.format(adjective=\"big\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "878bcde9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Give the antonym of every input\n",
      "\n",
      "Input: happy\n",
      "Output: sad\n",
      "\n",
      "Input: big and huge and massive and large and gigantic and tall and much much much much much bigger than everything else\n",
      "Output:\n"
     ]
    }
   ],
   "source": [
    "# An example with long input, so it selects only one example.\n",
    "long_string = \"big and huge and massive and large and gigantic and tall and much much much much much bigger than everything else\"\n",
    "print(dynamic_prompt.format(adjective=long_string))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e4bebcd9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Give the antonym of every input\n",
      "\n",
      "Input: happy\n",
      "Output: sad\n",
      "\n",
      "Input: tall\n",
      "Output: short\n",
      "\n",
      "Input: energetic\n",
      "Output: lethargic\n",
      "\n",
      "Input: sunny\n",
      "Output: gloomy\n",
      "\n",
      "Input: windy\n",
      "Output: calm\n",
      "\n",
      "Input: big\n",
      "Output: small\n",
      "\n",
      "Input: enthusiastic\n",
      "Output:\n"
     ]
    }
   ],
   "source": [
    "# You can add an example to an example selector as well.\n",
    "new_example = {\"input\": \"big\", \"output\": \"small\"}\n",
    "dynamic_prompt.example_selector.add_example(new_example)\n",
    "print(dynamic_prompt.format(adjective=\"enthusiastic\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39f30097",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}