1use crate::routing::classify::Lane;
5use crate::routing::executor::Completion;
6use crate::telemetry::GatewayMetrics;
7
8#[derive(Debug, Clone)]
11pub struct GenAiSpan {
12 pub system: &'static str, pub request_model: String,
14 pub response_model: String,
15 pub lane: Lane,
16 pub route: String,
17 pub tenant: String,
18 pub workspace: Option<String>,
19 pub legs_attempted: u32,
20 pub stream: bool,
21 pub input_tokens: u64,
22 pub output_tokens: u64,
23}
24
25pub fn system_for(provider: &str) -> &'static str {
27 match provider {
28 "vertex" => "vertexai",
29 "openai" => "openai",
30 "qwen" => "dashscope",
31 _ => "oai_compat",
32 }
33}
34
35impl GenAiSpan {
36 pub fn from_completion(
37 c: &Completion,
38 lane: Lane,
39 route: &str,
40 tenant: &str,
41 workspace: Option<&str>,
42 legs_attempted: u32,
43 stream: bool,
44 ) -> Self {
45 Self {
46 system: system_for(&c.provider),
47 request_model: route.to_string(),
48 response_model: c.model.clone(),
49 lane,
50 route: route.to_string(),
51 tenant: tenant.to_string(),
52 workspace: workspace.map(str::to_string),
53 legs_attempted,
54 stream,
55 input_tokens: c.input_tokens,
56 output_tokens: c.output_tokens,
57 }
58 }
59
60 pub fn emit_metrics(&self, metrics: &GatewayMetrics, latency_secs: f64) {
62 metrics.request(self, latency_secs);
63 }
64}
65
66#[cfg(test)]
67mod tests {
68 use super::*;
69
70 #[test]
71 fn maps_provider_to_gen_ai_system() {
72 assert_eq!(system_for("vertex"), "vertexai");
73 assert_eq!(system_for("qwen"), "dashscope");
74 assert_eq!(system_for("openai"), "openai");
75 assert_eq!(system_for("oai_compat"), "oai_compat");
76 }
77
78 #[test]
79 fn from_completion_carries_actual_served_model() {
80 let c = Completion {
81 provider: "vertex".into(),
82 model: "gemini-3-flash".into(),
83 content: "x".into(),
84 tool_calls: Vec::new(),
85 finish_reason: crate::routing::stream::FinishReason::Stop,
86 input_tokens: 1,
87 output_tokens: 2,
88 };
89 let span =
90 GenAiSpan::from_completion(&c, Lane::Standard, "fast", "acme", Some("ws1"), 2, false);
91 assert_eq!(span.system, "vertexai");
92 assert_eq!(span.response_model, "gemini-3-flash");
93 assert_eq!(span.workspace.as_deref(), Some("ws1"));
94 assert_eq!(span.input_tokens, 1);
95 }
96}