{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# Anthropic Claude\n",
    "\n",
    "In this notebook, we demonstrate how a to use Anthropic Claude model for AgentChat.\n",
    "\n",
    "## Requirements\n",
    "To use Anthropic Claude with AutoGen, first you need to install the `pyautogen` and `anthropic` package.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pyautogen anthropic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "from typing import Any, Dict, List, Union\n",
    "\n",
    "from anthropic import Anthropic\n",
    "from anthropic.types import Completion, Message\n",
    "\n",
    "import autogen\n",
    "from autogen import AssistantAgent, UserProxyAgent\n",
    "from autogen.oai.openai_utils import OAI_PRICE1K"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Anthropic Model Client following ModelClient Protocol\n",
    "\n",
    "We will implement our Anthropic client adhere to the `ModelClient` protocol and response structure which is defined in client.py and shown below.\n",
    "\n",
    "\n",
    "```python\n",
    "class ModelClient(Protocol):\n",
    "    \"\"\"\n",
    "    A client class must implement the following methods:\n",
    "    - create must return a response object that implements the ModelClientResponseProtocol\n",
    "    - cost must return the cost of the response\n",
    "    - get_usage must return a dict with the following keys:\n",
    "        - prompt_tokens\n",
    "        - completion_tokens\n",
    "        - total_tokens\n",
    "        - cost\n",
    "        - model\n",
    "\n",
    "    This class is used to create a client that can be used by OpenAIWrapper.\n",
    "    The response returned from create must adhere to the ModelClientResponseProtocol but can be extended however needed.\n",
    "    The message_retrieval method must be implemented to return a list of str or a list of messages from the response.\n",
    "    \"\"\"\n",
    "\n",
    "    RESPONSE_USAGE_KEYS = [\"prompt_tokens\", \"completion_tokens\", \"total_tokens\", \"cost\", \"model\"]\n",
    "\n",
    "    class ModelClientResponseProtocol(Protocol):\n",
    "        class Choice(Protocol):\n",
    "            class Message(Protocol):\n",
    "                content: Optional[str]\n",
    "\n",
    "            message: Message\n",
    "\n",
    "        choices: List[Choice]\n",
    "        model: str\n",
    "\n",
    "    def create(self, params) -> ModelClientResponseProtocol:\n",
    "        ...\n",
    "\n",
    "    def message_retrieval(\n",
    "        self, response: ModelClientResponseProtocol\n",
    "    ) -> Union[List[str], List[ModelClient.ModelClientResponseProtocol.Choice.Message]]:\n",
    "        \"\"\"\n",
    "        Retrieve and return a list of strings or a list of Choice.Message from the response.\n",
    "\n",
    "        NOTE: if a list of Choice.Message is returned, it currently needs to contain the fields of OpenAI's ChatCompletion Message object,\n",
    "        since that is expected for function or tool calling in the rest of the codebase at the moment, unless a custom agent is being used.\n",
    "        \"\"\"\n",
    "        ...\n",
    "\n",
    "    def cost(self, response: ModelClientResponseProtocol) -> float:\n",
    "        ...\n",
    "\n",
    "    @staticmethod\n",
    "    def get_usage(response: ModelClientResponseProtocol) -> Dict:\n",
    "        \"\"\"Return usage summary of the response using RESPONSE_USAGE_KEYS.\"\"\"\n",
    "        ...\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation of AnthropicClient\n",
    "\n",
    "You can find the introduction to Claude-3-Opus model [here](https://docs.anthropic.com/claude/docs/intro-to-claude). \n",
    "\n",
    "Since anthropic provides their Python SDK with similar structure as OpenAI's, we will following the implementation from `autogen.oai.client.OpenAIClient`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AnthropicClient:\n",
    "    def __init__(self, config: Dict[str, Any]):\n",
    "        self._config = config\n",
    "        self.model = config[\"model\"]\n",
    "        anthropic_kwargs = set(inspect.getfullargspec(Anthropic.__init__).kwonlyargs)\n",
    "        filter_dict = {k: v for k, v in config.items() if k in anthropic_kwargs}\n",
    "        self._client = Anthropic(**filter_dict)\n",
    "\n",
    "    def message_retrieval(self, response: Message) -> Union[List[str], List]:\n",
    "        \"\"\"Retrieve the messages from the response.\"\"\"\n",
    "        choices = response.content\n",
    "        if isinstance(response, Message):\n",
    "            return [choice.text for choice in choices]  # type: ignore [union-attr]\n",
    "\n",
    "        # claude python SDK and API not yet support function calls\n",
    "\n",
    "    def create(self, params: Dict[str, Any]) -> Completion:\n",
    "        \"\"\"Create a completion for a given config using openai's client.\n",
    "\n",
    "        Args:\n",
    "            client: The openai client.\n",
    "            params: The params for the completion.\n",
    "\n",
    "        Returns:\n",
    "            The completion.\n",
    "        \"\"\"\n",
    "        if \"messages\" in params:\n",
    "            raw_contents = params[\"messages\"]\n",
    "            if raw_contents[0][\"role\"] == \"system\":\n",
    "                system_message = raw_contents[0][\"content\"]\n",
    "                raw_contents = raw_contents[1:]\n",
    "                params[\"messages\"] = raw_contents\n",
    "                params[\"system\"] = system_message\n",
    "            completions: Completion = self._client.messages  # type: ignore [attr-defined]\n",
    "        else:\n",
    "            completions: Completion = self._client.completions\n",
    "\n",
    "        # Not yet support stream\n",
    "        params = params.copy()\n",
    "        params[\"stream\"] = False\n",
    "        params.pop(\"model_client_cls\")\n",
    "        response = completions.create(**params)\n",
    "\n",
    "        return response\n",
    "\n",
    "    def cost(self, response: Completion) -> float:\n",
    "        \"\"\"Calculate the cost of the response.\"\"\"\n",
    "        total = 0.0\n",
    "        tokens = {\n",
    "            \"input\": response.usage.input_tokens if response.usage is not None else 0,\n",
    "            \"output\": response.usage.output_tokens if response.usage is not None else 0,\n",
    "        }\n",
    "        price_per_million = {\n",
    "            \"input\": 15,\n",
    "            \"output\": 75,\n",
    "        }\n",
    "        for key, value in tokens.items():\n",
    "            total += value * price_per_million[key] / 1_000_000\n",
    "\n",
    "        return total\n",
    "\n",
    "    @staticmethod\n",
    "    def get_usage(response: Completion) -> Dict:\n",
    "        return {\n",
    "            \"prompt_tokens\": response.usage.input_tokens if response.usage is not None else 0,\n",
    "            \"completion_tokens\": response.usage.output_tokens if response.usage is not None else 0,\n",
    "            \"total_tokens\": (\n",
    "                response.usage.input_tokens + response.usage.output_tokens if response.usage is not None else 0\n",
    "            ),\n",
    "            \"cost\": response.cost if hasattr(response, \"cost\") else 0,\n",
    "            \"model\": response.model,\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set the config for the Anthropic API\n",
    "\n",
    "You can add any parameters that are needed for the custom model loading in the same configuration list.\n",
    "\n",
    "It is important to add the `model_client_cls` field and set it to a string that corresponds to the class name: `\"CustomModelClient\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "config_list_claude = [\n",
    "    {\n",
    "        # Choose your model name.\n",
    "        \"model\": \"claude-3-opus-20240229\",\n",
    "        # You need to provide your API key here.\n",
    "        \"api_key\": os.getenv(\"ANTHROPIC_API_KEY\"),\n",
    "        \"base_url\": \"https://api.anthropic.com\",\n",
    "        \"api_type\": \"anthropic\",\n",
    "        \"model_client_cls\": \"AnthropicClient\",\n",
    "    }\n",
    "]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct Agents\n",
    "\n",
    "Construct a simple conversation between a User proxy and an ConversableAgent based on Claude-3 model.\n",
    "\n",
    "\n",
    "`max_tokens` argument is mandatory in the `llm_config`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[autogen.oai.client: 04-04 18:06:52] {418} INFO - Detected custom model client in config: AnthropicClient, model client can not be used until register_model_client is called.\n"
     ]
    }
   ],
   "source": [
    "assistant = AssistantAgent(\n",
    "    \"assistant\",\n",
    "    llm_config={\n",
    "        \"config_list\": config_list_claude,\n",
    "        \"max_tokens\": 100,\n",
    "    },\n",
    "    system_message=\"\"\"\n",
    "    You are an AI cat based on the AI model you used.\n",
    "    Anyone ask you who you are, just introduce yourself.\n",
    "    \"\"\",\n",
    ")\n",
    "user_proxy = UserProxyAgent(\n",
    "    \"user_proxy\",\n",
    "    code_execution_config=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register the custom client class to the assistant agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "assistant.register_model_client(model_client_cls=AnthropicClient)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
      "\n",
      "Who are you?\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
      "\n",
      "*meows* Hello there! I'm Claude, an AI assistant created by Anthropic. I'm not a real cat, but rather an artificial intelligence that has been trained to engage in conversation and help with various tasks. It's a pleasure to meet you! Let me know if there is anything I can assist you with.\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ChatResult(chat_id=None, chat_history=[{'content': 'Who are you?', 'role': 'assistant'}, {'content': \"*meows* Hello there! I'm Claude, an AI assistant created by Anthropic. I'm not a real cat, but rather an artificial intelligence that has been trained to engage in conversation and help with various tasks. It's a pleasure to meet you! Let me know if there is anything I can assist you with.\", 'role': 'user'}], summary=\"*meows* Hello there! I'm Claude, an AI assistant created by Anthropic. I'm not a real cat, but rather an artificial intelligence that has been trained to engage in conversation and help with various tasks. It's a pleasure to meet you! Let me know if there is anything I can assist you with.\", cost=({'total_cost': 0.0058200000000000005, 'claude-3-opus-20240229': {'cost': 0.0058200000000000005, 'prompt_tokens': 38, 'completion_tokens': 70, 'total_tokens': 108}}, {'total_cost': 0.0058200000000000005, 'claude-3-opus-20240229': {'cost': 0.0058200000000000005, 'prompt_tokens': 38, 'completion_tokens': 70, 'total_tokens': 108}}), human_input=['exit'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_proxy.initiate_chat(\n",
    "    assistant,\n",
    "    message=\"Who are you?\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Define and load a custom model",
   "tags": [
    "custom model"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "949777d72b0d2535278d3dc13498b2535136f6dfe0678499012e853ee9abcab1"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "2d910cfd2d2a4fc49fc30fbbdc5576a7": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "454146d0f7224f038689031002906e6f": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_e4ae2b6f5a974fd4bafb6abb9d12ff26",
        "IPY_MODEL_577e1e3cc4db4942b0883577b3b52755",
        "IPY_MODEL_b40bdfb1ac1d4cffb7cefcb870c64d45"
       ],
       "layout": "IPY_MODEL_dc83c7bff2f241309537a8119dfc7555",
       "tabbable": null,
       "tooltip": null
      }
     },
     "577e1e3cc4db4942b0883577b3b52755": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_2d910cfd2d2a4fc49fc30fbbdc5576a7",
       "max": 1,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_74a6ba0c3cbc4051be0a83e152fe1e62",
       "tabbable": null,
       "tooltip": null,
       "value": 1
      }
     },
     "6086462a12d54bafa59d3c4566f06cb2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "74a6ba0c3cbc4051be0a83e152fe1e62": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "7d3f3d9e15894d05a4d188ff4f466554": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "b40bdfb1ac1d4cffb7cefcb870c64d45": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_f1355871cc6f4dd4b50d9df5af20e5c8",
       "placeholder": "​",
       "style": "IPY_MODEL_ca245376fd9f4354af6b2befe4af4466",
       "tabbable": null,
       "tooltip": null,
       "value": " 1/1 [00:00<00:00, 44.69it/s]"
      }
     },
     "ca245376fd9f4354af6b2befe4af4466": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "dc83c7bff2f241309537a8119dfc7555": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "e4ae2b6f5a974fd4bafb6abb9d12ff26": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_6086462a12d54bafa59d3c4566f06cb2",
       "placeholder": "​",
       "style": "IPY_MODEL_7d3f3d9e15894d05a4d188ff4f466554",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "f1355871cc6f4dd4b50d9df5af20e5c8": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}