{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "80f15d95-00d8-4c38-a291-07ff2233b4fd",
   "metadata": {},
   "source": [
    "# How to create a custom Output Parser\n",
    "\n",
    "In some situations you may want to implement a custom parser to structure the model output into a custom format.\n",
    "\n",
    "There are two ways to implement a custom parser:\n",
    "\n",
    "1. Using `RunnableLambda` or `RunnableGenerator` in LCEL -- we strongly recommend this for most use cases\n",
    "2. By inherting from one of the base classes for out parsing -- this is the hard way of doing things\n",
    "\n",
    "The difference between the two approaches are mostly superficial and are mainly in terms of which callbacks are triggered (e.g., `on_chain_start` vs. `on_parser_start`), and how a runnable lambda vs. a parser might be visualized in a tracing platform like LangSmith."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c651cc26-28cb-45d1-9969-d88deff8b819",
   "metadata": {},
   "source": [
    "## Runnable Lambdas and Generators\n",
    "\n",
    "The recommended way to parse is using **runnable lambdas** and **runnable generators**!\n",
    "\n",
    "Here, we will make a simple parse that inverts the case of the output from the model.\n",
    "\n",
    "For example, if the model outputs: \"Meow\", the parser will produce \"mEOW\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6cd7cc21-ec51-4e22-82d0-32c4401f5adc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'hELLO!'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Iterable\n",
    "\n",
    "from langchain_anthropic.chat_models import ChatAnthropic\n",
    "from langchain_core.messages import AIMessage, AIMessageChunk\n",
    "\n",
    "model = ChatAnthropic(model_name=\"claude-2.1\")\n",
    "\n",
    "\n",
    "def parse(ai_message: AIMessage) -> str:\n",
    "    \"\"\"Parse the AI message.\"\"\"\n",
    "    return ai_message.content.swapcase()\n",
    "\n",
    "\n",
    "chain = model | parse\n",
    "chain.invoke(\"hello\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed8baf2-f4c2-44c1-b47d-e9f560af6202",
   "metadata": {},
   "source": [
    ":::{.callout-tip}\n",
    "\n",
    "LCEL automatically upgrades the function `parse` to `RunnableLambda(parse)` when composed using a `|`  syntax.\n",
    "\n",
    "If you don't like that you can manually import `RunnableLambda` and then run`parse = RunnableLambda(parse)`.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "896f52ce-91e2-4c7c-bd62-1f901002ade2",
   "metadata": {},
   "source": [
    "Does streaming work?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4e35389a-caa5-4c0d-9d95-48648d0b8d4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i'M cLAUDE, AN ai ASSISTANT CREATED BY aNTHROPIC TO BE HELPFUL, HARMLESS, AND HONEST.|"
     ]
    }
   ],
   "source": [
    "for chunk in chain.stream(\"tell me about yourself in one sentence\"):\n",
    "    print(chunk, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11c486bb-b2d4-461b-8fd8-19b9e0472129",
   "metadata": {},
   "source": [
    "No, it doesn't because the parser aggregates the input before parsing the output.\n",
    "\n",
    "If we want to implement a streaming parser, we can have the parser accept an iterable over the input instead and yield\n",
    "the results as they're available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "930aa59e-82d0-447c-b711-b416d92a08b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableGenerator\n",
    "\n",
    "\n",
    "def streaming_parse(chunks: Iterable[AIMessageChunk]) -> Iterable[str]:\n",
    "    for chunk in chunks:\n",
    "        yield chunk.content.swapcase()\n",
    "\n",
    "\n",
    "streaming_parse = RunnableGenerator(streaming_parse)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62192808-c7e1-4b3a-85f4-b7901de7c0b8",
   "metadata": {},
   "source": [
    ":::{.callout-important}\n",
    "\n",
    "Please wrap the streaming parser in `RunnableGenerator` as we may stop automatically upgrading it with the `|` syntax.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c054d4da-66f3-4f11-8137-0734bb3de06c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'hELLO!'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = model | streaming_parse\n",
    "chain.invoke(\"hello\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d344ff2-5c93-49a9-af00-03856d2cbfdb",
   "metadata": {},
   "source": [
    "Let's confirm that streaming works!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "26d746ae-9c5a-4cda-a535-33f555e2e04a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i|'M| cLAUDE|,| AN| ai| ASSISTANT| CREATED| BY| aN|THROP|IC| TO| BE| HELPFUL|,| HARMLESS|,| AND| HONEST|.|"
     ]
    }
   ],
   "source": [
    "for chunk in chain.stream(\"tell me about yourself in one sentence\"):\n",
    "    print(chunk, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24067447-8a5a-4d6b-86a3-4b9cc4b4369b",
   "metadata": {},
   "source": [
    "## Inherting from Parsing Base Classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9713f547-b2e4-48eb-807f-a0f6f6d0e7e0",
   "metadata": {},
   "source": [
    "Another approach to implement a parser is by inherting from `BaseOutputParser`, `BaseGenerationOutputParser` or another one of the base parsers depending on what you need to do.\n",
    "\n",
    "In general, we **do not** recommend this approach for most use cases as it results in more code to write without significant benefits.\n",
    "\n",
    "The simplest kind of output parser extends the `BaseOutputParser` class and must implement the following methods:\n",
    "\n",
    "* `parse`: takes the string output from the model and parses it\n",
    "* (optional) `_type`: identifies the name of the parser.\n",
    "\n",
    "When the output from the chat model or LLM is malformed, the can throw an `OutputParserException` to indicate that parsing fails because of bad input. Using this exception allows code that utilizes the parser to handle the exceptions in a consistent manner.\n",
    "\n",
    ":::{.callout-tip} Parsers are Runnables! 🏃\n",
    "\n",
    "Because `BaseOutputParser` implements the `Runnable` interface, any custom parser you will create this way will become valid LangChain Runnables and will benefit from automatic async support, batch interface, logging support etc.\n",
    ":::\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e0f9c59-b5bd-4ed0-a187-ae514c203e80",
   "metadata": {},
   "source": [
    "### Simple Parser"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a96a846-1296-4d92-8e76-e29e583dee22",
   "metadata": {},
   "source": [
    "Here's a simple parser that can parse a **string** representation of a booealn (e.g., `YES` or `NO`) and convert it into the corresponding `boolean` type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "733a0c4f-471a-4161-ad3e-804f63053e6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.exceptions import OutputParserException\n",
    "from langchain_core.output_parsers import BaseOutputParser\n",
    "\n",
    "\n",
    "# The [bool] desribes a parameterization of a generic.\n",
    "# It's basically indicating what the return type of parse is\n",
    "# in this case the return type is either True or False\n",
    "class BooleanOutputParser(BaseOutputParser[bool]):\n",
    "    \"\"\"Custom boolean parser.\"\"\"\n",
    "\n",
    "    true_val: str = \"YES\"\n",
    "    false_val: str = \"NO\"\n",
    "\n",
    "    def parse(self, text: str) -> bool:\n",
    "        cleaned_text = text.strip().upper()\n",
    "        if cleaned_text not in (self.true_val.upper(), self.false_val.upper()):\n",
    "            raise OutputParserException(\n",
    "                f\"BooleanOutputParser expected output value to either be \"\n",
    "                f\"{self.true_val} or {self.false_val} (case-insensitive). \"\n",
    "                f\"Received {cleaned_text}.\"\n",
    "            )\n",
    "        return cleaned_text == self.true_val.upper()\n",
    "\n",
    "    @property\n",
    "    def _type(self) -> str:\n",
    "        return \"boolean_output_parser\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "101e54f0-12f1-4734-a80d-98e6f62644b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser = BooleanOutputParser()\n",
    "parser.invoke(\"YES\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "39ed9d84-16a1-4612-a1f7-13269b9f48e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Triggered an exception of type: <class 'langchain_core.exceptions.OutputParserException'>\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    parser.invoke(\"MEOW\")\n",
    "except Exception as e:\n",
    "    print(f\"Triggered an exception of type: {type(e)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c27da11a-2c64-4108-9a8a-38008d6041fc",
   "metadata": {},
   "source": [
    "Let's test changing the parameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2e94c0f4-f6c1-401b-8cee-2572a80846cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser = BooleanOutputParser(true_val=\"OKAY\")\n",
    "parser.invoke(\"OKAY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dac313d5-20c8-44a9-bfe9-c2b5020172e2",
   "metadata": {},
   "source": [
    "Let's confirm that other LCEL methods are present"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "97fb540f-83b2-46fd-a741-b200235f8f9e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, False]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parser.batch([\"OKAY\", \"NO\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "60cbdb2f-5538-4e74-ba03-53bc1bc4bb2f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[True, False]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await parser.abatch([\"OKAY\", \"NO\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6520dff0-259c-48e4-be69-829fb3275ac2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='OKAY')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_anthropic.chat_models import ChatAnthropic\n",
    "\n",
    "anthropic = ChatAnthropic(model_name=\"claude-2.1\")\n",
    "anthropic.invoke(\"say OKAY or NO\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12dc079e-c451-496c-953c-cba55ef26de8",
   "metadata": {},
   "source": [
    "Let's test that our parser works!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bb177c14-b1f5-474f-a1c8-5b32ae242259",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = anthropic | parser\n",
    "chain.invoke(\"say OKAY or NO\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18f83192-37e8-43f5-ab29-9568b1279f1b",
   "metadata": {},
   "source": [
    ":::{.callout-note}\n",
    "The parser will work with either the output from an LLM (a string) or the output from a chat model (an `AIMessage`)!\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ed063d3-3159-4f5b-8362-710956fc50bd",
   "metadata": {},
   "source": [
    "### Parsing Raw Model Outputs\n",
    "\n",
    "Sometimes there is additional metadata on the model output that is important besides the raw text. One example of this is tool calling, where arguments intended to be passed to called functions are returned in a separate property. If you need this finer-grained control, you can instead subclass the `BaseGenerationOutputParser` class. \n",
    "\n",
    "This class requires a single method `parse_result`. This method takes raw model output (e.g., list of `Generation` or `ChatGeneration`) and returns the parsed output.\n",
    "\n",
    "Supporting both `Generation` and `ChatGeneration` allows the parser to work with both regular LLMs as well as with Chat Models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0fd1f936-e77d-4602-921c-52a37e589e90",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_core.exceptions import OutputParserException\n",
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.output_parsers import BaseGenerationOutputParser\n",
    "from langchain_core.outputs import ChatGeneration, Generation\n",
    "\n",
    "\n",
    "class StrInvertCase(BaseGenerationOutputParser[str]):\n",
    "    \"\"\"An example parser that inverts the case of the characters in the message.\n",
    "\n",
    "    This is an example parse shown just for demonstration purposes and to keep\n",
    "    the example as simple as possible.\n",
    "    \"\"\"\n",
    "\n",
    "    def parse_result(self, result: List[Generation], *, partial: bool = False) -> str:\n",
    "        \"\"\"Parse a list of model Generations into a specific format.\n",
    "\n",
    "        Args:\n",
    "            result: A list of Generations to be parsed. The Generations are assumed\n",
    "                to be different candidate outputs for a single model input.\n",
    "                Many parsers assume that only a single generation is passed it in.\n",
    "                We will assert for that\n",
    "            partial: Whether to allow partial results. This is used for parsers\n",
    "                     that support streaming\n",
    "        \"\"\"\n",
    "        if len(result) != 1:\n",
    "            raise NotImplementedError(\n",
    "                \"This output parser can only be used with a single generation.\"\n",
    "            )\n",
    "        generation = result[0]\n",
    "        if not isinstance(generation, ChatGeneration):\n",
    "            # Say that this one only works with chat generations\n",
    "            raise OutputParserException(\n",
    "                \"This output parser can only be used with a chat generation.\"\n",
    "            )\n",
    "        return generation.message.content.swapcase()\n",
    "\n",
    "\n",
    "chain = anthropic | StrInvertCase()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accab8a3-6b0e-4ad0-89e6-1824ca20c726",
   "metadata": {},
   "source": [
    "Let's the new parser! It should be inverting the output from the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "568fae19-b09c-484f-8775-1c9a60aabdf4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'hELLO! mY NAME IS cLAUDE.'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke(\"Tell me a short sentence about yourself\")"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}