07_planning_tasks.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Planning and Executing a Query Plan\n",
                "\n",
                "This example demonstrates how to use the OpenAI Function Call ChatCompletion model to plan and execute a query plan in a question-answering system. By breaking down a complex question into smaller sub-questions with defined dependencies, the system can systematically gather the necessary information to answer the main question.\n",
                "\n",
                "## Motivation\n",
                "\n",
                "The goal of this example is to showcase how query planning can be used to handle complex questions, facilitate iterative information gathering, automate workflows, and optimize processes. By leveraging the OpenAI Function Call model, you can design and execute a structured plan to find answers effectively.\n",
                "\n",
                "**Use Cases:**\n",
                "- Complex question answering\n",
                "- Iterative information gathering\n",
                "- Workflow automation\n",
                "- Process optimization\n",
                "\n",
                "With the OpenAI Function Call model, you can customize the planning process and integrate it into your specific application to meet your unique requirements."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Imports"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [],
            "source": [
                "import enum\n",
                "from typing import List\n",
                "from pydantic import Field, BaseModel\n",
                "import instructor\n",
                "from openai import OpenAI"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {},
            "outputs": [],
            "source": [
                "from rich import pretty, print\n",
                "pretty.install()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/html": [
                            "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
                        ],
                        "text/plain": []
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "data": {
                        "text/plain": [
                            "\u001b[3;92mTrue\u001b[0m"
                        ]
                    },
                    "execution_count": 9,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "from dotenv import load_dotenv\n",
                "load_dotenv(\"../api_keys.env\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Defining the Structures\n",
                "\n",
                "Let's define the necessary Pydantic models to represent the query plan and the queries."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "metadata": {},
            "outputs": [],
            "source": [
                "class QueryType(str, enum.Enum):\n",
                "    \"\"\"Enumeration representing the types of queries that can be asked to a question answer system.\"\"\"\n",
                "\n",
                "    SINGLE_QUESTION = \"SINGLE\"\n",
                "    MERGE_MULTIPLE_RESPONSES = \"MERGE_MULTIPLE_RESPONSES\"\n",
                "\n",
                "\n",
                "class Query(BaseModel):\n",
                "    \"\"\"Class representing a single question in a query plan.\"\"\"\n",
                "\n",
                "    id: int = Field(..., description=\"Unique id of the query\")\n",
                "    question: str = Field(\n",
                "        ...,\n",
                "        description=\"Question asked using a question answering system\",\n",
                "    )\n",
                "    dependencies: List[int] = Field(\n",
                "        default_factory=list,\n",
                "        description=\"List of sub questions that need to be answered before asking this question\",\n",
                "    )\n",
                "    node_type: QueryType = Field(\n",
                "        default=QueryType.SINGLE_QUESTION,\n",
                "        description=\"Type of question, either a single question or a multi-question merge\",\n",
                "    )\n",
                "\n",
                "\n",
                "class QueryPlan(BaseModel):\n",
                "    \"\"\"Container class representing a tree of questions to ask a question answering system.\"\"\"\n",
                "\n",
                "    query_graph: List[Query] = Field(\n",
                "        ..., description=\"The query graph representing the plan\"\n",
                "    )\n",
                "\n",
                "    def _dependencies(self, ids: List[int]) -> List[Query]:\n",
                "        \"\"\"Returns the dependencies of a query given their ids.\"\"\"\n",
                "        return [q for q in self.query_graph if q.id in ids]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Client Initialization"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Apply the patch to the OpenAI client\n",
                "# enables response_model keyword\n",
                "client = instructor.from_openai(OpenAI())"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Planning a Query Plan\n",
                "\n",
                "Now, let's demonstrate how to plan and execute a query plan using the defined models and the OpenAI API."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/html": [
                            "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">{</span>\n",
                            "    <span style=\"color: #008000; text-decoration-color: #008000\">'query_graph'</span>: <span style=\"font-weight: bold\">[</span>\n",
                            "        <span style=\"font-weight: bold\">{</span>\n",
                            "            <span style=\"color: #008000; text-decoration-color: #008000\">'id'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>,\n",
                            "            <span style=\"color: #008000; text-decoration-color: #008000\">'question'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'What is the population of Canada?'</span>,\n",
                            "            <span style=\"color: #008000; text-decoration-color: #008000\">'dependencies'</span>: <span style=\"font-weight: bold\">[]</span>,\n",
                            "            <span style=\"color: #008000; text-decoration-color: #008000\">'node_type'</span>: <span style=\"font-weight: bold\">&lt;</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff; font-weight: bold\">QueryType.SINGLE_QUESTION:</span><span style=\"color: #000000; text-decoration-color: #000000\"> </span><span style=\"color: #008000; text-decoration-color: #008000\">'SINGLE'</span><span style=\"color: #000000; text-decoration-color: #000000\">&gt;</span>\n",
                            "<span style=\"color: #000000; text-decoration-color: #000000\">        </span><span style=\"color: #000000; text-decoration-color: #000000; font-weight: bold\">}</span><span style=\"color: #000000; text-decoration-color: #000000\">,</span>\n",
                            "<span style=\"color: #000000; text-decoration-color: #000000\">        </span><span style=\"color: #000000; text-decoration-color: #000000; font-weight: bold\">{</span>\n",
                            "<span style=\"color: #000000; text-decoration-color: #000000\">            </span><span style=\"color: #008000; text-decoration-color: #008000\">'id'</span><span style=\"color: #000000; text-decoration-color: #000000\">: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span><span style=\"color: #000000; text-decoration-color: #000000\">,</span>\n",
                            "<span style=\"color: #000000; text-decoration-color: #000000\">            </span><span style=\"color: #008000; text-decoration-color: #008000\">'question'</span><span style=\"color: #000000; text-decoration-color: #000000\">: </span><span style=\"color: #008000; text-decoration-color: #008000\">\"What is the population of Jason's home country?\"</span><span style=\"color: #000000; text-decoration-color: #000000\">,</span>\n",
                            "<span style=\"color: #000000; text-decoration-color: #000000\">            </span><span style=\"color: #008000; text-decoration-color: #008000\">'dependencies'</span><span style=\"color: #000000; text-decoration-color: #000000\">: </span><span style=\"color: #000000; text-decoration-color: #000000; font-weight: bold\">[]</span><span style=\"color: #000000; text-decoration-color: #000000\">,</span>\n",
                            "<span style=\"color: #000000; text-decoration-color: #000000\">            </span><span style=\"color: #008000; text-decoration-color: #008000\">'node_type'</span><span style=\"color: #000000; text-decoration-color: #000000\">: &lt;QueryType.SINGLE_QUESTION: </span><span style=\"color: #008000; text-decoration-color: #008000\">'SINGLE'</span><span style=\"font-weight: bold\">&gt;</span>\n",
                            "        <span style=\"font-weight: bold\">}</span>\n",
                            "    <span style=\"font-weight: bold\">]</span>\n",
                            "<span style=\"font-weight: bold\">}</span>\n",
                            "</pre>\n"
                        ],
                        "text/plain": [
                            "\u001b[1m{\u001b[0m\n",
                            "    \u001b[32m'query_graph'\u001b[0m: \u001b[1m[\u001b[0m\n",
                            "        \u001b[1m{\u001b[0m\n",
                            "            \u001b[32m'id'\u001b[0m: \u001b[1;36m1\u001b[0m,\n",
                            "            \u001b[32m'question'\u001b[0m: \u001b[32m'What is the population of Canada?'\u001b[0m,\n",
                            "            \u001b[32m'dependencies'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m,\n",
                            "            \u001b[32m'node_type'\u001b[0m: \u001b[1m<\u001b[0m\u001b[1;95mQueryType.SINGLE_QUESTION:\u001b[0m\u001b[39m \u001b[0m\u001b[32m'SINGLE'\u001b[0m\u001b[39m>\u001b[0m\n",
                            "\u001b[39m        \u001b[0m\u001b[1;39m}\u001b[0m\u001b[39m,\u001b[0m\n",
                            "\u001b[39m        \u001b[0m\u001b[1;39m{\u001b[0m\n",
                            "\u001b[39m            \u001b[0m\u001b[32m'id'\u001b[0m\u001b[39m: \u001b[0m\u001b[1;36m2\u001b[0m\u001b[39m,\u001b[0m\n",
                            "\u001b[39m            \u001b[0m\u001b[32m'question'\u001b[0m\u001b[39m: \u001b[0m\u001b[32m\"What is the population of Jason's home country?\"\u001b[0m\u001b[39m,\u001b[0m\n",
                            "\u001b[39m            \u001b[0m\u001b[32m'dependencies'\u001b[0m\u001b[39m: \u001b[0m\u001b[1;39m[\u001b[0m\u001b[1;39m]\u001b[0m\u001b[39m,\u001b[0m\n",
                            "\u001b[39m            \u001b[0m\u001b[32m'node_type'\u001b[0m\u001b[39m: <QueryType.SINGLE_QUESTION: \u001b[0m\u001b[32m'SINGLE'\u001b[0m\u001b[1m>\u001b[0m\n",
                            "        \u001b[1m}\u001b[0m\n",
                            "    \u001b[1m]\u001b[0m\n",
                            "\u001b[1m}\u001b[0m\n"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "def query_planner(question: str) -> QueryPlan:\n",
                "    PLANNING_MODEL = \"gpt-3.5-turbo\"\n",
                "\n",
                "    messages = [\n",
                "        {\n",
                "            \"role\": \"system\",\n",
                "            \"content\": \"You are a world class query planning algorithm capable of breaking apart questions into its dependency queries such that the answers can be used to inform the parent question. Do not answer the questions, simply provide a correct compute graph with good specific questions to ask and relevant dependencies. Before you call the function, think step-by-step to get a better understanding of the problem.\",\n",
                "        },\n",
                "        {\n",
                "            \"role\": \"user\",\n",
                "            \"content\": f\"Consider: {question}\\nGenerate the correct query plan.\",\n",
                "        },\n",
                "    ]\n",
                "\n",
                "    root = client.chat.completions.create(\n",
                "        model=PLANNING_MODEL,\n",
                "        temperature=0,\n",
                "        response_model=QueryPlan,\n",
                "        messages=messages,\n",
                "        max_tokens=1000,\n",
                "    )\n",
                "    return root\n",
                "\n",
                "# Example usage\n",
                "plan = query_planner(\n",
                "    \"What is the difference in populations of Canada and the Jason's home country?\"\n",
                ")\n",
                "print(plan.model_dump())"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Note on RAG\n",
                "\n",
                "While we build the query plan in this example, we do not propose a method to actually answer the question. You can implement your own answer function that perhaps makes a retrieval and calls OpenAI for retrieval augmented generation. That step would also make use of function calls but goes beyond the scope of this example.\n",
                "\n",
                "## Conclusion\n",
                "\n",
                "In this example, we demonstrated how to use the OpenAI Function Call ChatCompletion model to plan and execute a query plan using a question-answering system. We defined the necessary structures using Pydantic and created a query planner function.\n",
                "\n",
                "Feel free to modify the code to fit your specific use case and explore other possibilities of using Function Call models to plan and execute complex workflows."
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "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.12.4"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}