Skip to content

fix: [AIPLAT-1489]: stop double-counting LiteLLM cache tokens as gen_ai input #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shreyas-n-harness wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix: [AIPLAT-1489]: stop double-counting LiteLLM cache tokens as gen_ai input #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0486f3e
fix: [AIPLAT-1489]: emit cache-exclusive gen_ai.usage.input_tokens fo…
shreyas-n-harness Sep 2, 2026
7e53996
chore: drop explanatory comment from LiteLLM usage test
shreyas-n-harness Sep 2, 2026
0371d27
test: avoid LiteLLM ModelResponse on Python 3.10
shreyas-n-harness Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 19 additions & 6 deletions src/harness_sdk/instrumentation/litellm/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -438,8 +438,23 @@ def _set_response_attributes(
prompt_details = _get_value(usage, "prompt_tokens_details")
completion_details = _get_value(usage, "completion_tokens_details")

input_tokens = _get_value(usage, "prompt_tokens")
if input_tokens is None:
cache_read = _get_value(usage, "cache_read_input_tokens")
if cache_read is None:
cache_read = _get_value(prompt_details, "cached_tokens")

@t-santoshsahu t-santoshsahu Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe this should be handled at backend. SDK should generate the info that is needed.

cache_creation = _get_value(usage, "cache_creation_input_tokens")
if cache_creation is None:
cache_creation = _get_value(prompt_details, "cache_creation_tokens")

prompt_tokens = _get_value(usage, "prompt_tokens")
if prompt_tokens is not None:
# LiteLLM's OpenAI-shaped prompt_tokens is cache-inclusive. Emit
# disjoint gen_ai.usage buckets so input_tokens matches uncached input.
input_tokens = max(
int(prompt_tokens or 0) - int(cache_read or 0) - int(cache_creation or 0),
0,
)
else:
# Anthropic-shaped input_tokens is already cache-exclusive.
input_tokens = _get_value(usage, "input_tokens")

output_tokens = _get_value(usage, "completion_tokens")
Expand All @@ -458,15 +473,13 @@ def _set_response_attributes(
otel_logger,
span,
"gen_ai.usage.cache_read.input_tokens",
_get_value(usage, "cache_read_input_tokens")
or _get_value(prompt_details, "cached_tokens"),
cache_read,
)
_set_if_present(
otel_logger,
span,
"gen_ai.usage.cache_creation.input_tokens",
_get_value(usage, "cache_creation_input_tokens")
or _get_value(prompt_details, "cache_creation_tokens"),
cache_creation,
)
_set_if_present(
otel_logger,
Expand Down
72 changes: 66 additions & 6 deletions test/instrumentation/litellm/litellm_instrumentation_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,21 @@
from test.control_test_helpers import AlwaysBlockControlPlugin
"""Tests for LiteLLM instrumentation (gen_ai spans + evaluate_agent_span)."""

import sys
from types import SimpleNamespace
from unittest.mock import patch

import pytest

pytest.importorskip("litellm")

import litellm
from litellm.types.utils import EmbeddingResponse, ModelResponse
from litellm.types.utils import EmbeddingResponse

# LiteLLM's ModelResponse/Message models fail to rebuild on Python 3.10 with
# current pydantic (ChatCompletionReasoningSummaryTextBlock). Tests that hit
# real litellm.completion() (streaming / mock_response) skip on 3.10.
_SKIP_REAL_LITELLM_COMPLETION = sys.version_info < (3, 11)

from harness_sdk.plugins.control import ControlResult, get_control_registry
from harness_sdk.gen_ai.exceptions import ControlEvaluationBlocked
Expand All @@ -25,7 +32,7 @@ def litellm_instrumentor():


def _fake_model_response(*_args, **_kwargs):
return ModelResponse(
return SimpleNamespace(
id="chatcmpl-test",
choices=[
{
Expand All @@ -45,6 +52,7 @@ def _fake_model_response(*_args, **_kwargs):
},
"completion_tokens_details": {"reasoning_tokens": 1},
},
_hidden_params={},
)


Expand Down Expand Up @@ -88,7 +96,7 @@ def test_litellm_completion_span_has_gen_ai_attributes(agent, exporter, litellm_
assert attrs.get("gen_ai.response.model") == "gpt-4o-mini"
assert attrs.get("gen_ai.response.id") == "chatcmpl-test"
assert attrs.get("gen_ai.response.finish_reasons") == "['stop']"
assert attrs.get("gen_ai.usage.input_tokens") == 3
assert attrs.get("gen_ai.usage.input_tokens") == 0
assert attrs.get("gen_ai.usage.output_tokens") == 5
assert attrs.get("gen_ai.usage.total_tokens") == 8
assert attrs.get("gen_ai.usage.cache_read.input_tokens") == 1
Expand Down Expand Up @@ -198,11 +206,15 @@ async def _fake_async(*_args, **_kwargs):
assert attrs.get("gen_ai.response.model") == "gpt-4o-mini"
assert attrs.get("gen_ai.response.id") == "chatcmpl-test"
assert attrs.get("gen_ai.response.finish_reasons") == "['stop']"
assert attrs.get("gen_ai.usage.input_tokens") == 3
assert attrs.get("gen_ai.usage.input_tokens") == 0
assert attrs.get("gen_ai.usage.output_tokens") == 5
assert attrs.get("gen_ai.usage.total_tokens") == 8


@pytest.mark.skipif(
_SKIP_REAL_LITELLM_COMPLETION,
reason="LiteLLM ModelResponse is broken on Python 3.10 + current pydantic",
)
def test_litellm_streaming_span_defers_until_consumed(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument
litellm_instrumentor.instrument()
stream = litellm.completion(
Expand Down Expand Up @@ -231,6 +243,10 @@ def test_litellm_streaming_span_defers_until_consumed(agent, exporter, litellm_i


@pytest.mark.asyncio
@pytest.mark.skipif(
_SKIP_REAL_LITELLM_COMPLETION,
reason="LiteLLM ModelResponse is broken on Python 3.10 + current pydantic",
)
async def test_litellm_async_streaming_span_defers_until_consumed(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument
litellm_instrumentor.instrument()
stream = await litellm.acompletion(
Expand Down Expand Up @@ -288,7 +304,10 @@ async def test_litellm_async_embedding_emits_single_span(agent, exporter, litell

@pytest.mark.asyncio
async def test_litellm_async_completion_emits_single_span(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument
with patch("litellm.main.completion", new=_fake_model_response):
async def _fake_async(*_args, **_kwargs):
return _fake_model_response()

with patch("litellm.main.acompletion", new=_fake_async):
litellm_instrumentor.instrument()
await litellm.acompletion(
model="gpt-4o-mini",
Expand All @@ -300,10 +319,47 @@ async def test_litellm_async_completion_emits_single_span(agent, exporter, litel
llm_spans = _litellm_spans(spans)
assert len(llm_spans) == 1
attrs = llm_spans[0].attributes
assert attrs.get("gen_ai.usage.input_tokens") == 3
assert attrs.get("gen_ai.usage.input_tokens") == 0
assert attrs.get("gen_ai.usage.output_tokens") == 5


def test_litellm_anthropic_input_tokens_are_already_cache_exclusive(
agent, exporter, litellm_instrumentor
): # pylint: disable=unused-argument
def _anthropic_shaped(*_args, **_kwargs):
return {
"id": "msg-test",
"model": "claude-sonnet-4",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hello"},
"finish_reason": "stop",
}
],
"usage": {
"input_tokens": 15,
"output_tokens": 5,
"cache_read_input_tokens": 80,
"cache_creation_input_tokens": 5,
},
}

with patch("litellm.main.completion", new=_anthropic_shaped):
litellm_instrumentor.instrument()
litellm.completion(
model="claude-sonnet-4",
messages=[{"role": "user", "content": "hi"}],
)

spans = exporter.get_finished_spans()
exporter.clear()
attrs = _request_span(spans).attributes
assert attrs.get("gen_ai.usage.input_tokens") == 15
assert attrs.get("gen_ai.usage.cache_read.input_tokens") == 80
assert attrs.get("gen_ai.usage.cache_creation.input_tokens") == 5


def test_litellm_embedding_dict_response(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument
def _dict_embedding(*_args, **_kwargs):
return {
Expand Down Expand Up @@ -451,6 +507,10 @@ def counting_fake(*_a, **_k):
assert len(spans) == 1


@pytest.mark.skipif(
_SKIP_REAL_LITELLM_COMPLETION,
reason="LiteLLM ModelResponse is broken on Python 3.10 + current pydantic",
)
def test_litellm_mock_response_with_wrapper_enrichment(agent, exporter, litellm_instrumentor): # pylint: disable=unused-argument
litellm_instrumentor.instrument()
litellm.completion(
Expand Down
Loading