From 4e1cac78ce93479d112680e067e085969046d718 Mon Sep 17 00:00:00 2001 From: amabito Date: Mon, 9 Mar 2026 19:19:44 +0900 Subject: [PATCH] fix(cookbook): add async callback for acompletion support Extract shared _record_cost helper and implement both log_success_event and async_log_success_event so the budget cap works with litellm.acompletion as well. --- cookbook/LiteLLM_Veronica_Run_Budget.ipynb | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/cookbook/LiteLLM_Veronica_Run_Budget.ipynb b/cookbook/LiteLLM_Veronica_Run_Budget.ipynb index a18cdd46c5a..a87664ad719 100644 --- a/cookbook/LiteLLM_Veronica_Run_Budget.ipynb +++ b/cookbook/LiteLLM_Veronica_Run_Budget.ipynb @@ -49,13 +49,20 @@ "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", - " cost = kwargs.get(\"response_cost\", 0.0)\n", - " if cost > 0:\n", - " budget.spend(cost)\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",