notebooks/instructor.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U instructor\n",
    "# https://github.com/jxnl/instructor\n",
    "# https://python.useinstructor.com/examples/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import instructor\n",
    "from pydantic import BaseModel\n",
    "from openai import OpenAI\n",
    "\n",
    "\n",
    "# Define your desired output structure\n",
    "class UserInfo(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "# Patch the OpenAI client\n",
    "client = instructor.from_openai(OpenAI())\n",
    "\n",
    "# Extract structured data from natural language\n",
    "user_info = client.chat.completions.create(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "    response_model=UserInfo,\n",
    "    messages=[{\"role\": \"user\", \"content\": \"John Doe is 30 years old.\"}],\n",
    ")\n",
    "\n",
    "print(user_info.name)\n",
    "# > John Doe\n",
    "print(user_info.age)\n",
    "# > 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import instructor\n",
    "from anthropic import Anthropic\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "class User(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "client = instructor.from_anthropic(Anthropic())\n",
    "\n",
    "# note that client.chat.completions.create will also work\n",
    "resp = client.messages.create(\n",
    "    model=\"claude-3-opus-20240229\",\n",
    "    max_tokens=1024,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"Extract Jason is 25 years old.\",\n",
    "        }\n",
    "    ],\n",
    "    response_model=User,\n",
    ")\n",
    "\n",
    "assert isinstance(resp, User)\n",
    "assert resp.name == \"Jason\"\n",
    "assert resp.age == 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### CREATE\n",
    "import openai\n",
    "import instructor\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "class User(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "client = instructor.from_openai(openai.OpenAI())\n",
    "\n",
    "user = client.chat.completions.create(\n",
    "    model=\"gpt-4-turbo-preview\",\n",
    "    messages=[\n",
    "        {\"role\": \"user\", \"content\": \"Create a user\"},\n",
    "    ],\n",
    "    response_model=User,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import instructor\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "client = instructor.from_openai(openai.AsyncOpenAI())\n",
    "\n",
    "\n",
    "class User(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "async def extract():\n",
    "    return await client.chat.completions.create(\n",
    "        model=\"gpt-4-turbo-preview\",\n",
    "        messages=[\n",
    "            {\"role\": \"user\", \"content\": \"Create a user\"},\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import instructor\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "client = instructor.from_openai(openai.OpenAI())\n",
    "\n",
    "\n",
    "class User(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "user, completion = client.chat.completions.create_with_completion(\n",
    "    model=\"gpt-4-turbo-preview\",\n",
    "    messages=[\n",
    "        {\"role\": \"user\", \"content\": \"Create a user\"},\n",
    "    ],\n",
    "    response_model=User,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import instructor\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "client = instructor.from_openai(openai.OpenAI())\n",
    "\n",
    "\n",
    "class User(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "user_stream = client.chat.completions.create_partial(\n",
    "    model=\"gpt-4-turbo-preview\",\n",
    "    messages=[\n",
    "        {\"role\": \"user\", \"content\": \"Create a user\"},\n",
    "    ],\n",
    "    response_model=User,\n",
    ")\n",
    "\n",
    "for user in user_stream:\n",
    "    print(user)\n",
    "    # > name=None age=None\n",
    "    # > name=None age=None\n",
    "    # > name=None age=None\n",
    "    # > name=None age=None\n",
    "    # > name=None age=25\n",
    "    # > name=None age=25\n",
    "    # > name=None age=25\n",
    "    # > name=None age=25\n",
    "    # > name=None age=25\n",
    "    # > name=None age=25\n",
    "    # > name='John Doe' age=25\n",
    "    # name=None age=None\n",
    "    # name='' age=None\n",
    "    # name='John' age=None\n",
    "    # name='John Doe' age=None\n",
    "    # name='John Doe' age=30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import instructor\n",
    "from pydantic import BaseModel\n",
    "\n",
    "\n",
    "client = instructor.from_openai(openai.OpenAI())\n",
    "\n",
    "\n",
    "class User(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "users = client.chat.completions.create_iterable(\n",
    "    model=\"gpt-4-turbo-preview\",\n",
    "    messages=[\n",
    "        {\"role\": \"user\", \"content\": \"Create 2 users\"},\n",
    "    ],\n",
    "    response_model=User,\n",
    ")\n",
    "\n",
    "for user in users:\n",
    "    print(user)\n",
    "    # > name='John' age=30\n",
    "    # > name='Jane' age=25\n",
    "    # User(name='John Doe', age=30)\n",
    "    # User(name='Jane Smith', age=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import anthropic\n",
    "import instructor\n",
    "from pydantic import BaseModel, field_validator\n",
    "from typing import List, Literal\n",
    "from enum import Enum\n",
    "\n",
    "client = instructor.from_anthropic(\n",
    "    anthropic.Anthropic(), mode=instructor.Mode.ANTHROPIC_TOOLS\n",
    ")\n",
    "\n",
    "\n",
    "def test_simple():\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        age: int\n",
    "\n",
    "        @field_validator(\"name\")\n",
    "        def name_is_uppercase(cls, v: str):\n",
    "            assert v.isupper(), \"Name must be uppercase\"\n",
    "            return v\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=2,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Extract John is 18 years old.\",\n",
    "            }\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )  # type: ignore\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    assert resp.name == \"JOHN\"  # due to validation\n",
    "    assert resp.age == 18\n",
    "\n",
    "\n",
    "def test_nested_type():\n",
    "    class Address(BaseModel):\n",
    "        house_number: int\n",
    "        street_name: str\n",
    "\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        age: int\n",
    "        address: Address\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=0,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Extract John is 18 years old and lives at 123 First Avenue.\",\n",
    "            }\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )  # type: ignore\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    assert resp.name == \"John\"\n",
    "    assert resp.age == 18\n",
    "\n",
    "    assert isinstance(resp.address, Address)\n",
    "    assert resp.address.house_number == 123\n",
    "    assert resp.address.street_name == \"First Avenue\"\n",
    "\n",
    "\n",
    "def test_list_str():\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        age: int\n",
    "        family: List[str]\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=0,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Create a user for a model with a name, age, and family members.\",\n",
    "            }\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    assert isinstance(resp.family, List)\n",
    "    for member in resp.family:\n",
    "        assert isinstance(member, str)\n",
    "\n",
    "\n",
    "def test_enum():\n",
    "    class Role(str, Enum):\n",
    "        ADMIN = \"admin\"\n",
    "        USER = \"user\"\n",
    "\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        role: Role\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=0,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Create a user for a model with a name and role of admin.\",\n",
    "            }\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )  # type: ignore\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    assert resp.role == Role.ADMIN\n",
    "\n",
    "\n",
    "def test_literal():\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        role: Literal[\"admin\", \"user\"]\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=2,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Create a admin user for a model with a name and role.\",\n",
    "            }\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )  # type: ignore\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    assert resp.role == \"admin\"\n",
    "\n",
    "\n",
    "def test_nested_list():\n",
    "    class Properties(BaseModel):\n",
    "        key: str\n",
    "        value: str\n",
    "\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        age: int\n",
    "        properties: List[Properties]\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=0,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Create a user for a model with a name, age, and properties.\",\n",
    "            }\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )  # type: ignore\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    for property in resp.properties:\n",
    "        assert isinstance(property, Properties)\n",
    "\n",
    "\n",
    "def test_system_messages_allcaps():\n",
    "    class User(BaseModel):\n",
    "        name: str\n",
    "        age: int\n",
    "\n",
    "    resp = client.messages.create(\n",
    "        model=\"claude-3-haiku-20240307\",\n",
    "        max_tokens=1024,\n",
    "        max_retries=0,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"EVERYTHING MUST BE IN ALL CAPS\"},\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Create a user for a model with a name and age.\",\n",
    "            },\n",
    "        ],\n",
    "        response_model=User,\n",
    "    )  # type: ignore\n",
    "\n",
    "    assert isinstance(resp, User)\n",
    "    assert resp.name.isupper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import enum\n",
    "from itertools import product\n",
    "from typing import List\n",
    "\n",
    "import pytest\n",
    "import instructor\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from instructor.function_calls import Mode\n",
    "from ..util import models, modes\n",
    "\n",
    "\n",
    "class Labels(str, enum.Enum):\n",
    "    SPAM = \"spam\"\n",
    "    NOT_SPAM = \"not_spam\"\n",
    "\n",
    "\n",
    "class SinglePrediction(BaseModel):\n",
    "    \"\"\"\n",
    "    Correct class label for the given text\n",
    "    \"\"\"\n",
    "\n",
    "    class_label: Labels\n",
    "\n",
    "\n",
    "data = [\n",
    "    (\n",
    "        \"I am a spammer\",\n",
    "        Labels.SPAM,\n",
    "    ),\n",
    "    (\n",
    "        \"I am not a spammer\",\n",
    "        Labels.NOT_SPAM,\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"model, data, mode\", product(models, data, modes))\n",
    "def test_classification(model, data, mode, client):\n",
    "    client = instructor.patch(client, mode=mode)\n",
    "\n",
    "    if mode == instructor.Mode.JSON and model in {\"gpt-3.5-turbo\", \"gpt-4\"}:\n",
    "        pytest.skip(\n",
    "            \"JSON mode is not supported for gpt-3.5-turbo and gpt-4, skipping test\"\n",
    "        )\n",
    "\n",
    "    input, expected = data\n",
    "    resp = client.chat.completions.create(\n",
    "        model=model,\n",
    "        response_model=SinglePrediction,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": f\"Classify the following text: {input}\",\n",
    "            },\n",
    "        ],\n",
    "    )\n",
    "    assert resp.class_label == expected\n",
    "\n",
    "\n",
    "# Define new Enum class for multiple labels\n",
    "class MultiLabels(str, enum.Enum):\n",
    "    BILLING = \"billing\"\n",
    "    GENERAL_QUERY = \"general_query\"\n",
    "    HARDWARE = \"hardware\"\n",
    "\n",
    "\n",
    "# Adjust the prediction model to accommodate a list of labels\n",
    "class MultiClassPrediction(BaseModel):\n",
    "    predicted_labels: List[MultiLabels]\n",
    "\n",
    "\n",
    "data = [\n",
    "    (\n",
    "        \"I am having trouble with my billing\",\n",
    "        [MultiLabels.BILLING],\n",
    "    ),\n",
    "    (\n",
    "        \"I am having trouble with my hardware\",\n",
    "        [MultiLabels.HARDWARE],\n",
    "    ),\n",
    "    (\n",
    "        \"I have a general query and a billing issue\",\n",
    "        [MultiLabels.GENERAL_QUERY, MultiLabels.BILLING],\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"model, data, mode\", product(models, data, modes))\n",
    "def test_multi_classify(model, data, mode, client):\n",
    "    client = instructor.patch(client, mode=mode)\n",
    "\n",
    "    if (mode, model) in {\n",
    "        (Mode.JSON, \"gpt-3.5-turbo\"),\n",
    "        (Mode.JSON, \"gpt-4\"),\n",
    "    }:\n",
    "        pytest.skip(f\"{mode} mode is not supported for {model}, skipping test\")\n",
    "\n",
    "    input, expected = data\n",
    "\n",
    "    resp = client.chat.completions.create(\n",
    "        model=model,\n",
    "        response_model=MultiClassPrediction,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": f\"Classify the following support ticket: {input}\",\n",
    "            },\n",
    "        ],\n",
    "    )\n",
    "    assert set(resp.predicted_labels) == set(expected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "from typing import List, Literal\n",
    "\n",
    "import pytest\n",
    "import instructor\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from instructor.function_calls import Mode\n",
    "from ..util import models, modes\n",
    "\n",
    "\n",
    "class SinglePrediction(BaseModel):\n",
    "    \"\"\"\n",
    "    Correct class label for the given text\n",
    "    \"\"\"\n",
    "\n",
    "    class_label: Literal[\"spam\", \"not_spam\"]\n",
    "\n",
    "\n",
    "data = [\n",
    "    (\"I am a spammer\", \"spam\"),\n",
    "    (\"I am not a spammer\", \"not_spam\"),\n",
    "]\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"model, data, mode\", product(models, data, modes))\n",
    "@pytest.mark.asyncio\n",
    "async def test_classification(model, data, mode, aclient):\n",
    "    client = instructor.patch(aclient, mode=mode)\n",
    "\n",
    "    if mode == instructor.Mode.JSON and model in {\"gpt-3.5-turbo\", \"gpt-4\"}:\n",
    "        pytest.skip(\n",
    "            \"JSON mode is not supported for gpt-3.5-turbo and gpt-4, skipping test\"\n",
    "        )\n",
    "\n",
    "    input, expected = data\n",
    "    resp = await client.chat.completions.create(\n",
    "        model=model,\n",
    "        response_model=SinglePrediction,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": f\"Classify the following text: {input}\",\n",
    "            },\n",
    "        ],\n",
    "    )\n",
    "    assert resp.class_label == expected\n",
    "\n",
    "\n",
    "# Adjust the prediction model to accommodate a list of labels\n",
    "class MultiClassPrediction(BaseModel):\n",
    "    predicted_labels: List[Literal[\"billing\", \"general_query\", \"hardware\"]]\n",
    "\n",
    "\n",
    "data = [\n",
    "    (\n",
    "        \"I am having trouble with my billing\",\n",
    "        [\"billing\"],\n",
    "    ),\n",
    "    (\n",
    "        \"I am having trouble with my hardware\",\n",
    "        [\"hardware\"],\n",
    "    ),\n",
    "    (\n",
    "        \"I have a general query and a billing issue\",\n",
    "        [\"general_query\", \"billing\"],\n",
    "    ),\n",
    "]\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"model, data, mode\", product(models, data, modes))\n",
    "@pytest.mark.asyncio\n",
    "async def test_multi_classify(model, data, mode, aclient):\n",
    "    client = instructor.patch(aclient, mode=mode)\n",
    "\n",
    "    if (mode, model) in {\n",
    "        (Mode.JSON, \"gpt-3.5-turbo\"),\n",
    "        (Mode.JSON, \"gpt-4\"),\n",
    "    }:\n",
    "        pytest.skip(f\"{mode} mode is not supported for {model}, skipping test\")\n",
    "\n",
    "    input, expected = data\n",
    "\n",
    "    resp = await client.chat.completions.create(\n",
    "        model=model,\n",
    "        response_model=MultiClassPrediction,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": f\"Classify the following support ticket: {input}\",\n",
    "            },\n",
    "        ],\n",
    "    )\n",
    "    assert set(resp.predicted_labels) == set(expected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "from typing import List\n",
    "from pydantic import BaseModel, Field\n",
    "import pytest\n",
    "\n",
    "import instructor\n",
    "\n",
    "from instructor.function_calls import Mode\n",
    "from ..util import models, modes\n",
    "\n",
    "\n",
    "class Property(BaseModel):\n",
    "    key: str\n",
    "    value: str\n",
    "    resolved_absolute_value: str\n",
    "\n",
    "\n",
    "class Entity(BaseModel):\n",
    "    id: int = Field(\n",
    "        ...,\n",
    "        description=\"Unique identifier for the entity, used for deduplication, design a scheme allows multiple entities\",\n",
    "    )\n",
    "    subquote_string: List[str] = Field(\n",
    "        ...,\n",
    "        description=\"Correctly resolved value of the entity, if the entity is a reference to another entity, this should be the id of the referenced entity, include a few more words before and after the value to allow for some context to be used in the resolution\",\n",
    "    )\n",
    "    entity_title: str\n",
    "    properties: List[Property] = Field(\n",
    "        ..., description=\"List of properties of the entity\"\n",
    "    )\n",
    "    dependencies: List[int] = Field(\n",
    "        ...,\n",
    "        description=\"List of entity ids that this entity depends  or relies on to resolve it\",\n",
    "    )\n",
    "\n",
    "\n",
    "class DocumentExtraction(BaseModel):\n",
    "    entities: List[Entity] = Field(\n",
    "        ...,\n",
    "        description=\"Body of the answer, each fact should be its seperate object with a body and a list of sources\",\n",
    "    )\n",
    "\n",
    "\n",
    "def ask_ai(content, model, client) -> DocumentExtraction:\n",
    "    resp: DocumentExtraction = client.chat.completions.create(\n",
    "        model=model,\n",
    "        response_model=DocumentExtraction,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": \"You are a perfect entity resolution system that extracts facts from the document. Extract and resolve a list of entities from the following document:\",\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": content,\n",
    "            },\n",
    "        ],\n",
    "    )  # type: ignore\n",
    "    return resp\n",
    "\n",
    "\n",
    "content = \"\"\"\n",
    "Sample Legal Contract\n",
    "Agreement Contract\n",
    "\n",
    "This Agreement is made and entered into on 2020-01-01 by and between Company A (\"the Client\") and Company B (\"the Service Provider\").\n",
    "\n",
    "Article 1: Scope of Work\n",
    "\n",
    "The Service Provider will deliver the software product to the Client 30 days after the agreement date.\n",
    "\n",
    "Article 2: Payment Terms\n",
    "\n",
    "The total payment for the service is $50,000.\n",
    "An initial payment of $10,000 will be made within 7 days of the the signed date.\n",
    "The final payment will be due 45 days after [SignDate].\n",
    "\n",
    "Article 3: Confidentiality\n",
    "\n",
    "The parties agree not to disclose any confidential information received from the other party for 3 months after the final payment date.\n",
    "\n",
    "Article 4: Termination\n",
    "\n",
    "The contract can be terminated with a 30-day notice, unless there are outstanding obligations that must be fulfilled after the [DeliveryDate].\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"model, mode\", product(models, modes))\n",
    "def test_extract(model, mode, client):\n",
    "    client = instructor.patch(client, mode=mode)\n",
    "    if (mode, model) in {\n",
    "        (Mode.JSON, \"gpt-3.5-turbo\"),\n",
    "        (Mode.JSON, \"gpt-4\"),\n",
    "    }:\n",
    "        pytest.skip(f\"{mode} mode is not supported for {model}, skipping test\")\n",
    "\n",
    "    # Honestly, if there are no errors, then it's a pass\n",
    "    extract = ask_ai(content=content, model=model, client=client)\n",
    "    assert len(extract.entities) > 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytest\n",
    "from itertools import product\n",
    "from pydantic import BaseModel\n",
    "import instructor\n",
    "from instructor.function_calls import Mode\n",
    "from ..util import models, modes\n",
    "\n",
    "\n",
    "class UserDetails(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "\n",
    "\n",
    "# Lists for models, test data, and modes\n",
    "test_data = [\n",
    "    (\"Jason is 10\", \"Jason\", 10),\n",
    "    (\"Alice is 25\", \"Alice\", 25),\n",
    "    (\"Bob is 35\", \"Bob\", 35),\n",
    "]\n",
    "\n",
    "\n",
    "@pytest.mark.parametrize(\"model, data, mode\", product(models, test_data, modes))\n",
    "def test_extract(model, data, mode, client):\n",
    "    sample_data, expected_name, expected_age = data\n",
    "\n",
    "    if (mode, model) in {\n",
    "        (Mode.JSON, \"gpt-3.5-turbo\"),\n",
    "        (Mode.JSON, \"gpt-4\"),\n",
    "    }:\n",
    "        pytest.skip(f\"{mode} mode is not supported for {model}, skipping test\")\n",
    "\n",
    "    # Setting up the client with the instructor patch\n",
    "    client = instructor.patch(client, mode=mode)\n",
    "\n",
    "    # Calling the extract function with the provided model, sample data, and mode\n",
    "    response = client.chat.completions.create(\n",
    "        model=model,\n",
    "        response_model=UserDetails,\n",
    "        messages=[\n",
    "            {\"role\": \"user\", \"content\": sample_data},\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    # Assertions\n",
    "    assert (\n",
    "        response.name == expected_name\n",
    "    ), f\"Expected name {expected_name}, got {response.name}\"\n",
    "    assert (\n",
    "        response.age == expected_age\n",
    "    ), f\"Expected age {expected_age}, got {response.age}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "\n",
    "\n",
    "class PersonBirthday(BaseModel):\n",
    "    name: str\n",
    "    age: int\n",
    "    birthday: datetime.date\n",
    "\n",
    "\n",
    "schema = {\n",
    "    \"properties\": {\n",
    "        \"name\": {\"type\": \"string\"},\n",
    "        \"age\": {\"type\": \"integer\"},\n",
    "        \"birthday\": {\"type\": \"string\", \"format\": \"YYYY-MM-DD\"},\n",
    "    },\n",
    "    \"required\": [\"name\", \"age\"],\n",
    "    \"type\": \"object\",\n",
    "}\n",
    "resp = client.chat.completions.create(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": f\"Extract `Jason Liu is thirty years old his birthday is yesturday` into json today is {datetime.date.today()}\",\n",
    "        },\n",
    "    ],\n",
    "    functions=[{\"name\": \"Person\", \"parameters\": schema}],\n",
    "    function_call=\"auto\",\n",
    ")\n",
    "\n",
    "PersonBirthday.model_validate_json(\n",
    "    resp.choices[0].message.function_call.arguments\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import instructor\n",
    "from openai import OpenAI\n",
    "from enum import Enum\n",
    "from pydantic import BaseModel, Field\n",
    "from typing_extensions import Literal\n",
    "\n",
    "client = instructor.patch(\n",
    "    OpenAI()\n",
    ")  # Tip: Do not use auto() as they cast to 1,2,3,4\n",
    "\n",
    "\n",
    "class House(Enum):\n",
    "    Gryffindor = \"gryffindor\"\n",
    "    Hufflepuff = \"hufflepuff\"\n",
    "    Ravenclaw = \"ravenclaw\"\n",
    "    Slytherin = \"slytherin\"\n",
    "\n",
    "\n",
    "class Character(BaseModel):\n",
    "    age: int\n",
    "    name: str\n",
    "    house: House\n",
    "\n",
    "    def say_hello(self):\n",
    "        print(\n",
    "            f\"Hello, I'm {self.name}, I'm {self.age} years old and I'm from {self.house.value.title()}\"\n",
    "        )\n",
    "\n",
    "\n",
    "resp = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    messages=[{\"role\": \"user\", \"content\": \"Harry Potter\"}],\n",
    "    response_model=Character,\n",
    ")\n",
    "\n",
    "resp.model_dump()\n",
    "\n",
    "resp.say_hello()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Character(BaseModel):\n",
    "    age: int\n",
    "    name: str\n",
    "    house: Literal[\"Gryffindor\", \"Hufflepuff\", \"Ravenclaw\", \"Slytherin\"]\n",
    "\n",
    "\n",
    "resp = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    messages=[{\"role\": \"user\", \"content\": \"Harry Potter\"}],\n",
    "    response_model=Character,\n",
    ")\n",
    "\n",
    "resp.model_dump()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "\n",
    "class Property(BaseModel):\n",
    "    key: str = Field(description=\"Must be snake case\")\n",
    "    value: str\n",
    "\n",
    "\n",
    "class Character(BaseModel):\n",
    "    age: int\n",
    "    name: str\n",
    "    house: Literal[\"Gryffindor\", \"Hufflepuff\", \"Ravenclaw\", \"Slytherin\"]\n",
    "    properties: List[Property]\n",
    "\n",
    "\n",
    "resp = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    messages=[{\"role\": \"user\", \"content\": \"Snape from Harry Potter\"}],\n",
    "    response_model=Character,\n",
    ")\n",
    "\n",
    "resp.model_dump()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}