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