From b8f508b75b587395402ca004384a8dc0fc9e1b65 Mon Sep 17 00:00:00 2001
From: amabito <amabito@local>
Date: Thu, 12 Mar 2026 20:03:17 +0900
Subject: [PATCH] fix(cookbook): address review feedback -- append callback,
 add error handling

- Use litellm.callbacks.append() instead of assignment to preserve
  existing callbacks.
- Wrap litellm.completion() in try/except so API errors don't crash
  the loop before the final summary prints.
---
 cookbook/LiteLLM_Veronica_Run_Budget.ipynb | 53 +---------------------
 1 file changed, 2 insertions(+), 51 deletions(-)

diff --git a/cookbook/LiteLLM_Veronica_Run_Budget.ipynb b/cookbook/LiteLLM_Veronica_Run_Budget.ipynb
index a87664ad719..c63ee2f7e35 100644
--- a/cookbook/LiteLLM_Veronica_Run_Budget.ipynb
+++ b/cookbook/LiteLLM_Veronica_Run_Budget.ipynb
@@ -37,63 +37,14 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import os\n",
-    "import litellm\n",
-    "from litellm.integrations.custom_logger import CustomLogger\n",
-    "from veronica_core import BudgetEnforcer\n",
-    "\n",
-    "# Set your API key (or use any provider LiteLLM supports).\n",
-    "# os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
-    "\n",
-    "budget = BudgetEnforcer(limit_usd=0.05)\n",
-    "\n",
-    "\n",
-    "def _record_cost(kwargs):\n",
-    "    cost = kwargs.get(\"response_cost\", 0.0)\n",
-    "    if cost > 0:\n",
-    "        budget.spend(cost)\n",
-    "\n",
-    "\n",
-    "class RunBudgetCallback(CustomLogger):\n",
-    "    \"\"\"Record response_cost in BudgetEnforcer after each successful call.\"\"\"\n",
-    "\n",
-    "    def log_success_event(self, kwargs, response_obj, start_time, end_time):\n",
-    "        _record_cost(kwargs)\n",
-    "\n",
-    "    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):\n",
-    "        _record_cost(kwargs)\n",
-    "\n",
-    "\n",
-    "litellm.callbacks = [RunBudgetCallback()]\n",
-    "print(f\"Run budget: ${budget.limit_usd:.2f}\")"
-   ]
+   "source": "import os\nimport litellm\nfrom litellm.integrations.custom_logger import CustomLogger\nfrom veronica_core import BudgetEnforcer\n\n# Set your API key (or use any provider LiteLLM supports).\n# os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n\nbudget = BudgetEnforcer(limit_usd=0.05)\n\n\ndef _record_cost(kwargs):\n    cost = kwargs.get(\"response_cost\", 0.0)\n    if cost > 0:\n        budget.spend(cost)\n\n\nclass RunBudgetCallback(CustomLogger):\n    \"\"\"Record response_cost in BudgetEnforcer after each successful call.\"\"\"\n\n    def log_success_event(self, kwargs, response_obj, start_time, end_time):\n        _record_cost(kwargs)\n\n    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):\n        _record_cost(kwargs)\n\n\nlitellm.callbacks.append(RunBudgetCallback())\nprint(f\"Run budget: ${budget.limit_usd:.2f}\")"
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "MAX_CALLS = 20\n",
-    "\n",
-    "for i in range(MAX_CALLS):\n",
-    "    if budget.is_exceeded:\n",
-    "        print(f\"\\nStopped before call {i}: run budget exhausted.\")\n",
-    "        break\n",
-    "\n",
-    "    response = litellm.completion(\n",
-    "        model=\"gpt-4o-mini\",\n",
-    "        messages=[{\"role\": \"user\", \"content\": f\"Say 'hello' in one word. ({i})\"}],\n",
-    "        max_tokens=10,\n",
-    "    )\n",
-    "    print(f\"call {i:>2d}  spent=${budget.spent_usd:.4f}  \"\n",
-    "          f\"remaining=${budget.remaining_usd:.4f}  \"\n",
-    "          f\"calls={budget.call_count}\")\n",
-    "\n",
-    "print(f\"\\nFinal: ${budget.spent_usd:.4f} / ${budget.limit_usd:.2f}  \"\n",
-    "      f\"({budget.call_count} calls)\")"
-   ]
+   "source": "MAX_CALLS = 20\n\nfor i in range(MAX_CALLS):\n    if budget.is_exceeded:\n        print(f\"\\nStopped before call {i}: run budget exhausted.\")\n        break\n\n    try:\n        response = litellm.completion(\n            model=\"gpt-4o-mini\",\n            messages=[{\"role\": \"user\", \"content\": f\"Say 'hello' in one word. ({i})\"}],\n            max_tokens=10,\n        )\n    except Exception as e:\n        print(f\"call {i:>2d}  error: {type(e).__name__}\")\n        continue\n\n    print(f\"call {i:>2d}  spent=${budget.spent_usd:.4f}  \"\n          f\"remaining=${budget.remaining_usd:.4f}  \"\n          f\"calls={budget.call_count}\")\n\nprint(f\"\\nFinal: ${budget.spent_usd:.4f} / ${budget.limit_usd:.2f}  \"\n      f\"({budget.call_count} calls)\")"
   },
   {
    "cell_type": "markdown",