From 70e68aac4a0458a649bb26c87d092a7ddcf9b82d Mon Sep 17 00:00:00 2001 From: roroghost17 Date: Thu, 9 Jul 2026 15:17:45 +0530 Subject: [PATCH] feat: add connector test cases and missing attributes --- core/schemas/enrichment.go | 140 +++++++++++++++++++++++++++++++ core/schemas/enrichment_test.go | 30 +++++++ core/schemas/trace.go | 2 + framework/tracing/tracer.go | 17 ++++ plugins/telemetry/main.go | 48 +++++++---- plugins/telemetry/parity_test.go | 54 ++++++++++++ 6 files changed, 273 insertions(+), 18 deletions(-) create mode 100644 core/schemas/enrichment.go create mode 100644 core/schemas/enrichment_test.go create mode 100644 plugins/telemetry/parity_test.go diff --git a/core/schemas/enrichment.go b/core/schemas/enrichment.go new file mode 100644 index 00000000000..aadda3a150c --- /dev/null +++ b/core/schemas/enrichment.go @@ -0,0 +1,140 @@ +package schemas + +// EnrichmentDim describes one identity/context dimension that connectors attach +// to the telemetry they emit for a request (which team/customer/business unit/ +// virtual key/etc. the request belongs to). +// +// It is the single source of truth that keeps the CURATED emitters from drifting +// apart — the ones that hand-pick a dimension list: +// - Prometheus labels (OSS plugins/telemetry) +// - Datadog metric tags (enterprise plugins/datadog, buildMetricTags) +// - BigQuery columns (enterprise plugins/bigquery, traceColumns) +// +// Each of those derives its list from this registry, and a per-connector +// conformance test asserts the derived list matches — so adding a dimension in +// one place can't silently leave the others behind. +// +// The GENERIC emitters (otel, kafka, pubsub) project the entire span attribute +// map and therefore already carry every dimension; they need no derivation and +// no conformance test here. +// +// `alias` and `routing_engine_used` are in the metric tier but derived only +// post-response (once model resolution/routing has run). They are attached to +// the span in framework/tracing and carry a normal (non-empty) SpanAttr, so a +// record/trace-tier connector can read them like any other dimension. +type EnrichmentDim struct { + // Name is the canonical short identifier, used verbatim as the Prometheus + // label and the Datadog metric tag key. The BigQuery column name also equals + // it unless Column overrides (see below). + Name string + // Column is the BigQuery column name when it differs from Name. BigQuery + // predates the "method" naming and stores it as "request_type"; empty means + // the column name equals Name. + Column string + // SpanAttr is the canonical bifrost.* span-attribute key the dimension is + // stored under. It is what the record/trace-tier emitters read and what a + // connector derives its projection from. + SpanAttr string + // MetricSafe marks a LOW-cardinality dimension eligible to become a Prometheus + // label / Datadog metric tag. High-cardinality dims (per-user, arrays) are + // false and live only on records/traces (BigQuery columns, span attributes). + MetricSafe bool + // Multi marks an array-valued dimension — governance can attach several teams/ + // customers/business units to a single request. Array dims are never + // MetricSafe (they would explode metric series cardinality). + Multi bool +} + +// EnrichmentDims is the canonical, ordered registry of identity/context +// dimensions. Order is stable so derived lists (labels/tags/columns) are +// deterministic. Add a dimension here once and every curated connector picks it +// up via its derivation + conformance test. +var EnrichmentDims = []EnrichmentDim{ + // --- Metric tier: low-cardinality, safe as Prometheus labels / Datadog tags, + // and also present on records/traces. --- + {Name: "provider", SpanAttr: AttrBifrostProviderName, MetricSafe: true}, + {Name: "model", SpanAttr: AttrRequestModel, MetricSafe: true}, + {Name: "method", Column: "request_type", SpanAttr: AttrLegacyRequestType, MetricSafe: true}, + // alias and routing_engine_used are derived post-response and attached to the + // span in framework/tracing (they have no meaning until the model is resolved + // and routing has run), so connectors read them like any other dimension. + {Name: "alias", SpanAttr: AttrBifrostAlias, MetricSafe: true}, + {Name: "routing_engine_used", SpanAttr: AttrBifrostRoutingEngineUsed, MetricSafe: true}, + {Name: "virtual_key_id", SpanAttr: AttrBifrostVirtualKeyID, MetricSafe: true}, + {Name: "virtual_key_name", SpanAttr: AttrBifrostVirtualKeyName, MetricSafe: true}, + {Name: "selected_key_id", SpanAttr: AttrBifrostSelectedKeyID, MetricSafe: true}, + {Name: "selected_key_name", SpanAttr: AttrBifrostSelectedKeyName, MetricSafe: true}, + {Name: "routing_rule_id", SpanAttr: AttrBifrostRoutingRuleID, MetricSafe: true}, + {Name: "routing_rule_name", SpanAttr: AttrBifrostRoutingRuleName, MetricSafe: true}, + {Name: "team_id", SpanAttr: AttrBifrostTeamID, MetricSafe: true}, + {Name: "team_name", SpanAttr: AttrBifrostTeamName, MetricSafe: true}, + {Name: "customer_id", SpanAttr: AttrBifrostCustomerID, MetricSafe: true}, + {Name: "customer_name", SpanAttr: AttrBifrostCustomerName, MetricSafe: true}, + {Name: "business_unit_id", SpanAttr: AttrBifrostBusinessUnitID, MetricSafe: true}, + {Name: "business_unit_name", SpanAttr: AttrBifrostBusinessUnitName, MetricSafe: true}, + {Name: "fallback_index", SpanAttr: AttrBifrostFallbackIndex, MetricSafe: true}, + + // --- Record/trace tier only: high cardinality, NOT metric-safe. Present on + // BigQuery columns and span attributes, never as metric labels/tags. --- + {Name: "user_id", SpanAttr: AttrBifrostUserID}, + {Name: "user_name", SpanAttr: AttrBifrostUserName}, + {Name: "team_ids", SpanAttr: AttrBifrostTeamIDs, Multi: true}, + {Name: "team_names", SpanAttr: AttrBifrostTeamNames, Multi: true}, + {Name: "customer_ids", SpanAttr: AttrBifrostCustomerIDs, Multi: true}, + {Name: "customer_names", SpanAttr: AttrBifrostCustomerNames, Multi: true}, + {Name: "business_unit_ids", SpanAttr: AttrBifrostBusinessUnitIDs, Multi: true}, + {Name: "business_unit_names", SpanAttr: AttrBifrostBusinessUnitNames, Multi: true}, +} + +// ColumnName returns the BigQuery column name for the dimension — Column when +// set, otherwise Name. +func (d EnrichmentDim) ColumnName() string { + if d.Column != "" { + return d.Column + } + return d.Name +} + +// EnrichmentDimColumnNames returns the BigQuery column name for every dimension, +// in registry order (the record/trace-tier set). +func EnrichmentDimColumnNames() []string { + out := make([]string, len(EnrichmentDims)) + for i, d := range EnrichmentDims { + out[i] = d.ColumnName() + } + return out +} + +// MetricSafeEnrichmentDims returns the low-cardinality dimensions eligible to be +// Prometheus labels / Datadog metric tags, in registry order. +func MetricSafeEnrichmentDims() []EnrichmentDim { + out := make([]EnrichmentDim, 0, len(EnrichmentDims)) + for _, d := range EnrichmentDims { + if d.MetricSafe { + out = append(out, d) + } + } + return out +} + +// EnrichmentDimNames returns every dimension name, in registry order (the +// record/trace-tier set). +func EnrichmentDimNames() []string { + out := make([]string, len(EnrichmentDims)) + for i, d := range EnrichmentDims { + out[i] = d.Name + } + return out +} + +// MetricSafeEnrichmentDimNames returns the names of the metric-tier dimensions, +// in registry order. +func MetricSafeEnrichmentDimNames() []string { + out := make([]string, 0, len(EnrichmentDims)) + for _, d := range EnrichmentDims { + if d.MetricSafe { + out = append(out, d.Name) + } + } + return out +} diff --git a/core/schemas/enrichment_test.go b/core/schemas/enrichment_test.go new file mode 100644 index 00000000000..ab881448f5b --- /dev/null +++ b/core/schemas/enrichment_test.go @@ -0,0 +1,30 @@ +package schemas + +import "testing" + +// TestArrayDimsAreNeverMetricSafe is the structural guard that keeps array +// (Multi) dimensions out of the metric tier — i.e. out of Prometheus labels and +// Datadog metric tags. An array value like "team-a,team-b,team-c" would become a +// distinct label/tag value per team combination and explode series cardinality, +// so a Multi dimension must never be MetricSafe. The curated connectors derive +// their metric-tier lists from MetricSafeEnrichmentDims(), so this invariant is +// what actually prevents arrays from ever being shared as Prometheus labels. +func TestArrayDimsAreNeverMetricSafe(t *testing.T) { + for _, d := range EnrichmentDims { + if d.Multi && d.MetricSafe { + t.Errorf("dimension %q is Multi (array) AND MetricSafe — arrays must never be metric labels/tags (cardinality explosion)", d.Name) + } + } +} + +// TestEnrichmentDimNamesUnique guards against a copy-paste duplicate slipping into +// the registry, which would double-emit a label/column. +func TestEnrichmentDimNamesUnique(t *testing.T) { + seen := map[string]bool{} + for _, d := range EnrichmentDims { + if seen[d.Name] { + t.Errorf("duplicate enrichment dimension name %q", d.Name) + } + seen[d.Name] = true + } +} diff --git a/core/schemas/trace.go b/core/schemas/trace.go index c17c5e1026d..f578419250b 100644 --- a/core/schemas/trace.go +++ b/core/schemas/trace.go @@ -708,6 +708,8 @@ const ( AttrBifrostUserName = "bifrost.user.name" AttrBifrostRetries = "bifrost.retries" AttrBifrostFallbackIndex = "bifrost.fallback_index" + AttrBifrostAlias = "bifrost.alias" // original requested model when it differs from the resolved model + AttrBifrostRoutingEngineUsed = "bifrost.routing_engine_used" // comma-joined routing engines that handled the request AttrBifrostStopSequencesJoined = "bifrost.request.stop_sequences" // OTel general semconv (no gen_ai prefix). Emitted alongside the legacy diff --git a/framework/tracing/tracer.go b/framework/tracing/tracer.go index f4969f03153..21c423d6789 100644 --- a/framework/tracing/tracer.go +++ b/framework/tracing/tracer.go @@ -336,6 +336,23 @@ func (t *Tracer) PopulateLLMResponseAttributes(ctx *schemas.BifrostContext, hand for k, v := range PopulateErrorAttributes(err) { span.SetAttribute(k, v) } + + // Enrichment dimensions derivable only post-response, attached here so every + // connector reads them from one place (see core/schemas EnrichmentDims): + // - alias: the originally requested model when it differs from the resolved + // model (an alias was matched or a fallback swapped the model). + // - routing_engine_used: the comma-joined set of routing engines that handled + // the request; the context list is only complete once routing has run. + if resp != nil { + ef := resp.GetExtraFields() + if ef.ResolvedModelUsed != "" && ef.ResolvedModelUsed != ef.OriginalModelRequested && ef.OriginalModelRequested != "" { + span.SetAttribute(schemas.AttrBifrostAlias, ef.OriginalModelRequested) + } + } + if engines, ok := ctx.Value(schemas.BifrostContextKeyRoutingEnginesUsed).([]string); ok && len(engines) > 0 { + span.SetAttribute(schemas.AttrBifrostRoutingEngineUsed, strings.Join(engines, ",")) + } + // Populate cost attribute using pricing manager if t.pricingManager != nil && resp != nil { cost := t.pricingManager.CalculateCost(resp, modelcatalog.PricingLookupScopesFromContext(ctx, string(resp.GetExtraFields().Provider))) diff --git a/plugins/telemetry/main.go b/plugins/telemetry/main.go index 426edfb01d4..1cfb56c2a9e 100644 --- a/plugins/telemetry/main.go +++ b/plugins/telemetry/main.go @@ -227,6 +227,31 @@ var ( ) // Init creates a new PrometheusPlugin with initialized metrics. +// defaultBifrostLabelNames is the canonical set of Prometheus labels attached to +// bifrost.* metrics. It is a package var (not an Init local) so the connector- +// parity conformance test can assert it against the shared enrichment registry +// (core/schemas). Metric-tier dimensions only — no high-cardinality (user, arrays). +var defaultBifrostLabelNames = []string{ + "provider", + "model", + "alias", + "method", + "virtual_key_id", + "virtual_key_name", + "routing_engine_used", + "routing_rule_id", + "routing_rule_name", + "selected_key_id", + "selected_key_name", + "fallback_index", + "team_id", + "team_name", + "customer_id", + "customer_name", + "business_unit_id", + "business_unit_name", +} + func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger schemas.Logger) (*PrometheusPlugin, error) { if config == nil { return nil, fmt.Errorf("config is required") @@ -257,24 +282,7 @@ func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger sche } defaultHTTPLabels := []string{"path", "method", "status"} - defaultBifrostLabels := []string{ - "provider", - "model", - "alias", - "method", - "virtual_key_id", - "virtual_key_name", - "routing_engine_used", - "routing_rule_id", - "routing_rule_name", - "selected_key_id", - "selected_key_name", - "fallback_index", - "team_id", - "team_name", - "customer_id", - "customer_name", - } + defaultBifrostLabels := append([]string(nil), defaultBifrostLabelNames...) var filteredCustomLabels []string if len(config.CustomLabels) > 0 { @@ -707,6 +715,8 @@ func (p *PrometheusPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche teamName := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceTeamName) customerID := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceCustomerID) customerName := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceCustomerName) + businessUnitID := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceBusinessUnitID) + businessUnitName := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyGovernanceBusinessUnitName) // Extract ALL context values BEFORE spawning the goroutine. labelValues := map[string]string{ @@ -726,6 +736,8 @@ func (p *PrometheusPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche "team_name": teamName, "customer_id": customerID, "customer_name": customerName, + "business_unit_id": businessUnitID, + "business_unit_name": businessUnitName, } // Get all custom prometheus labels from context BEFORE the goroutine. diff --git a/plugins/telemetry/parity_test.go b/plugins/telemetry/parity_test.go new file mode 100644 index 00000000000..b34879d7c86 --- /dev/null +++ b/plugins/telemetry/parity_test.go @@ -0,0 +1,54 @@ +package telemetry + +import ( + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +// TestPrometheusLabelsMatchEnrichmentRegistry keeps the Prometheus bifrost label +// set in parity with the canonical enrichment registry (core/schemas): every +// metric-safe dimension must be a label, and no record-tier (high cardinality) +// dimension may be. Known divergences are enumerated in the allowlist below; the +// test fails on any new drift and when a listed entry no longer applies. +func TestPrometheusLabelsMatchEnrichmentRegistry(t *testing.T) { + labels := map[string]bool{} + for _, l := range defaultBifrostLabelNames { + labels[l] = true + } + + metricSafe := map[string]bool{} + for _, n := range schemas.MetricSafeEnrichmentDimNames() { + metricSafe[n] = true + } + allDims := map[string]bool{} + for _, n := range schemas.EnrichmentDimNames() { + allDims[n] = true + } + + // Metric-safe dims the telemetry plugin does not expose as labels. Empty: the + // Prometheus label set covers the full metric-safe registry. + knownMissing := map[string]string{} + + // 1) Every metric-safe dim must be a label, unless a known gap. + for n := range metricSafe { + if !labels[n] && knownMissing[n] == "" { + t.Errorf("Prometheus labels are missing registry metric-safe dimension %q", n) + } + } + // 2) Every enrichment-dim label must be metric-safe (no high-cardinality labels). + for l := range labels { + if !allDims[l] { + continue // non-enrichment label (e.g. status_code) — out of scope + } + if !metricSafe[l] { + t.Errorf("Prometheus exposes record-tier dimension %q as a label (cardinality risk)", l) + } + } + // 3) Keep the allowlist honest: a gap that has closed must be removed. + for n := range knownMissing { + if labels[n] { + t.Errorf("dimension %q is now a label — remove it from knownMissing", n) + } + } +}