docs/03_Pydantic.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: Could not find a version that satisfies the requirement agentai==0.0.4 (from versions: 0.0.1, 0.0.2, 0.0.3)\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[31mERROR: No matching distribution found for agentai==0.0.4\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install agentai==0.0.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "GPT_MODEL = \"gpt-3.5-turbo-0613\"\n",
    "EMBEDDING_MODEL = \"text-embedding-ada-002\"\n",
    "\n",
    "from agentai.api import chat_complete, chat_complete_execute_fn\n",
    "from agentai.annotations import tool, ToolRegistry\n",
    "from agentai.conversation import Conversation\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import List, Optional\n",
    "\n",
    "extraction_registry = ToolRegistry()  # Namespace for functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Medicines(BaseModel):\n",
    "    name: str = Field(..., description=\"The name of the medicine\")\n",
    "    dosage: Optional[str] = Field(..., description=\"The dosage of the medicine\")\n",
    "    frequency: Optional[str] = Field(..., description=\"The frequency of the medicine\")\n",
    "    effect: Optional[str] = Field(..., description=\"The response to existing medication\")\n",
    "\n",
    "\n",
    "class Symptoms(BaseModel):\n",
    "    description: str = Field(..., description=\"The description of the symptom\")\n",
    "    severity: Optional[str] = Field(\n",
    "        ...,\n",
    "        description=\"The severity of the symptom\",\n",
    "        enum=[\"Acute\", \"Mild\", \"Moderate\", \"Severe/Chronic\"],\n",
    "    )\n",
    "    duration: Optional[str] = Field(\n",
    "        ...,\n",
    "        description=\"The duration of the symptom\",\n",
    "    )\n",
    "\n",
    "\n",
    "class PatientInformation(BaseModel):\n",
    "    gender: str = Field(..., description=\"Gender of the patient\", enum=[\"Male\", \"Female\", \"Other\"])\n",
    "    age: int = Field(..., description=\"Age of the patient\")\n",
    "    symptoms: Optional[List[Symptoms]] = Field(\n",
    "        ...,\n",
    "        description=\"Symptoms that the patient is currently experiencing. Each symptom should be classified into separate item in the list.\",\n",
    "    )\n",
    "    medicines: Optional[List[Medicines]] = Field(\n",
    "        ...,\n",
    "        description=\"Medicines that the patient is currently taking. Each medicine should be classified into separate item in the list.\",\n",
    "    )\n",
    "\n",
    "\n",
    "@tool(extraction_registry)\n",
    "def save_patient_history(history: PatientInformation) -> None:\n",
    "    \"\"\"\n",
    "    Save the patient's history which consists of gender, age, symptoms and medicines\n",
    "\n",
    "    Args:\n",
    "        history (PatientInformation): The patient's history\n",
    "\n",
    "    Returns:\n",
    "        str: The current weather\n",
    "    \"\"\"\n",
    "    print(f\"Saving patient history: {history}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"role\": \"assistant\",\n",
      "  \"content\": null,\n",
      "  \"function_call\": {\n",
      "    \"name\": \"save_patient_history\",\n",
      "    \"arguments\": \"{\\n  \\\"history\\\": {\\n    \\\"gender\\\": \\\"Male\\\",\\n    \\\"age\\\": 49,\\n    \\\"symptoms\\\": [\\n      {\\n        \\\"description\\\": \\\"Chronic macular rash to face & hair\\\",\\n        \\\"severity\\\": \\\"Moderate\\\",\\n        \\\"duration\\\": \\\"Unknown\\\"\\n      },\\n      {\\n        \\\"description\\\": \\\"Itchy\\\",\\n        \\\"severity\\\": \\\"Mild\\\",\\n        \\\"duration\\\": \\\"Unknown\\\"\\n      },\\n      {\\n        \\\"description\\\": \\\"Flaky, slightly scaly\\\",\\n        \\\"severity\\\": \\\"Mild\\\",\\n        \\\"duration\\\": \\\"Unknown\\\"\\n      }\\n    ],\\n    \\\"medicines\\\": [\\n      {\\n        \\\"name\\\": \\\"OTC steroid cream\\\",\\n        \\\"dosage\\\": \\\"Unknown\\\",\\n        \\\"frequency\\\": \\\"Unknown\\\",\\n        \\\"effect\\\": \\\"Moderate response\\\"\\n      }\\n    ]\\n  }\\n}\"\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "convo = Conversation()\n",
    "convo.add_message(\n",
    "    role=\"system\",\n",
    "    content=\"\"\"You are an AI assistant.\n",
    "- Carefully think about functions to use.\n",
    "- DO NOT make assumptions.\n",
    "- Use `null` if you are not confident\"\"\",\n",
    ")\n",
    "\n",
    "convo.add_message(\n",
    "    role=\"user\",\n",
    "    content=\"\"\"\"49 y/o Male with chronic macular rash to face & hair, worse in beard, eyebrows & nares. \n",
    "Itchy, flaky, slightly scaly. Moderate response to OTC steroid cream\"\"\",\n",
    ")\n",
    "completion = chat_complete(convo, model=GPT_MODEL, tool_registry=extraction_registry)\n",
    "print(completion.choices[0].message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"history\": {\n",
      "    \"gender\": \"Male\",\n",
      "    \"age\": 49,\n",
      "    \"symptoms\": [\n",
      "      {\n",
      "        \"description\": \"Chronic macular rash to face & hair\",\n",
      "        \"severity\": \"Moderate\",\n",
      "        \"duration\": \"Unknown\"\n",
      "      },\n",
      "      {\n",
      "        \"description\": \"Itchy\",\n",
      "        \"severity\": \"Mild\",\n",
      "        \"duration\": \"Unknown\"\n",
      "      },\n",
      "      {\n",
      "        \"description\": \"Flaky, slightly scaly\",\n",
      "        \"severity\": \"Mild\",\n",
      "        \"duration\": \"Unknown\"\n",
      "      }\n",
      "    ],\n",
      "    \"medicines\": [\n",
      "      {\n",
      "        \"name\": \"OTC steroid cream\",\n",
      "        \"dosage\": \"Unknown\",\n",
      "        \"frequency\": \"Unknown\",\n",
      "        \"effect\": \"Moderate response\"\n",
      "      }\n",
      "    ]\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "print(completion.choices[0].message[\"function_call\"][\"arguments\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PatientInformation(gender='Male', age=49, symptoms=[Symptoms(description='Chronic macular rash to face & hair', severity='Moderate', duration='Unknown'), Symptoms(description='Itchy', severity='Mild', duration='Unknown'), Symptoms(description='Flaky, slightly scaly', severity='Mild', duration='Unknown')], medicines=[Medicines(name='OTC steroid cream', dosage='Unknown', frequency='Unknown', effect='Moderate response')])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "from pydantic import validate_arguments\n",
    "\n",
    "arguments = completion.choices[0].message[\"function_call\"][\"arguments\"]\n",
    "history = json.loads(arguments)[\"history\"]\n",
    "PatientInformation(**history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'save_patient_history',\n",
       " 'description': \"Save the patient's history which consists of gender, age, symptoms and medicines\",\n",
       " 'parameters': {'type': 'object',\n",
       "  'properties': {'history': {'title': 'PatientInformation',\n",
       "    'type': 'object',\n",
       "    'properties': {'gender': {'title': 'Gender',\n",
       "      'description': 'Gender of the patient',\n",
       "      'enum': ['Male', 'Female', 'Other'],\n",
       "      'type': 'string'},\n",
       "     'age': {'title': 'Age',\n",
       "      'description': 'Age of the patient',\n",
       "      'type': 'integer'},\n",
       "     'symptoms': {'title': 'Symptoms',\n",
       "      'description': 'Symptoms that the patient is currently experiencing. Each symptom should be classified into separate item in the list.',\n",
       "      'type': 'array',\n",
       "      'items': {'title': 'Symptoms',\n",
       "       'type': 'object',\n",
       "       'properties': {'description': {'title': 'Description',\n",
       "         'description': 'The description of the symptom',\n",
       "         'type': 'string'},\n",
       "        'severity': {'title': 'Severity',\n",
       "         'description': 'The severity of the symptom',\n",
       "         'enum': ['Acute', 'Mild', 'Moderate', 'Severe/Chronic'],\n",
       "         'type': 'string'},\n",
       "        'duration': {'title': 'Duration',\n",
       "         'description': 'The duration of the symptom',\n",
       "         'type': 'string'}},\n",
       "       'required': ['description', 'severity', 'duration']}},\n",
       "     'medicines': {'title': 'Medicines',\n",
       "      'description': 'Medicines that the patient is currently taking. Each medicine should be classified into separate item in the list.',\n",
       "      'type': 'array',\n",
       "      'items': {'title': 'Medicines',\n",
       "       'type': 'object',\n",
       "       'properties': {'name': {'title': 'Name',\n",
       "         'description': 'The name of the medicine',\n",
       "         'type': 'string'},\n",
       "        'dosage': {'title': 'Dosage',\n",
       "         'description': 'The dosage of the medicine',\n",
       "         'type': 'string'},\n",
       "        'frequency': {'title': 'Frequency',\n",
       "         'description': 'The frequency of the medicine',\n",
       "         'type': 'string'},\n",
       "        'effect': {'title': 'Effect',\n",
       "         'description': 'The response to existing medication',\n",
       "         'type': 'string'}},\n",
       "       'required': ['name', 'dosage', 'frequency', 'effect']}}},\n",
       "    'required': ['gender', 'age', 'symptoms', 'medicines'],\n",
       "    'description': \"The patient's history\"}},\n",
       "  'required': ['history']}}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "save_patient_history._schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}