function_calling_lc.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyO3aRPYbLiyOmZaM0X7JEEs",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/satvik314/ai_experiments/blob/main/function_calling_lc.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3O3biiV4JbAH",
        "outputId": "4dd08db7-f145-45b5-af65-70b78771426e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m24.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m73.6/73.6 kB\u001b[0m \u001b[31m7.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m50.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m90.0/90.0 kB\u001b[0m \u001b[31m11.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m114.5/114.5 kB\u001b[0m \u001b[31m13.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m26.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m149.6/149.6 kB\u001b[0m \u001b[31m11.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.1/49.1 kB\u001b[0m \u001b[31m5.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip install -qU langchain openai"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.chat_models import ChatOpenAI\n",
        "from langchain.chains import create_tagging_chain, create_tagging_chain_pydantic\n",
        "from langchain.prompts import ChatPromptTemplate"
      ],
      "metadata": {
        "id": "W3NhThEsJgAl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "from getpass import getpass\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = getpass()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2QWndidtKMNt",
        "outputId": "349a67fa-e9e4-4e52-b8f3-05be563549e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "··········\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "llm = ChatOpenAI(temperature = 0, model = \"gpt-3.5-turbo-0613\")"
      ],
      "metadata": {
        "id": "TG9WzdPmKWAn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "schema = {\n",
        "    \"properties\" : {\n",
        "        \"sentiment\" : {\"type\" : \"string\"},\n",
        "        \"stars\" : {\"type\" : \"integer\"},\n",
        "        \"language\" : {\"type\" : \"string\"}\n",
        "    }\n",
        "}"
      ],
      "metadata": {
        "id": "oRCmbpBiKjSA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain = create_tagging_chain(schema, llm)"
      ],
      "metadata": {
        "id": "WOF1vZ7VK9ow"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain.prompt.messages[0].prompt"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XzQDoy6RLEH_",
        "outputId": "3da2e4be-1209-4a26-a91c-8268151be93b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "PromptTemplate(input_variables=['input'], output_parser=None, partial_variables={}, template='Extract the desired information from the following passage.\\n\\nPassage:\\n{input}\\n', template_format='f-string', validate_template=True)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(chain.prompt.messages[0].prompt.template)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Z6qgjZFPLJYh",
        "outputId": "0f5e055f-56e2-4f61-f764-8584af4a13d0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extract the desired information from the following passage.\n",
            "\n",
            "Passage:\n",
            "{input}\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain.llm_kwargs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D8fFuCInLTUp",
        "outputId": "b4c467a0-bd59-4e7e-9a09-aeef50e58608"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'functions': [{'name': 'information_extraction',\n",
              "   'description': 'Extracts the relevant information from the passage.',\n",
              "   'parameters': {'type': 'object',\n",
              "    'properties': {'sentiment': {'title': 'sentiment', 'type': 'string'},\n",
              "     'stars': {'title': 'stars', 'type': 'integer'},\n",
              "     'language': {'title': 'language', 'type': 'string'}},\n",
              "    'required': []}}],\n",
              " 'function_call': {'name': 'information_extraction'}}"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain.output_parser"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KQrDHPpzLwaU",
        "outputId": "0737af09-0111-495f-d885-d88c085817a8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "JsonOutputFunctionsParser()"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "review_01 = \"Starts off kind of slow, but builds up quickly to give the reader a good understanding of how things unfolded. I’m anticipating a second book, and can’t wait for the rest of the story to be unveiled!\"\n",
        "chain.run(review_01)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u-I3jLVqMHQQ",
        "outputId": "89ffd301-2560-4013-af68-8d97e9047cdc"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'sentiment': 'positive', 'stars': 4}"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "review_02 = \"\"\"It's the BIGGEST CON \\n\n",
        "This book has NOT been security checked for correct information , I feel Meghan and Harry are laughing at me for buying this book , And putting my money in there pockets. Harry can't come up with any dates, says he can't remember the dates of anything , Said he was At college on a Hot summers day when the phone call came about his Great Great Grandmother passing , He was not at school he was Skiing in Switzerland with his Brother ( I remember this) And the papers have come back with the proof. Harry says that his mother bought him an XBox for his 13th birthday in advance of her death that her sister brought to the school for him, again not true the XBox didn't come out till 4 years later . There is so many holes in the book . Every other page has something on it that didn't happen. He takes NO reasonability for ANYthing even when the evidence is staring him in his face. This book is a BIG con, don't waste your money.\"\"\"\n",
        "\n",
        "chain.run(review_02)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fVeSzjjkMLae",
        "outputId": "eb92f6fa-1331-4041-91db-07eefa16c0f4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'sentiment': 'negative'}"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "schema = {\n",
        "    \"properties\" : {\n",
        "        \"sentiment\" : {\"type\" : \"string\", \"enum\" : [\"positive\", \"neutral\", \"negative\"]},\n",
        "        \"stars\" : {\n",
        "            \"type\" : \"integer\",\n",
        "            \"enum\" : [1,2,3,4,5]\n",
        "        },\n",
        "        \"language\" : {\n",
        "            \"type\" : \"string\",\n",
        "            \"enum\" : [\"spanish\", \"english\", \"french\", \"german\", \"italian\"]\n",
        "        }\n",
        "    },\n",
        "    \"required\" : [\"language\", \"sentiment\", \"stars\"]\n",
        "}"
      ],
      "metadata": {
        "id": "PkxJ9_m1MXyD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain = create_tagging_chain(schema, llm)"
      ],
      "metadata": {
        "id": "NNWDozjCNtFa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain.run(review_01)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HxKEkKnNODf4",
        "outputId": "c1dc6627-6e23-4ef1-d864-e20b76ab5213"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'sentiment': 'positive', 'stars': 4, 'language': 'english'}"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain.run(review_02)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LJk-TvX4OFTM",
        "outputId": "18c622c8-1a39-4c87-997e-8913866dd223"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'sentiment': 'negative', 'stars': 1, 'language': 'english'}"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Using Pydantic"
      ],
      "metadata": {
        "id": "9iypyDEbOLag"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from enum import Enum\n",
        "from pydantic import BaseModel, Field"
      ],
      "metadata": {
        "id": "xvYC40wqOHpM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class Tags(BaseModel):\n",
        "  sentiment : str = Field(..., enum = [\"positive\", \"neutral\", \"negative\"])\n",
        "  stars : int = Field(\n",
        "      ...,\n",
        "      description = \"describes the numbers of stars given by a reviewer on Amazon\",\n",
        "      enum = [1,2,3,4,5]\n",
        "  )\n",
        "  language : str = Field(\n",
        "      ...,\n",
        "      enum = ['spanish', 'english', 'french', 'german', 'italian']\n",
        "  )\n"
      ],
      "metadata": {
        "id": "cXroSw61PSDZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain = create_tagging_chain_pydantic(Tags, llm)"
      ],
      "metadata": {
        "id": "F-gu9rAtPfLG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res = chain.run(review_01)"
      ],
      "metadata": {
        "id": "ddyrSfH6Pjzm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(res)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tLg4BS3HQ0hz",
        "outputId": "048fc52a-2679-4018-9f79-025f0c025621"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "sentiment='positive' stars=4 language='english'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "res.sentiment"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "j3zWN2Y1Q2Ar",
        "outputId": "580836c2-fdcb-4c05-c8ec-75ef32a8f5c6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'positive'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Extraction, NER"
      ],
      "metadata": {
        "id": "ao-ZqeHPRDo2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.chat_models import ChatOpenAI\n",
        "from langchain.chains import create_extraction_chain, create_extraction_chain_pydantic\n",
        "from langchain.prompts import ChatPromptTemplate"
      ],
      "metadata": {
        "id": "A10B6M_SQ4Z_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "tc_01 = \"\"\"Reddit CEO Steve Huffman is not backing down amid protests against API changes made by the platform. In interviews with The Verge, NBC News and NPR, Huffman defended business decisions made by the company to charge third-party apps saying the API wasn’t designed to support these clients.\n",
        "\n",
        "The Reddit co-founder also talked about protesting moderators, changing site rules, and profitability in these interviews. The platform is facing one of the strongest backlashes from the community, but the CEO seemingly doesn’t want to budge.\n",
        "\n",
        "What’s happening at Reddit?\n",
        "In April, Reddit announced that it is going to charge for its API, but didn’t announce any pricing. Earlier this month, Christian Selig, the developer of a popular Reddit client for iOS called Apollo, posted that he had a call with Reddit. API pricing quoted by them could cost him $20 million a year to run the app. Selig later said that, because the social network is not ready to make any changes to the pricing structure, he is forced to shut down Apollo. Other third-party developers of clients like Reddit is Fun and Relay for Reddit also said that they will shut down their apps on June 30.\n",
        "\n",
        "The only exception Reddit made was to allow free access to the API to non-commercial apps providing accessibility features. The company has made deals with apps like RedReader, Dystopia, and Luna and given them exemptions from its “large-scale pricing terms.”\n",
        "\n",
        "Thousands of subreddits went dark starting June 12 to protest those changes — it caused a brief outage as well. Meanwhile, Huffman took a strong stance in his AMA and took a dig at Apollo and Selig. As moderators didn’t see anything changing, many subreddits decided to extend the blackout.\"\"\""
      ],
      "metadata": {
        "id": "vpfBMwceRg85"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "llm = ChatOpenAI(temperature = 0, model = \"gpt-3.5-turbo-0613\")"
      ],
      "metadata": {
        "id": "MRFJyvYWTAan"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "schema = {\n",
        "    \"properties\" : {\n",
        "        \"person_name\" : {\"type\" : \"string\"},\n",
        "        \"startup\" : {\"type\" : \"string\"},\n",
        "        \"news_outlet\" : {\"type\" : \"string\"},\n",
        "        \"app_name\" : {\"type\" : \"string\"},\n",
        "        \"month\" : {\"type\" : \"string\"}\n",
        "    },\n",
        "    \"required\" : [\"person_name\", \"startup\"]\n",
        "}"
      ],
      "metadata": {
        "id": "95F0EfBRTLDd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain = create_extraction_chain(schema, llm)"
      ],
      "metadata": {
        "id": "Dv8KeuCZT8Bo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain.prompt.messages[0]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZC6YRwk6UAuW",
        "outputId": "7392d2d3-fc3a-4000-dc3a-1fff05fe3184"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], output_parser=None, partial_variables={}, template='Extract and save the relevant entities mentioned in the following passage together with their properties.\\n\\nPassage:\\n{input}\\n', template_format='f-string', validate_template=True), additional_kwargs={})"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(chain.prompt.messages[0].prompt.template)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iYWtqqMdUEGw",
        "outputId": "50e762c3-4240-4c64-a2c8-50e63786393c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extract and save the relevant entities mentioned in the following passage together with their properties.\n",
            "\n",
            "Passage:\n",
            "{input}\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain.llm_kwargs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QozA4KNeUKua",
        "outputId": "94ca1d15-364e-41d8-fc62-4bc32d6e0776"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'functions': [{'name': 'information_extraction',\n",
              "   'description': 'Extracts the relevant information from the passage.',\n",
              "   'parameters': {'type': 'object',\n",
              "    'properties': {'info': {'type': 'array',\n",
              "      'items': {'type': 'object',\n",
              "       'properties': {'person_name': {'title': 'person_name',\n",
              "         'type': 'string'},\n",
              "        'startup': {'title': 'startup', 'type': 'string'},\n",
              "        'news_outlet': {'title': 'news_outlet', 'type': 'string'},\n",
              "        'app_name': {'title': 'app_name', 'type': 'string'},\n",
              "        'month': {'title': 'month', 'type': 'string'}},\n",
              "       'required': ['person_name', 'startup']}}},\n",
              "    'required': ['info']}}],\n",
              " 'function_call': {'name': 'information_extraction'}}"
            ]
          },
          "metadata": {},
          "execution_count": 38
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "chain.output_parser"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CYkdMpMkURmD",
        "outputId": "c41b12dd-6523-486b-d16e-f7284a87798e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "JsonKeyOutputFunctionsParser(key_name='info')"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "response = chain.run(tc_01)\n",
        "response"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nKkvcA0sUU9S",
        "outputId": "6de9c44b-22b9-479f-83c4-b8acce78e6d5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'person_name': 'Steve Huffman',\n",
              "  'startup': 'Reddit',\n",
              "  'news_outlet': 'The Verge',\n",
              "  'app_name': 'Apollo',\n",
              "  'month': 'April'},\n",
              " {'person_name': 'Christian Selig',\n",
              "  'startup': 'Apollo',\n",
              "  'news_outlet': 'N/A',\n",
              "  'app_name': 'Apollo',\n",
              "  'month': 'April'},\n",
              " {'person_name': 'Christian Selig',\n",
              "  'startup': 'Apollo',\n",
              "  'news_outlet': 'N/A',\n",
              "  'app_name': 'Apollo',\n",
              "  'month': 'June'},\n",
              " {'person_name': 'Reddit',\n",
              "  'startup': 'Reddit',\n",
              "  'news_outlet': 'N/A',\n",
              "  'app_name': 'Reddit is Fun',\n",
              "  'month': 'June'},\n",
              " {'person_name': 'Reddit',\n",
              "  'startup': 'Reddit',\n",
              "  'news_outlet': 'N/A',\n",
              "  'app_name': 'Relay for Reddit',\n",
              "  'month': 'June'}]"
            ]
          },
          "metadata": {},
          "execution_count": 40
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def sort_objects(obj_list):\n",
        "    # Initialize empty lists for each category\n",
        "    people = []\n",
        "    startups = []\n",
        "    news_outlets = []\n",
        "    apps = []\n",
        "    months = []\n",
        "\n",
        "    # Loop through each dictionary in the list\n",
        "    for obj in obj_list:\n",
        "        # Add each value to the appropriate list, if the key exists in the dictionary and the value is not already in the list\n",
        "        if 'person_name' in obj and obj['person_name'] not in people:\n",
        "            people.append(obj['person_name'])\n",
        "        if 'startup' in obj and obj['startup'] not in startups:\n",
        "            startups.append(obj['startup'])\n",
        "        if 'news_outlet' in obj and obj['news_outlet'] not in news_outlets:\n",
        "            news_outlets.append(obj['news_outlet'])\n",
        "        if 'app_name' in obj and obj['app_name'] not in apps:\n",
        "            apps.append(obj['app_name'])\n",
        "        if 'month' in obj and obj['month'] not in months:\n",
        "            months.append(obj['month'])\n",
        "\n",
        "    # Return the lists as a dictionary\n",
        "    return {'people': people, 'startups': startups, 'news_outlets': news_outlets, 'apps': apps, 'months': months}"
      ],
      "metadata": {
        "id": "tErcKJDbUaWA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sort_objects(response)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QBsFFx2ZUrzY",
        "outputId": "903ce3e6-a97a-4fba-8bc7-d2165ba56f2b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'people': ['Steve Huffman', 'Christian Selig', 'Reddit'],\n",
              " 'startups': ['Reddit', 'Apollo'],\n",
              " 'news_outlets': ['The Verge', 'N/A'],\n",
              " 'apps': ['Apollo', 'Reddit is Fun', 'Relay for Reddit'],\n",
              " 'months': ['April', 'June']}"
            ]
          },
          "metadata": {},
          "execution_count": 42
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pydantic import BaseModel\n",
        "\n",
        "class NER(BaseModel):\n",
        "  person_name : str\n",
        "  startup : str\n",
        "  news_outlet : str\n",
        "  app_name : str\n",
        "  month : str"
      ],
      "metadata": {
        "id": "c8G8aw7WUx20"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chain = create_extraction_chain_pydantic(NER, llm)"
      ],
      "metadata": {
        "id": "sI4qx-2MVv2M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "res = chain.run(tc_01)"
      ],
      "metadata": {
        "id": "hlyOTrZ8WtTV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(res)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U5ZGIuj0XEnv",
        "outputId": "22b3401f-8c53-4cbc-c45c-1a6dd2e73b29"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[NER(person_name='Steve Huffman', startup='Reddit', news_outlet='The Verge', app_name='Apollo', month='April'), NER(person_name='Christian Selig', startup='Apollo', news_outlet='N/A', app_name='Apollo', month='April'), NER(person_name='Christian Selig', startup='Apollo', news_outlet='N/A', app_name='Apollo', month='June'), NER(person_name='Reddit', startup='Reddit', news_outlet='N/A', app_name='Reddit is Fun', month='June'), NER(person_name='Reddit', startup='Reddit', news_outlet='N/A', app_name='Relay for Reddit', month='June')]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "res[0].person_name"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "_aVAhlnCXIU0",
        "outputId": "820f1f01-8a20-49f0-953d-21ef268c1fac"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'Steve Huffman'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 48
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "GCGlox6lXMi8"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}