From fe6edbef4cd9257df4d68e0f4523c1d4f5363759 Mon Sep 17 00:00:00 2001
From: Srivatsa03 <180591884+Srivatsa03@users.noreply.github.com>
Date: Sun, 11 Oct 2026 22:57:34 -0500
Subject: [PATCH] fix(proxy): tolerate a StandardLoggingPayload without
 total_tokens

Some call types (e.g. aretrieve_batch) build a payload without total_tokens,
so indexing it raised KeyError and tripped the cost-tracking error path. Read
it with a .get default of 0 like the DB-side write, suppressing just the one
overload-unknown diagnostic with a reason.
---
 litellm/proxy/hooks/proxy_track_cost_callback.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py
index 6144dc37e7a..4ebc44630ed 100644
--- a/litellm/proxy/hooks/proxy_track_cost_callback.py
+++ b/litellm/proxy/hooks/proxy_track_cost_callback.py
@@ -741,7 +741,9 @@ def _extract_total_tokens(sl_object: StandardLoggingPayload | None) -> int:
     """
     if sl_object is None:
         return 0
-    return sl_object["total_tokens"]
+    return sl_object.get(  # pyright: ignore[reportUnknownMemberType]  # StandardLoggingPayload.get overload is partially unknown; total_tokens is a required int
+        "total_tokens", 0
+    )
 
 
 async def _update_database_and_spend_counters(