Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
44 changes: 44 additions & 0 deletions tests/test_llm_attributes_extractors.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,17 @@
from types import SimpleNamespace
from unittest.mock import Mock, call

import pytest
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.adk.tools.function_tool import FunctionTool
from google.genai import types
from opentelemetry.sdk.trace import TracerProvider

from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extractors import (
llm_gen_ai_request_functions,
llm_gen_ai_usage_cache_creation_input_tokens,
llm_gen_ai_usage_cache_read_input_tokens,
llm_gen_ai_usage_output_tokens,
)
from veadk.tracing.telemetry.attributes.extractors.types import ExtractorResponse
Expand Down Expand Up @@ -118,6 +123,45 @@ def test_missing_output_token_count_is_not_written_to_span():
span.set_attribute.assert_not_called()


@pytest.mark.parametrize("cached_tokens", [600, 0, None])
def test_cached_tokens_are_only_reported_as_cache_reads(cached_tokens):
params = SimpleNamespace(
llm_response=LlmResponse(
usage_metadata=types.GenerateContentResponseUsageMetadata(
cached_content_token_count=cached_tokens,
)
)
)
provider = TracerProvider()
try:
with provider.get_tracer(__name__).start_as_current_span("cache-usage") as span:
ExtractorResponse.update_span(
span,
"gen_ai.usage.cache_read_input_tokens",
llm_gen_ai_usage_cache_read_input_tokens(params),
)
ExtractorResponse.update_span(
span,
"gen_ai.usage.cache_creation_input_tokens",
llm_gen_ai_usage_cache_creation_input_tokens(params),
)
expected = (
{"gen_ai.usage.cache_read_input_tokens": cached_tokens}
if cached_tokens is not None
else {}
)
assert dict(span.attributes) == expected
finally:
provider.shutdown()


def test_missing_usage_omits_cache_attributes():
params = SimpleNamespace(llm_response=LlmResponse())

assert llm_gen_ai_usage_cache_read_input_tokens(params).content is None
assert llm_gen_ai_usage_cache_creation_input_tokens(params).content is None


def test_falsy_attribute_values_are_written_to_span():
span = Mock()

Expand Down
77 changes: 77 additions & 0 deletions tests/test_portal_metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@
from types import SimpleNamespace

import pytest
from google.adk.agents.run_config import RunConfig, StreamingMode
from google.adk.models.llm_request import LlmRequest
from google.adk.models.llm_response import LlmResponse
from google.genai import types
from opentelemetry import metrics as metrics_api
from opentelemetry.metrics import _internal as metrics_internal
from opentelemetry.sdk import metrics as metrics_sdk
Expand Down Expand Up @@ -138,6 +142,79 @@ def test_proxy_instruments_follow_provider_installed_later(
]


@pytest.mark.parametrize(
"server_address",
[
"https://ark.cn-beijing.volces.com/api/v3/",
"https://ark.ap-southeast.bytepluses.com/api/v3",
],
ids=["volcengine", "byteplus"],
)
@pytest.mark.parametrize("streaming_mode", [StreamingMode.NONE, StreamingMode.SSE])
@pytest.mark.parametrize("cached_tokens", [600, 0, None])
def test_cache_read_token_metrics_preserve_input_usage(
fresh_global_meter_provider, server_address, streaming_mode, cached_tokens
):
reader = InMemoryMetricReader()
metrics_api.set_meter_provider(metrics_sdk.MeterProvider(metric_readers=[reader]))
recorder = PortalMetricRecorder(name="test-cache-read")
context = SimpleNamespace(
agent=SimpleNamespace(model_api_base=server_address),
run_config=RunConfig(streaming_mode=streaming_mode),
)
request = LlmRequest(model="test-model")

if streaming_mode == StreamingMode.SSE:
recorder.record_call_llm(context, "partial", request, LlmResponse(partial=True))

recorder.record_call_llm(
context,
"final",
request,
LlmResponse(
usage_metadata=types.GenerateContentResponseUsageMetadata(
prompt_token_count=1000,
candidates_token_count=200,
total_token_count=1200,
cached_content_token_count=cached_tokens,
),
),
)

