diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 9559a668..e6aeb86a 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -15,7 +15,7 @@ //! `"type": "not_implemented"`. use aisix_gateway::{BridgeContext, BridgeError, EmbeddingRequest}; -use aisix_obs::{AccessLog, RequestOutcome}; +use aisix_obs::{AccessLog, RequestOutcome, UsageEvent}; use axum::extract::State; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; @@ -72,25 +72,50 @@ pub async fn embeddings( let model_name = body.model.clone(); match dispatch(&state, &auth, body, &request_id).await { - Ok((resp, provider)) => { + Ok(success) => { let elapsed = started.elapsed(); let status = 200u16; emit_access_log( &model_name, - &provider, + &success.provider, &api_key_id, status, elapsed, &request_id, ); state.metrics.record_request( - &provider, + &success.provider, &model_name, status, RequestOutcome::Success, elapsed, ); - resp + // Issue #226: emit UsageEvent so cp-api's budget ledger + // and customer-facing /logs analytics see embeddings + // spend. Pre-#226 the embedding handler dropped the + // event entirely, so any /v1/embeddings traffic was + // invisible to budget enforcement and billing + // reconciliation. Skip when `upstream_called == false` + // — the dispatch's 501-NotImplemented path returns the + // false flag because no upstream call happened, and + // attributing a zero-everything event to the api_key + // would bloat /logs with noise. Distinguished from + // `prompt_tokens == 0` so a 200 with legitimately zero + // tokens (empty input, provider-specific billing + // convention) still emits. Same emit-on-success-only + // convention as chat.rs. + if success.upstream_called { + emit_usage_event( + &state, + &request_id, + &success.model_id, + &api_key_id, + status, + elapsed, + success.prompt_tokens, + ); + } + success.response } Err(err) => { let status = err.status().as_u16(); @@ -115,12 +140,35 @@ pub async fn embeddings( } } +/// Per-request payload from a successful dispatch — carries +/// everything the handler needs to emit access-log + UsageEvent + +/// the actual HTTP response. Pre-#226 the dispatch only returned +/// the response + provider label; expanding to a struct surfaces +/// the model_id + prompt_tokens the UsageEvent emission needs +/// without threading two more positional args through the handler. +struct EmbedDispatchSuccess { + response: Response, + provider: String, + model_id: String, + prompt_tokens: u32, + /// `true` when the dispatch produced a real 200 from the upstream + /// (we have authoritative usage data to attribute). `false` for the + /// 501-NotImplemented branch where no upstream call was made. + /// + /// Explicitly distinguished from `prompt_tokens == 0` so a future + /// provider that legitimately reports zero tokens on a 200 (empty + /// input, special billing convention) still gets emitted — the + /// "no upstream call" channel must not piggyback on a numeric + /// sentinel that real responses could also produce. + upstream_called: bool, +} + async fn dispatch( state: &ProxyState, auth: &AuthenticatedKey, body: EmbeddingRequestBody, request_id: &str, -) -> Result<(Response, String), ProxyError> { +) -> Result { let snapshot = state.snapshot.load(); let model_entry = snapshot @@ -174,18 +222,38 @@ async fn dispatch( // here even though other handlers commit 0. reservation.commit_tokens(embed_resp.usage.total_tokens as u64); let provider_label = provider.to_ascii_lowercase(); - Ok((Json(embed_resp).into_response(), provider_label)) + // Capture the prompt_tokens count BEFORE moving the + // embed_resp into the JSON response — the handler needs + // this for UsageEvent emission downstream (#226). + let prompt_tokens = embed_resp.usage.prompt_tokens; + Ok(EmbedDispatchSuccess { + response: Json(embed_resp).into_response(), + provider: provider_label, + model_id: model_entry.id.to_string(), + prompt_tokens, + upstream_called: true, + }) } Err(BridgeError::Config(msg)) if msg.contains("does not support embeddings") => { // Provider doesn't implement embed → 501 Not Implemented. // Drop the reservation without committing — the request - // didn't hit the upstream. + // didn't hit the upstream. No UsageEvent emission either + // (`upstream_called: false` → handler skips emit per the + // chat.rs convention that we only attribute usage on a + // real upstream completion). reservation.commit_tokens(0); let env = ErrorEnvelope::new(msg, "not_implemented"); - Ok(( - (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), - provider.to_ascii_lowercase(), - )) + Ok(EmbedDispatchSuccess { + response: (StatusCode::NOT_IMPLEMENTED, Json(env)).into_response(), + provider: provider.to_ascii_lowercase(), + model_id: model_entry.id.to_string(), + prompt_tokens: 0, + // No upstream call happened — the handler reads this + // and skips UsageEvent emission. Distinguished from + // `prompt_tokens == 0` so a 200 that legitimately + // reports zero tokens still emits. + upstream_called: false, + }) } Err(e) => { reservation.commit_tokens(0); @@ -227,6 +295,79 @@ fn emit_access_log( .emit(); } +/// Push one `UsageEvent` onto cp-api's telemetry sink **and** fan it +/// out to every per-env OTLP/HTTP exporter in the live snapshot. +/// Non-blocking on both legs: the CP sink drops on full queue, the +/// OTLP fan-out detaches a tokio task per exporter. Mirrors the +/// shape of chat.rs::emit_usage_event for the fields that matter +/// to /v1/embeddings: +/// +/// - `completion_tokens = 0` — embeddings have no completion side. +/// - `inbound_protocol = "openai"` — match chat.rs convention. +/// - No cache / streaming / reasoning / finish_reason metadata — +/// none of these concepts apply to the embeddings endpoint; +/// cp-api reads these UsageEvent fields with `omitempty`-equivalent +/// defaults so leaving them zero is the same as omitting. +/// - `cost_usd = 0.0` — cp-api computes cost server-side from +/// pricing catalog + token counts on ingestion (same convention +/// as every chat.rs emit site). +/// +/// Issue #226. /v1/embeddings is the first non-chat handler to gain +/// emission; follow-ups for completions / responses / rerank / +/// audio / images each get their own PR with the same shape. +fn emit_usage_event( + state: &ProxyState, + request_id: &str, + model_id: &str, + api_key_id: &str, + status_code: u16, + elapsed: Duration, + prompt_tokens: u32, +) { + // Only populate fields meaningful to /v1/embeddings; rely on + // UsageEvent's `#[derive(Default)]` for everything else. Wire-level + // empty / zero / false maps to NULL on cp-api via skip_serializing_if, + // identical to the legacy "field absent" semantics older DP images + // emitted. Specifically left at Default: + // - completion_tokens / cached_prompt_tokens / reasoning_tokens / + // cache_creation_tokens / cache_read_tokens — embeddings have + // no completion side and no cache token concepts + // - provider_request_id / provider_model_version / finish_reason + // — not exposed by the OpenAI embeddings response shape + // - cost_usd — cp-api computes server-side from pricing catalog + // - guardrail_blocked / guardrail_bypassed_reason — embeddings + // bypass guardrails today + // - cache_status / cache_hit_saved_* — no caching on embeddings + // - ttft_ms — embeddings are not streamed + // - served_by_model / routing_* — embeddings don't run routing + // - provider_kind / provider_featured / branded_provider / + // pk_label / byo_label — per-PK telemetry attribution is wired + // for chat completions only; tracked as a follow-up for the + // non-chat handlers (see #226 follow-up issues). + let event = UsageEvent { + request_id: request_id.to_string(), + // RFC 3339 UTC. cp-api parses with time.Parse(time.RFC3339, ...); + // chrono's `to_rfc3339_opts(Secs, true)` emits the trailing Z. + occurred_at: chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Secs, true), + model_id: model_id.to_string(), + api_key_id: api_key_id.to_string(), + prompt_tokens, + latency_ms: elapsed.as_millis().min(u32::MAX as u128) as u32, + status_code, + inbound_protocol: "openai".to_string(), + ..Default::default() + }; + state.usage_sink.try_emit(event.clone()); + // Per-env OTLP/HTTP fan-out — same shape as chat.rs:1334. The + // snapshot's exporter table is empty for envs that haven't + // configured any, so this is a cheap no-op on the common path. + let snap = state.snapshot.load(); + let exporters = snap.observability_exporters.entries(); + state + .otlp_fan_out + .fan_out(&event, exporters.iter().map(|e| &e.value)); +} + #[cfg(test)] mod tests { @@ -620,4 +761,165 @@ mod tests { .unwrap(); assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS); } + + /// Issue #226: a successful /v1/embeddings call must emit a + /// `UsageEvent` onto the `usage_sink`. Pre-#226 the embeddings + /// handler dropped the event entirely, so any /v1/embeddings + /// traffic was invisible to cp-api's budget ledger and the + /// customer-facing /logs analytics. This test pins the contract: + /// after a 200 response, exactly one event arrives with the + /// caller's prompt_tokens, model_id, status, and `inbound_protocol + /// = "openai"` (mirroring chat.rs's emission convention). + #[tokio::test] + async fn emits_usage_event_on_200_with_prompt_tokens_issue_226() { + use aisix_obs::UsageSink; + + let upstream = MockServer::start().await; + // Pin a specific upstream token count so a regression that + // hardcoded 0 (or swapped prompt/completion semantics) would + // surface here. + let upstream_body = serde_json::json!({ + "object": "list", + "data": [{ + "object": "embedding", + "index": 0, + "embedding": [0.1_f32, 0.2_f32] + }], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 42, "total_tokens": 42} + }); + Mock::given(method("POST")) + .and(path("/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(upstream_body)) + .mount(&upstream) + .await; + + let snap = new_snap(&upstream.uri()); + snap.models.insert(model_entry("my-embed")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let handle = SnapshotHandle::new(snap); + let state = crate::ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + let app = crate::build_router(state); + + let body = serde_json::json!({"model": "my-embed", "input": "hello world"}); + let resp = tower::ServiceExt::oneshot(app, make_req(body)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Event must arrive within a generous window — the emit is + // try_send (non-blocking) so a 500ms ceiling is well clear of + // schedule jitter. + let event = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect("UsageEvent was never emitted for /v1/embeddings 200") + .expect("usage_sink sender dropped before emission"); + + assert_eq!( + event.prompt_tokens, 42, + "UsageEvent.prompt_tokens must mirror upstream usage.prompt_tokens" + ); + assert_eq!( + event.completion_tokens, 0, + "embeddings have no completion side; completion_tokens must be zero" + ); + assert_eq!( + event.status_code, 200, + "status_code must reflect the response" + ); + assert_eq!( + event.api_key_id, "k-1", + "api_key_id must reflect the authenticated caller" + ); + assert_eq!( + event.model_id, "m-1", + "model_id must reflect the resolved model row" + ); + assert_eq!( + event.inbound_protocol, "openai", + "inbound_protocol must be \"openai\" for /v1/embeddings (matches chat.rs convention)" + ); + assert!( + !event.request_id.is_empty(), + "request_id must be set for join with x-aisix-call-id" + ); + assert!( + !event.occurred_at.is_empty(), + "occurred_at must be set (RFC 3339 UTC)" + ); + } + + /// Issue #226 audit M1: a 200 response with upstream-reported + /// `prompt_tokens = 0` MUST still emit a UsageEvent. Pre-fix the + /// handler gated emission on `prompt_tokens > 0`, conflating + /// "no upstream call" (501 NotImplemented) with "upstream returned + /// zero tokens" (legitimate 200, rare-but-possible for empty input + /// or provider-specific billing conventions). The fix introduces an + /// explicit `upstream_called: bool` flag on EmbedDispatchSuccess + /// so the channel is unambiguous; this test pins the post-fix + /// contract by driving a 200 with zero tokens and asserting the + /// event still arrives. + #[tokio::test] + async fn emits_usage_event_on_200_with_zero_prompt_tokens_audit_m1() { + use aisix_obs::UsageSink; + + let upstream = MockServer::start().await; + // 200 with usage.prompt_tokens=0. The contract: emit anyway — + // attribution belongs to the api_key for compliance / audit + // even when the billable count is zero. + let upstream_body = serde_json::json!({ + "object": "list", + "data": [{ + "object": "embedding", + "index": 0, + "embedding": [0.1_f32, 0.2_f32] + }], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 0, "total_tokens": 0} + }); + Mock::given(method("POST")) + .and(path("/embeddings")) + .respond_with(ResponseTemplate::new(200).set_body_json(upstream_body)) + .mount(&upstream) + .await; + + let snap = new_snap(&upstream.uri()); + snap.models.insert(model_entry("my-embed")); + snap.apikeys.insert(apikey_entry(&["*"])); + + let (tx, mut rx) = tokio::sync::mpsc::channel(8); + let hub = Arc::new(Hub::new()); + hub.register_specialized("openai", Arc::new(OpenAiBridge::new())); + let handle = SnapshotHandle::new(snap); + let state = crate::ProxyState::new(handle, hub, &cfg()) + .without_cache() + .with_usage_sink(UsageSink::new(tx)); + let app = crate::build_router(state); + + let body = serde_json::json!({"model": "my-embed", "input": ""}); + let resp = tower::ServiceExt::oneshot(app, make_req(body)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let event = tokio::time::timeout(std::time::Duration::from_millis(500), rx.recv()) + .await + .expect( + "UsageEvent must be emitted even when upstream reports prompt_tokens=0 \ + (audit M1 — emission keyed on upstream_called, not on token count)", + ) + .expect("usage_sink sender dropped before emission"); + + assert_eq!(event.prompt_tokens, 0); + assert_eq!(event.status_code, 200); + assert_eq!(event.api_key_id, "k-1"); + assert_eq!(event.model_id, "m-1"); + assert_eq!(event.inbound_protocol, "openai"); + } }