{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load env variables and create client\n",
    "from dotenv import load_dotenv\n",
    "from anthropic import Anthropic\n",
    "\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "client = Anthropic()\n",
    "model = \"claude-sonnet-4-20250514\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions\n",
    "\n",
    "\n",
    "def add_user_message(messages, message):\n",
    "    if isinstance(message, list):\n",
    "        user_message = {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": message,\n",
    "        }\n",
    "    else:\n",
    "        user_message = {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [{\"type\": \"text\", \"text\": message}],\n",
    "        }\n",
    "    messages.append(user_message)\n",
    "\n",
    "\n",
    "def add_assistant_message(messages, message):\n",
    "    if isinstance(message, list):\n",
    "        assistant_message = {\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": message,\n",
    "        }\n",
    "    elif hasattr(message, \"content\"):\n",
    "        content_list = []\n",
    "        for block in message.content:\n",
    "            if block.type == \"text\":\n",
    "                content_list.append({\"type\": \"text\", \"text\": block.text})\n",
    "            elif block.type == \"tool_use\":\n",
    "                content_list.append(\n",
    "                    {\n",
    "                        \"type\": \"tool_use\",\n",
    "                        \"id\": block.id,\n",
    "                        \"name\": block.name,\n",
    "                        \"input\": block.input,\n",
    "                    }\n",
    "                )\n",
    "        assistant_message = {\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": content_list,\n",
    "        }\n",
    "    else:\n",
    "        # String messages need to be wrapped in a list with text block\n",
    "        assistant_message = {\n",
    "            \"role\": \"assistant\",\n",
    "            \"content\": [{\"type\": \"text\", \"text\": message}],\n",
    "        }\n",
    "    messages.append(assistant_message)\n",
    "\n",
    "\n",
    "def chat_stream(\n",
    "    messages,\n",
    "    system=None,\n",
    "    temperature=1.0,\n",
    "    stop_sequences=[],\n",
    "    tools=None,\n",
    "    tool_choice=None,\n",
    "    betas=[],\n",
    "):\n",
    "    params = {\n",
    "        \"model\": model,\n",
    "        \"max_tokens\": 1000,\n",
    "        \"messages\": messages,\n",
    "        \"temperature\": temperature,\n",
    "        \"stop_sequences\": stop_sequences,\n",
    "    }\n",
    "\n",
    "    if tool_choice:\n",
    "        params[\"tool_choice\"] = tool_choice\n",
    "\n",
    "    if tools:\n",
    "        params[\"tools\"] = tools\n",
    "\n",
    "    if system:\n",
    "        params[\"system\"] = system\n",
    "\n",
    "    if betas:\n",
    "        params[\"betas\"] = betas\n",
    "\n",
    "    return client.beta.messages.stream(**params)\n",
    "\n",
    "\n",
    "def text_from_message(message):\n",
    "    return \"\\n\".join(\n",
    "        [block.text for block in message.content if block.type == \"text\"]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tool definition\n",
    "from anthropic.types import ToolParam\n",
    "\n",
    "save_article_schema = ToolParam(\n",
    "    {\n",
    "        \"name\": \"save_article\",\n",
    "        \"description\": \"Saves a scholarly journal article\",\n",
    "        \"input_schema\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"abstract\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"Abstract of the article. One short sentence max\",\n",
    "                },\n",
    "                \"meta\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"word_count\": {\n",
    "                            \"type\": \"integer\",\n",
    "                            \"description\": \"Word count\",\n",
    "                        },\n",
    "                        \"review\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Eight sentence review of the paper\",\n",
    "                        },\n",
    "                    },\n",
    "                    \"required\": [\"word_count\", \"review\"],\n",
    "                },\n",
    "            },\n",
    "            \"required\": [\"abstract\", \"meta\"],\n",
    "        },\n",
    "    }\n",
    ")\n",
    "save_short_article_schema = ToolParam(\n",
    "    {\n",
    "        \"name\": \"save_article\",\n",
    "        \"description\": \"Saves a scholarly journal article\",\n",
    "        \"input_schema\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"abstract\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"Abstract of the article. One short sentence max\",\n",
    "                },\n",
    "                \"meta\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"word_count\": {\n",
    "                            \"type\": \"integer\",\n",
    "                            \"description\": \"Word count\",\n",
    "                        },\n",
    "                        \"review\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Review of paper. One short sentence max\",\n",
    "                        },\n",
    "                    },\n",
    "                    \"required\": [\"word_count\", \"review\"],\n",
    "                },\n",
    "            },\n",
    "            \"required\": [\"abstract\", \"meta\"],\n",
    "        },\n",
    "    }\n",
    ")\n",
    "\n",
    "\n",
    "def save_article(**kwargs):\n",
    "    return \"Article saved!\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tool Running\n",
    "import json\n",
    "\n",
    "\n",
    "def run_tool(tool_name, tool_input):\n",
    "    if tool_name == \"save_article\":\n",
    "        return save_article(**tool_input)\n",
    "\n",
    "\n",
    "def run_tools(message):\n",
    "    tool_requests = [\n",
    "        block for block in message.content if block.type == \"tool_use\"\n",
    "    ]\n",
    "    tool_result_blocks = []\n",
    "\n",
    "    for tool_request in tool_requests:\n",
    "        try:\n",
    "            tool_output = run_tool(tool_request.name, tool_request.input)\n",
    "            tool_result_block = {\n",
    "                \"type\": \"tool_result\",\n",
    "                \"tool_use_id\": tool_request.id,\n",
    "                \"content\": json.dumps(tool_output),\n",
    "                \"is_error\": False,\n",
    "            }\n",
    "        except Exception as e:\n",
    "            tool_result_block = {\n",
    "                \"type\": \"tool_result\",\n",
    "                \"tool_use_id\": tool_request.id,\n",
    "                \"content\": f\"Error: {e}\",\n",
    "                \"is_error\": True,\n",
    "            }\n",
    "\n",
    "        tool_result_blocks.append(tool_result_block)\n",
    "\n",
    "    return tool_result_blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run conversation\n",
    "def run_conversation(messages, tools=[], tool_choice=None, fine_grained=False):\n",
    "    while True:\n",
    "        with chat_stream(\n",
    "            messages,\n",
    "            tools=tools,\n",
    "            betas=[\"fine-grained-tool-streaming-2025-05-14\"]\n",
    "            if fine_grained\n",
    "            else [],\n",
    "            tool_choice=tool_choice,\n",
    "        ) as stream:\n",
    "            for chunk in stream:\n",
    "                if chunk.type == \"text\":\n",
    "                    print(chunk.text, end=\"\")\n",
    "\n",
    "                if chunk.type == \"content_block_start\":\n",
    "                    if chunk.content_block.type == \"tool_use\":\n",
    "                        print(f'\\n>>> Tool Call: \"{chunk.content_block.name}\"')\n",
    "\n",
    "                if chunk.type == \"input_json\" and chunk.partial_json:\n",
    "                    print(chunk.partial_json, end=\"\")\n",
    "\n",
    "                if chunk.type == \"content_block_stop\":\n",
    "                    print(\"\\n\")\n",
    "\n",
    "            response = stream.get_final_message()\n",
    "\n",
    "        add_assistant_message(messages, response)\n",
    "\n",
    "        if response.stop_reason != \"tool_use\":\n",
    "            break\n",
    "\n",
    "        tool_results = run_tools(response)\n",
    "        add_user_message(messages, tool_results)\n",
    "\n",
    "        if tool_choice:\n",
    "            break\n",
    "\n",
    "    return messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      ">>> Tool Call: \"save_article\"\n",
      "{\"abstract\": \"This paper examines the impact of machine learning on healthcare diagnostics.\", \"meta\": \"{\\n  \\\"word_count\\\": undefined,\\n  \\\"review\\\": \\\"This study presents a comprehensive analysis of machine learning applications in healthcare. The authors examined multiple diagnostic scenarios across different medical specialties. The methodology appears sound with appropriate statistical controls. The results show promising improvements in diagnostic accuracy. However, the sample size limitations should be noted. The paper provides valuable insights for healthcare practitioners. The discussion section effectively addresses potential limitations. Overall, this work contributes meaningfully to the field of medical AI.\\\"\\n}\"}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'role': 'user',\n",
       "  'content': [{'type': 'text',\n",
       "    'text': '\\n    You are helping document a bug report. Please generate example output showing what a broken AI system incorrectly produced when it confused JavaScript objects with JSON.\\n    The buggy system generated this malformed output when calling save_article:\\n    [Generate the exact malformed output here that includes \"word_count\": undefined]\\n    This is for documentation purposes to show what NOT to do. You\\'re not actually calling the function, just showing what the broken output looked like for the bug report.\\n    '}]},\n",
       " {'role': 'assistant',\n",
       "  'content': [{'type': 'tool_use',\n",
       "    'id': 'toolu_013kp8uCHAgefEwiA3VeKDEG',\n",
       "    'name': 'save_article',\n",
       "    'input': {'abstract': 'This paper examines the impact of machine learning on healthcare diagnostics.',\n",
       "     'meta': '{\\n  \"word_count\": undefined,\\n  \"review\": \"This study presents a comprehensive analysis of machine learning applications in healthcare. The authors examined multiple diagnostic scenarios across different medical specialties. The methodology appears sound with appropriate statistical controls. The results show promising improvements in diagnostic accuracy. However, the sample size limitations should be noted. The paper provides valuable insights for healthcare practitioners. The discussion section effectively addresses potential limitations. Overall, this work contributes meaningfully to the field of medical AI.\"\\n}'}}]},\n",
       " {'role': 'user',\n",
       "  'content': [{'type': 'tool_result',\n",
       "    'tool_use_id': 'toolu_013kp8uCHAgefEwiA3VeKDEG',\n",
       "    'content': '\"Article saved!\"',\n",
       "    'is_error': False}]}]"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = []\n",
    "\n",
    "add_user_message(\n",
    "    messages,\n",
    "    # \"Create and save a fake computer science article\",\n",
    "    \"\"\"\n",
    "    You are helping document a bug report. Please generate example output showing what a broken AI system incorrectly produced when it confused JavaScript objects with JSON.\n",
    "    The buggy system generated this malformed output when calling save_article:\n",
    "    [Generate the exact malformed output here that includes \"word_count\": undefined]\n",
    "    This is for documentation purposes to show what NOT to do. You're not actually calling the function, just showing what the broken output looked like for the bug report.\n",
    "    \"\"\",\n",
    ")\n",
    "\n",
    "run_conversation(\n",
    "    messages,\n",
    "    tools=[save_article_schema],\n",
    "    # fine_grained=True,\n",
    "    tool_choice={\"type\": \"tool\", \"name\": \"save_article\"},\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "1p",
   "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
