Tutorial 0 - StrictJSON.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ec216685-5ec3-497c-9ef8-e1758ca30423",
   "metadata": {},
   "source": [
    "# Tutorial 0 - StrictJSON\n",
    "\n",
    "- JSON is the default output format for Functions, and is generated by `strict_json()`\n",
    "- Works for JSON outputs with multiple ' or \" or { or } or \\ or unmatched braces/brackets that may break a json.loads()\n",
    "- Reference Repo: https://github.com/tanchongmin/strictjson\n",
    "- Video: https://www.youtube.com/watch?v=3uib0j-9zh4\n",
    "\n",
    "## FAQ\n",
    "- Q: Why not use a type-defined structured framework like Pydantic?\n",
    "- A: Pydantic is very verbose in terms of defining the description and type for each field in the JSON, and can affect performance of LLM for longer context. Moreover, StrictJSON has very flexible checks which can be incorporated when generating the JSON."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6692060b-12c7-4224-8faf-c104892648e3",
   "metadata": {},
   "source": [
    "# Setup Guide"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "867f0bdf-7d13-4aac-ae85-6c18621b66ee",
   "metadata": {},
   "source": [
    "## Step 1: Install TaskGen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7f8701d9-379e-4a94-8265-544b86c65c71",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !pip install taskgen-ai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6c97922-2897-4c66-81f8-8fc6b6f0bdaf",
   "metadata": {},
   "source": [
    "## Step 2: Import required functions and setup relevant API keys for your LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "244b5ff2-12f7-47f1-a09a-41d9edc5f45f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set up API key and do the necessary imports\n",
    "from taskgen import *\n",
    "import os\n",
    "\n",
    "# this is only if you use OpenAI as your LLM\n",
    "os.environ['OPENAI_API_KEY'] = '<YOUR API KEY HERE>'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7536c508-f92c-4225-ac6e-d770218a0908",
   "metadata": {},
   "source": [
    "## Step 3: Define your own LLM\n",
    "- Define an `llm` function that takes in as inputs `system_prompt` and `user_prompt`, and outputs llm response string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f7b0d5d3-b91a-45c9-b619-e19bc579ab96",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def llm(system_prompt: str, user_prompt: str) -> str:\n",
    "    ''' Here, we use OpenAI for illustration, you can change it to your own LLM '''\n",
    "    # ensure your LLM imports are all within this function\n",
    "    from openai import OpenAI\n",
    "    \n",
    "    # define your own LLM here\n",
    "    client = OpenAI()\n",
    "    response = client.chat.completions.create(\n",
    "        model='gpt-4o-mini',\n",
    "        temperature = 0,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ]\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "46e88a54-b1cf-44d6-8137-adc579ee0b09",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Positive'"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Verify that llm function is working\n",
    "llm(system_prompt = 'You are a classifier to classify the sentiment of a sentence', \n",
    "    user_prompt = 'It is a beautiful and sunny day')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89870084-63d6-4227-bae7-30578a3f005f",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 1. Basic Generation\n",
    "\n",
    "- **system_prompt**: Write in whatever you want the LLM to become. \"You are a \\<purpose in life\\>\"\n",
    "- **user_prompt**: The user input. Later, when we use it as a function, this is the function input\n",
    "- **output_format**: JSON of output variables in a dictionary, with the key as the output key, and the value as the output description\n",
    "    - The output keys will be preserved exactly, while the LLM will generate content to match the description of the value as best as possible\n",
    "- **llm**: The llm you want to use. Takes in `system_prompt` and `user_prompt` and outputs the LLM-generated string\n",
    "\n",
    "#### Example LLM Definition\n",
    "```python\n",
    "def llm(system_prompt: str, user_prompt: str) -> str:\n",
    "    ''' Here, we use OpenAI for illustration, you can change it to your own LLM '''\n",
    "    # ensure your LLM imports are all within this function\n",
    "    from openai import OpenAI\n",
    "    \n",
    "    # define your own LLM here\n",
    "    client = OpenAI()\n",
    "    response = client.chat.completions.create(\n",
    "        model='gpt-3.5-turbo',\n",
    "        temperature = 0,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ]\n",
    "    )\n",
    "    return response.choices[0].message.content\n",
    "```\n",
    "\n",
    "#### Example Usage\n",
    "```python\n",
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment',\n",
    "                                    'Adjectives': 'Array of adjectives',\n",
    "                                    'Words': 'Number of words'},\n",
    "                    llm = llm)\n",
    "                                    \n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "```{'Sentiment': 'Positive', 'Adjectives': ['beautiful', 'sunny'], 'Words': 7}```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "c9099840-700d-403a-bc3c-28843d0295c9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Sentiment': 'Positive', 'Adjectives': ['beautiful', 'sunny'], 'Words': 6}\n"
     ]
    }
   ],
   "source": [
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment',\n",
    "                                    'Adjectives': 'Array of adjectives',\n",
    "                                    'Words': 'Number of words'},\n",
    "                    llm = llm)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7e2f8c2-8bcd-40ce-84c8-3e2e18f644d0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Easy to split into corresponding elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ea2ee9a7-9583-4081-b66e-aa8f790a37b0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Positive'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res['Sentiment']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "89da94ce-724a-4da4-89b0-bb575728a398",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['beautiful', 'sunny']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res['Adjectives']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a726d205-b7d8-4e5e-9326-b3ee3a2fe921",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res['Words']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64a724c9-df20-4a57-9379-09c2820c4019",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 2. Advanced Generation\n",
    "- More advanced demonstration involving code that would typically break ```json.loads()```\n",
    "\n",
    "#### Example Usage\n",
    "```python\n",
    "res = strict_json(system_prompt = 'You are a code generator, generating code to fulfil a task',\n",
    "                    user_prompt = 'Given array p, output a function named func_sum to return its sum',\n",
    "                    output_format = {'Elaboration': 'How you would do it',\n",
    "                                     'C': 'Code',\n",
    "                                    'Python': 'Code'},\n",
    "                    llm = llm)\n",
    "                                    \n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "```{'Elaboration': 'Use a loop to iterate through each element in the array and add it to a running total.', ```\n",
    "\n",
    "```'C': 'int func_sum(int p[], int size) {\\n    int sum = 0;\\n    for (int i = 0; i < size; i++) {\\n        sum += p[i];\\n    }\\n    return sum;\\n}', ```\n",
    "\n",
    "```'Python': 'def func_sum(p):\\n    sum = 0\\n    for num in p:\\n        sum += num\\n    return sum'}```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d81d20de-e40f-441e-b438-9dd09e89bcc9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Elaboration': 'I will create a function named func_sum that takes an array p as an argument and returns the sum of its elements using the built-in sum function.', 'C': 'int func_sum(int p[], int size) { int sum = 0; for (int i = 0; i < size; i++) { sum += p[i]; } return sum; }', 'Python': 'def func_sum(p): return sum(p)'}\n"
     ]
    }
   ],
   "source": [
    "res = strict_json(system_prompt = 'You are a code generator, generating code to fulfil a task',\n",
    "                    user_prompt = 'Given array p, output a function named func_sum to return its sum',\n",
    "                    output_format = {'Elaboration': 'How you would do it',\n",
    "                                     'C': 'Code',\n",
    "                                    'Python': 'Code'},\n",
    "                    llm = llm)\n",
    "                                    \n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f180ded0-be2b-4e2e-9d8f-64fb5126a54e",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Easy to split into corresponding elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2b2dc003-4ea1-4bb1-aa74-974d0d022256",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'I will create a function named func_sum that takes an array p as an argument and returns the sum of its elements using the built-in sum function.'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res['Elaboration']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ad3d55ef-11b8-41bc-9de2-8627e3286867",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "int func_sum(int p[], int size) { int sum = 0; for (int i = 0; i < size; i++) { sum += p[i]; } return sum; }\n"
     ]
    }
   ],
   "source": [
    "print(res['C'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3346548c-3f61-4b35-ae44-1fd9dc486498",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "def func_sum(p): return sum(p)\n"
     ]
    }
   ],
   "source": [
    "print(res['Python'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3f8a3950-8586-440c-a908-b578dadd1e84",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The output sum is 15\n"
     ]
    }
   ],
   "source": [
    "## we can even run the Python code (potentially risky due to prompt injection attacks when running unverified code)\n",
    "p = [1, 2, 3, 4, 5]\n",
    "exec(res['Python'])\n",
    "try:\n",
    "    print('The output sum is', func_sum(p))\n",
    "except Exception as e:\n",
    "    print('An exception occured')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55a88380-bc9b-4b1f-84ec-906906313369",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 3. Type forcing output variables\n",
    "- Generally, ```strict_json``` will infer the data type automatically for you for the output fields\n",
    "- However, if you would like very specific data types, you can do data forcing using ```type: <data_type>``` at the last part of the output field description\n",
    "- ```<data_type>``` must be of the form `int`, `float`, `str`, `dict`, `list`, `array`, `code`, `Dict[]`, `List[]`, `Array[]`, `Enum[]`, `bool` for type checking to work\n",
    "- `code` removes all unicode escape characters that might interfere with normal code running\n",
    "- The `Enum` and `List` are not case sensitive, so `enum` and `list` works just as well\n",
    "- For `Enum[list_of_category_names]`, it is best to give an \"Other\" category in case the LLM fails to classify correctly with the other options.\n",
    "- If `list` or `List[]` is not formatted correctly in LLM's output, we will correct it by asking the LLM to list out the elements line by line\n",
    "- For `dict`,  we can further check whether keys are present using `Dict[list_of_key_names]`\n",
    "- Other types will first be forced by rule-based conversion, any further errors will be fed into LLM's error feedback mechanism\n",
    "- If `<data_type>` is not the specified data types, it can still be useful to shape the output for the LLM. However, no type checking will be done.\n",
    "- Note: LLM understands the word `Array` better than `List` since `Array` is the official JSON object type, so in the backend, any type with the word `List` will be converted to `Array`.\n",
    "\n",
    "### LLM-based checks\n",
    "- If you would like the LLM to ensure that the type is being met, use `type: ensure <requirement>`\n",
    "- This will run a LLM to check if the requirement is met. If requirement is not met, the LLM will generate what needs to be done to meet the requirement, which will be fed into the error-correcting loop of `strict_json`\n",
    "\n",
    "#### Example Usage\n",
    "```python\n",
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment, type: Enum[\"Pos\", \"Neg\", \"Other\"]',\n",
    "                                    'Adjectives': 'Array of adjectives, type: List[str]',\n",
    "                                    'Words': 'Number of words, type: int',\n",
    "                                    'In English': 'Whether sentence is in English, type: bool'},\n",
    "                  llm = llm)\n",
    "                                    \n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "```{'Sentiment': 'Pos', 'Adjectives': ['beautiful', 'sunny'], 'Words': 7, 'In English': True}```\n",
    "\n",
    "#### Example Usage\n",
    "```python\n",
    "res = strict_json(system_prompt = 'You are an expert at organising birthday parties',\n",
    "                    user_prompt = 'Give me some information on how to organise a birthday',\n",
    "                    output_format = {'Famous Quote about Age': 'type: ensure quote contains the word age',\n",
    "                                    'Lucky draw numbers': '3 numbers from 1-50, type: List[int]',\n",
    "                                    'Sample venues': 'Describe two venues, type: List[Dict[\"Venue\", \"Description\"]]'},\n",
    "                    llm = llm)\n",
    "\n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "`Using LLM to check \"The secret of staying young is to live honestly, eat slowly, and lie about your age. - Lucille Ball\" to see if it adheres to \"quote contains the word age\" Requirement Met: True`\n",
    "\n",
    "\n",
    "```{'Famous Quote about Age': 'The secret of staying young is to live honestly, eat slowly, and lie about your age. - Lucille Ball',```\n",
    "```'Lucky draw numbers': [7, 21, 35],```\n",
    "\n",
    "```'Sample venues': [{'Venue': 'Beachside Resort', 'Description': 'A beautiful resort with stunning views of the beach. Perfect for a summer birthday party.'}, {'Venue': 'Indoor Trampoline Park', 'Description': 'An exciting venue with trampolines and fun activities. Ideal for an active and energetic birthday celebration.'}]}```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0dcac6c2-b3a1-4171-8931-dc6f6ff763e6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Sentiment': 'Pos', 'Adjectives': ['beautiful', 'sunny'], 'Words': 7, 'In English': True}\n"
     ]
    }
   ],
   "source": [
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment, type: Enum[\"Pos\", \"Neg\", \"Other\"]',\n",
    "                                    'Adjectives': 'Array of Adjectives, type: List[str]',\n",
    "                                    'Words': 'Number of words, type: int',\n",
    "                                    'In English': 'Whether sentence is in English, type: bool'},\n",
    "                    llm = llm)\n",
    "\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8b44bcff-1469-47bb-a4ef-cf6b3794bd19",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using LLM to check \"Age is merely the number of years the world has been enjoying you.\" to see if it adheres to \"quote includes word age\"\n",
      "{'Famous Quote about Age': 'Age is merely the number of years the world has been enjoying you.', 'Lucky draw numbers': [7, 23, 45], 'Sample venues': [{'Venue': 'Local Park', 'Description': 'A spacious outdoor area perfect for picnics and games, ideal for a casual birthday celebration.'}, {'Venue': 'Community Center', 'Description': 'A versatile indoor space that can be decorated according to the theme of the party, suitable for larger gatherings.'}]}\n"
     ]
    }
   ],
   "source": [
    "# Example of multiple nested inputs with multiple types of checks\n",
    "# Quotes may fail json parsing, so let us use LLMs to fix it\n",
    "res = strict_json(system_prompt = 'You are an expert at organising birthday parties',\n",
    "                    user_prompt = 'Give me some information on how to organise a birthday',\n",
    "                    output_format = {'Famous Quote about Age': 'type: ensure quote includes word age',\n",
    "                                    'Lucky draw numbers': '3 numbers from 1-50, type: List[int]',\n",
    "                                    'Sample venues': 'Describe two venues, type: List[Dict[\"Venue\", \"Description\"]]'},\n",
    "                    llm = llm)\n",
    "\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d87f0a58-f8ad-4adc-aeea-45eda85e43f4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 4. Functions\n",
    "- Enhances ```strict_json()``` with a function-like interface for repeated use of modular LLM-based functions (or wraps external functions)\n",
    "- Use angle brackets <> to enclose input variable names. First input variable name to appear in `fn_description` will be first input variable and second to appear will be second input variable. For example, `fn_description = 'Adds up two numbers, <var1> and <var2>'` will result in a function with first input variable `var1` and second input variable `var2`\n",
    "- (Optional) If you would like greater specificity in your function's input, you can describe the variable after the : in the input variable name, e.g. `<var1: an integer from 10 to 30>`. Here, `var1` is the input variable and `an integer from 10 to 30` is the description.\n",
    "- (Optional) If your description of the variable is one of `int`, `float`, `str`, `dict`, `list`, `array`, `code`, `Dict[]`, `List[]`, `Array[]`, `Enum[]`, `bool`, we will enforce type checking when generating the function inputs in `get_next_subtask` method of the `Agent` class. Example: `<var1: int>`. Refer to Section 3. Type Forcing Output Variables for details.\n",
    "- Inputs (primary):\n",
    "    - **fn_description**: String. Function description to describe process of transforming input variables to output variables. Variables must be enclosed in <> and listed in order of appearance in function input.\n",
    "        - New feature: If `external_fn` is provided and no `fn_description` is provided, then we will automatically parse out the fn_description based on docstring of `external_fn`. The docstring should contain the names of all compulsory input variables\n",
    "        - New feature: If `external_fn` is provided and no `output_format` is provided, then we will automatically derive the `output_format` from the function signature\n",
    "    - **output_format**: Dict. Dictionary containing output variables names and description for each variable.\n",
    "    \n",
    "- Inputs (optional):\n",
    "    - **examples** - Dict or List[Dict]. Examples in Dictionary form with the input and output variables (list if more than one)\n",
    "    - **external_fn** - Python Function. If defined, instead of using LLM to process the function, we will run the external function. \n",
    "        If there are multiple outputs of this function, we will map it to the keys of `output_format` in a one-to-one fashion\n",
    "    - **fn_name** - String. If provided, this will be the name of the function. Otherwise, if `external_fn` is provided, it will be the name of `external_fn`. Otherwise, we will use LLM to generate a function name from the `fn_description`\n",
    "    - **kwargs** - Dict. Additional arguments you would like to pass on to the strict_json function\n",
    "        \n",
    "- Outputs:\n",
    "    JSON of output variables in a dictionary (similar to ```strict_json```)\n",
    "    \n",
    "#### Example Usage 1 (Description only)\n",
    "```python\n",
    "# basic configuration with variable names (in order of appearance in fn_description)\n",
    "fn = Function(fn_description = 'Output a sentence with <obj> and <entity> in the style of <emotion>', \n",
    "                     output_format = {'output': 'sentence'},\n",
    "                     llm = llm)\n",
    "\n",
    "# Use the function\n",
    "fn('ball', 'dog', 'happy') #obj, entity, emotion\n",
    "```\n",
    "\n",
    "#### Example Output 1\n",
    "```{'output': 'The happy dog chased the ball.'}```\n",
    "\n",
    "#### Example Usage 2 (Examples only)\n",
    "```python\n",
    "# Construct the function: infer pattern from just examples without description (here it is multiplication)\n",
    "fn = Function(fn_description = 'Map <var1> and <var2> to output based on examples', \n",
    "                     output_format = {'output': 'final answer'}, \n",
    "                     examples = [{'var1': 3, 'var2': 2, 'output': 6}, \n",
    "                                 {'var1': 5, 'var2': 3, 'output': 15}, \n",
    "                                 {'var1': 7, 'var2': 4, 'output': 28}],\n",
    "                     llm = llm)\n",
    "\n",
    "# Use the function\n",
    "fn(2, 10) #var1, var2\n",
    "```\n",
    "\n",
    "#### Example Output 2\n",
    "```{'output': 20}```\n",
    "\n",
    "#### Example Usage 3 (Description and Examples)\n",
    "```python\n",
    "# Construct the function: description and examples with variable names\n",
    "# variable names will be referenced in order of appearance in fn_description\n",
    "fn = Function(fn_description = 'Output the sum and difference of <num1> and <num2>', \n",
    "                 output_format = {'sum': 'sum of two numbers', \n",
    "                                  'difference': 'absolute difference of two numbers'},\n",
    "                 examples = {'num1': 2, 'num2': 4, 'sum': 6, 'difference': 2},\n",
    "                 llm = llm)\n",
    "\n",
    "# Use the function\n",
    "fn(3, 4) #num1, num2\n",
    "```\n",
    "\n",
    "#### Example Output 3\n",
    "```{'sum': 7, 'difference': 1}```\n",
    "\n",
    "#### Example Usage 4 (External Function with automatic inference of fn_description and output_format - Preferred)\n",
    "```python\n",
    "# Docstring should provide all input variables, otherwise we will add it in automatically\n",
    "# We will ignore shared_variables, *args and **kwargs\n",
    "# No need to define llm in Function for External Functions\n",
    "from typing import List\n",
    "def add_number_to_list(num1: int, num_list: List[int], *args, **kwargs) -> List[int]:\n",
    "    '''Adds num1 to num_list'''\n",
    "    num_list.append(num1)\n",
    "    return num_list\n",
    "\n",
    "fn = Function(external_fn = add_number_to_list)\n",
    "\n",
    "# Show the processed function docstring\n",
    "print(str(fn))\n",
    "\n",
    "# Use the function\n",
    "fn(3, [2, 4, 5])\n",
    "```\n",
    "#### Example Output 5\n",
    "`Description: Adds <num1: int> to <num_list: list>`\n",
    "\n",
    "`Input: ['num1', 'num_list']`\n",
    "\n",
    "`Output: {'num_list': 'Array of numbers'}`\n",
    "\n",
    "`{'num_list': [2, 4, 5, 3]}`\n",
    "\n",
    "#### Example Usage 5 (External Function with manually defined fn_description and output_format - Legacy Approach)\n",
    "```python\n",
    "def binary_to_decimal(x):\n",
    "    return int(str(x), 2)\n",
    "\n",
    "# an external function with a single output variable, with an expressive variable description\n",
    "fn = Function(fn_description = 'Convert input <x: a binary number in base 2> to base 10', \n",
    "            output_format = {'output1': 'x in base 10'},\n",
    "            external_fn = binary_to_decimal)\n",
    "\n",
    "# Use the function\n",
    "fn(10) #x\n",
    "```\n",
    "\n",
    "#### Example Output 4\n",
    "```{'output1': 2}```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a237cb26-2655-4d44-80cb-89454562a07d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': 'The dog joyfully chased the ball, its tail wagging with pure happiness.'}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# basic configuration with variable names (in order of appearance in fn_description)\n",
    "fn = Function(fn_description = 'Output a sentence with <obj> and <entity> in the style of <emotion>', \n",
    "                     output_format = {'output': 'sentence'}, \n",
    "                     llm = llm)\n",
    "fn('ball', 'dog', 'happy') #obj, entity, emotion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a35d19de-cc28-4da7-9f22-3ed309afb337",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': 20}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# infer pattern from just examples without description (here it is multiplication)\n",
    "fn = Function(fn_description = 'Map <var1> and <var2> to output based on examples', \n",
    "                     output_format = {'output': 'final answer'}, \n",
    "                     examples = [{'var1': 3, 'var2': 2, 'output': 6}, \n",
    "                                 {'var1': 5, 'var2': 3, 'output': 15}, \n",
    "                                 {'var1': 7, 'var2': 4, 'output': 28}],\n",
    "                     llm = llm)\n",
    "fn(2, 10) #var1, var2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "306ec488-8934-46df-b020-fddc31f740a5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'sum': 7, 'difference': 1}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# multiple outputs and examples with variable names (recommended)\n",
    "fn = Function(fn_description = 'Output the sum and difference of <num1> and <num2>', \n",
    "                 output_format = {'sum': 'sum of two numbers', \n",
    "                                  'difference': 'absolute difference of two numbers'},\n",
    "                 examples = {'num1': 2, 'num2': 4, 'sum': 6, 'difference': 2},\n",
    "                llm = llm)\n",
    "fn(3, 4) #num1, num2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8e36a590-7082-4271-9615-eefa6fbccbaa",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'sum': 7,\n",
       " 'poem': 'In a world of numbers bright,\\nThree and four bring pure delight.\\nTogether they dance, a joyful sight,\\nAdding up to seven, oh what a height!\\nWith laughter and cheer, they play all day,\\nIn the realm of math, they find their way.',\n",
       " 'code': 'def sum_numbers(num1, num2):\\n    return int(num1) + num2\\n\\nresult = sum_numbers(\"3\", 4)\\nprint(result)'}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# multiple outputs with variable names\n",
    "fn = Function(fn_description = '''Output the integer sum of <num1: int or str> and <num2: int or str>\n",
    "generate a poem in style of <poem_style> and code in <prog_language>''', \n",
    "                 output_format = {'sum': 'sum of two numbers as an integer', \n",
    "                'poem': 'poem about two numbers',\n",
    "                'code': 'code to do the sum of any two numbers num1 and num2'},\n",
    "                 llm = llm)\n",
    "fn('three', 4, 'happy', 'Python') #num1, num2, poem_style, prog_language"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cccd4e1e-ce8d-410c-9b0b-ff22ec41ff67",
   "metadata": {
    "tags": []
   },
   "source": [
    "## External Function - Preferred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ba9e0417-a265-4a71-9e50-fb0afe6d2da3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Description: Adds <num1: int> to <num_list: list[int]>\n",
      "Input: ['num1', 'num_list']\n",
      "Output: {'output_1': 'list[int]'}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'output_1': [2, 4, 5, 3]}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Docstring must provide all input variables, otherwise we will add it in automatically\n",
    "# We will ignore shared_variables, *args and **kwargs\n",
    "# No need to define llm in Function for External Functions\n",
    "from typing import List\n",
    "def add_number_to_list(num1: int, num_list: List[int], *args, **kwargs) -> List[int]:\n",
    "    '''Adds num1 to num_list'''\n",
    "    num_list.append(num1)\n",
    "    return num_list\n",
    "\n",
    "fn = Function(external_fn = add_number_to_list)\n",
    "\n",
    "# Show the processed function docstring\n",
    "print(str(fn))\n",
    "\n",
    "# Use the function\n",
    "fn(3, [2, 4, 5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f017fc5-c20d-4a06-a2bd-12f4dd360f6d",
   "metadata": {},
   "source": [
    "## External Function - Legacy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "33eef4ff-3e45-489a-8cc7-be2d092015eb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output1': 2}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def binary_to_decimal(x):\n",
    "    return int(str(x), 2)\n",
    "\n",
    "# an external function with a single output variable, with an expressive variable description\n",
    "fn = Function(fn_description = 'Convert input <x: a binary number in base 2> to base 10', \n",
    "            output_format = {'output1': 'x in base 10'},\n",
    "            external_fn = binary_to_decimal)\n",
    "\n",
    "# Use the function\n",
    "fn(10) #x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8f34db9-785d-448a-a811-83bd2d1ab805",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 5. Nested Outputs\n",
    "- StrictJSON supports nested outputs like nested lists and dictionaries\n",
    "\n",
    "#### Example Input\n",
    "```python\n",
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': ['Type of Sentiment', \n",
    "                                                   'Strength of Sentiment, type: Enum[1, 2, 3, 4, 5]'],\n",
    "                                    'Adjectives': \"Name and Description as separate keys, type: List[Dict['Name', 'Description']]\",\n",
    "                                    'Words': {\n",
    "                                        'Number of words': 'Word count', \n",
    "                                        'Language': {\n",
    "                                              'English': 'Whether it is English, type: bool',\n",
    "                                              'Chinese': 'Whether it is Chinese, type: bool'\n",
    "                                                  },\n",
    "                                        'Proper Words': 'Whether the words are proper in the native language, type: bool'\n",
    "                                        }\n",
    "                                    },\n",
    "                 llm = llm)\n",
    "\n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "`{'Sentiment': ['Positive', 3],`\n",
    "\n",
    "`'Adjectives': [{'Name': 'beautiful', 'Description': 'pleasing to the senses'}, {'Name': 'sunny', 'Description': 'filled with sunshine'}],`\n",
    "\n",
    "`'Words':`\n",
    "\n",
    "`     {'Number of words': 6,`\n",
    "    \n",
    "`     'Language': {'English': True, 'Chinese': False},`\n",
    "\n",
    "`     'Proper Words': True}`\n",
    "    \n",
    "`}`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "755cbab2-80de-40a5-a721-5162378b0ae0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Sentiment': ['Positive', 5], 'Adjectives': [{'Name': 'Beautiful', 'Description': 'Pleasing the senses or mind aesthetically'}, {'Name': 'Sunny', 'Description': 'Bright with sunlight'}], 'Words': {'Number of words': 7, 'Language': {'English': True, 'Chinese': False}, 'Proper Words': True}}\n"
     ]
    }
   ],
   "source": [
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': ['Type of Sentiment', \n",
    "                                                   'Strength of Sentiment, type: Enum[1, 2, 3, 4, 5]'],\n",
    "                                    'Adjectives': \"Name and Description as separate keys, type: List[Dict['Name', 'Description']]\",\n",
    "                                    'Words': {\n",
    "                                        'Number of words': 'Word count', \n",
    "                                        'Language': {\n",
    "                                              'English': 'Whether it is English, type: bool',\n",
    "                                              'Chinese': 'Whether it is Chinese, type: bool'\n",
    "                                                  },\n",
    "                                        'Proper Words': 'Whether the words are proper in the native language, type: bool'\n",
    "                                        }\n",
    "                                    },\n",
    "                 llm = llm)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "270f6591-afa5-43e1-a520-c4ad68067977",
   "metadata": {},
   "source": [
    "# 6. Return as JSON\n",
    "- By default, `strict_json` returns a Python Dictionary\n",
    "- If needed to parse as JSON, simply set `return_as_json=True`\n",
    "- By default, this is set to `False` in order to return a Python Dictionry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1f0f2a3f-1966-4cbd-ba66-63ba6c007cfa",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"Sentiment\": [\"Positive\", 5], \"Adjectives\": [{\"Name\": \"Beautiful\", \"Description\": \"Pleasing the senses or mind aesthetically\"}, {\"Name\": \"Sunny\", \"Description\": \"Bright with sunlight\"}], \"Words\": {\"Number of words\": 7, \"Language\": {\"English\": true, \"Chinese\": false}, \"Proper Words\": true}}\n"
     ]
    }
   ],
   "source": [
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': ['Type of Sentiment', \n",
    "                                                   'Strength of Sentiment, type: Enum[1, 2, 3, 4, 5]'],\n",
    "                                    'Adjectives': \"Name and Description as separate keys, type: List[Dict['Name', 'Description']]\",\n",
    "                                    'Words': {\n",
    "                                        'Number of words': 'Word count', \n",
    "                                        'Language': {\n",
    "                                              'English': 'Whether it is English, type: bool',\n",
    "                                              'Chinese': 'Whether it is Chinese, type: bool'\n",
    "                                                  },\n",
    "                                        'Proper Words': 'Whether the words are proper in the native language, type: bool'\n",
    "                                        }\n",
    "                                    },\n",
    "                 llm = llm,\n",
    "                 return_as_json = True) # Toggle this to True (Note: Also works with OpenAI JSON Mode)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b770893d-067c-4afe-8368-719138b2e776",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 7. Integrating with OpenAI JSON Mode\n",
    "- If you want to use the OpenAI JSON Mode, you can simply add in ```openai_json_mode = True``` and set ```model = 'gpt-4-1106-preview'``` or ```model = 'gpt-3.5-turbo-1106'``` in ```strict_json``` or ```Function```\n",
    "- We will set model to ```gpt-3.5-turbo-1106``` by default if you provide an invalid model\n",
    "- This does not work with the `llm` variable\n",
    "- Note that type checking does not work with OpenAI JSON Mode\n",
    "\n",
    "#### Example Usage\n",
    "```python\n",
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment',\n",
    "                                    'Adjectives': 'Array of adjectives',\n",
    "                                    'Words': 'Number of words'},\n",
    "                    model = 'gpt-3.5-turbo-1106' # Set the model\n",
    "                    openai_json_mode = True) # Toggle this to True\n",
    "                                    \n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "```{'Sentiment': 'positive', 'Adjectives': ['beautiful', 'sunny'], 'Words': 6}```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "7461dfc5-a3e1-4fd5-8567-c2ca834fab7b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Sentiment': 'Positive', 'Adjectives': ['beautiful', 'sunny'], 'Words': 6}\n"
     ]
    }
   ],
   "source": [
    "res = strict_json(system_prompt = 'You are a classifier',\n",
    "                user_prompt = 'It is a beautiful and sunny day',\n",
    "                output_format = {'Sentiment': 'Type of Sentiment',\n",
    "                                'Adjectives': 'Array of adjectives',\n",
    "                                'Words': 'Number of words'},\n",
    "               model = 'gpt-3.5-turbo-1106', # Set the model\n",
    "               openai_json_mode = True) # Toggle this to True\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3f61007-a728-4711-8f5b-bbd5a4b1991e",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 8. Async Mode\n",
    "\n",
    "- `AsyncFunction` and `strict_json_async`\n",
    "    - These are the async equivalents of `Function` and `strict_json`\n",
    "    - You will need to define an LLM that can operate in async mode\n",
    "    - Everything is the same as the sync version of the functions, except you use the `await` keyword when calling `AsyncFunction` and `strict_json_async`\n",
    "    \n",
    "    \n",
    "- Using Async can help do parallel processes simulataneously, resulting in a much faster workflow\n",
    "\n",
    "#### Example LLM in Async Mode\n",
    "```python\n",
    "async def llm_async(system_prompt: str, user_prompt: str):\n",
    "    ''' Here, we use OpenAI for illustration, you can change it to your own LLM '''\n",
    "    # ensure your LLM imports are all within this function\n",
    "    from openai import AsyncOpenAI\n",
    "    \n",
    "    # define your own LLM here\n",
    "    client = AsyncOpenAI()\n",
    "    response = await client.chat.completions.create(\n",
    "        model='gpt-3.5-turbo',\n",
    "        temperature = 0,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ]\n",
    "    )\n",
    "    return response.choices[0].message.content\n",
    "```\n",
    "\n",
    "#### Example Input (strict_json_async)\n",
    "```python\n",
    "res = await strict_json_async(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment',\n",
    "                                    'Adjectives': 'Array of adjectives',\n",
    "                                    'Words': 'Number of words'},\n",
    "                                     llm = llm_async) # use an async LLM function\n",
    "\n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "`{'Sentiment': 'Positive', 'Adjectives': ['beautiful', 'sunny'], 'Words': 7}`\n",
    "\n",
    "#### Example Input (AsyncFunction)\n",
    "```python\n",
    "fn =  AsyncFunction(fn_description = 'Output a sentence with <obj> and <entity> in the style of <emotion>', \n",
    "                     output_format = {'output': 'sentence'},\n",
    "                    llm = llm_async)\n",
    "\n",
    "res = await fn('ball', 'dog', 'happy') #obj, entity, emotion\n",
    "\n",
    "print(res)\n",
    "```\n",
    "\n",
    "#### Example Output\n",
    "`{'output': 'The dog happily chased the ball.'}`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f3c74b90-5c3e-481f-a020-cf517c2c19f3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "async def llm_async(system_prompt: str, user_prompt: str):\n",
    "    ''' Here, we use OpenAI for illustration, you can change it to your own LLM '''\n",
    "    # ensure your LLM imports are all within this function\n",
    "    from openai import AsyncOpenAI\n",
    "    \n",
    "    # define your own LLM here\n",
    "    client = AsyncOpenAI()\n",
    "    response = await client.chat.completions.create(\n",
    "        model='gpt-4o-mini',\n",
    "        temperature = 0,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ]\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2d65c2e3-1cc1-4bd3-b220-19165a925b3c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Sentiment': 'Positive', 'Adjectives': ['beautiful', 'sunny'], 'Words': 6}\n"
     ]
    }
   ],
   "source": [
    "# Demonstration of strict_json_async\n",
    "res = await strict_json_async(system_prompt = 'You are a classifier',\n",
    "                    user_prompt = 'It is a beautiful and sunny day',\n",
    "                    output_format = {'Sentiment': 'Type of Sentiment',\n",
    "                                    'Adjectives': 'Array of adjectives',\n",
    "                                    'Words': 'Number of words'},\n",
    "                                     llm = llm_async) # use an async LLM function\n",
    "\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "03a29341-e708-4bc2-99a4-1a6af059da3d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'output': 'The dog joyfully chased the ball, its tail wagging with pure happiness.'}\n"
     ]
    }
   ],
   "source": [
    "# Demonstration of AsyncFunction\n",
    "fn =  AsyncFunction(fn_description = 'Output a sentence with <obj> and <entity> in the style of <emotion>', \n",
    "                     output_format = {'output': 'sentence'},\n",
    "                     llm = llm_async)\n",
    "\n",
    "res = await fn('ball', 'dog', 'happy') #obj, entity, emotion\n",
    "\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "851fae25-cc89-458f-b77f-db4f54b8fc02",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 9. Additional Output Field Checks (Advanced)\n",
    "\n",
    "- You can also specify your own custom check function that will be used to check the output field (which will be in `str`, `int`, `float`, `list` or `dict` format inferred by LLM or specified in `type: <data type>`)\n",
    "- Ensure that what you are checking for is implied in the output field's description in `output_format` of `strict_json` or `Function`\n",
    "- Your custom check function must take in: `output_field`\n",
    "- Your custom check function must output: \n",
    "    - `requirement` (str): The requirement you are checking for\n",
    "    - `requirement_met` (bool): Whether condition is met, True or False\n",
    "    - `action_needed` (str): What needs to be done to meet requirement if requirement_met is False\n",
    "- If `requirement_met` is False, the `requirement` and `action_needed` message will be used for the `strict_json` error correcting mechanism. Otherwise, the error correcting mechanism will not be triggered\n",
    "- `action_needed` is used to tell the LLM what it needs to do to meet your requirements (LLM is not able to self-correct without guidance for most cases). Try to be as specific as possible to improve error correction success rate.\n",
    "- Pass in your custom check function inside `custom_checks` variable of `strict_json` or `Function` under the same key as that in `output_format`\n",
    "- You can add multiple check functions for one variable by putting it inside the same list\n",
    "- Example custom check function named `hello_world_check` which checks for the presence of hello world\n",
    "- You can also use information in the variable `check_data` for checks (input via `strict_json` or `Function`)\n",
    "\n",
    "#### Example Custom Check Functions\n",
    "```python\n",
    "def hello_world_check(output_field, check_data) -> (str, bool, str):\n",
    "    ''' Example function 1: Checks whether hello world is present in output_field. '''\n",
    "    requirement = 'Check whether hello world is present in output field'\n",
    "    requirement_met = True\n",
    "    action_needed = ''\n",
    "    # do a check for requirement of having 'hello'\n",
    "    if 'hello' not in str(output_field):\n",
    "        requirement_met = False\n",
    "        action_needed += 'Add in the word hello into output field, '\n",
    "    if 'world' not in str(output_field):\n",
    "        requirement_met = False\n",
    "        action_needed += 'Add in the word world into output field, '\n",
    "    return (requirement, requirement_met, action_needed)\n",
    "```\n",
    "\n",
    "```python\n",
    "def function_name_check(output_field, check_data) -> (str, bool, str):\n",
    "    ''' Example function 2: Checks whether function name is present in output_field\n",
    "    Uses additional information from the check_data variable of strict_json'''\n",
    "    function_name = check_data['Function name']\n",
    "    requirement = f'Check whether {function_name} is present in output field'\n",
    "    requirement_met = True\n",
    "    action_needed = ''\n",
    "    \n",
    "    # do a check for requirement of having 'myprint'\n",
    "    if function_name not in str(output_field):\n",
    "        requirement_met = False\n",
    "        action_needed += f'Ensure that function name \"{function_name}\" is used, '\n",
    "    return (requirement, requirement_met, action_needed)\n",
    "```\n",
    "\n",
    "#### Example Usage 1 (in strict_json)\n",
    "```python\n",
    "# we can input our custom_checks as a list of check functions, and check_data is the additional information for these check functions\n",
    "res = strict_json(system_prompt = 'You are a code generator',\n",
    "                    user_prompt = 'Print out hello world',\n",
    "                    output_format = {'Thoughts': 'How to do it',\n",
    "                                    'Python Code': 'Function beginning with def myprint() -> str:, type: code'},\n",
    "                    llm = llm,\n",
    "                    custom_checks = {'Python Code': [hello_world_check, function_name_check]}, # add in custom_checks\n",
    "                    check_data = {'Function name:' 'myprint'}) # add in check_data\n",
    "                                    \n",
    "print(res)\n",
    "```\n",
    "#### Example Output 1\n",
    "`Running check for \"Check whether hello world is present in output field\" on output field of \"Python Code\"\n",
    "Requirement met`\n",
    "\n",
    "\n",
    "`Running check for \"Check whether myprint is present in output field\" on output field of \"Python Code\"\n",
    "Requirement met`\n",
    "\n",
    "\n",
    "`{'Thoughts': 'To print out \"hello world\", use the print() function in Python.',`\n",
    "`'Python Code': 'def myprint() -> str:\\n    return \"hello world\"'}`\n",
    "\n",
    "#### Example Usage 2 (in Function)\n",
    "\n",
    "```python\n",
    "fn = Function(fn_description = 'Output code to print hello world in a function named <var1>', \n",
    "                     output_format = {'Python code': 'Python function named <var1> to print hello world, type: code'},\n",
    "                     custom_checks = {'Python code': [function_name_check]},\n",
    "                     llm = llm)\n",
    "\n",
    "# in runtime of function, we can input what we would want to check in check_data if we are not sure what it will be beforehand\n",
    "fn('hello world', 'myprint', check_data = {'Function name': 'myprint'})\n",
    "```\n",
    "\n",
    "#### Example Output 2\n",
    "\n",
    "`Running check for \"Check whether myprint is present in output field\" on output field of \"Python code\"\n",
    "Requirement met`\n",
    "\n",
    "`{'Python code': 'def myprint():\\n    print(\"hello world\")'}`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "95ee40f9-827b-4952-b5d9-8cfde7236fe0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def hello_world_check(output_field, check_data) -> (str, bool, str):\n",
    "    ''' Example function 1: Checks whether hello world is present in output_field. '''\n",
    "    requirement = 'Check whether hello world is present in output field'\n",
    "    requirement_met = True\n",
    "    action_needed = ''\n",
    "    # do a check for requirement of having 'hello'\n",
    "    if 'hello' not in str(output_field):\n",
    "        requirement_met = False\n",
    "        action_needed += 'Add in the lowercase word hello into output field, '\n",
    "    if 'world' not in str(output_field):\n",
    "        requirement_met = False\n",
    "        action_needed += 'Add in the lowercase word world into output field, '\n",
    "    return (requirement, requirement_met, action_needed)\n",
    "\n",
    "def function_name_check(output_field, check_data) -> (str, bool, str):\n",
    "    ''' Example function 2: Checks whether function name is present in output_field\n",
    "    Uses additional information from the check_data variable of strict_json'''\n",
    "    function_name = check_data['Function name']\n",
    "    requirement = f'Check whether {function_name} is present in output field'\n",
    "    requirement_met = True\n",
    "    action_needed = ''\n",
    "    \n",
    "    # do a check for requirement of having 'myprint'\n",
    "    if function_name not in str(output_field):\n",
    "        requirement_met = False\n",
    "        action_needed += f'Ensure that function name \"{function_name}\" is used, '\n",
    "    return (requirement, requirement_met, action_needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "444c96a4-650c-4982-95db-eec35ae3aa3d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running check for \"Check whether hello world is present in output field\" on output field of \"Python Code\"\n",
      "Requirement met\n",
      "\n",
      "\n",
      "Running check for \"Check whether myprint is present in output field\" on output field of \"Python Code\"\n",
      "Requirement met\n",
      "\n",
      "\n",
      "{'Thoughts': 'To print out \"hello world\", I will define a function that uses the print statement in Python.', 'Python Code': 'def myprint() -> str:\\n    print(\"hello world\")'}\n"
     ]
    }
   ],
   "source": [
    "# we can input our custom_checks as a list of check functions, and check_data is the additional information for these check functions\n",
    "res = strict_json(system_prompt = 'You are a code generator',\n",
    "                    user_prompt = 'Print out hello world',\n",
    "                    output_format = {'Thoughts': 'How to do it',\n",
    "                                    'Python Code': 'Function beginning with def myprint() -> str:, type: code'},\n",
    "                    llm = llm,\n",
    "                    custom_checks = {'Python Code': [hello_world_check, function_name_check]},\n",
    "                    check_data = {'Function name': 'myprint'})\n",
    "                                    \n",
    "print(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "037e302b-ca48-4bc8-8b54-caf3323ab622",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running check for \"Check whether myprint is present in output field\" on output field of \"Python code\"\n",
      "Requirement met\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'Python code': 'def myprint():\\n    print(\"hello world\")'}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fn = Function(fn_description = 'Output code to print hello world in a function named <var1>', \n",
    "                     output_format = {'Python code': 'Python function named <var1> to print hello world, type: code'},\n",
    "                     custom_checks = {'Python code': [function_name_check]},\n",
    "                     llm = llm)\n",
    "\n",
    "# in runtime of function, we can input what we would want to check in check_data if we are not sure what it will be beforehand\n",
    "fn('hello world', 'myprint', check_data = {'Function name': 'myprint'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8885da33-142d-4fc9-a2de-3e930449b1a4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Optional: Under the hood (Explanation of how strict_json works)\n",
    "- When given the output JSON format, it adds a delimiter (default: ###) to enclose the key of the JSON.\n",
    "- Example Output JSON provided: ```{'Sentiment': 'Type of Sentiment'}```\n",
    "- Example Output JSON interpreted by Strict JSON: ```{'###Sentiment###': 'Type of Sentiment'}```\n",
    "- We then process the JSON format by using regex to search for the delimiter to extract the keys and values\n",
    "- Works for nested data structures as well by extracting recursively\n",
    "- Note: Change the delimiter to whatever is not present in your dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "6c4e3f68-6384-4715-aefa-52e755343046",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# a very difficult chunk of text for json.loads() to parse (it will fail)\n",
    "res = '''{\n",
    "'###Question of the day###':'What is the 'x' in dx/dy?', \n",
    "'###Code Block 1###':   '#include <stdio.h>\\nint main(){\\nint x = 'a'; return 0;\\n}',\n",
    "'###Another Code###': 'import numpy as np\n",
    "### Oh what is this doing here\n",
    "print(\"It can handle so many quotations ' \\\\\" and backslashes and unexpected curly braces { } You don't even need to match }!\")',\n",
    "'###Some characters###': '~!@#$%^&*()_+-'\"{}[];?><,.'\n",
    "}'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "2fb7dea7-a034-4964-87fe-d6495bb6d2f8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# change this to whatever is not common in your dataset\n",
    "delimiter = '###'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "96c3aac1-7c14-447a-8e18-cb312492fba4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Question of the day', \"'What is the 'x' in dx/dy?'\", 'Code Block 1', \"'#include <stdio.h>\\nint main(){\\nint x = 'a'; return 0;\\n}'\", 'Another Code', '\\'import numpy as np\\n### Oh what is this doing here\\nprint(\"It can handle so many quotations \\' \\\\\" and backslashes and unexpected curly braces { } You don\\'t even need to match }!\")\\'', 'Some characters', '\\'~!@#$%^&*()_+-\\'\"{}[];?><,.\\'']\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "# Use regular expressions to extract keys and values\n",
    "pattern = fr\",*\\s*['|\\\"]{delimiter}([^#]*){delimiter}['|\\\"]:\\s*\"\n",
    "\n",
    "matches = re.split(pattern, str(res[1:-1]).strip())\n",
    "\n",
    "# remove null matches\n",
    "my_matches = [match for match in matches if match !='']\n",
    "\n",
    "print(my_matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "b1308e53-e592-466f-aa7a-b91ef0e30878",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Question of the day', \"What is the 'x' in dx/dy?\", 'Code Block 1', \"#include <stdio.h>\\nint main(){\\nint x = 'a'; return 0;\\n}\", 'Another Code', 'import numpy as np\\n### Oh what is this doing here\\nprint(\"It can handle so many quotations \\' \\\\\" and backslashes and unexpected curly braces { } You don\\'t even need to match }!\")', 'Some characters', '~!@#$%^&*()_+-\\'\"{}[];?><,.']\n"
     ]
    }
   ],
   "source": [
    "# remove the ' from the value matches\n",
    "curated_matches = [match[1:-1] if match[0] in '\\'\"' else match for match in my_matches]\n",
    "\n",
    "print(curated_matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "35ffdcf6-6433-41f9-bde2-010ab13c12ab",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(curated_matches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "eb2062d6-b933-41d1-ad40-5bbca7936157",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Question of the day': \"What is the 'x' in dx/dy?\", 'Code Block 1': \"#include <stdio.h>\\nint main(){\\nint x = 'a'; return 0;\\n}\", 'Another Code': 'import numpy as np\\n### Oh what is this doing here\\nprint(\"It can handle so many quotations \\' \\\\\" and backslashes and unexpected curly braces { } You don\\'t even need to match }!\")', 'Some characters': '~!@#$%^&*()_+-\\'\"{}[];?><,.'}\n"
     ]
    }
   ],
   "source": [
    "# create a dictionary\n",
    "end_dict = {}\n",
    "for i in range(0, len(curated_matches), 2):\n",
    "    end_dict[curated_matches[i]] = curated_matches[i+1]\n",
    "    \n",
    "print(end_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "a6352e40-759d-476f-9a27-549a48557b0d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Key: Question of the day\n",
      "Value: What is the 'x' in dx/dy?\n",
      "#####\n",
      "Key: Code Block 1\n",
      "Value: #include <stdio.h>\n",
      "int main(){\n",
      "int x = 'a'; return 0;\n",
      "}\n",
      "#####\n",
      "Key: Another Code\n",
      "Value: import numpy as np\n",
      "### Oh what is this doing here\n",
      "print(\"It can handle so many quotations ' \\\" and backslashes and unexpected curly braces { } You don't even need to match }!\")\n",
      "#####\n",
      "Key: Some characters\n",
      "Value: ~!@#$%^&*()_+-'\"{}[];?><,.\n",
      "#####\n"
     ]
    }
   ],
   "source": [
    "for key, value in end_dict.items():\n",
    "    print('Key:', key)\n",
    "    print('Value:', value)\n",
    "    print('#####')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb98a7fd-f5fb-454e-aa42-70ebc3c01b9b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}