py/promptlib/notebooks/basics_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AI Controller: PromptLib\n",
    "\n",
    "PromptLib is a library for creating AST runner programs for execution by the AICI AST VM. In addition to wrapping AICI AST functionality, PromptLib also adds higher level functionality that it compiles down to AST runner's primitives.\n",
    "\n",
    "In this tutorial notebook, we'll show how to go over the basic usage of PromptLib, including prompting and generation, constrained generation with regexes and context free grammars,  branches, variables, and control flow.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Please see the AICI README to install and run an AICI-enabled LLM service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import promptlib as pl\n",
    "import json\n",
    "\n",
    "AICI_API_BASE=\"\"\n",
    "\n",
    "# PromptLib currently uses the DeclCtrl, the Declarative Controller\n",
    "aici = pl.AICI(base_url=AICI_API_BASE, wasm_runner_id=\"declctrl-latest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll begin by demonstrating how to run a simple PromptLib text generation program.\n",
    "\n",
    "We'll create a PromptProgram object, create a prompt by appending static strings to it, and then we'll append instructions to generate text.\n",
    "\n",
    "When we run the program, we'll see the log of running the program, followed by '[DONE]'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please answer the following simple question:\n",
      "Q: What is the meaning of life?\n",
      "A: The meaning of life is to find your purpose.\n",
      "[DONE]\n"
     ]
    }
   ],
   "source": [
    "# Create a new program and set its endpoint to our AICI node\n",
    "program = pl.PromptProgram(aici)\n",
    "\n",
    "# A PromptLib program is built up by appending fixed text, generation\n",
    "# requests, and other commands\n",
    "p = program + \"Please answer the following simple question:\\n\"\n",
    "p += \"Q: What is the meaning of life?\\nA:\"\n",
    "p = p.gen(max_tokens = 20, stop_at=\"\\n\")\n",
    "\n",
    "# Once we've finished building the program, we can run it. \n",
    "result = program.run()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Results are returned in a structured format, so we can inspect the result variable to see the generated text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Please answer the following simple question:\\nQ: What is the meaning of life?\\nA: The meaning of life is to find your purpose.\\n']\n"
     ]
    }
   ],
   "source": [
    "# result['text'] is the full generated text\n",
    "print(result['text'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition to the generated text, the result also includes the original request sent to the server, the detailed response, logs, and other information. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The original request is:\n",
      "\n",
      "{\n",
      "    \"model\": \"\",\n",
      "    \"prompt\": \"\",\n",
      "    \"max_tokens\": 200,\n",
      "    \"n\": 1,\n",
      "    \"temperature\": 0,\n",
      "    \"stream\": true,\n",
      "    \"aici_module\": \"declctrl-latest\",\n",
      "    \"aici_arg\": {\n",
      "        \"steps\": [\n",
      "            {\n",
      "                \"Fixed\": {\n",
      "                    \"text\": {\n",
      "                        \"String\": {\n",
      "                            \"str\": \"Please answer the following simple question:\\n\"\n",
      "                        }\n",
      "                    },\n",
      "                    \"tag\": \"_1\",\n",
      "                    \"following\": null,\n",
      "                    \"label\": \"_1\"\n",
      "                }\n",
      "            },\n",
      "            {\n",
      "                \"Fixed\": {\n",
      "                    \"text\": {\n",
      "                        \"String\": {\n",
      "                            \"str\": \"Q: What is the meaning of life?\\nA:\"\n",
      "                        }\n",
      "                    },\n",
      "                    \"tag\": \"_2\",\n",
      "                    \"following\": null,\n",
      "                    \"label\": \"_2\"\n",
      "                }\n",
      "            },\n",
      "            {\n",
      "                \"Gen\": {\n",
      "                    \"max_tokens\": 20,\n",
      "                    \"stop_at\": \"\\n\",\n",
      "                    \"stmts\": []\n",
      "                }\n",
      "            }\n",
      "        ]\n",
      "    },\n",
      "    \"ignore_eos\": false\n",
      "}\n",
      "\n",
      "\n",
      "The full response is:\n",
      "\n",
      "[\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \"Please answer the following simple question:\\n\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"dfa: 144 bytes\\n[0] Fixed(\\\"Please answer the following simple question:\\\\n\\\", tag:TagName(\\\"_1\\\"), label:_1) [Fixed()] tok:0/inf\\n[1] Fixed(\\\"Q: What is the meaning of life?\\\\nA:\\\", tag:TagName(\\\"_2\\\"), label:_2) [Fixed()] tok:0/inf\\n[2] Gen(stop_at:\\\"\\\\n\\\", max_tokens:20, ) [Gen()] tok:0/20\\n[3] Stop [Stop] tok:0/inf\\nprompt: [1]\\ntokenize_bytes: \\\"Please answer the following simple question:\\\\n\\\" -> [12148, 1234, 278, 1494, 2560, 1139, 29901, 13]\\nfinish: [Fixed()] tok:8/inf \\\"Please answer the following simple question:\\\\n\\\"\\n\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 1,\n",
      "            \"prompt_tokens\": 1,\n",
      "            \"total_tokens\": 2,\n",
      "            \"fuel_tokens\": 3\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \"Q: What is the meaning of life?\\nA:\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"tokenize_bytes: \\\"Q: What is the meaning of life?\\\\nA:\\\" -> [29984, 29901, 1724, 338, 278, 6593, 310, 2834, 29973, 13, 29909, 29901]\\nfinish: [Fixed()] tok:12/inf \\\"Q: What is the meaning of life?\\\\nA:\\\"\\n\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 2,\n",
      "            \"prompt_tokens\": 9,\n",
      "            \"total_tokens\": 11,\n",
      "            \"fuel_tokens\": 13\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" The\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 3,\n",
      "            \"prompt_tokens\": 21,\n",
      "            \"total_tokens\": 24,\n",
      "            \"fuel_tokens\": 27\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" meaning\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 4,\n",
      "            \"prompt_tokens\": 22,\n",
      "            \"total_tokens\": 26,\n",
      "            \"fuel_tokens\": 30\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" of\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 5,\n",
      "            \"prompt_tokens\": 23,\n",
      "            \"total_tokens\": 28,\n",
      "            \"fuel_tokens\": 33\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" life\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 6,\n",
      "            \"prompt_tokens\": 24,\n",
      "            \"total_tokens\": 30,\n",
      "            \"fuel_tokens\": 36\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" is\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 7,\n",
      "            \"prompt_tokens\": 25,\n",
      "            \"total_tokens\": 32,\n",
      "            \"fuel_tokens\": 39\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" to\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 8,\n",
      "            \"prompt_tokens\": 26,\n",
      "            \"total_tokens\": 34,\n",
      "            \"fuel_tokens\": 42\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" find\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 9,\n",
      "            \"prompt_tokens\": 27,\n",
      "            \"total_tokens\": 36,\n",
      "            \"fuel_tokens\": 45\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" your\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 10,\n",
      "            \"prompt_tokens\": 28,\n",
      "            \"total_tokens\": 38,\n",
      "            \"fuel_tokens\": 48\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \" purpose\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 11,\n",
      "            \"prompt_tokens\": 29,\n",
      "            \"total_tokens\": 40,\n",
      "            \"fuel_tokens\": 51\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": null,\n",
      "                \"text\": \".\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 12,\n",
      "            \"prompt_tokens\": 30,\n",
      "            \"total_tokens\": 42,\n",
      "            \"fuel_tokens\": 54\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": \"aici-stop\",\n",
      "                \"text\": \"\\n\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"finish: [Gen()] tok:11/20 \\\" The meaning of life is to find your purpose.\\\\n\\\"\\n\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 13,\n",
      "            \"prompt_tokens\": 31,\n",
      "            \"total_tokens\": 44,\n",
      "            \"fuel_tokens\": 57\n",
      "        }\n",
      "    },\n",
      "    {\n",
      "        \"object\": \"text_completion\",\n",
      "        \"id\": \"cmpl-774c4b94-780c-4270-ab93-fd1f69249ded\",\n",
      "        \"model\": \"codellama/CodeLlama-13b-Instruct-hf\",\n",
      "        \"created\": 1706608472,\n",
      "        \"choices\": [\n",
      "            {\n",
      "                \"index\": 0,\n",
      "                \"finish_reason\": \"aici-stop\",\n",
      "                \"text\": \"\",\n",
      "                \"error\": \"\",\n",
      "                \"logs\": \"\",\n",
      "                \"storage\": []\n",
      "            }\n",
      "        ],\n",
      "        \"usage\": {\n",
      "            \"completion_tokens\": 13,\n",
      "            \"prompt_tokens\": 31,\n",
      "            \"total_tokens\": 44,\n",
      "            \"fuel_tokens\": 57\n",
      "        }\n",
      "    }\n",
      "]\n",
      "\n",
      "\n",
      "Other sections of information in the response include: dict_keys(['request', 'response', 'text', 'logs', 'raw_storage', 'error', 'usage', 'storage'])\n"
     ]
    }
   ],
   "source": [
    "print(\"The original request is:\\n\")\n",
    "print(json.dumps(result['request'], indent=4))\n",
    "\n",
    "print(\"\\n\\nThe full response is:\\n\")\n",
    "print(json.dumps(result['response'], indent=4))\n",
    "\n",
    "print(\"\\n\\nOther sections of information in the response include: \" + str(result.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of generating arbitrary text, we can also restrict the generation in various ways.  The simplest is to force the language model to choose between two or more options.  The model will then choose the most likely of the options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please answer the following simple question:\n",
      "Q: What is the best television program ever made?\n",
      "She-Ra[DONE]\n"
     ]
    }
   ],
   "source": [
    "program = pl.PromptProgram(aici)\n",
    "\n",
    "p = program + \"Please answer the following simple question:\\n\"\n",
    "p += \"Q: What is the best television program ever made?\\n\"\n",
    "\n",
    "# Force the LLM to choose between the listed options\n",
    "p = p.choose([\"MacGyver\", \"Star Trek\", \"She-Ra\"])\n",
    "\n",
    "result = program.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can save the selected choice to a variable.  In fact, we can store the text generated by any gen() command to a variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please answer the following simple question:\n",
      "Q: What is the best television program ever made?\n",
      "She-Ra[DONE]\n",
      "The answer to the question was: She-Ra\n"
     ]
    }
   ],
   "source": [
    "program = pl.PromptProgram(aici)\n",
    "\n",
    "p = program + \"Please answer the following simple question:\\n\"\n",
    "p += \"Q: What is the best television program ever made?\\n\"\n",
    "\n",
    "# We'll add a second argument that will copy the chosen text into a variable for easy access\n",
    "p = p.choose([\"MacGyver\", \"Star Trek\", \"She-Ra\"], set_var=\"best_tv_show\")\n",
    "\n",
    "result = program.run()\n",
    "\n",
    "# Results are returned in a structured representation\n",
    "print(\"The answer to the question was: \" + result['storage']['best_tv_show'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constrained Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PromptLib supports several ways to constrain text generation, using regular expressions and context free grammars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please guess a number: '10'[DONE]\n",
      "The LLM guessed '10'\n"
     ]
    }
   ],
   "source": [
    "# fixed generation with regex constraints (and see more detailed regex notebook)\n",
    "\n",
    "program = pl.PromptProgram(aici)\n",
    "\n",
    "p = program + \"Please guess a number: \"\n",
    "\n",
    "# pass in a regular expression (rx) that forces the LLM to generate a number\n",
    "p = p.gen(rx = r\"'\\d+'\", set_var=\"number\")\n",
    "\n",
    "result = program.run()\n",
    "\n",
    "print(\"The LLM guessed \" + result['storage']['number'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can combine fixed text with choices and regexes to generate text that matches a specific pattern, like JSON."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please summarize and extract key factors from the following text:\n",
      "<text>\n",
      "The 3 astronauts who flew to the International Space Station reached their destination today (May 31). The\n",
      "crew successfully docked at 10:16 a.m. EDT, about 19 hours after blasting off from\n",
      "NASA's Kennedy Space Center. The docking went smoothly, with the Crew Dragon's soft capture docking\n",
      "mechanism attaching automatically at 10:16 a.m. EDT.</text>\n",
      "Return as JSON:{shorttitle: 'Astronauts Reach Space Station', who: '3 astronauts', what: 'reached their destination', when: 'today (May 31)', where: 'International Space Station', excitement_level: '10', quality_level: 'high'}\n",
      "[DONE]\n",
      "{shorttitle: 'Astronauts Reach Space Station', who: '3 astronauts', what: 'reached their destination', when: 'today (May 31)', where: 'International Space Station', excitement_level: '10', quality_level: 'high'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "program = pl.PromptProgram(aici)\n",
    "\n",
    "note = \"\"\"\n",
    "The 3 astronauts who flew to the International Space Station reached their destination today (May 31). The\n",
    "crew successfully docked at 10:16 a.m. EDT, about 19 hours after blasting off from\n",
    "NASA's Kennedy Space Center. The docking went smoothly, with the Crew Dragon's soft capture docking\n",
    "mechanism attaching automatically at 10:16 a.m. EDT.\"\"\"\n",
    "\n",
    "p = program + \"Please summarize and extract key factors from the following text:\\n\"\n",
    "p += \"<text>\" + note + \"</text>\\n\"\n",
    "p += \"Return as JSON:\"\n",
    "p += \"{shorttitle: '\"\n",
    "p = p.gen(set_var=\"title\") + \"', who: '\"\n",
    "p = p.gen(set_var=\"who\") + \"', what: '\"\n",
    "p = p.gen(set_var=\"what\") + \"', when: '\"\n",
    "p = p.gen(set_var=\"when\") + \"', where: '\"\n",
    "p = p.gen(set_var=\"where\") + \"', excitement_level: '\"\n",
    "p = p.gen(rx = r\"\\d+\", set_var=\"excitement_level\") + \"', quality_level: '\"\n",
    "p = p.choose([\"high\", \"medium\", \"low\"], set_var=\"quality_level\") + \"'}\\n\"\n",
    "\n",
    "result = program.run()\n",
    "\n",
    "# print everything in generated text after the 'JSON:'\n",
    "print(result['text'][0].split('JSON:')[1])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a simple constrained generation example using a context-free grammar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please write a C program that recursively calculates the Fibonnaci sequence to N digits.\n",
      "#include<stdio.h>\n",
      "int fib(int n)\n",
      "{\n",
      "    if(n==0)\n",
      "        return 0;\n",
      "    else if(n==1)\n",
      "        return 1;\n",
      "    else\n",
      "        return fib(n-1)+fib(n-2);\n",
      "}\n",
      "int main()\n",
      "{\n",
      "    int n;\n",
      "    printf(\"Enter the number of digits: \");\n",
      "    scanf(\"%d\",&n);\n",
      "    printf(\"The Fibonnaci sequence is: %d\",fib(n));\n",
      "    return 0;\n",
      "}\n",
      "\n",
      "[DONE]\n",
      "int fib(int n)\n",
      "{\n",
      "    if(n==0)\n",
      "        return 0;\n",
      "    else if(n==1)\n",
      "        return 1;\n",
      "    else\n",
      "        return fib(n-1)+fib(n-2);\n",
      "}\n",
      "int main()\n",
      "{\n",
      "    int n;\n",
      "    printf(\"Enter the number of digits: \");\n",
      "    scanf(\"%d\",&n);\n",
      "    printf(\"The Fibonnaci sequence is: %d\",fib(n));\n",
      "    return 0;\n",
      "}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "program = pl.PromptProgram(aici)\n",
    "\n",
    "# load a grammar file\n",
    "c_grammar = open(\"aici_abi/grammars/c.y\", \"r\").read()\n",
    "\n",
    "p = program + \"Please write a C program that recursively calculates the Fibonnaci sequence to N digits.\\n\"\n",
    "p += \"#include<stdio.h>\\n\"\n",
    "p = p.gen(max_tokens = 200, yacc=c_grammar, set_var=\"code\", stop_at='\\n\\n')\n",
    "\n",
    "result = program.run()\n",
    "\n",
    "print(result['storage']['code'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For more examples of regular expressions and context free grammars see the detailed constraints tutorials."
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}