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