fix(telemetry): per-PK attribution tags on all non-chat usage events

[jarvis9443]

Problem

Provider attribution in the dashboard Logs — 服务商类型 (provider_kind), 品牌服务商 (branded_provider), PK 标签 (pk_label), plus provider_featured/byo_label — was wired into /v1/chat/completions and /v1/messages, and recently /v1/responses (#646 / AISIX-Cloud#867). The remaining billing endpoints left all five UsageEvent fields at their defaults, so their Logs rows showed no upstream vendor or PK label:

• /v1/completions
• /v1/embeddings
• /v1/rerank
• /v1/audio/transcriptions, /translations, /speech
• /v1/images/generations

This is the same oversight #867 fixed for /v1/responses, replicated across the rest of the handler family (each handler's emit even carried a "wired for chat only / follow-up for the non-chat handlers" comment).

Fix

• New shared helper usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id) — one source of truth for both the snapshot telemetry_tags lookup AND the five-field wire mapping (with sanitize_tag), so the handler family can't drift apart again.
• Each handler already resolves the target ProviderKey for dispatch; thread its id (pk_entry.id) through the dispatch-success struct to the emitter and call the shared helper. Empty/unknown id → all-empty tags → wire NULL, matching the chat/messages contract.
• Point the existing /v1/responses lookup at the shared helper too (removes the duplicate copy added in fix(responses): attribute per-PK telemetry tags on /v1/responses usage events #646).

No wire/schema change — the cp-api columns already exist; they were simply never filled by these endpoints.

Tests

Per-handler integration tests (real router + bridge + mock upstream + usage sink) assert the emitted UsageEvent carries the resolved PK's tags. Verified each fails before the fix (empty tags) and passes after. Full aisix-proxy suite green (484), clippy + fmt clean.

E2E note: these attribution tags don't traverse the OTLP fan-out the standalone DP E2E harness observes (the OTLP encoder is an intentional allowlist, shared across endpoints, that excludes them) — they only reach cp-api via the usage-sink POST. Same testing approach as the original chat/messages attribution feature and #646.

Origin: surfaced by a cross-API consistency audit after AISIX-Cloud#867.

Summary by CodeRabbit

• New Features

  • Usage reporting now includes provider attribution details for more proxy endpoints, improving reporting consistency across audio, completions, embeddings, images, and rerank requests.
  • Telemetry emitted for successful requests can now carry richer provider labels and metadata.

• Bug Fixes

  • Fixed missing attribution in usage events so successful requests are tracked with the correct provider-related details.
  • Improved consistency for responses that don’t support a requested operation.

[coderabbitai]

📝 Walkthrough

Walkthrough

The proxy now resolves provider-key telemetry tags through a shared helper and threads provider_key_id through audio, completions, embeddings, images, and rerank usage-event paths. Each endpoint’s emitted UsageEvent can now be stamped with provider telemetry fields from the resolved ProviderKey.

Changes

Provider-key telemetry attribution

Layer / File(s)	Summary
Shared attribution helper crates/aisix-proxy/src/lib.rs, crates/aisix-proxy/src/usage_attr.rs, crates/aisix-proxy/src/responses.rs	usage_attr is added as a crate module, defines shared ProviderKey telemetry lookup and event tagging helpers, and responses.rs now imports the shared lookup instead of keeping a local copy.
Audio usage attribution crates/aisix-proxy/src/audio.rs	AudioDispatchSuccess, speech_dispatch, multipart_dispatch, and /v1/audio/speech now carry provider_key_id through emit_usage_event, which applies per-provider-key telemetry tags before emission; new test fixtures and a regression test cover the emitted audio telemetry fields.
Completions usage attribution crates/aisix-proxy/src/completions.rs	CompletionDispatchSuccess, the upstream-success and 501 NotImplemented dispatch arms, and emit_usage_event now pass provider_key_id and apply per-PK telemetry tags before emission; the test module adds tagged fixtures and a regression test.
Embeddings usage attribution crates/aisix-proxy/src/embeddings.rs	EmbedDispatchSuccess, both dispatch branches, and emit_usage_event now thread provider_key_id into apply_pk_telemetry; tagged snapshot helpers and a regression test assert the emitted UsageEvent fields.
Images usage attribution crates/aisix-proxy/src/images.rs	ImageDispatchSuccess, the success and 501 NotImplemented paths, and emit_usage_event now carry provider_key_id and apply per-PK telemetry tags; tagged snapshot helpers and a regression test validate the emitted fields.
Rerank usage attribution crates/aisix-proxy/src/rerank.rs	RerankDispatchSuccess, the success path, and emit_usage_event now pass provider_key_id into apply_pk_telemetry; tagged snapshot helpers and a regression test cover the emitted UsageEvent fields.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• api7/aisix#646: Changes responses.rs to resolve and apply provider telemetry tags for usage events; this PR moves that lookup into usage_attr and extends the same attribution path to additional endpoints.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
E2e Test Quality Review	⚠️ Warning	Most endpoints have real-router integration coverage, but /v1/audio/speech never asserts the new PK tags, so the changed attribution path isn't fully pinned.	Add a tagged /v1/audio/speech end-to-end test (and ideally direct empty/unknown provider-key cases for usage_attr) to cover the new attribution path.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: per-provider-key telemetry attribution for non-chat usage events.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Check	✅ Passed	Only sanitized provider telemetry tags are added to UsageEvent; no new secret logging, storage, auth, or ownership logic was introduced.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/non-chat-telemetry-parity

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

crates/aisix-proxy/src/responses.rs (1)

1631-1664: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consider using apply_pk_telemetry to avoid re-duplicating the 5-field mapping.

This path now imports provider_telemetry_tags for the lookup but still hand-rolls the identical provider_kind / provider_featured / branded_provider / pk_label / byo_label mapping (Lines 1655-1659). That is precisely the per-field drift the usage_attr module set out to eliminate (usage_attr.rs lines 6-10: "Centralising the snapshot lookup AND the wire-field mapping here keeps the handler family ... from drifting apart again"). Building the event first and then stamping it keeps /v1/responses on the same single source of truth as the non-chat handlers.

♻️ Proposed refactor

     let snap = state.snapshot.load();
-    let tags = provider_telemetry_tags(&snap, provider_key_id);
-    let event = UsageEvent {
+    let mut event = UsageEvent {
         request_id: request_id.to_string(),
         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(),
         requested_model: requested_model.to_string(),
         prompt_tokens: usage.prompt_tokens,
         completion_tokens: usage.completion_tokens,
         cached_prompt_tokens: usage.cached_prompt_tokens,
         reasoning_tokens: usage.reasoning_tokens,
         cache_creation_tokens: usage.cache_creation_tokens,
         cache_read_tokens: usage.cache_read_tokens,
         latency_ms: elapsed.as_millis().min(u32::MAX as u128) as u32,
         status_code,
         inbound_protocol: "openai".to_string(),
         attempt_index: attempt.index,
         attempt_kind: attempt.kind,
         attempt_model: attempt.model,
         error_class: attempt.error_class,
         error_message: attempt.error_message,
-        provider_kind: sanitize_tag(tags.kind.map(|k| k.as_str().to_owned()).unwrap_or_default()),
-        provider_featured: tags.featured,
-        branded_provider: sanitize_tag(tags.branded_provider.unwrap_or_default()),
-        pk_label: sanitize_tag(tags.pk_label.unwrap_or_default()),
-        byo_label: sanitize_tag(tags.byo_label.unwrap_or_default()),
         client_source_ip: client.source_ip.clone(),
         client_user_agent: client.user_agent.clone(),
         guardrail_blocked,
         ..Default::default()
     };
+    crate::usage_attr::apply_pk_telemetry(&mut event, &snap, provider_key_id);

This would also let the provider_telemetry_tags import at Line 36 be dropped in favor of apply_pk_telemetry.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/aisix-proxy/src/responses.rs` around lines 1631 - 1664, The
`/v1/responses` event assembly is still duplicating the provider telemetry field
mapping instead of using the shared helper. Update the `UsageEvent` construction
in `responses.rs` to build the event first, then call `apply_pk_telemetry` so
the `provider_kind`, `provider_featured`, `branded_provider`, `pk_label`, and
`byo_label` fields come from the same single source of truth as the other
handlers. This should remove the manual `provider_telemetry_tags` lookup and
keep the `UsageEvent` wiring aligned with the `usage_attr` module.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/aisix-proxy/src/responses.rs`:
- Around line 1631-1664: The `/v1/responses` event assembly is still duplicating
the provider telemetry field mapping instead of using the shared helper. Update
the `UsageEvent` construction in `responses.rs` to build the event first, then
call `apply_pk_telemetry` so the `provider_kind`, `provider_featured`,
`branded_provider`, `pk_label`, and `byo_label` fields come from the same single
source of truth as the other handlers. This should remove the manual
`provider_telemetry_tags` lookup and keep the `UsageEvent` wiring aligned with
the `usage_attr` module.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 6724764a-1238-480a-a627-5c1fb8cffefd

📥 Commits

Reviewing files that changed from the base of the PR and between 8f758c4 and ca92d33.

📒 Files selected for processing (8)

• crates/aisix-proxy/src/audio.rs
• crates/aisix-proxy/src/completions.rs
• crates/aisix-proxy/src/embeddings.rs
• crates/aisix-proxy/src/images.rs
• crates/aisix-proxy/src/lib.rs
• crates/aisix-proxy/src/rerank.rs
• crates/aisix-proxy/src/responses.rs
• crates/aisix-proxy/src/usage_attr.rs