braintrust_sdk_rust/
types.rs

1use std::collections::HashMap;
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use serde_json::{Map, Value};
6use serde_repr::{Deserialize_repr, Serialize_repr};
7use std::fmt;
8
9pub const LOGS_API_VERSION: u8 = 2;
10
11/// The type of span object, serialized as its integer representation for wire compatibility.
12#[derive(Debug, Clone, Copy, Serialize_repr, Deserialize_repr, PartialEq, Eq)]
13#[repr(u8)]
14pub enum SpanObjectType {
15    Experiment = 1,
16    ProjectLogs = 2,
17    PlaygroundLogs = 3,
18}
19
20impl SpanObjectType {
21    pub fn as_str(self) -> &'static str {
22        match self {
23            SpanObjectType::Experiment => "experiment",
24            SpanObjectType::ProjectLogs => "project_logs",
25            SpanObjectType::PlaygroundLogs => "playground_logs",
26        }
27    }
28}
29
30impl fmt::Display for SpanObjectType {
31    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
32        f.write_str(self.as_str())
33    }
34}
35
36/// Error returned when an invalid u8 value is converted to SpanObjectType.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct InvalidSpanObjectType(pub u8);
39
40impl fmt::Display for InvalidSpanObjectType {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        write!(f, "invalid SpanObjectType value: {}", self.0)
43    }
44}
45
46impl std::error::Error for InvalidSpanObjectType {}
47
48impl TryFrom<u8> for SpanObjectType {
49    type Error = InvalidSpanObjectType;
50
51    fn try_from(value: u8) -> Result<Self, Self::Error> {
52        match value {
53            1 => Ok(SpanObjectType::Experiment),
54            2 => Ok(SpanObjectType::ProjectLogs),
55            3 => Ok(SpanObjectType::PlaygroundLogs),
56            _ => Err(InvalidSpanObjectType(value)),
57        }
58    }
59}
60
61/// The type of span.
62#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
63#[serde(rename_all = "snake_case")]
64pub enum SpanType {
65    #[default]
66    Llm,
67    Score,
68    Function,
69    Eval,
70    Task,
71    Tool,
72    Automation,
73    Facet,
74    Preprocessor,
75}
76
77/// Span attributes with typed known fields and passthrough for extras.
78///
79/// Uses `#[serde(flatten)]` to serialize/deserialize unknown fields into the
80/// `extra` map, matching the TS SDK's `.passthrough()` behavior.
81#[derive(Debug, Clone, Default, Serialize, Deserialize)]
82pub(crate) struct SpanAttributes {
83    #[serde(skip_serializing_if = "Option::is_none")]
84    pub name: Option<String>,
85    #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
86    pub span_type: Option<SpanType>,
87    #[serde(skip_serializing_if = "Option::is_none")]
88    pub purpose: Option<String>,
89    /// Catch-all for additional fields (passthrough behavior).
90    #[serde(flatten, skip_serializing_if = "HashMap::is_empty")]
91    pub extra: HashMap<String, Value>,
92}
93
94/// The destination for a log row. Each variant represents a mutually exclusive
95/// target: an experiment, project logs, or playground logs.
96///
97/// NOTE: The `untagged` attribute is only safe if the field sets are also
98/// mutually exclusive.
99#[derive(Debug, Clone, Serialize, Deserialize)]
100#[serde(untagged)]
101pub(crate) enum LogDestination {
102    /// Log to an experiment (for evaluation runs).
103    Experiment { experiment_id: String },
104    /// Log to project logs (general observability).
105    ProjectLogs { project_id: String, log_id: String },
106    /// Log to playground logs (interactive sessions).
107    PlaygroundLogs {
108        prompt_session_id: String,
109        log_id: String,
110    },
111}
112
113impl LogDestination {
114    /// Create a new experiment destination.
115    pub fn experiment(experiment_id: impl Into<String>) -> Self {
116        Self::Experiment {
117            experiment_id: experiment_id.into(),
118        }
119    }
120
121    /// Create a new project logs destination.
122    pub fn project_logs(project_id: impl Into<String>) -> Self {
123        Self::ProjectLogs {
124            project_id: project_id.into(),
125            log_id: "g".to_string(),
126        }
127    }
128
129    /// Create a new playground logs destination.
130    pub fn playground_logs(prompt_session_id: impl Into<String>) -> Self {
131        Self::PlaygroundLogs {
132            prompt_session_id: prompt_session_id.into(),
133            log_id: "x".to_string(),
134        }
135    }
136}
137
138#[derive(Debug, Clone, Serialize)]
139pub(crate) struct Logs3Request {
140    pub rows: Vec<Logs3Row>,
141    pub api_version: u8,
142}
143
144#[derive(Debug, Clone, Serialize)]
145pub(crate) struct Logs3Row {
146    pub id: String,
147    #[serde(rename = "_is_merge", skip_serializing_if = "Option::is_none")]
148    pub is_merge: Option<bool>,
149    pub span_id: String,
150    pub root_span_id: String,
151    #[serde(skip_serializing_if = "Option::is_none")]
152    pub span_parents: Option<Vec<String>>,
153    #[serde(flatten)]
154    pub destination: LogDestination,
155    pub org_id: String,
156    #[serde(skip_serializing_if = "Option::is_none")]
157    pub org_name: Option<String>,
158    #[serde(skip_serializing_if = "Option::is_none")]
159    pub input: Option<Value>,
160    #[serde(skip_serializing_if = "Option::is_none")]
161    pub output: Option<Value>,
162    #[serde(skip_serializing_if = "Option::is_none")]
163    pub metadata: Option<Map<String, Value>>,
164    #[serde(skip_serializing_if = "Option::is_none")]
165    pub metrics: Option<HashMap<String, f64>>,
166    #[serde(skip_serializing_if = "Option::is_none")]
167    pub span_attributes: Option<SpanAttributes>,
168    pub created: DateTime<Utc>,
169}
170
171#[derive(Debug, Clone)]
172pub(crate) struct SpanPayload {
173    pub row_id: String,
174    pub span_id: String,
175    pub is_merge: bool,
176    pub org_id: String,
177    pub org_name: Option<String>,
178    pub project_name: Option<String>,
179    pub input: Option<Value>,
180    pub output: Option<Value>,
181    pub metadata: Option<Map<String, Value>>,
182    pub metrics: Option<HashMap<String, f64>>,
183    pub span_attributes: Option<SpanAttributes>,
184}
185
186#[derive(Debug, Clone, Serialize, Deserialize)]
187pub enum ParentSpanInfo {
188    Experiment {
189        object_id: String,
190    },
191    ProjectLogs {
192        object_id: String,
193    },
194    ProjectName {
195        project_name: String,
196    },
197    PlaygroundLogs {
198        object_id: String,
199    },
200    FullSpan {
201        object_type: SpanObjectType,
202        object_id: String,
203        span_id: String,
204        root_span_id: String,
205    },
206}
207
208#[derive(Debug, Clone, Default, Serialize, Deserialize)]
209pub struct PromptTokensDetails {
210    pub audio_tokens: Option<u32>,
211    pub cached_tokens: Option<u32>,
212    pub cache_creation_tokens: Option<u32>,
213}
214
215#[derive(Debug, Clone, Default, Serialize, Deserialize)]
216pub struct CompletionTokensDetails {
217    pub audio_tokens: Option<u32>,
218    pub reasoning_tokens: Option<u32>,
219    pub accepted_prediction_tokens: Option<u32>,
220    pub rejected_prediction_tokens: Option<u32>,
221}
222
223#[derive(Debug, Clone, Default)]
224pub struct UsageMetrics {
225    pub prompt_tokens: Option<u32>,
226    pub completion_tokens: Option<u32>,
227    pub total_tokens: Option<u32>,
228    pub reasoning_tokens: Option<u32>,
229    pub prompt_cached_tokens: Option<u32>,
230    pub prompt_cache_creation_tokens: Option<u32>,
231    pub completion_reasoning_tokens: Option<u32>,
232    pub prompt_tokens_details: Option<PromptTokensDetails>,
233    pub completion_tokens_details: Option<CompletionTokensDetails>,
234}
235
236/// Usage statistics that can deserialize from both OpenAI and Anthropic formats.
237///
238/// OpenAI uses `prompt_tokens`/`completion_tokens`, while Anthropic uses
239/// `input_tokens`/`output_tokens`. The serde aliases allow this struct to
240/// deserialize from either format automatically.
241#[derive(Debug, Clone, Serialize, Deserialize, Default)]
242pub struct Usage {
243    #[serde(default, alias = "input_tokens")]
244    pub prompt_tokens: u32,
245    #[serde(default, alias = "output_tokens")]
246    pub completion_tokens: u32,
247    #[serde(default)]
248    pub total_tokens: u32,
249    #[serde(default)]
250    pub reasoning_tokens: Option<u32>,
251    #[serde(
252        default,
253        skip_serializing_if = "Option::is_none",
254        alias = "cache_read_input_tokens"
255    )]
256    pub prompt_cached_tokens: Option<u32>,
257    #[serde(
258        default,
259        skip_serializing_if = "Option::is_none",
260        alias = "cache_creation_input_tokens"
261    )]
262    pub prompt_cache_creation_tokens: Option<u32>,
263    #[serde(default, skip_serializing_if = "Option::is_none")]
264    pub completion_reasoning_tokens: Option<u32>,
265    #[serde(default, skip_serializing_if = "Option::is_none")]
266    pub prompt_tokens_details: Option<PromptTokensDetails>,
267    #[serde(default, skip_serializing_if = "Option::is_none")]
268    pub completion_tokens_details: Option<CompletionTokensDetails>,
269}
270
271impl Usage {
272    /// Create a Usage from UsageMetrics, returning None if no metrics are present.
273    pub fn from_metrics(metrics: UsageMetrics) -> Option<Self> {
274        let has_metrics = metrics.prompt_tokens.is_some()
275            || metrics.completion_tokens.is_some()
276            || metrics.total_tokens.is_some()
277            || metrics.reasoning_tokens.is_some()
278            || metrics.prompt_cached_tokens.is_some()
279            || metrics.prompt_cache_creation_tokens.is_some()
280            || metrics.completion_reasoning_tokens.is_some();
281
282        if !has_metrics {
283            return None;
284        }
285
286        let prompt = metrics.prompt_tokens.unwrap_or_default();
287        let completion = metrics.completion_tokens.unwrap_or_default();
288        let total = metrics
289            .total_tokens
290            .or_else(|| {
291                if prompt != 0 || completion != 0 {
292                    Some(prompt + completion)
293                } else {
294                    None
295                }
296            })
297            .unwrap_or_default();
298        let prompt_details = metrics.prompt_tokens_details.clone();
299        let completion_details = metrics.completion_tokens_details.clone();
300
301        Some(Self {
302            prompt_tokens: prompt,
303            completion_tokens: completion,
304            total_tokens: total,
305            reasoning_tokens: metrics.reasoning_tokens,
306            prompt_cached_tokens: metrics.prompt_cached_tokens.or_else(|| {
307                prompt_details
308                    .as_ref()
309                    .and_then(|details| details.cached_tokens)
310            }),
311            prompt_cache_creation_tokens: metrics.prompt_cache_creation_tokens.or_else(|| {
312                prompt_details
313                    .as_ref()
314                    .and_then(|details| details.cache_creation_tokens)
315            }),
316            completion_reasoning_tokens: metrics.completion_reasoning_tokens.or_else(|| {
317                completion_details
318                    .as_ref()
319                    .and_then(|details| details.reasoning_tokens)
320            }),
321            prompt_tokens_details: prompt_details,
322            completion_tokens_details: completion_details,
323        })
324    }
325}
326
327pub fn usage_metrics_to_map(usage: UsageMetrics) -> HashMap<String, f64> {
328    let mut metrics = HashMap::new();
329    insert_metric(&mut metrics, "prompt_tokens", usage.prompt_tokens);
330    insert_metric(&mut metrics, "completion_tokens", usage.completion_tokens);
331    insert_metric(&mut metrics, "tokens", usage.total_tokens);
332    insert_metric(&mut metrics, "reasoning_tokens", usage.reasoning_tokens);
333    insert_metric(
334        &mut metrics,
335        "completion_reasoning_tokens",
336        usage.completion_reasoning_tokens,
337    );
338    insert_metric(
339        &mut metrics,
340        "prompt_cached_tokens",
341        usage.prompt_cached_tokens,
342    );
343    insert_metric(
344        &mut metrics,
345        "prompt_cache_creation_tokens",
346        usage.prompt_cache_creation_tokens,
347    );
348
349    if let Some(details) = usage.prompt_tokens_details {
350        insert_metric(&mut metrics, "prompt_audio_tokens", details.audio_tokens);
351        if usage.prompt_cached_tokens.is_none() {
352            insert_metric(&mut metrics, "prompt_cached_tokens", details.cached_tokens);
353        }
354        if usage.prompt_cache_creation_tokens.is_none() {
355            insert_metric(
356                &mut metrics,
357                "prompt_cache_creation_tokens",
358                details.cache_creation_tokens,
359            );
360        }
361    }
362
363    if let Some(details) = usage.completion_tokens_details {
364        insert_metric(
365            &mut metrics,
366            "completion_audio_tokens",
367            details.audio_tokens,
368        );
369        if usage.completion_reasoning_tokens.is_none() {
370            insert_metric(
371                &mut metrics,
372                "completion_reasoning_tokens",
373                details.reasoning_tokens,
374            );
375        }
376        insert_metric(
377            &mut metrics,
378            "completion_accepted_prediction_tokens",
379            details.accepted_prediction_tokens,
380        );
381        insert_metric(
382            &mut metrics,
383            "completion_rejected_prediction_tokens",
384            details.rejected_prediction_tokens,
385        );
386    }
387
388    metrics
389}
390
391fn insert_metric(metrics: &mut HashMap<String, f64>, key: &str, value: Option<u32>) {
392    if let Some(value) = value {
393        metrics.insert(key.to_string(), value as f64);
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use serde_json::json;
401
402    #[test]
403    fn log_destination_experiment_serializes_flat() {
404        let dest = LogDestination::experiment("exp-123");
405        let json = serde_json::to_value(&dest).unwrap();
406
407        assert_eq!(json, json!({"experiment_id": "exp-123"}));
408        // No log_id field for experiments
409        assert!(json.get("log_id").is_none());
410    }
411
412    #[test]
413    fn log_destination_project_logs_serializes_with_log_id() {
414        let dest = LogDestination::project_logs("proj-456");
415        let json = serde_json::to_value(&dest).unwrap();
416
417        assert_eq!(json.get("project_id").unwrap(), "proj-456");
418        assert_eq!(json.get("log_id").unwrap(), "g");
419    }
420
421    #[test]
422    fn log_destination_playground_serializes_with_log_id() {
423        let dest = LogDestination::playground_logs("session-789");
424        let json = serde_json::to_value(&dest).unwrap();
425
426        assert_eq!(json.get("prompt_session_id").unwrap(), "session-789");
427        assert_eq!(json.get("log_id").unwrap(), "x");
428    }
429
430    #[test]
431    fn log_destination_deserializes_experiment() {
432        let json = json!({"experiment_id": "exp-123"});
433        let dest: LogDestination = serde_json::from_value(json).unwrap();
434
435        assert!(
436            matches!(dest, LogDestination::Experiment { experiment_id } if experiment_id == "exp-123")
437        );
438    }
439
440    #[test]
441    fn log_destination_deserializes_project_logs() {
442        let json = json!({"project_id": "proj-456", "log_id": "g"});
443        let dest: LogDestination = serde_json::from_value(json).unwrap();
444
445        assert!(
446            matches!(dest, LogDestination::ProjectLogs { project_id, log_id }
447            if project_id == "proj-456" && log_id == "g")
448        );
449    }
450
451    #[test]
452    fn log_destination_deserializes_playground() {
453        let json = json!({"prompt_session_id": "session-789", "log_id": "x"});
454        let dest: LogDestination = serde_json::from_value(json).unwrap();
455
456        assert!(
457            matches!(dest, LogDestination::PlaygroundLogs { prompt_session_id, log_id }
458            if prompt_session_id == "session-789" && log_id == "x")
459        );
460    }
461
462    #[test]
463    fn log_destination_rejects_empty_object() {
464        let json = json!({});
465        let result: Result<LogDestination, _> = serde_json::from_value(json);
466
467        assert!(result.is_err());
468    }
469
470    #[test]
471    fn log_destination_rejects_missing_required_fields() {
472        // project_id without log_id should fail to match ProjectLogs,
473        // and won't match other variants either
474        let json = json!({"project_id": "proj-456"});
475        let result: Result<LogDestination, _> = serde_json::from_value(json);
476
477        assert!(result.is_err());
478    }
479
480    #[test]
481    fn logs3_row_flattens_destination() {
482        let row = Logs3Row {
483            id: "row-1".to_string(),
484            is_merge: None,
485            span_id: "span-1".to_string(),
486            root_span_id: "span-1".to_string(),
487            span_parents: None,
488            destination: LogDestination::experiment("exp-123"),
489            org_id: "org-1".to_string(),
490            org_name: None,
491            input: None,
492            output: None,
493            metadata: None,
494            metrics: None,
495            span_attributes: None,
496            created: Utc::now(),
497        };
498
499        let json = serde_json::to_value(&row).unwrap();
500
501        // experiment_id should be at top level, not nested
502        assert_eq!(json.get("experiment_id").unwrap(), "exp-123");
503        assert!(json.get("destination").is_none());
504        // No log_id for experiments
505        assert!(json.get("log_id").is_none());
506        // org_id and created are always present
507        assert!(json.get("org_id").is_some());
508        assert!(json.get("created").is_some());
509    }
510
511    #[test]
512    fn parent_span_info_full_span_serializes_object_type_as_u8() {
513        let parent = ParentSpanInfo::FullSpan {
514            object_type: SpanObjectType::Experiment,
515            object_id: "exp-123".to_string(),
516            span_id: "span-1".to_string(),
517            root_span_id: "root-1".to_string(),
518        };
519
520        let json = serde_json::to_value(&parent).unwrap();
521        let obj = json.get("FullSpan").unwrap();
522
523        // SpanObjectType serializes as u8 for wire compatibility
524        assert_eq!(obj.get("object_type").unwrap(), 1);
525    }
526
527    #[test]
528    fn parent_span_info_deserializes_with_typed_object_type() {
529        // Deserializes from integer (wire format) into SpanObjectType
530        let json = json!({
531            "FullSpan": {
532                "object_type": 1,
533                "object_id": "exp-123",
534                "span_id": "span-1",
535                "root_span_id": "root-1"
536            }
537        });
538
539        let parent: ParentSpanInfo = serde_json::from_value(json).unwrap();
540
541        match parent {
542            ParentSpanInfo::FullSpan { object_type, .. } => {
543                assert_eq!(object_type, SpanObjectType::Experiment);
544            }
545            _ => panic!("Expected FullSpan variant"),
546        }
547    }
548
549    #[test]
550    fn parent_span_info_rejects_invalid_object_type() {
551        // Invalid integer value should fail
552        let json = json!({
553            "FullSpan": {
554                "object_type": 99,
555                "object_id": "exp-123",
556                "span_id": "span-1",
557                "root_span_id": "root-1"
558            }
559        });
560
561        let result: Result<ParentSpanInfo, _> = serde_json::from_value(json);
562        assert!(result.is_err());
563    }
564
565    #[test]
566    fn parent_span_info_rejects_string_object_type() {
567        // String value should fail (must be integer)
568        let json = json!({
569            "FullSpan": {
570                "object_type": "Experiment",
571                "object_id": "exp-123",
572                "span_id": "span-1",
573                "root_span_id": "root-1"
574            }
575        });
576
577        let result: Result<ParentSpanInfo, _> = serde_json::from_value(json);
578        assert!(result.is_err());
579    }
580
581    #[test]
582    fn span_object_type_try_from_u8() {
583        assert_eq!(SpanObjectType::try_from(1), Ok(SpanObjectType::Experiment));
584        assert_eq!(SpanObjectType::try_from(2), Ok(SpanObjectType::ProjectLogs));
585        assert_eq!(
586            SpanObjectType::try_from(3),
587            Ok(SpanObjectType::PlaygroundLogs)
588        );
589        assert_eq!(SpanObjectType::try_from(0), Err(InvalidSpanObjectType(0)));
590        assert_eq!(SpanObjectType::try_from(99), Err(InvalidSpanObjectType(99)));
591    }
592
593    #[test]
594    fn span_type_serializes_as_snake_case() {
595        assert_eq!(serde_json::to_value(SpanType::Llm).unwrap(), json!("llm"));
596        assert_eq!(
597            serde_json::to_value(SpanType::Score).unwrap(),
598            json!("score")
599        );
600        assert_eq!(
601            serde_json::to_value(SpanType::Function).unwrap(),
602            json!("function")
603        );
604        assert_eq!(
605            serde_json::to_value(SpanType::Automation).unwrap(),
606            json!("automation")
607        );
608        assert_eq!(
609            serde_json::to_value(SpanType::Facet).unwrap(),
610            json!("facet")
611        );
612        assert_eq!(
613            serde_json::to_value(SpanType::Preprocessor).unwrap(),
614            json!("preprocessor")
615        );
616    }
617
618    #[test]
619    fn span_type_deserializes_from_snake_case() {
620        let llm: SpanType = serde_json::from_value(json!("llm")).unwrap();
621        assert_eq!(llm, SpanType::Llm);
622
623        let tool: SpanType = serde_json::from_value(json!("tool")).unwrap();
624        assert_eq!(tool, SpanType::Tool);
625
626        let automation: SpanType = serde_json::from_value(json!("automation")).unwrap();
627        assert_eq!(automation, SpanType::Automation);
628
629        let facet: SpanType = serde_json::from_value(json!("facet")).unwrap();
630        assert_eq!(facet, SpanType::Facet);
631
632        let preprocessor: SpanType = serde_json::from_value(json!("preprocessor")).unwrap();
633        assert_eq!(preprocessor, SpanType::Preprocessor);
634    }
635
636    #[test]
637    fn span_attributes_serializes_flat_with_extras() {
638        let attrs = SpanAttributes {
639            name: Some("my-span".to_string()),
640            span_type: Some(SpanType::Llm),
641            purpose: None,
642            extra: [("custom_field".to_string(), json!(42))]
643                .into_iter()
644                .collect(),
645        };
646
647        let json = serde_json::to_value(&attrs).unwrap();
648
649        // All fields should be at top level (flat, not nested "extra")
650        assert_eq!(json.get("name").unwrap(), "my-span");
651        assert_eq!(json.get("type").unwrap(), "llm");
652        assert_eq!(json.get("custom_field").unwrap(), 42);
653        // No "extra" key in output
654        assert!(json.get("extra").is_none());
655        // purpose is None so should be omitted
656        assert!(json.get("purpose").is_none());
657    }
658
659    #[test]
660    fn span_attributes_deserializes_with_passthrough() {
661        let json = json!({
662            "name": "test-span",
663            "type": "score",
664            "purpose": "scorer",
665            "exec_counter": 5,
666            "unknown_field": "hello"
667        });
668
669        let attrs: SpanAttributes = serde_json::from_value(json).unwrap();
670
671        assert_eq!(attrs.name, Some("test-span".to_string()));
672        assert_eq!(attrs.span_type, Some(SpanType::Score));
673        assert_eq!(attrs.purpose, Some("scorer".to_string()));
674        // Unknown fields captured in extra
675        assert_eq!(attrs.extra.get("exec_counter").unwrap(), &json!(5));
676        assert_eq!(attrs.extra.get("unknown_field").unwrap(), &json!("hello"));
677    }
678
679    #[test]
680    fn span_attributes_empty_serializes_to_empty_object() {
681        let attrs = SpanAttributes::default();
682        let json = serde_json::to_value(&attrs).unwrap();
683
684        assert_eq!(json, json!({}));
685    }
686}
braintrust_sdk_rust/types.rs

braintrust_sdk_rust/
types.rs
