diff --git a/framework/tracing/llmspan.go b/framework/tracing/llmspan.go index a61afbe59c2..477ef3f4717 100644 --- a/framework/tracing/llmspan.go +++ b/framework/tracing/llmspan.go @@ -148,6 +148,60 @@ func PopulateErrorAttributes(err *schemas.BifrostError) map[string]any { attrs[schemas.AttrHTTPResponseStatusCode] = *err.StatusCode } + // Usage the provider billed us for even though the request failed or was + // cancelled (see BifrostError.ExtraFields.BilledUsage). Governance and the + // logging plugin already charge for this; without emitting it here every + // span-based consumer — the otel plugin and the BigQuery / Datadog / Kafka / + // Pub-Sub connectors — records zero tokens for the same request. + if u := err.ExtraFields.BilledUsage; u != nil { + // Everything below is gated on > 0: attachBilledUsageFromContext + // (core/providers/utils) attaches a BilledUsage when *any* of tokens, + // details or cost is set, so a details-only usage reaches here with + // zero totals; emitting those would stamp explicit zeros on the span + // where the success paths only ever see provider-reported values. + if u.PromptTokens > 0 { + attrs[schemas.AttrInputTokens] = u.PromptTokens + } + if u.CompletionTokens > 0 { + attrs[schemas.AttrOutputTokens] = u.CompletionTokens + } + if u.TotalTokens > 0 { + attrs[schemas.AttrTotalTokens] = u.TotalTokens + } + + if d := u.PromptTokensDetails; d != nil { + // The nested cached-write detail keys are namespaced per request + // type on the success paths (input_token_details.* for Responses, + // prompt_token_details.* for chat); mirror that here instead of + // emitting both families, which would break the otel plugin's + // assumption that the two namespaces are mutually exclusive. + isResponses := err.ExtraFields.RequestType == schemas.ResponsesRequest || + err.ExtraFields.RequestType == schemas.ResponsesStreamRequest + if d.CachedReadTokens > 0 { + attrs[schemas.AttrUsageCacheReadInputTokens] = d.CachedReadTokens + } + if d.CachedWriteTokens > 0 { + attrs[schemas.AttrUsageCacheCreationInputTokens] = d.CachedWriteTokens + } + if wd := d.CachedWriteTokenDetails; wd != nil { + if wd.CachedWriteTokens5m > 0 { + if isResponses { + attrs[schemas.AttrInputTokenDetailsCachedWrite5m] = wd.CachedWriteTokens5m + } else { + attrs[schemas.AttrPromptTokenDetailsCachedWrite5m] = wd.CachedWriteTokens5m + } + } + if wd.CachedWriteTokens1h > 0 { + if isResponses { + attrs[schemas.AttrInputTokenDetailsCachedWrite1h] = wd.CachedWriteTokens1h + } else { + attrs[schemas.AttrPromptTokenDetailsCachedWrite1h] = wd.CachedWriteTokens1h + } + } + } + } + } + return attrs } diff --git a/framework/tracing/llmspan_test.go b/framework/tracing/llmspan_test.go index 8e6d70d5155..72e4697c653 100644 --- a/framework/tracing/llmspan_test.go +++ b/framework/tracing/llmspan_test.go @@ -176,3 +176,162 @@ func TestPopulateRequestExtraParamsSerializesStructuredValues(t *testing.T) { }) } } + +func TestPopulateErrorAttributesEmitsBilledUsage(t *testing.T) { + msg := "stream cancelled by client" + bifrostErr := &schemas.BifrostError{ + Error: &schemas.ErrorField{Message: msg}, + } + bifrostErr.ExtraFields.RequestType = schemas.ChatCompletionStreamRequest + bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{ + PromptTokens: 1200, + CompletionTokens: 34, + TotalTokens: 1234, + PromptTokensDetails: &schemas.ChatPromptTokensDetails{ + CachedReadTokens: 1000, + CachedWriteTokens: 200, + CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ + CachedWriteTokens5m: 120, + CachedWriteTokens1h: 80, + }, + }, + } + + attrs := PopulateErrorAttributes(bifrostErr) + + for key, want := range map[string]any{ + schemas.AttrInputTokens: 1200, + schemas.AttrOutputTokens: 34, + schemas.AttrTotalTokens: 1234, + schemas.AttrUsageCacheReadInputTokens: 1000, + schemas.AttrUsageCacheCreationInputTokens: 200, + schemas.AttrPromptTokenDetailsCachedWrite5m: 120, + schemas.AttrPromptTokenDetailsCachedWrite1h: 80, + } { + if got := attrs[key]; got != want { + t.Errorf("attribute %s = %v, want %v", key, got, want) + } + } + // A failed chat span must not carry the Responses namespace: the otel + // plugin treats the two 5m/1h families as mutually exclusive per request. + for _, key := range []string{ + schemas.AttrInputTokenDetailsCachedWrite5m, + schemas.AttrInputTokenDetailsCachedWrite1h, + } { + if _, ok := attrs[key]; ok { + t.Errorf("Responses-namespace attribute %s present on a chat span", key) + } + } +} + +func TestPopulateErrorAttributesUsesResponsesNamespace(t *testing.T) { + bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "responses stream cancelled"}} + bifrostErr.ExtraFields.RequestType = schemas.ResponsesStreamRequest + bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{ + PromptTokensDetails: &schemas.ChatPromptTokensDetails{ + CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ + CachedWriteTokens5m: 120, + CachedWriteTokens1h: 80, + }, + }, + } + + attrs := PopulateErrorAttributes(bifrostErr) + + for key, want := range map[string]any{ + schemas.AttrInputTokenDetailsCachedWrite5m: 120, + schemas.AttrInputTokenDetailsCachedWrite1h: 80, + } { + if got := attrs[key]; got != want { + t.Errorf("attribute %s = %v, want %v", key, got, want) + } + } + for _, key := range []string{ + schemas.AttrPromptTokenDetailsCachedWrite5m, + schemas.AttrPromptTokenDetailsCachedWrite1h, + } { + if _, ok := attrs[key]; ok { + t.Errorf("chat-namespace attribute %s present on a Responses span", key) + } + } +} + +func TestPopulateErrorAttributesWithoutBilledUsageEmitsNoTokens(t *testing.T) { + msg := "401 before the model ran" + bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: msg}} + + attrs := PopulateErrorAttributes(bifrostErr) + + for _, key := range []string{schemas.AttrInputTokens, schemas.AttrOutputTokens, schemas.AttrTotalTokens} { + if _, ok := attrs[key]; ok { + t.Errorf("attribute %s present for a request that consumed no tokens", key) + } + } +} + +func TestPopulateErrorAttributesEmitsCacheWriteDetailsWithoutAggregate(t *testing.T) { + bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "stream failed during cache creation"}} + bifrostErr.ExtraFields.RequestType = schemas.ChatCompletionStreamRequest + bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{ + PromptTokensDetails: &schemas.ChatPromptTokensDetails{ + CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ + CachedWriteTokens5m: 120, + CachedWriteTokens1h: 80, + }, + }, + } + + attrs := PopulateErrorAttributes(bifrostErr) + + for key, want := range map[string]any{ + schemas.AttrPromptTokenDetailsCachedWrite5m: 120, + schemas.AttrPromptTokenDetailsCachedWrite1h: 80, + } { + if got := attrs[key]; got != want { + t.Errorf("attribute %s = %v, want %v", key, got, want) + } + } + // Zero-valued aggregates and totals stay absent: this BilledUsage carries + // only cache-write details, so emitting the totals would stamp explicit + // zeros on the span. + for _, key := range []string{ + schemas.AttrUsageCacheCreationInputTokens, + schemas.AttrInputTokens, + schemas.AttrOutputTokens, + schemas.AttrTotalTokens, + } { + if _, ok := attrs[key]; ok { + t.Errorf("zero-valued attribute %s is present", key) + } + } +} + +// A cancelled stream reaches PopulateLLMResponseAttributes with BOTH a non-nil +// accumulated response and a non-nil error (see core/providers/utils). The +// accumulated response is missing the final usage chunk, so the error's +// BilledUsage must win. This mirrors the merge order in Tracer. +func TestErrorAttributesOverrideAccumulatedResponseTokens(t *testing.T) { + partial := &schemas.BifrostResponse{ + ChatResponse: &schemas.BifrostChatResponse{ + Usage: &schemas.BifrostLLMUsage{PromptTokens: 0, CompletionTokens: 0, TotalTokens: 0}, + }, + } + bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "client cancelled the stream"}} + bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{ + PromptTokens: 4096, + CompletionTokens: 128, + TotalTokens: 4224, + } + + attrs := PopulateResponseAttributes(partial) + for k, v := range PopulateErrorAttributes(bifrostErr) { + attrs[k] = v + } + + if got := attrs[schemas.AttrInputTokens]; got != 4096 { + t.Errorf("%s = %v, want 4096 from BilledUsage", schemas.AttrInputTokens, got) + } + if got := attrs[schemas.AttrTotalTokens]; got != 4224 { + t.Errorf("%s = %v, want 4224 from BilledUsage", schemas.AttrTotalTokens, got) + } +} diff --git a/framework/tracing/tracer.go b/framework/tracing/tracer.go index da466d5d2cb..816d819fe30 100644 --- a/framework/tracing/tracer.go +++ b/framework/tracing/tracer.go @@ -451,7 +451,19 @@ func (t *Tracer) PopulateLLMResponseAttributes(ctx *schemas.BifrostContext, hand return } respAttrs := PopulateResponseAttributes(resp) + // A cancelled stream arrives here with an accumulated response whose usage + // is missing the final chunk, so its aggregate token counts read zero. When + // the error carries the authoritative BilledUsage, drop those zeros from + // the response side: PopulateErrorAttributes gates its own emissions on + // > 0, so a zero stamped here would survive the merge and turn "not + // recorded" into a false zero on the span. + billed := err != nil && err.ExtraFields.BilledUsage != nil for k, v := range respAttrs { + if billed && (k == schemas.AttrInputTokens || k == schemas.AttrOutputTokens || k == schemas.AttrTotalTokens) { + if n, ok := v.(int); ok && n == 0 { + continue + } + } if k == schemas.AttrFinishReasons { // Spec: gen_ai.response.finish_reasons (string[]) belongs on the GenAI (llm.call) span. span.SetAttribute(schemas.AttrFinishReasons, v) @@ -481,8 +493,40 @@ func (t *Tracer) PopulateLLMResponseAttributes(ctx *schemas.BifrostContext, hand span.SetAttribute(schemas.AttrBifrostRoutingEngineUsed, strings.Join(engines, ",")) } - // Populate cost attribute using pricing manager - if t.pricingManager != nil && resp != nil { + // Populate cost attribute using pricing manager. BilledUsage wins when it is + // present: it is what the provider actually charged for a failed or cancelled + // turn. A cancelled stream still yields a non-nil accumulated response (see + // providers/utils, which passes both accumulatedResp and err), but that + // response is missing the final usage chunk, so pricing it would report 0. + if t.pricingManager != nil && err != nil && err.ExtraFields.BilledUsage != nil { + // Core calls BifrostError.PopulateExtraFields around RunPostLLMHooks, so + // Provider / RequestType / the model fields are always set here. + ef := err.ExtraFields + model := ef.ResolvedModelUsed + if model == "" { + model = ef.OriginalModelRequested + } + cost := t.pricingManager.CalculateCostForUsage( + ef.BilledUsage, + ef.Provider, + model, + ef.RequestType, + modelcatalog.PricingLookupScopesFromContext(ctx, string(ef.Provider)), + ) + // When the catalog cannot price the model, fall back to the cost the + // provider itself reported (deep-copied into BilledUsage by + // attachBilledUsageFromContext) rather than discarding it. + if cost == 0 && ef.BilledUsage.Cost != nil { + cost = ef.BilledUsage.Cost.TotalCost + } + // Guarded write: a resp == nil failure emitted no cost attribute before + // this path existed, and consumers rely on distinguishing "no cost + // recorded" from a genuine zero (see plugins/logging, which guards the + // same way). + if cost > 0 { + span.SetAttribute(schemas.AttrUsageCost, cost) + } + } else if t.pricingManager != nil && resp != nil { cost := t.pricingManager.CalculateCost(resp, modelcatalog.PricingLookupScopesFromContext(ctx, string(resp.GetExtraFields().Provider))) span.SetAttribute(schemas.AttrUsageCost, cost) } diff --git a/framework/tracing/tracer_test.go b/framework/tracing/tracer_test.go index a36b3a264b2..2b5c2d336ad 100644 --- a/framework/tracing/tracer_test.go +++ b/framework/tracing/tracer_test.go @@ -417,6 +417,50 @@ func TestTracer_SetAttribute(t *testing.T) { } } +// A cancelled stream reaches PopulateLLMResponseAttributes with an accumulated +// response whose usage exists but reads zero (the final usage chunk never +// arrived) and an error carrying the authoritative BilledUsage. The response +// side's zero aggregates must not survive onto the span: PopulateErrorAttributes +// gates its emissions on > 0, so a details-only BilledUsage would otherwise +// leave a false zero in gen_ai.usage.*. +func TestTracer_PopulateLLMResponseAttributesDropsZeroAggregatesWhenBilled(t *testing.T) { + store := NewTraceStore(5*time.Minute, nil) + defer store.Stop() + + tracer := NewTracer(store, nil, nil) + defer tracer.Stop() + + traceID := tracer.CreateTrace("") + ctx := context.WithValue(context.Background(), schemas.BifrostContextKeyTraceID, traceID) + _, handle := tracer.StartSpan(ctx, "llm.call", schemas.SpanKindLLMCall) + + resp := &schemas.BifrostResponse{ + ChatResponse: &schemas.BifrostChatResponse{ + Usage: &schemas.BifrostLLMUsage{PromptTokens: 0, CompletionTokens: 0, TotalTokens: 0}, + }, + } + bifrostErr := &schemas.BifrostError{Error: &schemas.ErrorField{Message: "client cancelled the stream"}} + bifrostErr.ExtraFields.RequestType = schemas.ChatCompletionStreamRequest + bifrostErr.ExtraFields.BilledUsage = &schemas.BifrostLLMUsage{ + PromptTokensDetails: &schemas.ChatPromptTokensDetails{ + CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{CachedWriteTokens5m: 120}, + }, + } + + bctx := schemas.NewBifrostContext(context.Background(), time.Time{}) + tracer.PopulateLLMResponseAttributes(bctx, handle, resp, bifrostErr) + + span := store.GetTrace(traceID).RootSpan + for _, key := range []string{schemas.AttrInputTokens, schemas.AttrOutputTokens, schemas.AttrTotalTokens} { + if v, ok := span.Attributes[key]; ok { + t.Errorf("zero-valued response aggregate %s = %v survived onto the billed failed span", key, v) + } + } + if got := span.Attributes[schemas.AttrPromptTokenDetailsCachedWrite5m]; got != 120 { + t.Errorf("attribute %s = %v, want 120", schemas.AttrPromptTokenDetailsCachedWrite5m, got) + } +} + func TestTracer_GetSpanHandleByID_RootSpan(t *testing.T) { store := NewTraceStore(5*time.Minute, nil) defer store.Stop()