cookbooks/0-rhubarb-cookbook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is Rhubarb?\n",
    "---\n",
    "\n",
    "Rhubarb is a light-weight Python framework that makes it easy to build document understanding applications using Multi-modal Large Language Models (LLMs). Rhubarb is created from the ground up to work with Amazon Bedrock and supports multiple foundation models including Anthropic Claude Sonnet, Haiku, and Opus models, as well as Amazon Nova Pro and Nova Lite models for document understanding and analysis.\n",
    "\n",
    "\n",
    "## What can I do with Rhubarb?\n",
    "---\n",
    "\n",
    "Rhubarb can do multiple document processing tasks such as\n",
    "\n",
    "- ✅ Document Q&A\n",
    "- ✅ Streaming chat with documents (Q&A)\n",
    "- ✅ Document Summarization\n",
    "  - 🚀 Page level summaries\n",
    "  - 🚀 Full summaries\n",
    "  - 🚀 Summaries of specific pages\n",
    "  - 🚀 Streaming Summaries\n",
    "- ✅ Extraction based on a JSON schema\n",
    "  - 🚀 Key-value extractions\n",
    "  - 🚀 Table extractions\n",
    "- ✅ Named entity recognition (NER) \n",
    "  - 🚀 With 50 built-in common entities\n",
    "- ✅ PII recognition with built-in entities\n",
    "- ✅ Figure and image understanding from documents\n",
    "- ✅ Document classification with Multi-modal Language models\n",
    "- ✅ Document classification with vector sampling using Multi-modal embedding models\n",
    "\n",
    "Rhubarb comes with built-in system prompts that makes it easy to use it for a number of different document understanding use-cases. You can customize Rhubarb by passing in your own system prompts. It supports exact JSON schema based output generation which makes it easy to integrate into downstream applications.\n",
    "\n",
    "- Supports PDF, TIFF, DOCX, PNG, JPG files\n",
    "- Performs document to image conversion internally to work with the multi-modal models\n",
    "- Works on local files or files stored in S3\n",
    "- Supports specifying page numbers for multi-page documents\n",
    "- Supports chat-history based chat for documents\n",
    "- Supports streaming and non-streaming mode\n",
    "- Supports Converse API \n",
    "- Supports Cross-Region Inference\n",
    "\n",
    "## How do I use Rhubarb?\n",
    "---\n",
    "\n",
    "Start by installing Rhubarb using `pip`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m pip install pyrhubarb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize Boto3 Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "session = boto3.Session(profile_name=\"anjanavb+demo1-Admin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Usage - Q&A with local file\n",
    "---\n",
    "\n",
    "Initiaize `DocAnalysis` with a local file and `bedrock` boto3 client and call the `run` method to get response back. In it's default form, it uses a default system prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Martha C Rivera'}],\n",
       " 'token_usage': {'input_tokens': 5093,\n",
       "  'output_tokens': 45,\n",
       "  'total_tokens': 5138}}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,)\n",
    "resp = da.run(message=\"What is the employee's name?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basic Usage - Q&A with S3 file\n",
    "---\n",
    "\n",
    "Initiaize `DocAnalysis` with a file in S3, and boto3 session and call the `run` method to get response back. In it's default form, it uses a default system prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Martha C Rivera'},\n",
       "  {'page': 2,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Answer not found'},\n",
       "  {'page': 3,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Answer not found'}],\n",
       " 'token_usage': {'input_tokens': 5093,\n",
       "  'output_tokens': 113,\n",
       "  'total_tokens': 5206}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"s3://<your-bucket>/<prefix>/employee_enrollment.pdf\", \n",
    "                 boto3_session=session)\n",
    "resp = da.run(message=\"What is the employee's name?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Default Model\n",
    "---\n",
    "By default Rhubarb uses Claude Sonnet model, however you can also use Haiku, Sonnet 3.5 or Opus (when available)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': \"The employee's name is Martha.\",\n",
       " 'token_usage': {'input_tokens': 5093,\n",
       "  'output_tokens': 10,\n",
       "  'total_tokens': 5103}}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, LanguageModels\n",
    "\n",
    "da = DocAnalysis(file_path=\"s3://<your-bucket>/<prefix>/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 modelId=LanguageModels.CLAUDE_HAIKU_V1)\n",
    "resp = da.run(message=\"What is the employee's name?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Martha C Rivera'}],\n",
       " 'token_usage': {'input_tokens': 5092,\n",
       "  'output_tokens': 51,\n",
       "  'total_tokens': 5143}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, LanguageModels\n",
    "\n",
    "da = DocAnalysis(file_path=\"s3://<your-bucket>/<prefix>/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 enable_cri=True,\n",
    "                 modelId=LanguageModels.CLAUDE_OPUS_V1)\n",
    "resp = da.run(message=\"What is the employee name?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Converse API\n",
    "---\n",
    "\n",
    "Rhubarb supports streaming responses using the Bedrock converse API. This feature enables real-time streaming of responses as they are generated. Here's how to use it:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "import boto3\n",
    "\n",
    "# Initialize a boto3 session\n",
    "session = boto3.Session()\n",
    "\n",
    "# Create a DocAnalysis instance with converse API enabled\n",
    "da = DocAnalysis(\n",
    "    file_path=\"./test_docs/employee_enrollment.pdf\",\n",
    "    boto3_session=session,\n",
    "    use_converse_api=True,  # Enable converse API\n",
    "    system_prompt=SystemPrompts().SummarySysPrompt\n",
    ")\n",
    "\n",
    "# Stream the response\n",
    "for resp in da.run_stream(message=\"Give me a brief summary of this document.\"):\n",
    "    if isinstance(resp, str):\n",
    "        print(resp, end='')\n",
    "    else:\n",
    "        print(\"\\n\")\n",
    "        print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Cross-Region Inference\n",
    "\n",
    "Rhubarb supports cross-region inference capabilities, allowing you to process documents using models deployed in different AWS regions. This feature can help optimize latency and provide regional failover support. Here's how to enable it:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "import boto3\n",
    "\n",
    "# Initialize a boto3 session\n",
    "session = boto3.Session()\n",
    "\n",
    "# Create a DocAnalysis instance with cross-region inference enabled\n",
    "da = DocAnalysis(\n",
    "    file_path=\"./test_docs/employee_enrollment.pdf\",\n",
    "    boto3_session=session,\n",
    "    enable_cri=True,  # Enable cross-region inference\n",
    "    system_prompt=SystemPrompts().SummarySysPrompt\n",
    ")\n",
    "\n",
    "# Run document analysis\n",
    "resp = da.run(message=\"Give me a brief summary of this document.\")\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Q&A with specific pages\n",
    "---\n",
    "\n",
    "Initiaize `DocAnalysis` with a file and page numbers, and boto3 session and call the `run` method to get response back. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 3,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Pat Rivera'}],\n",
       " 'token_usage': {'input_tokens': 2004,\n",
       "  'output_tokens': 44,\n",
       "  'total_tokens': 2048}}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 pages=[3])\n",
    "resp = da.run(message=\"For beneficiary type 'Secondary', what is the full name?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or specify multiple pages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'Martha C Rivera'},\n",
       "  {'page': 3,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': \"Employee's name: Martha C Rivera, Spouse's name: Mateo Rivera\"}],\n",
       " 'token_usage': {'input_tokens': 3552,\n",
       "  'output_tokens': 93,\n",
       "  'total_tokens': 3645}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 pages=[1,3])\n",
    "resp = da.run(message=\"What is the employee's name and what is the spouse's name?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Document Classification\n",
    "---\n",
    "You can classify the pages of a document using Rhubarb by using either the `ClassificationSysPrompt` system prompt for single class classification or `MultiClassificationSysPrompt` system prompt for multi-class classification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1, 'class': 'BANK_STATEMENT'},\n",
       "  {'page': 2, 'class': 'RECEIPT'},\n",
       "  {'page': 3, 'class': 'DRIVERS_LICENSE'},\n",
       "  {'page': 4, 'class': 'INSURANCE_ID'},\n",
       "  {'page': 5, 'class': 'W2'},\n",
       "  {'page': 6, 'class': 'MOM'}],\n",
       " 'token_usage': {'input_tokens': 8803, 'output_tokens': 156}}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/Sample1.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().ClassificationSysPrompt)\n",
    "resp = da.run(message=\"\"\"Given the document, classify the pages into the following classes\n",
    "                        <classes>\n",
    "                        DRIVERS_LICENSE  # a driver's license\n",
    "                        INSURANCE_ID     # a medical insurance ID card\n",
    "                        RECEIPT          # a store receipt\n",
    "                        BANK_STATEMENT   # a bank statement\n",
    "                        W2               # a W2 tax document\n",
    "                        MOM              # a minutes of meeting or meeting notes\n",
    "                        </classes>\"\"\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or multi-class classification. Note that in Multi-class classification it is helpful to clarify the hierarchy of classes to the model in two different list of classes. This should typically match with your document taxonomy such as\n",
    "\n",
    "```\n",
    "FINANCIAL           (Level-2)\n",
    "├── BANK_STATEMENT  (Level-1 leaf)\n",
    "└── W2              (Level-1 leaf)\n",
    "\n",
    "IDENTIFICATION      (Level-2)\n",
    "├── DRIVERS_LICENSE (Level-1 leaf)\n",
    "└── INSURANCE_ID    (Level-1 leaf)\n",
    "```\n",
    "\n",
    "And so on"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1, 'class': ['BANK_STATEMENT', 'FINANCIAL']},\n",
       "  {'page': 2, 'class': ['RECEIPT', 'GENERAL']},\n",
       "  {'page': 3, 'class': ['DRIVERS_LICENSE', 'IDENTIFICATION']},\n",
       "  {'page': 4, 'class': ['INSURANCE_ID', 'IDENTIFICATION']},\n",
       "  {'page': 5, 'class': ['W2', 'FINANCIAL']},\n",
       "  {'page': 6, 'class': ['MOM', 'GENERAL']}],\n",
       " 'token_usage': {'input_tokens': 8925, 'output_tokens': 228}}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/Sample1.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().MultiClassificationSysPrompt)\n",
    "resp = da.run(message=\"\"\"Given the document, classify the pages into the following classes\n",
    "                        <classes_level1>\n",
    "                        DRIVERS_LICENSE  # a driver's license\n",
    "                        INSURANCE_ID     # a medical insurance ID card\n",
    "                        RECEIPT          # a store receipt\n",
    "                        BANK_STATEMENT   # a bank statement\n",
    "                        W2               # a W2 tax document\n",
    "                        MOM              # a minutes of meeting or meeting notes\n",
    "                        <classes_level1>\n",
    "                        <classes_level2>\n",
    "                        FINANCIAL        # a document related to finances of a person\n",
    "                        IDENTIFICATION   # a personal document such as ID, membership cards, etc.\n",
    "                        GENERAL          # any other general document\n",
    "                        </classes_level2>\"\"\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform Named Entity Recognition \n",
    "---\n",
    "Rhubarb comes with 50 built-in entities which includes common entities such as LOCATION, EVENT etc. and PII entities such as NAME, SSN, ADDRESS and so on. Entities are available via the `Entities` class. You can pick and choose which entities to detect and then pass them onto the `run_entity` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'entities': [{'PERSON': 'Martha C Rivera'},\n",
       "    {'ADDRESS': '5005 ANY AVENUE, NEW YORK, NY- 10021'},\n",
       "    {'ADDRESS': '8 Any Plaza, 21 Street, Any City, CA 90210'}]},\n",
       "  {'page': 3,\n",
       "   'entities': [{'PERSON': 'Mateo Rivera'},\n",
       "    {'PERSON': 'Pat Rivera'},\n",
       "    {'ADDRESS': '8 Any Plaza, 21 Street, Any City, CA 90210'}]}],\n",
       " 'token_usage': {'input_tokens': 3523,\n",
       "  'output_tokens': 196,\n",
       "  'total_tokens': 3719}}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, Entities\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 pages=[1,3])\n",
    "resp = da.run_entity(message=\"Extract all the specified entities from this document.\", \n",
    "                     entities=[Entities.PERSON, Entities.ADDRESS])\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform PII Recognition \n",
    "---\n",
    "You can use the same `run_entity` method with PII entities available via `Entities`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'entities': [{'SSN': '376 12 1987'},\n",
       "    {'ADDRESS': '8 Any Plaza, 21 Street'}]},\n",
       "  {'page': 3,\n",
       "   'entities': [{'SSN': '791 36 9771'},\n",
       "    {'ADDRESS': '8 Any Plaza, 21 Street'},\n",
       "    {'SSN': '824 26 2211'},\n",
       "    {'ADDRESS': '8 Any Plaza, 21 Street'}]}],\n",
       " 'token_usage': {'input_tokens': 3534, 'output_tokens': 183}}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, Entities\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 pages=[1,3])\n",
    "resp = da.run_entity(message=\"Extract all the specified entities from this document.\", \n",
    "                     entities=[Entities.SSN, Entities.ADDRESS])\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform key-value extraction using custom JSON schema\n",
    "---\n",
    "Rhubarb supports extraction of key values using JSON Schema. You can pass in a valid JSON schema to extract specific data out of your document. Let's define a custom JSON schema appropriate for our document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"employee_name\": {\n",
    "            \"description\": \"Employee's Name\",\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"employee_ssn\": {\n",
    "            \"description\": \"Employee's social security number\",\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"employee_address\": {\n",
    "            \"description\": \"Employee's mailing address\",\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"employee_dob\": {\n",
    "            \"description\": \"Employee's date of birth\",\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"employee_gender\": {\n",
    "            \"description\": \"Employee's gender\",\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"male\":{\n",
    "                    \"description\": \"Whether the employee gender is Male\",\n",
    "                    \"type\": \"boolean\"\n",
    "                },\n",
    "                \"female\":{\n",
    "                    \"description\": \"Whether the employee gender is Female\",\n",
    "                    \"type\": \"boolean\"\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"male\", \"female\"]\n",
    "        },\n",
    "        \"employee_hire_date\": {\n",
    "            \"description\": \"Employee's hire date\",\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"employer_no\": {\n",
    "            \"description\": \"Employer number\",\n",
    "            \"type\": \"string\"\n",
    "        },\n",
    "        \"employment_status\": {\n",
    "            \"type\": \"object\",\n",
    "            \"description\": \"Employment status\",\n",
    "            \"properties\": {\n",
    "                \"full_time\":{\n",
    "                    \"description\": \"Whether employee is full-time\",\n",
    "                    \"type\": \"boolean\"\n",
    "                },\n",
    "                \"part_time\": {\n",
    "                    \"description\": \"Whether employee is part-time\",\n",
    "                    \"type\": \"boolean\"\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"full_time\", \"part_time\"]\n",
    "        },\n",
    "        \"employee_salary_rate\":{\n",
    "            \"description\": \"The dollar value of employee's salary\",\n",
    "            \"type\": \"integer\"\n",
    "        },\n",
    "        \"employee_salary_frequency\":{\n",
    "            \"type\": \"object\",\n",
    "            \"description\": \"Salary rate of the employee\",\n",
    "            \"properties\": {\n",
    "                \"annual\":{\n",
    "                    \"description\": \"Whether salary rate is monthly\",\n",
    "                    \"type\": \"boolean\"\n",
    "                },\n",
    "                \"monthly\": {\n",
    "                    \"description\": \"Whether salary rate is monthly\",\n",
    "                    \"type\": \"boolean\"\n",
    "                },\n",
    "                \"semi_monthly\": {\n",
    "                    \"description\": \"Whether salary rate is semi_monthly\",\n",
    "                    \"type\": \"boolean\"\n",
    "                },\n",
    "                \"bi_weekly\": {\n",
    "                    \"description\": \"Whether salary rate is bi_weekly\",\n",
    "                    \"type\": \"boolean\"\n",
    "                },\n",
    "                \"weekly\": {\n",
    "                    \"description\": \"Whether salary rate is weekly\",\n",
    "                    \"type\": \"boolean\"\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"annual\", \"monthly\", \"semi_monthly\",\"bi_weekly\",\"weekly\"]\n",
    "        }\n",
    "    },\n",
    "    \"required\": [\"employee_name\",\"employee_hire_date\", \"employer_no\", \"employment_status\"]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': {'employee_name': 'Martha C Rivera',\n",
       "  'employee_ssn': '376 12 1987',\n",
       "  'employee_address': '8 Any Plaza, 21 Street, Any City, CA 90210',\n",
       "  'employee_dob': '09/19/80',\n",
       "  'employee_gender': {'male': False, 'female': True},\n",
       "  'employee_hire_date': '07/19/2023',\n",
       "  'employer_no': '784371',\n",
       "  'employment_status': {'full_time': True, 'part_time': False},\n",
       "  'employee_salary_rate': 79930,\n",
       "  'employee_salary_frequency': {'annual': True,\n",
       "   'monthly': False,\n",
       "   'semi_monthly': False,\n",
       "   'bi_weekly': False,\n",
       "   'weekly': False}},\n",
       " 'token_usage': {'input_tokens': 5375,\n",
       "  'output_tokens': 224,\n",
       "  'total_tokens': 5599}}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session)\n",
    "resp = da.run(message=\"Give me the output based on the provided schema.\", output_schema=schema)\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform table extraction using custom JSON schema\n",
    "---\n",
    "You can also perform table extraction using custom JSON schema. In this case we will use a rather complex table from an AMZN 10-k filing document and attempt to extract the data from it. Here's what a JSON schema might look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_schema = {\n",
    "  \"additionalProperties\": {\n",
    "    \"type\": \"object\",\n",
    "    \"patternProperties\": {\n",
    "      \"^(2022|2023)$\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "          \"Net Sales\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "              \"North America\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"International\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"AWS\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"Consolidated\": {\n",
    "                \"type\": \"number\"\n",
    "              }\n",
    "            },\n",
    "            \"required\": [\"North America\", \"International\", \"AWS\", \"Consolidated\"]\n",
    "          },\n",
    "          \"Year-over-year Percentage Growth (Decline)\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "              \"North America\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"International\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"AWS\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"Consolidated\": {\n",
    "                \"type\": \"number\"\n",
    "              }\n",
    "            },\n",
    "            \"required\": [\"North America\", \"International\", \"AWS\", \"Consolidated\"]\n",
    "          },\n",
    "          \"Year-over-year Percentage Growth, excluding the effect of foreign exchange rates\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "              \"North America\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"International\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"AWS\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"Consolidated\": {\n",
    "                \"type\": \"number\"\n",
    "              }\n",
    "            },\n",
    "            \"required\": [\"North America\", \"International\", \"AWS\", \"Consolidated\"]\n",
    "          },\n",
    "          \"Net Sales Mix\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "              \"North America\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"International\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"AWS\": {\n",
    "                \"type\": \"number\"\n",
    "              },\n",
    "              \"Consolidated\": {\n",
    "                \"type\": \"number\"\n",
    "              }\n",
    "            },\n",
    "            \"required\": [\"North America\", \"International\", \"AWS\", \"Consolidated\"]\n",
    "          }\n",
    "        },\n",
    "        \"required\": [\"Net Sales\", \"Year-over-year Percentage Growth (Decline)\", \"Year-over-year Percentage Growth, excluding the effect of foreign exchange rates\", \"Net Sales Mix\"]\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are only interested in the results of operation which we know is in the first page, we will call Rhubarb with just the first page in this case to save costs. However, in situations where the table's exact location isn't known, the full document can be passed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': {'2022': {'Net Sales': {'North America': 315880,\n",
       "    'International': 118007,\n",
       "    'AWS': 80096,\n",
       "    'Consolidated': 513983},\n",
       "   'Year-over-year Percentage Growth (Decline)': {'North America': 13,\n",
       "    'International': -8,\n",
       "    'AWS': 29,\n",
       "    'Consolidated': 9},\n",
       "   'Year-over-year Percentage Growth, excluding the effect of foreign exchange rates': {'North America': 13,\n",
       "    'International': 4,\n",
       "    'AWS': 29,\n",
       "    'Consolidated': 13},\n",
       "   'Net Sales Mix': {'North America': 61,\n",
       "    'International': 23,\n",
       "    'AWS': 16,\n",
       "    'Consolidated': 100}},\n",
       "  '2023': {'Net Sales': {'North America': 352828,\n",
       "    'International': 131200,\n",
       "    'AWS': 90757,\n",
       "    'Consolidated': 574785},\n",
       "   'Year-over-year Percentage Growth (Decline)': {'North America': 12,\n",
       "    'International': 11,\n",
       "    'AWS': 13,\n",
       "    'Consolidated': 12},\n",
       "   'Year-over-year Percentage Growth, excluding the effect of foreign exchange rates': {'North America': 12,\n",
       "    'International': 11,\n",
       "    'AWS': 13,\n",
       "    'Consolidated': 12},\n",
       "   'Net Sales Mix': {'North America': 61,\n",
       "    'International': 23,\n",
       "    'AWS': 16,\n",
       "    'Consolidated': 100}}},\n",
       " 'token_usage': {'input_tokens': 2177,\n",
       "  'output_tokens': 433,\n",
       "  'total_tokens': 2610}}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/amzn-10k.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 pages=[1])\n",
    "resp = da.run(message=\"Give me data in the results of operation table from this 10-K SEC filing document. Use the schema provided.\", \n",
    "              output_schema=table_schema)\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Schema creation assistant\n",
    "---\n",
    "Rhubarb can also help create accurate JSON schemas from plain text prompts. You can provide a document and ask it to extract certain values from the document and it will respond back with a JSON schema. You can then use the JSON schema with the `output_schema` as shown above, or you can tweak and modify it to fit your need further. You do this using the `generate_schema` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'type': 'object',\n",
       " 'description': 'Employee enrollment form information',\n",
       " 'properties': {'employee_name': {'type': 'object',\n",
       "   'description': \"The employee's full name\",\n",
       "   'properties': {'first': {'type': 'string',\n",
       "     'description': \"Employee's first name\"},\n",
       "    'initial': {'type': 'string', 'description': \"Employee's middle initial\"},\n",
       "    'last': {'type': 'string', 'description': \"Employee's last name\"}},\n",
       "   'required': ['first', 'last']},\n",
       "  'ssn': {'type': 'string',\n",
       "   'description': \"Employee's Social Security Number\"},\n",
       "  'address': {'type': 'object',\n",
       "   'description': \"Employee's mailing address\",\n",
       "   'properties': {'street': {'type': 'string',\n",
       "     'description': 'Street address'},\n",
       "    'city': {'type': 'string', 'description': 'City'},\n",
       "    'state': {'type': 'string', 'description': 'State'},\n",
       "    'zip_code': {'type': 'string', 'description': 'ZIP code'}},\n",
       "   'required': ['street', 'city', 'state', 'zip_code']},\n",
       "  'date_of_birth': {'type': 'string',\n",
       "   'description': \"Employee's date of birth in MM/DD/YY format\"},\n",
       "  'phone_number': {'type': 'string',\n",
       "   'description': \"Employee's home phone number\"}},\n",
       " 'required': ['employee_name',\n",
       "  'ssn',\n",
       "  'address',\n",
       "  'date_of_birth',\n",
       "  'phone_number']}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 pages=[1])\n",
    "resp = da.generate_schema(message=\"I want to extract the employee name, employee SSN, employee address, date of birth and phone number from this document.\")\n",
    "resp['output']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then use this schema to perform extraction on the same document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': {'employee_name': {'first': 'Martha',\n",
       "   'initial': 'C',\n",
       "   'last': 'Rivera'},\n",
       "  'ssn': '376 12 1987',\n",
       "  'address': {'street': '8 Any Plaza, 21 Street',\n",
       "   'city': 'Any City',\n",
       "   'state': 'CA',\n",
       "   'zip_code': '90210'},\n",
       "  'date_of_birth': '09/19/80',\n",
       "  'phone_number': '(888) 555-0100'},\n",
       " 'token_usage': {'input_tokens': 2113,\n",
       "  'output_tokens': 145,\n",
       "  'total_tokens': 2258}}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_schema = resp['output']\n",
    "resp = da.run(message=\"I want to extract the employee name, employee SSN, employee address, date of birth and phone number from this document. Use the schema provided.\", \n",
    "              output_schema=output_schema)\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Schema creation assistance with question rephrase\n",
    "---\n",
    "In many cases you may want to quickly get started with creating a JSON Schema for your document wihtout spending too much time crafting a proper prompt for the document. For example, in a birth certificate you could be vague in asking a question such as \"_I want to get the child's, the mother's and father's details from the given document_\". In such cases Rhubarb can help rephrasing the question and create an appropriate rephrased question based on the document and generate a subsequent schema for it which can directly be used to extract the data. For this, you use the `assistive_rephrase` parameter in your call to `generate_schema` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'rephrased_question': \"Extract the child's, mother's, and father's details from the given birth certificate document.\",\n",
       " 'output_schema': {'type': 'object',\n",
       "  'properties': {'child': {'type': 'object',\n",
       "    'properties': {'name': {'type': 'object',\n",
       "      'properties': {'first': {'type': 'string'},\n",
       "       'middle': {'type': 'string'},\n",
       "       'last': {'type': 'string'}},\n",
       "      'required': ['first', 'middle', 'last']},\n",
       "     'sex': {'type': 'string'},\n",
       "     'date_of_birth': {'type': 'string'},\n",
       "     'time_of_birth': {'type': 'string'},\n",
       "     'place_of_birth': {'type': 'object',\n",
       "      'properties': {'hospital': {'type': 'string'},\n",
       "       'city': {'type': 'string'},\n",
       "       'county': {'type': 'string'}},\n",
       "      'required': ['hospital', 'city', 'county']}},\n",
       "    'required': ['name',\n",
       "     'sex',\n",
       "     'date_of_birth',\n",
       "     'time_of_birth',\n",
       "     'place_of_birth']},\n",
       "   'mother': {'type': 'object',\n",
       "    'properties': {'name': {'type': 'object',\n",
       "      'properties': {'first': {'type': 'string'},\n",
       "       'middle': {'type': 'string'},\n",
       "       'last': {'type': 'string'}},\n",
       "      'required': ['first', 'middle', 'last']},\n",
       "     'age': {'type': 'integer'},\n",
       "     'birthplace': {'type': 'string'},\n",
       "     'residence': {'type': 'object',\n",
       "      'properties': {'state': {'type': 'string'},\n",
       "       'county': {'type': 'string'},\n",
       "       'city': {'type': 'string'},\n",
       "       'address': {'type': 'string'}},\n",
       "      'required': ['state', 'county', 'city', 'address']}},\n",
       "    'required': ['name', 'age', 'birthplace', 'residence']},\n",
       "   'father': {'type': 'object',\n",
       "    'properties': {'name': {'type': 'object',\n",
       "      'properties': {'first': {'type': 'string'},\n",
       "       'middle': {'type': 'string'},\n",
       "       'last': {'type': 'string'}},\n",
       "      'required': ['first', 'middle', 'last']},\n",
       "     'age': {'type': 'integer'},\n",
       "     'birthplace': {'type': 'string'}},\n",
       "    'required': ['name', 'age', 'birthplace']}},\n",
       "  'required': ['child', 'mother', 'father']}}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/birth_cert.jpeg\",\n",
    "                 boto3_session=session)\n",
    "resp = da.generate_schema(message=\"I want to get the child's, the mother's and father's details from the given document\",\n",
    "                          assistive_rephrase=True)\n",
    "resp['output']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': {'child': {'name': {'first': 'PAULO',\n",
       "    'middle': 'SOUZA',\n",
       "    'last': 'SANTOS'},\n",
       "   'sex': 'MALE',\n",
       "   'date_of_birth': 'MARCH 23,1981',\n",
       "   'time_of_birth': '7:52A',\n",
       "   'place_of_birth': {'hospital': 'ANYGOVERNMENT MEMORIAL HOSPITAL',\n",
       "    'city': 'ANY TOWN',\n",
       "    'county': 'ANY COUNTY'}},\n",
       "  'mother': {'name': {'first': 'MARIA',\n",
       "    'middle': 'OLIVERIA',\n",
       "    'last': 'GARCIA'},\n",
       "   'age': 29,\n",
       "   'birthplace': 'SWITZERLAND',\n",
       "   'residence': {'state': 'FLORIDA',\n",
       "    'county': 'ANY COUNTY',\n",
       "    'city': 'ANY TOWN',\n",
       "    'address': '543 ANYSTREET DR'}},\n",
       "  'father': {'name': {'first': 'DIEGO', 'middle': '', 'last': 'RAMIREZ'},\n",
       "   'age': 31,\n",
       "   'birthplace': 'ILLINOIS'}},\n",
       " 'token_usage': {'input_tokens': 2301,\n",
       "  'output_tokens': 325,\n",
       "  'total_tokens': 2626}}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_schema = resp['output']['output_schema']\n",
    "question = resp['output']['rephrased_question']\n",
    "resp = da.run(message = question,\n",
    "              output_schema = output_schema)\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform page level summarization\n",
    "---\n",
    "Rhubarb can generate sumarries of every page in the document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': \"This page contains an employee enrollment form for a 401(k) plan with designated Roth contributions from Anycompany of America Life Insurance Co. It includes employer information, employee details such as name (Martha C Rivera), address, social security number, salary, and employment status. The employee's hire date is 07/19/2023, with a full-time status and an annual salary of $79,930. The form also includes sections for prior tax-exempt service and contribution details.\"},\n",
       "  {'page': 2,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': 'This page shows Section 2 - Allocation of Contributions for the 401(k) plan. It provides options for allocating contributions between the Interest Accumulation Account and various Separate Account Investment Funds. The employee has allocated 50% to the AnyCompany of America Interest Accumulation Account. The page also includes Section 3 - Beneficiary Designations, which explains the rules for naming beneficiaries for married and unmarried participants.'},\n",
       "  {'page': 3,\n",
       "   'detected_languages': ['English'],\n",
       "   'content': \"This page contains beneficiary designation forms. The primary beneficiary is listed as Mateo Rivera (spouse) with a 50% benefit, and the secondary beneficiary is Pat Rivera (child) also with a 50% benefit. The employee, Martha, is indicated as married. The page includes a Spouse's Waiver section, signed by Mateo Rivera on 8/25/2022. The bottom of the page shows Martha's signature dated 8/25/2022 in the Statement and Signature section.\"}],\n",
       " 'token_usage': {'input_tokens': 5095,\n",
       "  'output_tokens': 408,\n",
       "  'total_tokens': 5503}}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session)\n",
    "resp = da.run(message=\"Give me a brief summary for each page.\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform full summarization\n",
    "---\n",
    "Or you can generate an overall summary of the entire document. In this case, we will override the default System Prompt which breaks down the response per page. Rhubarb comes with a Summary specific System Prompt for the model, available via `SystemPrompts`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': \"This document is an Employee Enrollment Form for a 401(k) Plan with Designated Roth Contributions from Anycompany of America Life Insurance Co. The key details are:\\n\\n1. The employee, Martha C. Rivera, is enrolling in the 401(k) plan.\\n2. She was hired on 07/19/2023 as a full-time employee with an annual salary of $79,930.\\n3. The form includes personal information such as address, Social Security number, and date of birth.\\n4. There's a section for allocation of contributions to various investment funds.\\n5. The beneficiary designation section shows:\\n   - Primary beneficiary: Mateo Rivera (spouse)\\n   - Secondary beneficiary: Pat Rivera (child)\\n6. Both beneficiaries are allocated 50% of the benefits.\\n7. The employee is married, and there's a spouse's waiver section signed by Mateo Rivera.\\n8. The form is signed and dated by Martha on 8/25/2022.\\n\\nThis document serves as an official record of the employee's enrollment in the company's 401(k) plan and their choices regarding contributions and beneficiaries.\",\n",
       " 'token_usage': {'input_tokens': 4725,\n",
       "  'output_tokens': 262,\n",
       "  'total_tokens': 4987}}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().SummarySysPrompt)\n",
    "resp = da.run(message=\"Give me a brief summary of this document.\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform summarization of specific pages\n",
    "---\n",
    "You can also perform summarization of specific pages using the `pages` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': \"## Summary\\n\\nThis document is an employee enrollment form for a 401(k) plan with designated Roth contributions from AnyCompany of America Life Insurance Co. The key details are:\\n\\n- The employee is Martha C. Rivera, hired on 07/19/2023 as a full-time employee with an annual salary of $79,930. \\n- Her employer is AnyCompany Constructions Inc. with the employer number 784371.\\n- Traditional pre-tax contributions are set at 10% of salary effective 08/19/23 and employer matching contributions effective 08/19/23.\\n- Martha has designated her spouse Mateo Rivera as the primary beneficiary at 50% and her child Pat Rivera as the secondary beneficiary at 50%.\\n- Mateo has signed a spouse's waiver allowing Martha to receive the death benefit after his death.\\n- Martha has signed the form on 3/25/2022, indicating she has read the plan materials and wishes to participate in the Thrift Plan.\",\n",
       " 'token_usage': {'input_tokens': 3207, 'output_tokens': 228}}"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().SummarySysPrompt,\n",
    "                 pages=[1,3])\n",
    "resp = da.run(message=\"Give me a brief summary of this document.\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Streaming summaries\n",
    "---\n",
    "In some cases, you may want to stream the summaries for example let's say a real time chat application. You can easily do that using the `run_stream` method. Let's generate the full summary and stream it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|-START-|\n",
      "This document is an Employee Enrollment Form for 401(k) Plans with Designated Roth Contributions from AnyCompany of America Life Insurance Co. It contains the following key information:\n",
      "\n",
      "- Employee details: A female employee named Martha C. Rivera, hired on 07/19/2023 as a full-time employee with an annual salary of $79,930. Her personal information like address, date of birth, and contact details are provided.\n",
      "\n",
      "- Contribution details: The employee is enrolling in the company's 401(k) plan with 10% of salary allocated to Traditional Pre-Tax Contributions effective 08/19/23 and 08/19/23 for employer matching and non-matching contributions respectively. No Designated Roth Contributions are specified.\n",
      "\n",
      "- Investment allocation: 50% of contributions will be placed in the AnyCompany of America Interest Accumulation Account, while the remaining options for separate investment funds like equity, real estate, retirement, balanced, asset allocation, and fixed income funds are listed.\n",
      "\n",
      "- Beneficiary designations: Martha has designated her spouse (Mateo Rivera) as the primary beneficiary at 50% and her son (Pat Rivera) as the secondary beneficiary at 50%. The spouse has signed a waiver allowing this arrangement.\n",
      "\n",
      "- Employee signature: Martha has signed the form on 3/25/2022, indicating she has reviewed the plan materials and finds it suitable for her financial needs.\n",
      "|-END-|\n",
      "\n",
      "{'input_tokens': 4755, 'output_tokens': 317}\n"
     ]
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().SummarySysPrompt)\n",
    "for resp in da.run_stream(message=\"Give me a brief summary of this document.\"):\n",
    "    if isinstance(resp, str):\n",
    "        print(resp,end='')\n",
    "    else:\n",
    "        print(\"\\n\")\n",
    "        print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### With `converse` API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|-START-|\n",
      "This document is an Employee Enrollment Form for a 401(k) Plan with Designated Roth Contributions from Anycompany of America Life Insurance Co. The key details are:\n",
      "\n",
      "1. The employee, Martha C. Rivera, is enrolling in the 401(k) plan.\n",
      "2. She was hired on 07/19/2023 as a full-time employee with an annual salary of $79,930.\n",
      "3. The form includes personal information such as address, Social Security number, and date of birth.\n",
      "4. There's a section for allocation of contributions, with 50% allocated to the Interest Accumulation Account.\n",
      "5. The beneficiary designation section shows:\n",
      "   - Primary beneficiary: Mateo Rivera (spouse)\n",
      "   - Secondary beneficiary: Pat Rivera (child)\n",
      "6. Both beneficiaries are set to receive 50% of the benefits.\n",
      "7. The employee is married, and there's a spouse's waiver section signed by Mateo Rivera.\n",
      "8. The form is signed and dated by Martha on 8/25/2022.\n",
      "\n",
      "This document represents a standard 401(k) enrollment process, including personal information, contribution allocations, and beneficiary designations.\n",
      "|-END-|\n",
      "\n",
      "{'input_tokens': 4725, 'output_tokens': 267, 'total_tokens': 4992, 'latency': 7771}\n"
     ]
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/employee_enrollment.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 use_converse_api=True,\n",
    "                 system_prompt=SystemPrompts().SummarySysPrompt)\n",
    "for resp in da.run_stream(message=\"Give me a brief summary of this document.\"):\n",
    "    if isinstance(resp, str):\n",
    "        print(resp,end='')\n",
    "    else:\n",
    "        print(\"\\n\")\n",
    "        print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Text streaming starts with a `|-START-|` marker and the end of streaming is marked with an `|-END-|` marker. This is to make sure that the application client recieving the stream has a clear demarcation of when the streaming starts and ends. Currently Rhubarb doesn't support `stop_words` kwargs during model invocation but that is coming soon.\n",
    "\n",
    "One thing to keep in mind is that streaming only makes sense for a couple of use cases\n",
    "- One where you have a real time chat interface where a lot of text (like summary) is expected from the model\n",
    "- A real time conversational chat interface a.k.a. ChatBot\n",
    "\n",
    "As such, Rhubarb supports streaming for only Summary System Prompt and Chat System Prompt as we will see in the next section. You can also view the chat history by accessing the `history` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da.history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversational Chat with Documents (no streaming)\n",
    "---\n",
    "\n",
    "You can chat with your documents with Rhubarb using the Chat System Prompt `ChatSysPrompt` available via `SystemPrompts` class. Here's an example of a (non-streaming) chat. Note, that internally Rhubarb does support chat history, however, much of the history implementation is left to the developer. This coupled with the fact that Claude models only accept 20 images per invocation, makes it a little complicated to properly implement chat history based conversational system. We are working on simplifying this more, and will add support in upcoming releases.\n",
    "\n",
    "Also note that there isn't a lot of fundamental difference between the default invocation using Rhubarb vs. invocation using the `ChatSysPrompt`. The difference is in the response structure. While default is capable of giving you responses to your question on a per page basis, chat gives you the output (i.e. response to your chat) and optionally cites the source (i.e. the page number). So in conclusion, the only difference is the response payload structure.\n",
    "\n",
    "For non streaming chat responses, below is the output structure-\n",
    "\n",
    "```json\n",
    "{\n",
    "    'output': \n",
    "    {\n",
    "        'text': \"Response\", \n",
    "        'sources': [1, 2, 3],  // these are the page numbers\n",
    "        'quotes': ['quote 1', 'quote 2', 'quote 3'] // these are the verbatim quotes from the document\n",
    "    }, \n",
    "    'token_usage': {'input_tokens': 5029, 'output_tokens': 215, 'total_tokens': 5244}\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': {'text': \"This document appears to be a financial report or annual report for a company, likely Amazon, based on the references to business segments like North America, International, and AWS (Amazon Web Services). It provides details on the company's results of operations, net sales, operating income/loss, operating expenses, cost of sales, and fulfillment costs across these segments for the years 2022 and 2023.\",\n",
       "  'sources': [1, 2, 3],\n",
       "  'quotes': ['We have organized our operations into three segments: North America, International, and AWS. These segments reflect the way the Company evaluates its business performance and manages its operations.',\n",
       "   'Operating income (loss) by segment is as follows (in millions):',\n",
       "   'Cost of sales primarily consists of the purchase price of consumer products, inbound and outbound shipping costs, including costs related to sortation and delivery centers and Where we are the transportation service provider, and digital media content costs where we record revenue gross, including video and music.']},\n",
       " 'token_usage': {'input_tokens': 5060, 'output_tokens': 237}}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/amzn-10k.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().ChatSysPrompt)\n",
    "resp = da.run(message=\"What is this document about?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversational Chat with Documents (streaming)\n",
    "---\n",
    "Let's perform a streaming chat. Note `SystemPrompts(streaming=True)`, this tells Rhubarb not to respond back with JSON output since JSON output in streaming text is hard to parse and re-construct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|-START-|\n",
      "Based on the content in these pages, this document appears to be Amazon's annual report or financial statements, specifically discussing the company's results of operations for the fiscal year 2023 compared to 2022.\n",
      "\n",
      "The key points covered include:\n",
      "\n",
      "- Net sales breakdown by segment - North America, International, and AWS (Amazon Web Services) (Page 1)\n",
      "- Operating income/(loss) by segment (Page 2)\n",
      "- Detailed operating expenses like cost of sales, fulfillment, technology/infrastructure, marketing etc. (Page 3)\n",
      "- Explanations for changes in net sales, operating income, and operating expenses across the different segments (Pages 1-3)\n",
      "\n",
      "This report provides an overview of Amazon's financial performance, growth rates, and underlying factors driving the results in its core business segments for the fiscal year.\n",
      "|-END-|\n",
      "\n",
      "{'input_tokens': 4933, 'output_tokens': 178}\n"
     ]
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/amzn-10k.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts(streaming=True).ChatSysPrompt)\n",
    "for resp in da.run_stream(message=\"What is this document about?\"):\n",
    "    if isinstance(resp, str):\n",
    "        print(resp,end='')\n",
    "    else:\n",
    "        print(\"\\n\")\n",
    "        print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reasoning and explaining figures, charts etc.\n",
    "---\n",
    "Rhubarb can also help perform explanation and reasoning on images, charts, graphs within documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'figure_analysis': \"The bar chart in this document is labeled 'Figure 2: Normalized purchases of the same type, as function of the days from the request, in late shopping stages.' It shows three different categories: 'Same Product', 'Same Type', and 'Unrelated Purchases' over a period of 0 to 15 days after the request. The y-axis represents 'Normalized number of purchases' with values ranging from 0 to 30. The 'Same Product' category has the highest bars, particularly in the first few days, followed by 'Same Type'. 'Unrelated Purchases' shows the lowest values consistently across all days.\",\n",
       "   'figure_description': 'Figure 2: Normalized purchases of the same type, as function of the days from the request, in late shopping stages.',\n",
       "   'reasoning': \"1. Identified the chart as a bar chart from visual inspection.\\n2. Located the title of the chart in the caption below it.\\n3. Observed the x-axis showing days from 0 to 15.\\n4. Noted the y-axis label 'Normalized number of purchases' with values up to 30.\\n5. Identified three categories represented by different colored bars.\\n6. Analyzed the relative heights of the bars for each category across the time period.\\n7. Noted that 'Same Product' has the highest bars, especially in the early days.\\n8. Observed that 'Same Type' has the second highest bars.\\n9. Recognized that 'Unrelated Purchases' consistently shows the lowest values.\"}],\n",
       " 'token_usage': {'input_tokens': 1910,\n",
       "  'output_tokens': 381,\n",
       "  'total_tokens': 2291}}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/scientific_paper.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().FigureSysPrompt)\n",
    "resp = da.run(message=\"Explain the bar chart in this document.\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reasoning with tables (experimental)\n",
    "---\n",
    "We can also perform reasoning with tables using the Figure System Prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'figure_analysis': 'The Net Sales for North America increased from $315,880 million in 2022 to $352,828 million in 2023, a difference of $36,948 million.',\n",
       "   'figure_description': 'A table showing Net Sales broken down by segment (North America, International, AWS) for the years ended December 31, 2022 and 2023.',\n",
       "   'reasoning': \"In the 'Net Sales' table on page 1, the value for 'North America' in the 2022 column is $315,880 million. The value for 'North America' in the 2023 column is $352,828 million. The difference between these two values is $352,828 million - $315,880 million = $36,948 million.\"}],\n",
       " 'token_usage': {'input_tokens': 4960, 'output_tokens': 209}}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/amzn-10k.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().FigureSysPrompt)\n",
    "resp = da.run(message=\"What is the dollar value difference of Net Sales between 2022 and 2023 for North America in the given table. Explain your answer.\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-lingual documents\n",
    "---\n",
    "Analyze multi-lingual documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'detected_languages': ['English', 'Spanish'],\n",
       "   'content': 'El Departamento de Vivienda de la Ciudad de Phoenix (City of Phoenix Housing Department, COPHD) aceptará solicitudes previas para su lista de espera del Programa de vales de elección de vivienda (Housing Choice Voucher, HCV) de la Sección 8 a partir del martes 12 de septiembre a las 8 a. m. hasta el martes 26 de septiembre de 2023 a las 7 p. m. hora de Arizona (AZ).'},\n",
       "  {'page': 2,\n",
       "   'detected_languages': ['English', 'Spanish'],\n",
       "   'content': 'No hay costo para solicitar o recibir un vale de elección de vivienda de la Sección 8.'},\n",
       "  {'page': 3,\n",
       "   'detected_languages': ['English', 'Spanish'],\n",
       "   'content': 'Compromiso con la equidad en la vivienda y la no discriminación La ciudad de Phoenix no discrimina por motivos de raza, etnia, sexo, identidad de género, color, religión, estado civil, estado familiar, país de origen, edad, discapacidad, ascendencia, fuente de ingresos u orientación sexual en el acceso, admisión o empleo en programas o actividades de vivienda.'}],\n",
       " 'token_usage': {'input_tokens': 4962, 'output_tokens': 384}}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rhubarb import DocAnalysis\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/public-notice-spanish.pdf\", \n",
    "                 boto3_session=session)\n",
    "resp = da.run(message=\"Which entity or organization issued this notice?\")\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate ALT Text from images in a document\n",
    "---\n",
    "We will generate alt texts for images in a document in order to digitize the books and make them available in web format. Alt text helps with accessibility and screenreaders for people who rely on them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'output': [{'page': 1,\n",
       "   'figure_analysis': \"The image depicts the initial phase of the hydrological cycle, where rainwater infiltrates the topsoil layer ('T') and percolates through subsurface layers, recharging aquifers ('R'). The colors used are blue for water, brown for soil/rock layers, and green for vegetation.\",\n",
       "   'figure_description': 'Figure 1: The Surface Connection',\n",
       "   'reasoning': \"The image illustrates the process of rainwater infiltration into the topsoil ('T') and subsequent percolation through subsurface layers to recharge aquifers ('R'). The colors blue, brown, and green represent water, soil/rock layers, and vegetation, respectively.\"},\n",
       "  {'page': 2,\n",
       "   'figure_analysis': \"The image shows a cross-sectional view of an aquifer system, with 'Y' representing wells and springs that tap into the aquifer, 'R' depicting the aquifer itself, and 'K' representing the bedrock layer beneath. The colors used are green for the surface, blue for water, and brown/gray for rock layers.\",\n",
       "   'figure_description': 'Figure 2: Subsurface Dynamics',\n",
       "   'reasoning': \"The image illustrates the subsurface dynamics of an aquifer system, with wells ('Y') and springs ('J') accessing the aquifer ('R'), which is situated above the bedrock layer ('K'). The colors green, blue, and brown/gray represent the surface, water, and rock layers, respectively.\"},\n",
       "  {'page': 3,\n",
       "   'figure_analysis': \"The image depicts the interaction between a river ('N') and groundwater ('D'), with the blue color representing water and the brown/green colors representing the land surface. The river's flow is shown as a dynamic artery on the landscape, while the groundwater is an unseen layer beneath.\",\n",
       "   'figure_description': 'Figure 3: Riverbed Interactions',\n",
       "   'reasoning': \"The image illustrates the hydrological symbiosis between surface water (the river 'N') and groundwater ('D'), with the blue color representing water and the brown/green colors representing the land surface. The river's flow and the underlying groundwater are shown as interconnected components of the ecosystem.\"},\n",
       "  {'page': 4,\n",
       "   'figure_analysis': \"The image depicts the construction of wells ('F' and 'P') that tap into different layers of an aquifer system, including a confining layer ('L') and a permeable layer ('K'). The colors used are green for the surface, blue for water, and brown/gray for rock layers.\",\n",
       "   'figure_description': 'No figure description provided.',\n",
       "   'reasoning': \"The image illustrates the human intervention of well-drilling to access groundwater from different layers of an aquifer system, including a confining layer ('L') that restricts water flow and a permeable layer ('K') that facilitates water movement. The colors green, blue, and brown/gray represent the surface, water, and rock layers, respectively.\"}],\n",
       " 'token_usage': {'input_tokens': 6507, 'output_tokens': 715}}"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#./test_docs/Anatomy_and_Physiology_2e_small.pdf\n",
    "\n",
    "from rhubarb import DocAnalysis, SystemPrompts\n",
    "\n",
    "da = DocAnalysis(file_path=\"./test_docs/aquifers.pdf\", \n",
    "                 boto3_session=session,\n",
    "                 system_prompt=SystemPrompts().FigureSysPrompt,\n",
    "                 max_tokens=4096)\n",
    "resp = da.run(message=\"You will generate alt-texts for images found in the pages of this document. Make sure to be short but descriptive and explain the colors since alt text is used by screen-readers thus enabling accessibility.\")\n",
    "resp"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rhubarbenv",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}