{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = []\n",
    "\n",
    "add_user_message(\n",
    "    messages,\n",
    "    \"Create and save a fake computer science article\",\n",
    ")\n",
    "\n",
    "run_conversation(\n",
    "    messages,\n",
    "    tools=[save_article_schema],\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
}