metrics_data = reader.get_metrics_data()
recorded_metrics = {
metric.name: list(metric.data.data_points)
for resource_metrics in metrics_data.resource_metrics
for scope_metrics in resource_metrics.scope_metrics
for metric in scope_metrics.metrics
}
points = recorded_metrics["gen_ai.client.token.usage"]
expected = {"input": 1000, "output": 200}
if cached_tokens is not None:
expected["cache_read"] = cached_tokens
assert {point.attributes["gen_ai_token_type"]: point.sum for point in points} == (
expected
)
assert len(points) == len(expected)
for point in points:
assert point.count == 1
assert point.attributes["server_address"] == server_address
assert point.attributes["stream"] == (streaming_mode == StreamingMode.SSE)
assert point.attributes["gen_ai_response_model"] == "test-model"
assert [point.value for point in recorded_metrics["gen_ai.chat.count"]] == [1]


def test_missing_usage_does_not_record_token_metrics(fresh_global_meter_provider):
reader = InMemoryMetricReader()
metrics_api.set_meter_provider(metrics_sdk.MeterProvider(metric_readers=[reader]))
recorder = PortalMetricRecorder(name="test-missing-usage")
context = SimpleNamespace(agent=SimpleNamespace(), run_config=None)

recorder.record_call_llm(context, "event", LlmRequest(), LlmResponse())

assert reader.get_metrics_data() is None


def test_apmplus_reuses_preconfigured_global_provider(
fresh_global_meter_provider,
monkeypatch,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -205,29 +205,16 @@ def llm_gen_ai_usage_total_tokens(params: LLMAttributesParams) -> ExtractorRespo
return ExtractorResponse(content=None)


# FIXME
def llm_gen_ai_usage_cache_creation_input_tokens(
params: LLMAttributesParams,
) -> ExtractorResponse:
"""Extract the number of tokens used for cache creation.
"""Omit cache creation usage until response metadata provides it

Provides the count of tokens used for creating cached content,
which affects cost calculation in caching-enabled models.

Args:
params: LLM execution parameters containing response metadata

Returns:
ExtractorResponse: Response containing cache creation token count or None
cached_content_token_count measures cache reads, not cache creation
"""
if params.llm_response.usage_metadata:
return ExtractorResponse(
content=params.llm_response.usage_metadata.cached_content_token_count,
)
return ExtractorResponse(content=None)


# FIXME
def llm_gen_ai_usage_cache_read_input_tokens(
params: LLMAttributesParams,
) -> ExtractorResponse:
Expand Down
9 changes: 7 additions & 2 deletions veadk/tracing/telemetry/portal_metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ class PortalMetricRecorder:

Metrics Collected:
- LLM invocation counts and frequencies
- Token consumption (input/output) with histogram distribution
- Token consumption (input/output/cache_read) with histogram distribution
- Operation latency with performance bucket analysis
- Error rates and exception details
- Span-level performance metrics for APMPlus dashboards
Expand Down Expand Up @@ -258,7 +258,7 @@ def record_call_llm(

Metrics Recorded:
- Invocation count with model and operation attributes
- Input/output token usage with separate tracking
- Input/output token usage and the cached subset of input tokens
- Operation duration from span timing data
- Error counts and exception details
- Span latency for performance analysis
Expand Down Expand Up @@ -290,13 +290,18 @@ def record_call_llm(
# upload token usage
input_token = llm_response.usage_metadata.prompt_token_count
output_token = llm_response.usage_metadata.candidates_token_count
cache_read_token = llm_response.usage_metadata.cached_content_token_count

if input_token:
token_attributes = {**attributes, "gen_ai_token_type": "input"}
self.token_usage.record(input_token, attributes=token_attributes)
if output_token:
token_attributes = {**attributes, "gen_ai_token_type": "output"}
self.token_usage.record(output_token, attributes=token_attributes)
# Cache reads are already included in input tokens, not extra usage
if cache_read_token is not None:
token_attributes = {**attributes, "gen_ai_token_type": "cache_read"}
self.token_usage.record(cache_read_token, attributes=token_attributes)

# Get llm duration
span = trace.get_current_span()
Expand Down
Loading