tutorials/28_Structured_Output_With_Loop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AVBtOVlNJ51C"
   },
   "source": [
    "# Tutorial: Generating Structured Output with Loop-Based Auto-Correction\n",
    "\n",
    "- **Level**: Intermediate\n",
    "- **Time to complete**: 15 minutes\n",
    "- **Prerequisites**: You must have an API key from an active OpenAI account as this tutorial is using the gpt-4o-mini model by OpenAI.\n",
    "- **Components Used**: `PromptBuilder`, `OpenAIChatGenerator`, `OutputValidator` (Custom component)\n",
    "- **Goal**: After completing this tutorial, you will have built a system that extracts unstructured data, puts it in a JSON schema, and automatically corrects errors in the JSON output from a large language model (LLM) to make sure it follows the specified structure.\n",
    "\n",
    "> This tutorial uses the latest version of Haystack 2.x (`haystack-ai`). For more information on Haystack 2.0, read the [Haystack 2.0 announcement](https://haystack.deepset.ai/blog/haystack-2-release) or visit the [Haystack Documentation](https://docs.haystack.deepset.ai/docs/intro).\n",
    "\n",
    "## Overview\n",
    "This tutorial demonstrates how to use Haystack 2.0's advanced [looping pipelines](https://docs.haystack.deepset.ai/docs/pipelines#loops) with LLMs for more dynamic and flexible data processing. You'll learn how to extract structured data from unstructured data using an LLM, and to validate the generated output against a predefined schema.\n",
    "\n",
    "This tutorial uses `gpt-4o-mini` to change unstructured passages into JSON outputs that follow the [Pydantic](https://github.com/pydantic/pydantic) schema. It uses a custom OutputValidator component to validate the JSON and loop back to make corrections, if necessary."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jmiAHh1oGsKI"
   },
   "source": [
    "## Preparing the Colab Environment\n",
    "\n",
    "Enable the debug mode of logging:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "Vor9IHuNRvEh"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logging.basicConfig()\n",
    "logging.getLogger(\"canals.pipeline.pipeline\").setLevel(logging.DEBUG)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ljbWiyJkKiPw"
   },
   "source": [
    "## Installing Dependencies\n",
    "Install Haystack and [colorama](https://pypi.org/project/colorama/) with pip:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kcc1AlLQd_jI",
    "outputId": "efc4bbab-a9fe-46ee-d8af-9d86edacaf04"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "pip install haystack-ai\n",
    "pip install colorama"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nTA5fdvCLMKD"
   },
   "source": [
    "### Enabling Telemetry\n",
    "\n",
    "Enable telemetry to let us know you're using this tutorial. (You can always opt out by commenting out this line). For details, see [Telemetry](https://docs.haystack.deepset.ai/docs/telemetry)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "Apay3QSQLKdM"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/amna.mubashar/Library/Python/3.9/lib/python/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from haystack.telemetry import tutorial_running\n",
    "\n",
    "tutorial_running(28)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cmjfa8CiCeFl"
   },
   "source": [
    "## Defining a Schema to Parse the JSON Object\n",
    "\n",
    "Define a simple JSON schema for the data you want to extract from a text passsage using the LLM. As the first step, define two [Pydantic models](https://docs.pydantic.dev/1.10/usage/models/), `City` and `CitiesData`, with suitable fields and types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "xwKrDOOGdaAz"
   },
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "class City(BaseModel):\n",
    "    name: str\n",
    "    country: str\n",
    "    population: int\n",
    "\n",
    "\n",
    "class CitiesData(BaseModel):\n",
    "    cities: List[City]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zv-6-l_PCeFl"
   },
   "source": [
    "> You can change these models according to the format you wish to extract from the text."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ouk1mAOUCeFl"
   },
   "source": [
    "Then, generate a JSON schema from Pydantic models using `schema_json()`. You will later on use this schema in the prompt to instruct the LLM.\n",
    "\n",
    "To learn more about the JSON schemas, visit [Pydantic Schema](https://docs.pydantic.dev/1.10/usage/schema/).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "8Lg9_72jCeFl"
   },
   "outputs": [],
   "source": [
    "json_schema = CitiesData.schema_json(indent=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KvNhg0bP7kfg"
   },
   "source": [
    "## Creating a Custom Component: OutputValidator\n",
    "\n",
    "`OutputValidator` is a custom component that validates if the JSON object the LLM generates complies with the provided [Pydantic model](https://docs.pydantic.dev/1.10/usage/models/). If it doesn't, OutputValidator returns an error message along with the incorrect JSON object to get it fixed in the next loop.\n",
    "\n",
    "For more details about custom components, see [Creating Custom Components](https://docs.haystack.deepset.ai/docs/custom-components)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yr6D8RN2d7Vy"
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "import pydantic\n",
    "from pydantic import ValidationError\n",
    "from typing import Optional, List\n",
    "from colorama import Fore\n",
    "from haystack import component\n",
    "from haystack.dataclasses import ChatMessage\n",
    "\n",
    "\n",
    "# Define the component input parameters\n",
    "@component\n",
    "class OutputValidator:\n",
    "    def __init__(self, pydantic_model: pydantic.BaseModel):\n",
    "        self.pydantic_model = pydantic_model\n",
    "        self.iteration_counter = 0\n",
    "\n",
    "    # Define the component output\n",
    "    @component.output_types(valid_replies=List[str], invalid_replies=Optional[List[str]], error_message=Optional[str])\n",
    "    def run(self, replies: List[ChatMessage]):\n",
    "\n",
    "        self.iteration_counter += 1\n",
    "\n",
    "        ## Try to parse the LLM's reply ##\n",
    "        # If the LLM's reply is a valid object, return `\"valid_replies\"`\n",
    "        try:\n",
    "            output_dict = json.loads(replies[0].text)\n",
    "            self.pydantic_model.parse_obj(output_dict)\n",
    "            print(\n",
    "                Fore.GREEN\n",
    "                + f\"OutputValidator at Iteration {self.iteration_counter}: Valid JSON from LLM - No need for looping: {replies[0]}\"\n",
    "            )\n",
    "            return {\"valid_replies\": replies}\n",
    "\n",
    "        # If the LLM's reply is corrupted or not valid, return \"invalid_replies\" and the \"error_message\" for LLM to try again\n",
    "        except (ValueError, ValidationError) as e:\n",
    "            print(\n",
    "                Fore.RED\n",
    "                + f\"OutputValidator at Iteration {self.iteration_counter}: Invalid JSON from LLM - Let's try again.\\n\"\n",
    "                f\"Output from LLM:\\n {replies[0]} \\n\"\n",
    "                f\"Error from OutputValidator: {e}\"\n",
    "            )\n",
    "            return {\"invalid_replies\": replies, \"error_message\": str(e)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vQ_TfSBkCeFm"
   },
   "source": [
    "Then, create an OutputValidator instance with `CitiesData` that you have created before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "bhPCLCBCCeFm"
   },
   "outputs": [],
   "source": [
    "output_validator = OutputValidator(pydantic_model=CitiesData)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xcIWKjW4k42r"
   },
   "source": [
    "## Creating the Prompt\n",
    "\n",
    "Write instructions for the LLM for converting a passage into a JSON format. Ensure the instructions explain how to identify and correct errors if the JSON doesn't match the required schema. Once you create the prompt, initialize PromptBuilder to use it.  \n",
    "\n",
    "For information about Jinja2 template and ChatPromptBuilder, see [ChatPromptBuilder](https://docs.haystack.deepset.ai/docs/chatpromptbuilder)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "ohPpNALjdVKt"
   },
   "outputs": [],
   "source": [
    "from haystack.components.builders import ChatPromptBuilder\n",
    "\n",
    "\n",
    "prompt_template = [\n",
    "    ChatMessage.from_user(\n",
    "        \"\"\"\n",
    "Create a JSON object from the information present in this passage: {{passage}}.\n",
    "Only use information that is present in the passage. Follow this JSON schema, but only return the actual instances without any additional schema definition:\n",
    "{{schema}}\n",
    "Make sure your response is a dict and not a list.\n",
    "{% if invalid_replies and error_message %}\n",
    "  You already created the following output in a previous attempt: {{invalid_replies}}\n",
    "  However, this doesn't comply with the format requirements from above and triggered this Python exception: {{error_message}}\n",
    "  Correct the output and try again. Just return the corrected output without any extra explanations.\n",
    "{% endif %}\n",
    "\"\"\"\n",
    "    )\n",
    "]\n",
    "prompt_builder = ChatPromptBuilder(template=prompt_template)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KM9-Zq2FL7Nn"
   },
   "source": [
    "## Initalizing the ChatGenerator\n",
    "\n",
    "[OpenAIChatGenerator](https://docs.haystack.deepset.ai/docs/openaichatgenerator) generates\n",
    "text using OpenAI's `gpt-4o-mini` model by default. Set the `OPENAI_API_KEY` variable and provide a model name to the ChatGenerator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "Z4cQteIgunUR"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "from haystack.components.generators.chat import OpenAIChatGenerator\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter OpenAI API key:\")\n",
    "chat_generator = OpenAIChatGenerator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zbotIOgXHkC5"
   },
   "source": [
    "## Building the Pipeline\n",
    "\n",
    "Add all components to your pipeline and connect them. Add connections from `output_validator` back to the `prompt_builder` for cases where the produced JSON doesn't comply with the JSON schema. Set `max_runs_per_component` to avoid infinite looping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "eFglN9YEv-1W"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<haystack.core.pipeline.pipeline.Pipeline object at 0x304c658e0>\n",
       "🚅 Components\n",
       "  - prompt_builder: ChatPromptBuilder\n",
       "  - llm: OpenAIChatGenerator\n",
       "  - output_validator: OutputValidator\n",
       "🛤️ Connections\n",
       "  - prompt_builder.prompt -> llm.messages (List[ChatMessage])\n",
       "  - llm.replies -> output_validator.replies (List[ChatMessage])\n",
       "  - output_validator.invalid_replies -> prompt_builder.invalid_replies (Optional[List[str]])\n",
       "  - output_validator.error_message -> prompt_builder.error_message (Optional[str])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from haystack import Pipeline\n",
    "\n",
    "pipeline = Pipeline(max_runs_per_component=5)\n",
    "\n",
    "# Add components to your pipeline\n",
    "pipeline.add_component(instance=prompt_builder, name=\"prompt_builder\")\n",
    "pipeline.add_component(instance=chat_generator, name=\"llm\")\n",
    "pipeline.add_component(instance=output_validator, name=\"output_validator\")\n",
    "\n",
    "# Now, connect the components to each other\n",
    "pipeline.connect(\"prompt_builder.prompt\", \"llm.messages\")\n",
    "pipeline.connect(\"llm.replies\", \"output_validator\")\n",
    "# If a component has more than one output or input, explicitly specify the connections:\n",
    "pipeline.connect(\"output_validator.invalid_replies\", \"prompt_builder.invalid_replies\")\n",
    "pipeline.connect(\"output_validator.error_message\", \"prompt_builder.error_message\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-UKW5wtIIT7w"
   },
   "source": [
    "### Visualize the Pipeline\n",
    "\n",
    "Draw the pipeline with the [`draw()`](https://docs.haystack.deepset.ai/docs/drawing-pipeline-graphs) method to confirm the connections are correct. You can find the diagram in the Files section of this Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "RZJg6YHId300"
   },
   "outputs": [],
   "source": [
    "pipeline.draw(\"auto-correct-pipeline.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kV_kexTjImpo"
   },
   "source": [
    "## Testing the Pipeline\n",
    "\n",
    "Run the pipeline with an example passage that you want to convert into a JSON format and the `json_schema` you have created for `CitiesData`. For the given example passage, the generated JSON object should be like:\n",
    "```json\n",
    "{\n",
    "  \"cities\": [\n",
    "    {\n",
    "      \"name\": \"Berlin\",\n",
    "      \"country\": \"Germany\",\n",
    "      \"population\": 3850809\n",
    "    },\n",
    "    {\n",
    "      \"name\": \"Paris\",\n",
    "      \"country\": \"France\",\n",
    "      \"population\": 2161000\n",
    "    },\n",
    "    {\n",
    "      \"name\": \"Lisbon\",\n",
    "      \"country\": \"Portugal\",\n",
    "      \"population\": 504718\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```\n",
    "The output of the LLM should be compliant with the `json_schema`. If the LLM doesn't generate the correct JSON object, it will loop back and try again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yIoMedb6eKia",
    "outputId": "4a9ef924-cf26-4908-d83f-b0bc0dc03b54"
   },
   "outputs": [],
   "source": [
    "passage = \"Berlin is the capital of Germany. It has a population of 3,850,809. Paris, France's capital, has 2.161 million residents. Lisbon is the capital and the largest city of Portugal with the population of 504,718.\"\n",
    "result = pipeline.run({\"prompt_builder\": {\"passage\": passage, \"schema\": json_schema}})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WWxmPgADS_Fa"
   },
   "source": [
    "> If you encounter `PipelineMaxLoops: Maximum loops count (5) exceeded for component 'prompt_builder'.` error, consider increasing the maximum loop count or simply rerun the pipeline."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eWPawSjgSJAM"
   },
   "source": [
    "### Print the Correct JSON\n",
    "If you didn't get any error, you can now print the corrected JSON."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BVO47gXQQnDC",
    "outputId": "460a10d4-a69a-49cd-bbb2-fc4980907299"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'cities': [{'name': 'Berlin', 'country': 'Germany', 'population': 3850809}, {'name': 'Paris', 'country': 'France', 'population': 2161000}, {'name': 'Lisbon', 'country': 'Portugal', 'population': 504718}]}\n"
     ]
    }
   ],
   "source": [
    "valid_reply = result[\"output_validator\"][\"valid_replies\"][0].text\n",
    "valid_json = json.loads(valid_reply)\n",
    "print(valid_json)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Egz_4h2vI_QL"
   },
   "source": [
    "## What's next\n",
    "\n",
    "🎉 Congratulations! You've built a system that generates structured JSON out of unstructured text passages, and auto-corrects it by using the looping functionality of Haystack pipelines.\n",
    "\n",
    "To stay up to date on the latest Haystack developments, you can [subscribe to our newsletter](https://landing.deepset.ai/haystack-community-updates) and [join Haystack discord community](https://discord.gg/haystack).\n",
    "\n",
    "Thanks for reading!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}