{
 "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",
    "from pathlib import Path\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "client = Anthropic(\n",
    "    default_headers={\n",
    "        \"anthropic-beta\": \"code-execution-2025-05-22, files-api-2025-04-14\"\n",
    "    }\n",
    ")\n",
    "model = \"claude-sonnet-4-20250514\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions\n",
    "from anthropic.types import Message\n",
    "\n",
    "\n",
    "def add_user_message(messages, message):\n",
    "    user_message = {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": message.content if isinstance(message, Message) else message,\n",
    "    }\n",
    "    messages.append(user_message)\n",
    "\n",
    "\n",
    "def add_assistant_message(messages, message):\n",
    "    assistant_message = {\n",
    "        \"role\": \"assistant\",\n",
    "        \"content\": message.content if isinstance(message, Message) else message,\n",
    "    }\n",
    "    messages.append(assistant_message)\n",
    "\n",
    "\n",
    "def chat(\n",
    "    messages,\n",
    "    system=None,\n",
    "    temperature=1.0,\n",
    "    stop_sequences=[],\n",
    "    tools=None,\n",
    "    thinking=False,\n",
    "    thinking_budget=2000,\n",
    "):\n",
    "    params = {\n",
    "        \"model\": model,\n",
    "        \"max_tokens\": 10000,\n",
    "        \"messages\": messages,\n",
    "        \"temperature\": temperature,\n",
    "        \"stop_sequences\": stop_sequences,\n",
    "    }\n",
    "\n",
    "    if thinking:\n",
    "        params[\"thinking\"] = {\n",
    "            \"type\": \"enabled\",\n",
    "            \"budget_tokens\": thinking_budget,\n",
    "        }\n",
    "\n",
    "    if tools:\n",
    "        params[\"tools\"] = tools\n",
    "\n",
    "    if system:\n",
    "        params[\"system\"] = system\n",
    "\n",
    "    message = client.messages.create(**params)\n",
    "    return message\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",
    "    )\n",
    "\n",
    "\n",
    "def upload(file_path):\n",
    "    path = Path(file_path)\n",
    "    extension = path.suffix.lower()\n",
    "\n",
    "    mime_type_map = {\n",
    "        \".pdf\": \"application/pdf\",\n",
    "        \".txt\": \"text/plain\",\n",
    "        \".md\": \"text/plain\",\n",
    "        \".py\": \"text/plain\",\n",
    "        \".js\": \"text/plain\",\n",
    "        \".html\": \"text/plain\",\n",
    "        \".css\": \"text/plain\",\n",
    "        \".csv\": \"text/csv\",\n",
    "        \".json\": \"application/json\",\n",
    "        \".xml\": \"application/xml\",\n",
    "        \".xlsx\": \"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet\",\n",
    "        \".xls\": \"application/vnd.ms-excel\",\n",
    "        \".jpeg\": \"image/jpeg\",\n",
    "        \".jpg\": \"image/jpeg\",\n",
    "        \".png\": \"image/png\",\n",
    "        \".gif\": \"image/gif\",\n",
    "        \".webp\": \"image/webp\",\n",
    "    }\n",
    "\n",
    "    mime_type = mime_type_map.get(extension)\n",
    "\n",
    "    if not mime_type:\n",
    "        raise ValueError(f\"Unknown mimetype for extension: {extension}\")\n",
    "    filename = path.name\n",
    "\n",
    "    with open(file_path, \"rb\") as file:\n",
    "        return client.beta.files.upload(file=(filename, file, mime_type))\n",
    "\n",
    "\n",
    "def list_files():\n",
    "    return client.beta.files.list()\n",
    "\n",
    "\n",
    "def delete_file(id):\n",
    "    return client.beta.files.delete(id)\n",
    "\n",
    "\n",
    "def download_file(id, filename=None):\n",
    "    file_content = client.beta.files.download(id)\n",
    "\n",
    "    if not filename:\n",
    "        file_metadata = get_metadata(id)\n",
    "        file_content.write_to_file(file_metadata.filename)\n",
    "    else:\n",
    "        file_content.write_to_file(filename)\n",
    "\n",
    "\n",
    "def get_metadata(id):\n",
    "    return client.beta.files.retrieve_metadata(id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_metadata = upload(\"streaming.csv\")\n",
    "file_metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = []\n",
    "\n",
    "add_user_message(\n",
    "    messages,\n",
    "    [\n",
    "        {\n",
    "            \"type\": \"text\",\n",
    "            \"text\": \"\"\"\n",
    "Run a detailed analysis to determine major drivers of churn.\n",
    "Your final output should include at least one detailed plot summarizing your findings.\n",
    "\n",
    "Critical note: Every time you execute code, you're starting with a completely clean slate. \n",
    "No variables or library imports from previous executions exist. You need to redeclare/reimport all variables/libraries.\n",
    "            \"\"\",\n",
    "        },\n",
    "        {\"type\": \"container_upload\", \"file_id\": file_metadata.id},\n",
    "    ],\n",
    ")\n",
    "\n",
    "chat(\n",
    "    messages,\n",
    "    tools=[{\"type\": \"code_execution_20250522\", \"name\": \"code_execution\"}]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_file(\"file_011CPYZqxoMSsfbrSzFw8j9X\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}