diff --git a/src/rius/generation.py b/src/rius/generation.py index 61c9bf9..be1bdc6 100644 --- a/src/rius/generation.py +++ b/src/rius/generation.py @@ -27,6 +27,8 @@ GEN_AI_REQUEST_PREFIX, GEN_AI_RESPONSE_FINISH_REASONS, GEN_AI_RESPONSE_MODEL, + GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, + GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, TRACER_NAME, @@ -151,21 +153,46 @@ def set_response_model(self, response_model: str) -> None: self._span.set_attribute(GEN_AI_RESPONSE_MODEL, response_model) def set_usage( - self, *, input_tokens: int | None = None, output_tokens: int | None = None + self, + *, + input_tokens: int | None = None, + output_tokens: int | None = None, + cache_read_input_tokens: int | None = None, + cache_creation_input_tokens: int | None = None, ) -> None: - """Record token usage (``gen_ai.usage.input_tokens`` / ``output_tokens``). + """Record token usage (``gen_ai.usage.*`` attributes). Send token counts, never cost: cost is computed server-side from model pricing. + Pass provider-reported values as-is. Per the GenAI conventions, + ``input_tokens`` is the total including cached tokens (the cache + counts are subsets of it); some providers instead report an + ``input_tokens`` that excludes cache tokens (e.g. Anthropic, whose + OpenAI-style counterpart already includes them). The backend detects + and normalizes the exclusive case, so no client-side arithmetic is + needed. + Args: input_tokens: Prompt tokens consumed, when known. output_tokens: Completion tokens produced, when known. + cache_read_input_tokens: Input tokens served from a + provider-managed prompt cache + (``gen_ai.usage.cache_read.input_tokens``). + cache_creation_input_tokens: Input tokens written to a + provider-managed prompt cache + (``gen_ai.usage.cache_creation.input_tokens``). """ if input_tokens is not None: self._span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, input_tokens) if output_tokens is not None: self._span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) + if cache_read_input_tokens is not None: + self._span.set_attribute(GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read_input_tokens) + if cache_creation_input_tokens is not None: + self._span.set_attribute( + GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cache_creation_input_tokens + ) def record_first_token(self) -> None: """Mark the arrival of the first streamed token (``gen_ai.first_token`` event). diff --git a/src/rius/semconv.py b/src/rius/semconv.py index 7be0df0..112c683 100644 --- a/src/rius/semconv.py +++ b/src/rius/semconv.py @@ -40,6 +40,8 @@ GEN_AI_RESPONSE_MODEL = "gen_ai.response.model" GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens" GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens" +GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read.input_tokens" +GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS = "gen_ai.usage.cache_creation.input_tokens" GEN_AI_INPUT_MESSAGES = "gen_ai.input.messages" GEN_AI_OUTPUT_MESSAGES = "gen_ai.output.messages" GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons" diff --git a/tests/test_generation.py b/tests/test_generation.py index db46132..8ed0dab 100644 --- a/tests/test_generation.py +++ b/tests/test_generation.py @@ -52,6 +52,35 @@ def test_cm_output_and_usage(exported_spans: InMemorySpanExporter) -> None: assert attrs["gen_ai.response.model"] == "gpt-4o-2026-05" +def test_cm_usage_cache_tokens(exported_spans: InMemorySpanExporter) -> None: + with start_as_current_generation("chat") as gen: + gen.set_usage( + input_tokens=10, + output_tokens=202, + cache_read_input_tokens=11579, + cache_creation_input_tokens=12694, + ) + attrs = exported_spans.get_finished_spans()[0].attributes + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 11579 + assert attrs["gen_ai.usage.cache_creation.input_tokens"] == 12694 + + +def test_cm_usage_cache_tokens_omitted_are_absent(exported_spans: InMemorySpanExporter) -> None: + with start_as_current_generation("chat") as gen: + gen.set_usage(input_tokens=10, output_tokens=5) + attrs = exported_spans.get_finished_spans()[0].attributes + assert "gen_ai.usage.cache_read.input_tokens" not in attrs + assert "gen_ai.usage.cache_creation.input_tokens" not in attrs + + +def test_cm_usage_cache_tokens_zero_recorded(exported_spans: InMemorySpanExporter) -> None: + with start_as_current_generation("chat") as gen: + gen.set_usage(cache_read_input_tokens=0, cache_creation_input_tokens=0) + attrs = exported_spans.get_finished_spans()[0].attributes + assert attrs["gen_ai.usage.cache_read.input_tokens"] == 0 + assert attrs["gen_ai.usage.cache_creation.input_tokens"] == 0 + + def test_cm_finish_reasons_list(exported_spans: InMemorySpanExporter) -> None: with start_as_current_generation("chat") as gen: gen.set_finish_reasons(["stop", "length"])