Claude new API invocation bedrock demo

claude_messages_api_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "region='us-east-1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "bedrock_runtime = boto3.client('bedrock-runtime', region_name=region)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Class declarations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "\n",
    "from pydantic import BaseModel, Field, computed_field, model_validator\n",
    "\n",
    "\n",
    "class ContentSource(BaseModel):\n",
    "    \"\"\"A content source.\"\"\"\n",
    "\n",
    "    type: str\n",
    "    media_type: str\n",
    "    data: str\n",
    "\n",
    "\n",
    "class ContentBlock(BaseModel):\n",
    "    \"\"\"A content block.\"\"\"\n",
    "\n",
    "    type: str\n",
    "    text: str | None = None\n",
    "    source: ContentSource | None = None\n",
    "\n",
    "    @model_validator(mode='after')\n",
    "    def validate_content(self) -> 'ContentBlock':\n",
    "        \"\"\"Ensure that only one of text or source is provided.\"\"\"\n",
    "        if self.text is None and self.source is None:\n",
    "            raise ValueError('One of text or source must be provided')\n",
    "        if self.text is not None and self.source is not None:\n",
    "            raise ValueError('Only one of text or source can be provided')\n",
    "        return self\n",
    "\n",
    "\n",
    "class ClaudeRole(str, Enum):\n",
    "    \"\"\"Roles for Claude.\"\"\"\n",
    "\n",
    "    USER = 'user'\n",
    "    SYSTEM = 'system'\n",
    "    ASSISTANT = 'assistant'\n",
    "\n",
    "\n",
    "class ClaudeMessage(BaseModel):\n",
    "    \"\"\"A message from Claude.\"\"\"\n",
    "\n",
    "    role: ClaudeRole\n",
    "    content: list[ContentBlock] | str\n",
    "\n",
    "\n",
    "class LLMResponse(BaseModel):\n",
    "    \"\"\"Response from the language model.\"\"\"\n",
    "\n",
    "    contents: list[ContentBlock] = []\n",
    "    input_token_count: int | None = None\n",
    "    output_token_count: int | None = None\n",
    "    stop_reason: str\n",
    "    stop_sequence: str | None = None\n",
    "\n",
    "\n",
    "class LLMInput(BaseModel):\n",
    "    \"\"\"Input to the language model.\"\"\"\n",
    "\n",
    "    stop_sequences: list[str] | None = None\n",
    "\n",
    "    temperature: float | None = None\n",
    "    max_tokens: int = 200\n",
    "    top_p: float| None = None\n",
    "    top_k: int| None = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AnthropicInput(LLMInput):\n",
    "    \"\"\"Input to the Anthropic model.\"\"\"\n",
    "\n",
    "    anthropic_version: str = 'bedrock-2023-05-31'\n",
    "\n",
    "    inputs: list[ClaudeMessage] = Field(exclude=True)\n",
    "\n",
    "\n",
    "    @computed_field\n",
    "    def system(self) -> str | None:\n",
    "        \"\"\"Get the system message.\"\"\"\n",
    "        for message in self.inputs:\n",
    "            if message.role == ClaudeRole.SYSTEM:\n",
    "                return message.content if isinstance(message.content, str) else message.content[0].text\n",
    "        return None\n",
    "\n",
    "    @computed_field\n",
    "    def messages(self) ->list[ClaudeMessage]:\n",
    "        \"\"\"Get the user and assistant messages.\"\"\"\n",
    "        return [\n",
    "            message\n",
    "            for message in self.inputs\n",
    "            if message.role in (ClaudeRole.USER, ClaudeRole.ASSISTANT)\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AnthropicResponseUsage(BaseModel):\n",
    "    \"\"\"Usage information for the Anthropic model.\"\"\"\n",
    "\n",
    "    input_tokens: int\n",
    "    output_tokens: int\n",
    "\n",
    "class AnthropicResponse(BaseModel):\n",
    "    \"\"\"Response from the Anthropic model.\"\"\"\n",
    "\n",
    "    type: str\n",
    "    role: ClaudeRole\n",
    "    content: list[ContentBlock]\n",
    "\n",
    "    stop_reason: str\n",
    "    stop_sequence: str | None\n",
    "    usage: AnthropicResponseUsage\n",
    "\n",
    "    def to_llm_response(self) -> LLMResponse:\n",
    "        \"\"\"Convert to an LLMResponse.\"\"\"\n",
    "        return LLMResponse(\n",
    "            contents=self.content,\n",
    "            input_token_count=self.usage.input_tokens,\n",
    "            output_token_count=self.usage.output_tokens,\n",
    "            stop_reason=self.stop_reason,\n",
    "            stop_sequence=self.stop_sequence,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def invoke_llm(model_name: str, inp: LLMInput) -> LLMResponse:\n",
    "    \"\"\"Invoke the language model.\"\"\"\n",
    "    body = inp.json(exclude_none=True)\n",
    "\n",
    "    response = bedrock_runtime.invoke_model(\n",
    "        body=body,\n",
    "        modelId=model_name,\n",
    "        contentType='application/json',\n",
    "        accept = '*/*',\n",
    "    )\n",
    "\n",
    "    return AnthropicResponse.model_validate_json(response['body'].read())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LLMResponse(contents=[ContentBlock(type='text', text='The capital of ', source=None)], input_token_count=22, output_token_count=4, stop_reason='stop_sequence', stop_sequence='France')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k=AnthropicInput(\n",
    "        inputs=[\n",
    "            ClaudeMessage(\n",
    "                role=ClaudeRole.SYSTEM,\n",
    "                content=[\n",
    "                    ContentBlock(\n",
    "                        type='text',\n",
    "                        text='You will be acting as a tour guide',\n",
    "                    ),\n",
    "                ],\n",
    "            ),\n",
    "            ClaudeMessage(\n",
    "                role=ClaudeRole.USER,\n",
    "                content=[\n",
    "                    ContentBlock(\n",
    "                        type='text',\n",
    "                        text='What is the capital of France?',\n",
    "                    ),\n",
    "                ],\n",
    "            ),\n",
    "        ],\n",
    "        stop_sequences=['France'],\n",
    "    )\n",
    "\n",
    "invoke_llm(\n",
    "    'anthropic.claude-3-sonnet-20240229-v1:0',k).to_llm_response()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai-developer-assist-F2jItxF0-py3.11",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}

use pydantic v2