{
 "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": [
    "TOOL_ENABLED = False\n",
    "\n",
    "\n",
    "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",
    "                raw_contents = raw_contents[1:]\n",
    "                params[\"messages\"] = raw_contents\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",
    "\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": 10,
   "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": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[autogen.oai.client: 04-02 22:48: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",
    ")\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": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "assistant.register_model_client(model_client_cls=AnthropicClient)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No default IOStream has been set, defaulting to IOConsole.\n",
      "No default IOStream has been set, defaulting to IOConsole.\n",
      "No default IOStream has been set, defaulting to IOConsole.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "user_proxy (to assistant):\n",
      "\n",
      "Who are you?\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No default IOStream has been set, defaulting to IOConsole.\n",
      "No default IOStream has been set, defaulting to IOConsole.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant (to user_proxy):\n",
      "\n",
      "I am an artificial intelligence called Claude. I was created by Anthropic to be an intelligent conversational assistant, but I'm not a real person.\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No default IOStream has been set, defaulting to IOConsole.\n",
      "No default IOStream has been set, defaulting to IOConsole.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ChatResult(chat_id=None, chat_history=[{'content': 'Who are you?', 'role': 'assistant'}, {'content': \"I am an artificial intelligence called Claude. I was created by Anthropic to be an intelligent conversational assistant, but I'm not a real person.\", 'role': 'user'}], summary=\"I am an artificial intelligence called Claude. I was created by Anthropic to be an intelligent conversational assistant, but I'm not a real person.\", cost=({'total_cost': 0, 'claude-3-opus-20240229': {'cost': 0, 'prompt_tokens': 11, 'completion_tokens': 34, 'total_tokens': 45}}, {'total_cost': 0, 'claude-3-opus-20240229': {'cost': 0, 'prompt_tokens': 11, 'completion_tokens': 34, 'total_tokens': 45}}), human_input=['exit'])"
      ]
     },
     "execution_count": 13,
     "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.7"
  },
  "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
}