From 0486f3e438fb4d0e6baa5e581f2afa5212d7ab5f Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Wed, 2 Sep 2026 12:16:40 +0530 Subject: [PATCH 1/3] fix: [AIPLAT-1489]: emit cache-exclusive gen_ai.usage.input_tokens for LiteLLM LiteLLM prompt_tokens is cache-inclusive. Copying it onto gen_ai.usage.input_tokens while also exporting cache read/write double-counted prompt cache against GCP-style disjoint buckets. Anthropic-shaped input_tokens is already exclusive and is left unchanged. Co-authored-by: Cursor --- .../instrumentation/litellm/__init__.py | 25 ++++++++--- .../litellm/litellm_instrumentation_test.py | 45 +++++++++++++++++-- 2 files changed, 61 insertions(+), 9 deletions(-) diff --git a/src/harness_sdk/instrumentation/litellm/__init__.py b/src/harness_sdk/instrumentation/litellm/__init__.py index d0220b8..f182c88 100644 --- a/src/harness_sdk/instrumentation/litellm/__init__.py +++ b/src/harness_sdk/instrumentation/litellm/__init__.py @@ -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") + 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") @@ -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, diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index e921a73..8ba5aa6 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -88,7 +88,9 @@ 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 + # LiteLLM prompt_tokens is cache-inclusive (3 = 0 uncached + 1 read + 2 write). + # gen_ai.usage.input_tokens must be the uncached remainder so buckets are disjoint. + 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 @@ -198,7 +200,7 @@ 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 @@ -300,10 +302,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 { From 7e53996f24b38d9db1f4ebc1d88a125e3236932a Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Wed, 2 Sep 2026 12:19:30 +0530 Subject: [PATCH 2/3] chore: drop explanatory comment from LiteLLM usage test Co-authored-by: Cursor --- test/instrumentation/litellm/litellm_instrumentation_test.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index 8ba5aa6..ec625f8 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -88,8 +88,6 @@ 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']" - # LiteLLM prompt_tokens is cache-inclusive (3 = 0 uncached + 1 read + 2 write). - # gen_ai.usage.input_tokens must be the uncached remainder so buckets are disjoint. 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 From 0371d27ddf1687926ad99863373afe3346067c61 Mon Sep 17 00:00:00 2001 From: Shreyas Nagaraj Date: Wed, 2 Sep 2026 18:36:29 +0530 Subject: [PATCH 3/3] test: avoid LiteLLM ModelResponse on Python 3.10 Current litellm + pydantic cannot rebuild Message on 3.10, which broke CI. Use a SimpleNamespace fake for patched completions and skip tests that call real completion/mock_response on 3.10. Co-authored-by: Cursor --- .../litellm/litellm_instrumentation_test.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/test/instrumentation/litellm/litellm_instrumentation_test.py b/test/instrumentation/litellm/litellm_instrumentation_test.py index ec625f8..c65b483 100644 --- a/test/instrumentation/litellm/litellm_instrumentation_test.py +++ b/test/instrumentation/litellm/litellm_instrumentation_test.py @@ -1,6 +1,8 @@ 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 @@ -8,7 +10,12 @@ 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 @@ -25,7 +32,7 @@ def litellm_instrumentor(): def _fake_model_response(*_args, **_kwargs): - return ModelResponse( + return SimpleNamespace( id="chatcmpl-test", choices=[ { @@ -45,6 +52,7 @@ def _fake_model_response(*_args, **_kwargs): }, "completion_tokens_details": {"reasoning_tokens": 1}, }, + _hidden_params={}, ) @@ -203,6 +211,10 @@ async def _fake_async(*_args, **_kwargs): 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( @@ -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( @@ -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", @@ -488,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(