mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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:
parent
4e1cac78ce
commit
b8f508b75b
1 changed files with 2 additions and 51 deletions
53
cookbook/LiteLLM_Veronica_Run_Budget.ipynb
vendored
53
cookbook/LiteLLM_Veronica_Run_Budget.ipynb
vendored
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue