xtrace_client/
lib.rs

1#[cfg(feature = "tracing")]
2pub mod layer;
3#[cfg(feature = "tracing")]
4pub use layer::current_trace_id;
5#[cfg(feature = "tracing")]
6pub use layer::XtraceLayer;
7
8use chrono::{DateTime, Utc};
9use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
10use serde::{Deserialize, Serialize};
11use serde_json::Value as JsonValue;
12use std::collections::HashMap;
13use std::time::Duration;
14use url::Url;
15use uuid::Uuid;
16
17#[derive(Debug, thiserror::Error)]
18pub enum Error {
19    #[error("invalid base url: {0}")]
20    InvalidBaseUrl(#[from] url::ParseError),
21
22    #[error("http error: {0}")]
23    Http(#[from] reqwest::Error),
24}
25
26#[derive(Clone)]
27pub struct Client {
28    base_url: Url,
29    http: reqwest::Client,
30}
31
32impl Client {
33    pub fn new(base_url: &str, bearer_token: &str) -> Result<Self, Error> {
34        let base_url = Url::parse(base_url)?;
35
36        let mut headers = HeaderMap::new();
37        headers.insert(
38            AUTHORIZATION,
39            HeaderValue::from_str(&format!("Bearer {}", bearer_token)).unwrap(),
40        );
41        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
42
43        let http = reqwest::Client::builder()
44            .default_headers(headers)
45            .timeout(Duration::from_secs(30))
46            .build()?;
47
48        Ok(Self { base_url, http })
49    }
50
51    pub async fn healthz(&self) -> Result<(), Error> {
52        let url = self.base_url.join("healthz")?;
53        self.http.get(url).send().await?.error_for_status()?;
54        Ok(())
55    }
56
57    pub async fn ingest_batch(
58        &self,
59        req: &BatchIngestRequest,
60    ) -> Result<ApiResponse<JsonValue>, Error> {
61        let url = self.base_url.join("v1/l/batch")?;
62        let res = self
63            .http
64            .post(url)
65            .json(req)
66            .send()
67            .await?
68            .error_for_status()?;
69        Ok(res.json::<ApiResponse<JsonValue>>().await?)
70    }
71
72    pub async fn list_traces(&self, q: &TraceListQuery) -> Result<PagedData<TraceListItem>, Error> {
73        let mut url = self.base_url.join("api/public/traces")?;
74        {
75            let mut pairs = url.query_pairs_mut();
76            if let Some(v) = q.page {
77                pairs.append_pair("page", &v.to_string());
78            }
79            if let Some(v) = q.limit {
80                pairs.append_pair("limit", &v.to_string());
81            }
82            if let Some(v) = q.user_id.as_deref() {
83                pairs.append_pair("userId", v);
84            }
85            if let Some(v) = q.name.as_deref() {
86                pairs.append_pair("name", v);
87            }
88            if let Some(v) = q.session_id.as_deref() {
89                pairs.append_pair("sessionId", v);
90            }
91            if let Some(v) = q.from_timestamp.as_ref() {
92                pairs.append_pair("fromTimestamp", &v.to_rfc3339());
93            }
94            if let Some(v) = q.to_timestamp.as_ref() {
95                pairs.append_pair("toTimestamp", &v.to_rfc3339());
96            }
97            if let Some(v) = q.order_by.as_deref() {
98                pairs.append_pair("orderBy", v);
99            }
100            for tag in &q.tags {
101                pairs.append_pair("tags", tag);
102            }
103            if let Some(v) = q.version.as_deref() {
104                pairs.append_pair("version", v);
105            }
106            if let Some(v) = q.release.as_deref() {
107                pairs.append_pair("release", v);
108            }
109            for env in &q.environment {
110                pairs.append_pair("environment", env);
111            }
112            if let Some(v) = q.fields.as_deref() {
113                pairs.append_pair("fields", v);
114            }
115        }
116
117        let res = self.http.get(url).send().await?.error_for_status()?;
118        Ok(res.json::<PagedData<TraceListItem>>().await?)
119    }
120
121    pub async fn get_trace(&self, trace_id: Uuid) -> Result<TraceDetailDto, Error> {
122        let url = self
123            .base_url
124            .join(&format!("api/public/traces/{}", trace_id))?;
125        let res = self.http.get(url).send().await?.error_for_status()?;
126        Ok(res.json::<TraceDetailDto>().await?)
127    }
128
129    pub async fn metrics_daily(
130        &self,
131        q: &MetricsDailyQuery,
132    ) -> Result<PagedData<MetricsDailyItem>, Error> {
133        let mut url = self.base_url.join("api/public/metrics/daily")?;
134        {
135            let mut pairs = url.query_pairs_mut();
136            if let Some(v) = q.page {
137                pairs.append_pair("page", &v.to_string());
138            }
139            if let Some(v) = q.limit {
140                pairs.append_pair("limit", &v.to_string());
141            }
142            if let Some(v) = q.trace_name.as_deref() {
143                pairs.append_pair("traceName", v);
144            }
145            if let Some(v) = q.user_id.as_deref() {
146                pairs.append_pair("userId", v);
147            }
148            for tag in &q.tags {
149                pairs.append_pair("tags", tag);
150            }
151            if let Some(v) = q.from_timestamp.as_ref() {
152                pairs.append_pair("fromTimestamp", &v.to_rfc3339());
153            }
154            if let Some(v) = q.to_timestamp.as_ref() {
155                pairs.append_pair("toTimestamp", &v.to_rfc3339());
156            }
157            if let Some(v) = q.version.as_deref() {
158                pairs.append_pair("version", v);
159            }
160            if let Some(v) = q.release.as_deref() {
161                pairs.append_pair("release", v);
162            }
163        }
164
165        let res = self.http.get(url).send().await?.error_for_status()?;
166        Ok(res.json::<PagedData<MetricsDailyItem>>().await?)
167    }
168
169    pub async fn push_metrics(
170        &self,
171        metrics: &[MetricPoint],
172    ) -> Result<ApiResponse<JsonValue>, Error> {
173        let url = self.base_url.join("v1/metrics/batch")?;
174        let req = MetricsBatchRequest {
175            metrics: metrics.to_vec(),
176        };
177        let res = self
178            .http
179            .post(url)
180            .json(&req)
181            .send()
182            .await?
183            .error_for_status()?;
184        Ok(res.json::<ApiResponse<JsonValue>>().await?)
185    }
186
187    /// Query time-series metrics with optional downsampling and aggregation.
188    pub async fn query_metrics(
189        &self,
190        q: &MetricsQueryParams,
191    ) -> Result<MetricsQueryResponse, Error> {
192        let mut url = self.base_url.join("api/public/metrics/query")?;
193        {
194            let mut pairs = url.query_pairs_mut();
195            pairs.append_pair("name", &q.name);
196            if let Some(v) = q.from.as_ref() {
197                pairs.append_pair("from", &v.to_rfc3339());
198            }
199            if let Some(v) = q.to.as_ref() {
200                pairs.append_pair("to", &v.to_rfc3339());
201            }
202            if let Some(v) = q.labels.as_ref() {
203                pairs.append_pair("labels", &serde_json::to_string(v).unwrap_or_default());
204            }
205            if let Some(v) = q.step.as_deref() {
206                pairs.append_pair("step", v);
207            }
208            if let Some(v) = q.agg.as_deref() {
209                pairs.append_pair("agg", v);
210            }
211            if let Some(v) = q.group_by.as_deref() {
212                pairs.append_pair("group_by", v);
213            }
214        }
215
216        let res = self.http.get(url).send().await?.error_for_status()?;
217        Ok(res.json::<MetricsQueryResponse>().await?)
218    }
219
220    /// List all available metric names.
221    pub async fn list_metric_names(&self) -> Result<Vec<String>, Error> {
222        let url = self.base_url.join("api/public/metrics/names")?;
223        let res = self.http.get(url).send().await?.error_for_status()?;
224        let wrapper = res.json::<MetricNamesResponse>().await?;
225        Ok(wrapper.data)
226    }
227}
228
229#[derive(Debug, Serialize, Deserialize, Clone)]
230pub struct MetricPoint {
231    pub name: String,
232    #[serde(default)]
233    pub labels: HashMap<String, String>,
234    pub value: f64,
235    pub timestamp: DateTime<Utc>,
236}
237
238#[derive(Debug, Serialize)]
239struct MetricsBatchRequest {
240    metrics: Vec<MetricPoint>,
241}
242
243/// Parameters for `query_metrics`.
244#[derive(Debug, Default, Serialize, Deserialize)]
245pub struct MetricsQueryParams {
246    /// Required. Metric name (e.g. `pending_requests`, `kv_cache_usage`).
247    pub name: String,
248    /// Start time (inclusive). Defaults to 1 hour before `to`.
249    #[serde(default)]
250    pub from: Option<DateTime<Utc>>,
251    /// End time (inclusive). Defaults to now.
252    #[serde(default)]
253    pub to: Option<DateTime<Utc>>,
254    /// Label filter as a JSON object (JSONB containment).
255    #[serde(default)]
256    pub labels: Option<HashMap<String, String>>,
257    /// Downsample step: `1m`, `5m`, `1h`, `1d`. Default `1m`.
258    #[serde(default)]
259    pub step: Option<String>,
260    /// Aggregation: `avg`, `max`, `min`, `sum`, `last`, `p50`, `p90`, `p99`. Default `avg`.
261    #[serde(default)]
262    pub agg: Option<String>,
263    /// Group results by a specific label key instead of the full label set.
264    #[serde(default)]
265    pub group_by: Option<String>,
266}
267
268/// A single time-series data point.
269#[derive(Debug, Deserialize, Clone)]
270pub struct MetricValuePoint {
271    pub timestamp: String,
272    pub value: f64,
273}
274
275/// A time-series for one unique label combination.
276#[derive(Debug, Deserialize, Clone)]
277pub struct MetricsSeries {
278    pub labels: JsonValue,
279    pub values: Vec<MetricValuePoint>,
280}
281
282/// Metadata about the metrics query result.
283#[derive(Debug, Deserialize, Clone)]
284pub struct MetricsQueryMeta {
285    /// Timestamp (RFC3339 UTC) of the most recent data point. Absent when no data.
286    #[serde(default)]
287    pub latest_ts: Option<String>,
288    /// Number of distinct series returned.
289    pub series_count: usize,
290    /// `true` when results were truncated due to server limits.
291    pub truncated: bool,
292}
293
294/// Response from `GET /api/public/metrics/query`.
295#[derive(Debug, Deserialize, Clone)]
296pub struct MetricsQueryResponse {
297    pub data: Vec<MetricsSeries>,
298    pub meta: MetricsQueryMeta,
299}
300
301#[derive(Debug, Deserialize)]
302struct MetricNamesResponse {
303    data: Vec<String>,
304}
305
306#[derive(Debug, Deserialize)]
307pub struct ApiResponse<T> {
308    pub message: String,
309    /// Machine-readable error code. Present only on error responses.
310    /// Possible values: `UNAUTHORIZED`, `BAD_REQUEST`, `TOO_MANY_REQUESTS`,
311    /// `INTERNAL_ERROR`, `SERVICE_UNAVAILABLE`, `NOT_FOUND`.
312    #[serde(default)]
313    pub code: Option<String>,
314    #[serde(default)]
315    pub data: Option<T>,
316}
317
318#[derive(Debug, Deserialize)]
319#[serde(rename_all = "camelCase")]
320pub struct PageMeta {
321    pub page: i64,
322    pub limit: i64,
323    pub total_items: i64,
324    pub total_pages: i64,
325}
326
327#[derive(Debug, Deserialize)]
328pub struct PagedData<T> {
329    pub data: Vec<T>,
330    pub meta: PageMeta,
331}
332
333#[derive(Debug, Serialize, Deserialize, Default)]
334pub struct BatchIngestRequest {
335    #[serde(default)]
336    pub trace: Option<TraceIngest>,
337    #[serde(default)]
338    pub observations: Vec<ObservationIngest>,
339}
340
341#[derive(Debug, Serialize, Deserialize)]
342#[serde(rename_all = "camelCase")]
343pub struct TraceIngest {
344    pub id: Uuid,
345    #[serde(default)]
346    pub timestamp: Option<DateTime<Utc>>,
347
348    #[serde(default)]
349    pub name: Option<String>,
350    #[serde(default)]
351    pub input: Option<JsonValue>,
352    #[serde(default)]
353    pub output: Option<JsonValue>,
354    #[serde(default)]
355    pub session_id: Option<String>,
356    #[serde(default)]
357    pub release: Option<String>,
358    #[serde(default)]
359    pub version: Option<String>,
360    #[serde(default, rename = "userId")]
361    pub user_id: Option<String>,
362    #[serde(default)]
363    pub metadata: Option<JsonValue>,
364    #[serde(default)]
365    pub tags: Vec<String>,
366    #[serde(default)]
367    pub public: Option<bool>,
368    #[serde(default)]
369    pub environment: Option<String>,
370    #[serde(default)]
371    pub external_id: Option<String>,
372    #[serde(default)]
373    pub bookmarked: Option<bool>,
374
375    #[serde(default)]
376    pub latency: Option<f64>,
377    #[serde(default, rename = "totalCost")]
378    pub total_cost: Option<f64>,
379
380    #[serde(default, rename = "projectId")]
381    pub project_id: Option<String>,
382}
383
384impl TraceIngest {
385    pub fn new(id: Uuid) -> Self {
386        Self {
387            id,
388            timestamp: None,
389            name: None,
390            input: None,
391            output: None,
392            session_id: None,
393            release: None,
394            version: None,
395            user_id: None,
396            metadata: None,
397            tags: vec![],
398            public: None,
399            environment: None,
400            external_id: None,
401            bookmarked: None,
402            latency: None,
403            total_cost: None,
404            project_id: None,
405        }
406    }
407
408    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
409        self.session_id = Some(session_id.into());
410        self
411    }
412
413    pub fn with_metadata_field(mut self, key: &str, value: impl Serialize) -> Self {
414        let mut meta = match self.metadata {
415            Some(JsonValue::Object(map)) => map,
416            _ => serde_json::Map::new(),
417        };
418        if let Ok(v) = serde_json::to_value(value) {
419            meta.insert(key.to_string(), v);
420        }
421        self.metadata = Some(JsonValue::Object(meta));
422        self
423    }
424
425    pub fn with_turn_id(self, turn_id: impl Into<String>) -> Self {
426        self.with_metadata_field("turn_id", turn_id.into())
427    }
428
429    pub fn with_run_id(self, run_id: impl Into<String>) -> Self {
430        self.with_metadata_field("run_id", run_id.into())
431    }
432}
433
434#[derive(Debug, Serialize, Deserialize)]
435#[serde(rename_all = "camelCase")]
436pub struct ObservationIngest {
437    pub id: Uuid,
438    #[serde(rename = "traceId")]
439    pub trace_id: Uuid,
440
441    #[serde(default, rename = "type")]
442    pub r#type: Option<String>,
443    #[serde(default)]
444    pub name: Option<String>,
445
446    #[serde(default)]
447    pub start_time: Option<DateTime<Utc>>,
448    #[serde(default)]
449    pub end_time: Option<DateTime<Utc>>,
450    #[serde(default)]
451    pub completion_start_time: Option<DateTime<Utc>>,
452
453    #[serde(default)]
454    pub model: Option<String>,
455    #[serde(default)]
456    pub model_parameters: Option<JsonValue>,
457
458    #[serde(default)]
459    pub input: Option<JsonValue>,
460    #[serde(default)]
461    pub output: Option<JsonValue>,
462
463    #[serde(default)]
464    pub usage: Option<JsonValue>,
465
466    #[serde(default)]
467    pub level: Option<String>,
468    #[serde(default)]
469    pub status_message: Option<String>,
470    #[serde(default)]
471    pub parent_observation_id: Option<Uuid>,
472
473    #[serde(default)]
474    pub prompt_id: Option<String>,
475    #[serde(default)]
476    pub prompt_name: Option<String>,
477    #[serde(default)]
478    pub prompt_version: Option<String>,
479
480    #[serde(default)]
481    pub model_id: Option<String>,
482
483    #[serde(default)]
484    pub input_price: Option<f64>,
485    #[serde(default)]
486    pub output_price: Option<f64>,
487    #[serde(default)]
488    pub total_price: Option<f64>,
489
490    #[serde(default)]
491    pub calculated_input_cost: Option<f64>,
492    #[serde(default)]
493    pub calculated_output_cost: Option<f64>,
494    #[serde(default)]
495    pub calculated_total_cost: Option<f64>,
496
497    #[serde(default)]
498    pub latency: Option<f64>,
499    #[serde(default)]
500    pub time_to_first_token: Option<f64>,
501
502    #[serde(default)]
503    pub completion_tokens: Option<i64>,
504    #[serde(default)]
505    pub prompt_tokens: Option<i64>,
506    #[serde(default)]
507    pub total_tokens: Option<i64>,
508    #[serde(default)]
509    pub unit: Option<String>,
510
511    #[serde(default)]
512    pub metadata: Option<JsonValue>,
513
514    #[serde(default)]
515    pub environment: Option<String>,
516
517    #[serde(default, rename = "projectId")]
518    pub project_id: Option<String>,
519}
520
521impl ObservationIngest {
522    pub fn new(id: Uuid, trace_id: Uuid) -> Self {
523        Self {
524            id,
525            trace_id,
526            r#type: None,
527            name: None,
528            start_time: None,
529            end_time: None,
530            completion_start_time: None,
531            model: None,
532            model_parameters: None,
533            input: None,
534            output: None,
535            usage: None,
536            level: None,
537            status_message: None,
538            parent_observation_id: None,
539            prompt_id: None,
540            prompt_name: None,
541            prompt_version: None,
542            model_id: None,
543            input_price: None,
544            output_price: None,
545            total_price: None,
546            calculated_input_cost: None,
547            calculated_output_cost: None,
548            calculated_total_cost: None,
549            latency: None,
550            time_to_first_token: None,
551            completion_tokens: None,
552            prompt_tokens: None,
553            total_tokens: None,
554            unit: None,
555            metadata: None,
556            environment: None,
557            project_id: None,
558        }
559    }
560
561    pub fn with_metadata_field(mut self, key: &str, value: impl Serialize) -> Self {
562        let mut meta = match self.metadata {
563            Some(JsonValue::Object(map)) => map,
564            _ => serde_json::Map::new(),
565        };
566        if let Ok(v) = serde_json::to_value(value) {
567            meta.insert(key.to_string(), v);
568        }
569        self.metadata = Some(JsonValue::Object(meta));
570        self
571    }
572
573    pub fn with_step_id(self, step_id: impl Into<String>) -> Self {
574        self.with_metadata_field("step_id", step_id.into())
575    }
576
577    pub fn with_parent_step_id(self, parent_step_id: impl Into<String>) -> Self {
578        self.with_metadata_field("parent_step_id", parent_step_id.into())
579    }
580
581    pub fn with_step_type(self, step_type: impl Into<String>) -> Self {
582        self.with_metadata_field("step_type", step_type.into())
583    }
584}
585
586#[derive(Debug, Default, Serialize, Deserialize)]
587#[serde(rename_all = "camelCase")]
588pub struct TraceListQuery {
589    #[serde(default)]
590    pub page: Option<i64>,
591    #[serde(default)]
592    pub limit: Option<i64>,
593
594    #[serde(default, rename = "userId")]
595    pub user_id: Option<String>,
596    #[serde(default)]
597    pub name: Option<String>,
598    #[serde(default, rename = "sessionId")]
599    pub session_id: Option<String>,
600
601    #[serde(default, rename = "fromTimestamp")]
602    pub from_timestamp: Option<DateTime<Utc>>,
603    #[serde(default, rename = "toTimestamp")]
604    pub to_timestamp: Option<DateTime<Utc>>,
605
606    #[serde(default, rename = "orderBy")]
607    pub order_by: Option<String>,
608
609    #[serde(default)]
610    pub tags: Vec<String>,
611
612    #[serde(default)]
613    pub version: Option<String>,
614    #[serde(default)]
615    pub release: Option<String>,
616    #[serde(default)]
617    pub environment: Vec<String>,
618
619    #[serde(default)]
620    pub fields: Option<String>,
621}
622
623#[derive(Debug, Deserialize)]
624#[serde(rename_all = "camelCase")]
625pub struct TraceListItem {
626    pub id: Uuid,
627    pub timestamp: DateTime<Utc>,
628    pub name: Option<String>,
629    #[serde(default)]
630    pub input: Option<JsonValue>,
631    #[serde(default)]
632    pub output: Option<JsonValue>,
633    pub session_id: Option<String>,
634    pub release: Option<String>,
635    pub version: Option<String>,
636    pub user_id: Option<String>,
637    #[serde(default)]
638    pub metadata: Option<JsonValue>,
639    pub tags: Vec<String>,
640    pub public: bool,
641    pub environment: String,
642    pub html_path: String,
643    pub latency: Option<f64>,
644    pub total_cost: Option<f64>,
645    pub observations: Vec<String>,
646    pub scores: Vec<String>,
647}
648
649#[derive(Debug, Default, Serialize, Deserialize)]
650#[serde(rename_all = "camelCase")]
651pub struct MetricsDailyQuery {
652    #[serde(default)]
653    pub page: Option<i64>,
654    #[serde(default)]
655    pub limit: Option<i64>,
656
657    #[serde(default, rename = "traceName")]
658    pub trace_name: Option<String>,
659    #[serde(default, rename = "userId")]
660    pub user_id: Option<String>,
661    #[serde(default)]
662    pub tags: Vec<String>,
663
664    #[serde(default, rename = "fromTimestamp")]
665    pub from_timestamp: Option<DateTime<Utc>>,
666    #[serde(default, rename = "toTimestamp")]
667    pub to_timestamp: Option<DateTime<Utc>>,
668
669    #[serde(default)]
670    pub version: Option<String>,
671    #[serde(default)]
672    pub release: Option<String>,
673}
674
675#[derive(Debug, Deserialize)]
676#[serde(rename_all = "camelCase")]
677pub struct MetricsDailyItem {
678    pub date: String,
679    pub count_traces: i64,
680    pub count_observations: i64,
681    pub total_cost: f64,
682    pub usage: JsonValue,
683}
684
685#[derive(Debug, Deserialize)]
686#[serde(rename_all = "camelCase")]
687pub struct TraceDetailDto {
688    pub id: Uuid,
689    pub timestamp: DateTime<Utc>,
690    pub name: Option<String>,
691    pub input: JsonValue,
692    pub output: JsonValue,
693    pub session_id: Option<String>,
694    pub release: Option<String>,
695    pub version: Option<String>,
696    pub user_id: Option<String>,
697    pub metadata: JsonValue,
698    pub tags: Vec<String>,
699    pub public: bool,
700    pub environment: String,
701    pub html_path: String,
702    pub latency: Option<f64>,
703    pub total_cost: Option<f64>,
704    pub observations: Vec<JsonValue>,
705    pub scores: Vec<JsonValue>,
706}
xtrace_client/lib.rs

xtrace_client/
lib.rs
