From 8cca2c3d2c259b254f8470bab7105b474cd62252 Mon Sep 17 00:00:00 2001
From: Artem Rootman <4586640+artemrootman@users.noreply.github.com>
Date: Thu, 6 Aug 2026 16:59:22 +0300
Subject: [PATCH] feat(litellm): return token usage and cost with the response
 (#6678)

---
 .../litellm_provider/litellm_provider.py      |  8 +++
 .../test_litellm_response_parsing.py          | 49 ++++++++++++++++---
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/keep/providers/litellm_provider/litellm_provider.py b/keep/providers/litellm_provider/litellm_provider.py
index 32b22a4fc5..935332f58e 100644
--- a/keep/providers/litellm_provider/litellm_provider.py
+++ b/keep/providers/litellm_provider/litellm_provider.py
@@ -139,8 +139,16 @@ def _query(
                         f"Failed to parse generated text as JSON: {generated_text}. Model not following the structured output format. Response: {result}"
                     )
 
+            # Surface usage so a workflow can report what a run cost. Absent
+            # on some OpenAI-compatible backends, hence the .get() chain.
+            usage = result.get("usage") or {}
             return {
                 "response": generated_text,
+                "model": result.get("model", model),
+                "prompt_tokens": usage.get("prompt_tokens"),
+                "completion_tokens": usage.get("completion_tokens"),
+                "total_tokens": usage.get("total_tokens"),
+                "cost": usage.get("cost"),
             }
 
         except requests.exceptions.RequestException as e:
diff --git a/tests/providers/litellm_provider/test_litellm_response_parsing.py b/tests/providers/litellm_provider/test_litellm_response_parsing.py
index 6dd2f5bec3..c1f9ac30fb 100644
--- a/tests/providers/litellm_provider/test_litellm_response_parsing.py
+++ b/tests/providers/litellm_provider/test_litellm_response_parsing.py
@@ -34,16 +34,17 @@ def _build_provider() -> LitellmProvider:
     return LitellmProvider(ContextManager(tenant_id="test"), "litellm-test", config)
 
 
-def _response(content, finish_reason="stop"):
+def _response(content, finish_reason="stop", usage=None, model=None):
     response = MagicMock()
     response.raise_for_status = MagicMock()
-    response.json = MagicMock(
-        return_value={
-            "choices": [
-                {"message": {"content": content}, "finish_reason": finish_reason}
-            ]
-        }
-    )
+    payload = {
+        "choices": [{"message": {"content": content}, "finish_reason": finish_reason}]
+    }
+    if usage is not None:
+        payload["usage"] = usage
+    if model is not None:
+        payload["model"] = model
+    response.json = MagicMock(return_value=payload)
     return response
 
 
@@ -88,3 +89,35 @@ def test_plain_text_response_still_works():
     provider = _build_provider()
     with patch("requests.post", return_value=_response("plain answer")):
         assert provider._query(prompt="hi")["response"] == "plain answer"
+
+
+def test_usage_is_returned_alongside_the_response():
+    """A workflow needs the token count and cost to report what a run spent."""
+    provider = _build_provider()
+    usage = {
+        "prompt_tokens": 256,
+        "completion_tokens": 98,
+        "total_tokens": 354,
+        "cost": 0.00042,
+    }
+    with patch(
+        "requests.post", return_value=_response("hi", usage=usage, model="gpt-4o")
+    ):
+        result = provider._query(prompt="hi")
+    assert result["response"] == "hi"
+    assert result["prompt_tokens"] == 256
+    assert result["completion_tokens"] == 98
+    assert result["total_tokens"] == 354
+    assert result["cost"] == 0.00042
+    assert result["model"] == "gpt-4o"
+
+
+def test_missing_usage_does_not_break_the_response():
+    """Not every OpenAI-compatible backend reports usage."""
+    provider = _build_provider()
+    with patch("requests.post", return_value=_response("hi")):
+        result = provider._query(prompt="hi", model="local-model")
+    assert result["response"] == "hi"
+    assert result["prompt_tokens"] is None
+    assert result["cost"] is None
+    assert result["model"] == "local-model"