From 89fb1fc5731ba4185bb19e248e101a598a9752b1 Mon Sep 17 00:00:00 2001 From: Jacob Sznajdman Date: Tue, 2 Jun 2026 11:49:42 +0200 Subject: [PATCH] feat(llm): expose cached_tokens from Vertex AI context caching in LLMUsage and OTel traces - Add `cached_tokens` field to `LLMUsage` for prompt tokens served from Vertex AI's implicit context cache (cached_content_token_count). - Update `VertexAILLM._parse_content_response` to populate `cached_tokens` and, when non-zero, set the `llm.token_count.prompt_details.cache_read` span attribute on the currently active OpenTelemetry span. The OTel attribute is set inside `_parse_content_response`, which is called while the openinference-instrumentation-vertexai span is still active, so the attribute appears on the correct LLM span in Cloud Trace. Co-Authored-By: Claude Sonnet 4.6 --- src/neo4j_graphrag/llm/types.py | 4 ++++ src/neo4j_graphrag/llm/vertexai_llm.py | 10 ++++++++++ 2 files changed, 14 insertions(+) diff --git a/src/neo4j_graphrag/llm/types.py b/src/neo4j_graphrag/llm/types.py index d199919ec..84eb116fa 100644 --- a/src/neo4j_graphrag/llm/types.py +++ b/src/neo4j_graphrag/llm/types.py @@ -27,11 +27,15 @@ class LLMUsage(BaseModel): ``None`` when not reported by the provider. total_tokens (Optional[int]): Total tokens consumed by the call. ``None`` when not reported by the provider. + cached_tokens (Optional[int]): Number of prompt tokens served from the + provider's context cache (e.g. Vertex AI implicit caching). + ``None`` when not reported by the provider. """ request_tokens: Optional[int] = None response_tokens: Optional[int] = None total_tokens: Optional[int] = None + cached_tokens: Optional[int] = None class LLMResponse(BaseModel): diff --git a/src/neo4j_graphrag/llm/vertexai_llm.py b/src/neo4j_graphrag/llm/vertexai_llm.py index f9ee804be..313899e7e 100644 --- a/src/neo4j_graphrag/llm/vertexai_llm.py +++ b/src/neo4j_graphrag/llm/vertexai_llm.py @@ -593,9 +593,19 @@ def _parse_content_response(self, response: GenerationResponse) -> LLMResponse: usage = None metadata = response.usage_metadata if metadata: + cached = metadata.cached_content_token_count or None usage = LLMUsage( request_tokens=metadata.prompt_token_count, response_tokens=metadata.candidates_token_count, total_tokens=metadata.total_token_count, + cached_tokens=cached, ) + if cached: + from openinference.semconv.trace import SpanAttributes + from opentelemetry import trace as otel_trace + + otel_trace.get_current_span().set_attribute( + SpanAttributes.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ, + cached, + ) return LLMResponse(content=response.text, usage=usage)