rag-query-planning.ipynb

rag-query-planning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyM9AyH/ZUQ6JgjWb8dVNAVJ",
      "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/gist/virattt/692d31c71ae871b2d3fc21557ff9dd68/rag-query-planning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 0.  Installation"
      ],
      "metadata": {
        "id": "gVVb7vgPjWgZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PXjMEX2qgzii"
      },
      "outputs": [],
      "source": [
        "!pip install openai"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 1. Initialize OpenAI client"
      ],
      "metadata": {
        "id": "d6IcA-zzjh73"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from openai import OpenAI\n",
        "client = OpenAI(api_key=\"YOUR_OPENAI_KEY\")"
      ],
      "metadata": {
        "id": "ZEc4wAjEg2l_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Step 2. Call OpenAI\n",
        "We instruct GPT-4 to break our query down into smaller sub-queries."
      ],
      "metadata": {
        "id": "6pe-VoIYjmUQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# User query\n",
        "query = \"What is Costco's competitive advantage and what is Walmart's competitive advantage?\"\n",
        "\n",
        "# Call OpenAI\n",
        "response = client.chat.completions.create(\n",
        "  model=\"gpt-4-1106-preview\",\n",
        "  response_format={ \"type\": \"json_object\" },\n",
        "  messages=[\n",
        "    {\"role\": \"system\", \"content\": \"You are a helpful assistant that decomposes a user query into 2 simpler sub-queries per unique entity in the query in a JSON for retrieval-assisted generation.  The JSON response is a list of `queries`, where each query has a `sub_query` field.\"},\n",
        "    {\"role\": \"user\", \"content\": query}\n",
        "  ]\n",
        ")"
      ],
      "metadata": {
        "id": "aOR8N9CHjbSX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Print sub-queries\n",
        "print(response.choices[0].message.content)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WjuVzCQ0htA9",
        "outputId": "86c70c58-d897-43c7-c9a3-2a03e6f739dd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{\n",
            "  \"queries\": [\n",
            "    {\n",
            "      \"sub_query\": \"What is Costco's competitive advantage?\"\n",
            "    },\n",
            "    {\n",
            "      \"sub_query\": \"What are the key business strategies of Costco?\"\n",
            "    },\n",
            "    {\n",
            "      \"sub_query\": \"What is Walmart's competitive advantage?\"\n",
            "    },\n",
            "    {\n",
            "      \"sub_query\": \"What are the key business strategies of Walmart?\"\n",
            "    }\n",
            "  ]\n",
            "}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "EGb8kxoDiAH1"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}