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.
This commit is contained in:
amabito 2026-03-12 20:03:17 +09:00
parent 4e1cac78ce
commit b8f508b75b

View file

@ -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",