Skip to main content

xtrace_client/
lib.rs

1use chrono::{DateTime, Utc};
2use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};
3use serde::{Deserialize, Serialize};
4use serde_json::Value as JsonValue;
5use std::collections::HashMap;
6use std::time::Duration;
7use url::Url;
8use uuid::Uuid;
9
10#[derive(Debug, thiserror::Error)]
11pub enum Error {
12    #[error("invalid base url: {0}")]
13    InvalidBaseUrl(#[from] url::ParseError),
14
15    #[error("http error: {0}")]
16    Http(#[from] reqwest::Error),
17}
18
19#[derive(Clone)]
20pub struct Client {
21    base_url: Url,
22    http: reqwest::Client,
23}
24
25impl Client {
26    pub fn new(base_url: &str, bearer_token: &str) -> Result<Self, Error> {
27        let base_url = Url::parse(base_url)?;
28
29        let mut headers = HeaderMap::new();
30        headers.insert(
31            AUTHORIZATION,
32            HeaderValue::from_str(&format!("Bearer {}", bearer_token)).unwrap(),
33        );
34        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
35
36        let http = reqwest::Client::builder()
37            .default_headers(headers)
38            .timeout(Duration::from_secs(30))
39            .build()?;
40
41        Ok(Self { base_url, http })
42    }
43
44    pub async fn healthz(&self) -> Result<(), Error> {
45        let url = self.base_url.join("healthz")?;
46        self.http.get(url).send().await?.error_for_status()?;
47        Ok(())
48    }
49
50    pub async fn ingest_batch(
51        &self,
52        req: &BatchIngestRequest,
53    ) -> Result<ApiResponse<JsonValue>, Error> {
54        let url = self.base_url.join("v1/l/batch")?;
55        let res = self
56            .http
57            .post(url)
58            .json(req)
59            .send()
60            .await?
61            .error_for_status()?;
62        Ok(res.json::<ApiResponse<JsonValue>>().await?)
63    }
64
65    pub async fn list_traces(&self, q: &TraceListQuery) -> Result<PagedData<TraceListItem>, Error> {
66        let mut url = self.base_url.join("api/public/traces")?;
67        {
68            let mut pairs = url.query_pairs_mut();
69            if let Some(v) = q.page {
70                pairs.append_pair("page", &v.to_string());
71            }
72            if let Some(v) = q.limit {
73                pairs.append_pair("limit", &v.to_string());
74            }
75            if let Some(v) = q.user_id.as_deref() {
76                pairs.append_pair("userId", v);
77            }
78            if let Some(v) = q.name.as_deref() {
79                pairs.append_pair("name", v);
80            }
81            if let Some(v) = q.session_id.as_deref() {
82                pairs.append_pair("sessionId", v);
83            }
84            if let Some(v) = q.from_timestamp.as_ref() {
85                pairs.append_pair("fromTimestamp", &v.to_rfc3339());
86            }
87            if let Some(v) = q.to_timestamp.as_ref() {
88                pairs.append_pair("toTimestamp", &v.to_rfc3339());
89            }
90            if let Some(v) = q.order_by.as_deref() {
91                pairs.append_pair("orderBy", v);
92            }
93            for tag in &q.tags {
94                pairs.append_pair("tags", tag);
95            }
96            if let Some(v) = q.version.as_deref() {
97                pairs.append_pair("version", v);
98            }
99            if let Some(v) = q.release.as_deref() {
100                pairs.append_pair("release", v);
101            }
102            for env in &q.environment {
103                pairs.append_pair("environment", env);
104            }
105            if let Some(v) = q.fields.as_deref() {
106                pairs.append_pair("fields", v);
107            }
108        }
109
110        let res = self.http.get(url).send().await?.error_for_status()?;
111        Ok(res.json::<PagedData<TraceListItem>>().await?)
112    }
113
114    pub async fn get_trace(&self, trace_id: Uuid) -> Result<TraceDetailDto, Error> {
115        let url = self
116            .base_url
117            .join(&format!("api/public/traces/{}", trace_id))?;
118        let res = self.http.get(url).send().await?.error_for_status()?;
119        Ok(res.json::<TraceDetailDto>().await?)
120    }
121
122    pub async fn metrics_daily(
123        &self,
124        q: &MetricsDailyQuery,
125    ) -> Result<PagedData<MetricsDailyItem>, Error> {
126        let mut url = self.base_url.join("api/public/metrics/daily")?;
127        {
128            let mut pairs = url.query_pairs_mut();
129            if let Some(v) = q.page {
130                pairs.append_pair("page", &v.to_string());
131            }
132            if let Some(v) = q.limit {
133                pairs.append_pair("limit", &v.to_string());
134            }
135            if let Some(v) = q.trace_name.as_deref() {
136                pairs.append_pair("traceName", v);
137            }
138            if let Some(v) = q.user_id.as_deref() {
139                pairs.append_pair("userId", v);
140            }
141            for tag in &q.tags {
142                pairs.append_pair("tags", tag);
143            }
144            if let Some(v) = q.from_timestamp.as_ref() {
145                pairs.append_pair("fromTimestamp", &v.to_rfc3339());
146            }
147            if let Some(v) = q.to_timestamp.as_ref() {
148                pairs.append_pair("toTimestamp", &v.to_rfc3339());
149            }
150        }
151
152        let res = self.http.get(url).send().await?.error_for_status()?;
153        Ok(res.json::<PagedData<MetricsDailyItem>>().await?)
154    }
155
156    pub async fn push_metrics(
157        &self,
158        metrics: &[MetricPoint],
159    ) -> Result<ApiResponse<JsonValue>, Error> {
160        let url = self.base_url.join("v1/metrics/batch")?;
161        let req = MetricsBatchRequest {
162            metrics: metrics.to_vec(),
163        };
164        let res = self
165            .http
166            .post(url)
167            .json(&req)
168            .send()
169            .await?
170            .error_for_status()?;
171        Ok(res.json::<ApiResponse<JsonValue>>().await?)
172    }
173
174    /// Query time-series metrics from xtrace.
175    /// Returns series grouped by labels.
176    pub async fn query_metrics(
177        &self,
178        q: &MetricsQueryParams,
179    ) -> Result<MetricsQueryResponse, Error> {
180        let mut url = self.base_url.join("api/public/metrics/query")?;
181        {
182            let mut pairs = url.query_pairs_mut();
183            pairs.append_pair("name", &q.name);
184            if let Some(v) = q.from.as_ref() {
185                pairs.append_pair("from", &v.to_rfc3339());
186            }
187            if let Some(v) = q.to.as_ref() {
188                pairs.append_pair("to", &v.to_rfc3339());
189            }
190            if let Some(v) = q.labels.as_deref() {
191                pairs.append_pair("labels", v);
192            }
193            if let Some(v) = q.step.as_deref() {
194                pairs.append_pair("step", v);
195            }
196            if let Some(v) = q.agg.as_deref() {
197                pairs.append_pair("agg", v);
198            }
199        }
200        let res = self.http.get(url).send().await?.error_for_status()?;
201        Ok(res.json::<MetricsQueryResponse>().await?)
202    }
203}
204
205#[derive(Debug, Serialize, Deserialize, Clone)]
206pub struct MetricPoint {
207    pub name: String,
208    #[serde(default)]
209    pub labels: HashMap<String, String>,
210    pub value: f64,
211    pub timestamp: DateTime<Utc>,
212}
213
214#[derive(Debug, Serialize)]
215struct MetricsBatchRequest {
216    metrics: Vec<MetricPoint>,
217}
218
219#[derive(Debug, Deserialize)]
220pub struct ApiResponse<T> {
221    pub message: String,
222    #[serde(default)]
223    pub data: Option<T>,
224}
225
226#[derive(Debug, Deserialize)]
227#[serde(rename_all = "camelCase")]
228pub struct PageMeta {
229    pub page: i64,
230    pub limit: i64,
231    pub total_items: i64,
232    pub total_pages: i64,
233}
234
235#[derive(Debug, Deserialize)]
236pub struct PagedData<T> {
237    pub data: Vec<T>,
238    pub meta: PageMeta,
239}
240
241#[derive(Debug, Serialize, Deserialize, Default)]
242pub struct BatchIngestRequest {
243    #[serde(default)]
244    pub trace: Option<TraceIngest>,
245    #[serde(default)]
246    pub observations: Vec<ObservationIngest>,
247}
248
249#[derive(Debug, Serialize, Deserialize)]
250#[serde(rename_all = "camelCase")]
251pub struct TraceIngest {
252    pub id: Uuid,
253    #[serde(default)]
254    pub timestamp: Option<DateTime<Utc>>,
255
256    #[serde(default)]
257    pub name: Option<String>,
258    #[serde(default)]
259    pub input: Option<JsonValue>,
260    #[serde(default)]
261    pub output: Option<JsonValue>,
262    #[serde(default)]
263    pub session_id: Option<String>,
264    #[serde(default)]
265    pub release: Option<String>,
266    #[serde(default)]
267    pub version: Option<String>,
268    #[serde(default, rename = "userId")]
269    pub user_id: Option<String>,
270    #[serde(default)]
271    pub metadata: Option<JsonValue>,
272    #[serde(default)]
273    pub tags: Vec<String>,
274    #[serde(default)]
275    pub public: Option<bool>,
276    #[serde(default)]
277    pub environment: Option<String>,
278    #[serde(default)]
279    pub external_id: Option<String>,
280    #[serde(default)]
281    pub bookmarked: Option<bool>,
282
283    #[serde(default)]
284    pub latency: Option<f64>,
285    #[serde(default, rename = "totalCost")]
286    pub total_cost: Option<f64>,
287
288    #[serde(default, rename = "projectId")]
289    pub project_id: Option<String>,
290}
291
292#[derive(Debug, Serialize, Deserialize)]
293#[serde(rename_all = "camelCase")]
294pub struct ObservationIngest {
295    pub id: Uuid,
296    #[serde(rename = "traceId")]
297    pub trace_id: Uuid,
298
299    #[serde(default, rename = "type")]
300    pub r#type: Option<String>,
301    #[serde(default)]
302    pub name: Option<String>,
303
304    #[serde(default)]
305    pub start_time: Option<DateTime<Utc>>,
306    #[serde(default)]
307    pub end_time: Option<DateTime<Utc>>,
308    #[serde(default)]
309    pub completion_start_time: Option<DateTime<Utc>>,
310
311    #[serde(default)]
312    pub model: Option<String>,
313    #[serde(default)]
314    pub model_parameters: Option<JsonValue>,
315
316    #[serde(default)]
317    pub input: Option<JsonValue>,
318    #[serde(default)]
319    pub output: Option<JsonValue>,
320
321    #[serde(default)]
322    pub usage: Option<JsonValue>,
323
324    #[serde(default)]
325    pub level: Option<String>,
326    #[serde(default)]
327    pub status_message: Option<String>,
328    #[serde(default)]
329    pub parent_observation_id: Option<Uuid>,
330
331    #[serde(default)]
332    pub prompt_id: Option<String>,
333    #[serde(default)]
334    pub prompt_name: Option<String>,
335    #[serde(default)]
336    pub prompt_version: Option<String>,
337
338    #[serde(default)]
339    pub model_id: Option<String>,
340
341    #[serde(default)]
342    pub input_price: Option<f64>,
343    #[serde(default)]
344    pub output_price: Option<f64>,
345    #[serde(default)]
346    pub total_price: Option<f64>,
347
348    #[serde(default)]
349    pub calculated_input_cost: Option<f64>,
350    #[serde(default)]
351    pub calculated_output_cost: Option<f64>,
352    #[serde(default)]
353    pub calculated_total_cost: Option<f64>,
354
355    #[serde(default)]
356    pub latency: Option<f64>,
357    #[serde(default)]
358    pub time_to_first_token: Option<f64>,
359
360    #[serde(default)]
361    pub completion_tokens: Option<i64>,
362    #[serde(default)]
363    pub prompt_tokens: Option<i64>,
364    #[serde(default)]
365    pub total_tokens: Option<i64>,
366    #[serde(default)]
367    pub unit: Option<String>,
368
369    #[serde(default)]
370    pub metadata: Option<JsonValue>,
371
372    #[serde(default)]
373    pub environment: Option<String>,
374
375    #[serde(default, rename = "projectId")]
376    pub project_id: Option<String>,
377}
378
379#[derive(Debug, Default, Serialize, Deserialize)]
380#[serde(rename_all = "camelCase")]
381pub struct TraceListQuery {
382    #[serde(default)]
383    pub page: Option<i64>,
384    #[serde(default)]
385    pub limit: Option<i64>,
386
387    #[serde(default, rename = "userId")]
388    pub user_id: Option<String>,
389    #[serde(default)]
390    pub name: Option<String>,
391    #[serde(default, rename = "sessionId")]
392    pub session_id: Option<String>,
393
394    #[serde(default, rename = "fromTimestamp")]
395    pub from_timestamp: Option<DateTime<Utc>>,
396    #[serde(default, rename = "toTimestamp")]
397    pub to_timestamp: Option<DateTime<Utc>>,
398
399    #[serde(default, rename = "orderBy")]
400    pub order_by: Option<String>,
401
402    #[serde(default)]
403    pub tags: Vec<String>,
404
405    #[serde(default)]
406    pub version: Option<String>,
407    #[serde(default)]
408    pub release: Option<String>,
409    #[serde(default)]
410    pub environment: Vec<String>,
411
412    #[serde(default)]
413    pub fields: Option<String>,
414}
415
416#[derive(Debug, Deserialize)]
417#[serde(rename_all = "camelCase")]
418pub struct TraceListItem {
419    pub id: Uuid,
420    pub timestamp: DateTime<Utc>,
421    pub name: Option<String>,
422    #[serde(default)]
423    pub input: Option<JsonValue>,
424    #[serde(default)]
425    pub output: Option<JsonValue>,
426    pub session_id: Option<String>,
427    pub release: Option<String>,
428    pub version: Option<String>,
429    pub user_id: Option<String>,
430    #[serde(default)]
431    pub metadata: Option<JsonValue>,
432    pub tags: Vec<String>,
433    pub public: bool,
434    pub environment: String,
435    pub html_path: String,
436    pub latency: Option<f64>,
437    pub total_cost: Option<f64>,
438    pub observations: Vec<String>,
439    pub scores: Vec<String>,
440}
441
442#[derive(Debug, Default, Serialize, Deserialize)]
443#[serde(rename_all = "camelCase")]
444pub struct MetricsDailyQuery {
445    #[serde(default)]
446    pub page: Option<i64>,
447    #[serde(default)]
448    pub limit: Option<i64>,
449
450    #[serde(default, rename = "traceName")]
451    pub trace_name: Option<String>,
452    #[serde(default, rename = "userId")]
453    pub user_id: Option<String>,
454    #[serde(default)]
455    pub tags: Vec<String>,
456
457    #[serde(default, rename = "fromTimestamp")]
458    pub from_timestamp: Option<DateTime<Utc>>,
459    #[serde(default, rename = "toTimestamp")]
460    pub to_timestamp: Option<DateTime<Utc>>,
461}
462
463#[derive(Debug, Deserialize)]
464#[serde(rename_all = "camelCase")]
465pub struct MetricsDailyItem {
466    pub date: String,
467    pub count_traces: i64,
468    pub count_observations: i64,
469    pub total_cost: f64,
470    pub usage: JsonValue,
471}
472
473#[derive(Debug, Deserialize)]
474#[serde(rename_all = "camelCase")]
475pub struct TraceDetailDto {
476    pub id: Uuid,
477    pub timestamp: DateTime<Utc>,
478    pub name: Option<String>,
479    pub input: JsonValue,
480    pub output: JsonValue,
481    pub session_id: Option<String>,
482    pub release: Option<String>,
483    pub version: Option<String>,
484    pub user_id: Option<String>,
485    pub metadata: JsonValue,
486    pub tags: Vec<String>,
487    pub public: bool,
488    pub environment: String,
489    pub html_path: String,
490    pub latency: Option<f64>,
491    pub total_cost: Option<f64>,
492    pub observations: Vec<JsonValue>,
493    pub scores: Vec<JsonValue>,
494}
495
496#[derive(Debug, Default, Serialize, Deserialize)]
497pub struct MetricsQueryParams {
498    pub name: String,
499    #[serde(default)]
500    pub from: Option<DateTime<Utc>>,
501    #[serde(default)]
502    pub to: Option<DateTime<Utc>>,
503    #[serde(default)]
504    pub labels: Option<String>,
505    #[serde(default)]
506    pub step: Option<String>,
507    #[serde(default)]
508    pub agg: Option<String>,
509}
510
511#[derive(Debug, Deserialize)]
512pub struct MetricsQueryResponse {
513    pub data: Vec<MetricsSeriesItem>,
514}
515
516#[derive(Debug, Deserialize)]
517pub struct MetricsSeriesItem {
518    pub labels: JsonValue,
519    pub values: Vec<MetricsValuePoint>,
520}
521
522#[derive(Debug, Deserialize)]
523pub struct MetricsValuePoint {
524    pub timestamp: String,
525    pub value: f64,
526}