Skip to content

fix: emit gen_ai.usage.cache_write.input_tokens per the semconv rename #69

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
PabloPardoGarcia wants to merge 1 commit into
base: main
Choose a base branch
from
+15 −12
Open

fix: emit gen_ai.usage.cache_write.input_tokens per the semconv rename #69

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
13 changes: 7 additions & 6 deletions src/rius/generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +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_CACHE_WRITE_INPUT_TOKENS,
GEN_AI_USAGE_INPUT_TOKENS,
GEN_AI_USAGE_OUTPUT_TOKENS,
TRACER_NAME,
Expand Down Expand Up @@ -158,7 +158,7 @@ def set_usage(
input_tokens: int | None = None,
output_tokens: int | None = None,
cache_read_input_tokens: int | None = None,
cache_creation_input_tokens: int | None = None,
cache_write_input_tokens: int | None = None,
) -> None:
"""Record token usage (``gen_ai.usage.*`` attributes).

Expand All @@ -179,19 +179,20 @@ def set_usage(
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
cache_write_input_tokens: Input tokens written to a
provider-managed prompt cache
(``gen_ai.usage.cache_creation.input_tokens``).
(``gen_ai.usage.cache_write.input_tokens``, called
"cache creation" by Anthropic).
"""
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:
if cache_write_input_tokens is not None:
self._span.set_attribute(
GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cache_creation_input_tokens
GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS, cache_write_input_tokens
)

def record_first_token(self) -> None:
Expand Down
4 changes: 3 additions & 1 deletion src/rius/semconv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,9 @@
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"
# semconv-genai renamed cache_creation -> cache_write (PR #440) before our
# cache fields first shipped; SDKs 0.12.x emitted the pre-rename name.
GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS = "gen_ai.usage.cache_write.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"
Expand Down
10 changes: 5 additions & 5 deletions tests/test_generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,27 +58,27 @@ def test_cm_usage_cache_tokens(exported_spans: InMemorySpanExporter) -> None:
input_tokens=10,
output_tokens=202,
cache_read_input_tokens=11579,
cache_creation_input_tokens=12694,
cache_write_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
assert attrs["gen_ai.usage.cache_write.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
assert "gen_ai.usage.cache_write.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)
gen.set_usage(cache_read_input_tokens=0, cache_write_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
assert attrs["gen_ai.usage.cache_write.input_tokens"] == 0


def test_cm_finish_reasons_list(exported_spans: InMemorySpanExporter) -> None:
Expand Down