cognis_trace/
span.rs

1//! Core span / generation / score types. Aligned with Langfuse's
2//! observation model so the Langfuse exporter is a 1:1 mapping. See
3//! the spec, §4.
4
5use std::collections::HashMap;
6use std::time::SystemTime;
7
8use serde::{Deserialize, Serialize};
9use uuid::Uuid;
10
11/// Type of an observation. Aligned with Langfuse's `ObservationType`
12/// enum so values cross-walk directly. See spec §4.1.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "UPPERCASE")]
15pub enum SpanKind {
16    /// Generic span (chains, custom blocks).
17    Span,
18    /// LLM generation. Has a `Some(generation)` payload.
19    Generation,
20    /// Discrete event (no duration semantics).
21    Event,
22    /// Agent-loop span.
23    Agent,
24    /// Tool execution.
25    Tool,
26    /// Logical chain block.
27    Chain,
28    /// Retriever step.
29    Retriever,
30    /// Embedding step.
31    Embedding,
32    /// Guardrail / safety check.
33    Guardrail,
34}
35
36/// Observation severity, matching Langfuse's `ObservationLevel`.
37#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
38#[serde(rename_all = "UPPERCASE")]
39pub enum ObservationLevel {
40    /// Normal completion.
41    #[default]
42    Default,
43    /// Diagnostic.
44    Debug,
45    /// Recoverable issue.
46    Warning,
47    /// Failure.
48    Error,
49}
50
51/// One node in the trace tree.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct Span {
54    /// This observation's id.
55    pub run_id: Uuid,
56    /// Parent observation id, when nested.
57    pub parent_run_id: Option<Uuid>,
58    /// Root run_id of the trace this span belongs to.
59    pub trace_id: Uuid,
60    /// Observation type.
61    pub kind: SpanKind,
62    /// Friendly name (e.g. "openai.gpt-4o", "search_tool", "agent_node").
63    pub name: String,
64    /// Wall-clock start.
65    pub started_at: SystemTime,
66    /// Wall-clock end (None until span closes).
67    pub ended_at: Option<SystemTime>,
68    /// Severity.
69    pub level: ObservationLevel,
70    /// Set when level == Warning or Error.
71    pub status_message: Option<String>,
72    /// Input payload (Value to be backend-agnostic).
73    pub input: Option<serde_json::Value>,
74    /// Output payload.
75    pub output: Option<serde_json::Value>,
76    /// Set only on the trace root.
77    pub session_id: Option<String>,
78    /// Set only on the trace root.
79    pub user_id: Option<String>,
80    /// Set only on the trace root.
81    pub tags: Vec<String>,
82    /// Free-form metadata (excluding well-known fields above).
83    pub metadata: HashMap<String, serde_json::Value>,
84    /// Some iff `kind == Generation`.
85    pub generation: Option<Generation>,
86}
87
88/// LLM-specific data carried on a `Generation` span.
89#[derive(Debug, Clone, Default, Serialize, Deserialize)]
90pub struct Generation {
91    /// Concrete model id (e.g. "gpt-4o-2024-08-06").
92    pub model: String,
93    /// Provider tag (e.g. "openai", "anthropic").
94    pub provider: String,
95    /// Provider-supplied request parameters (temperature, max_tokens, …).
96    #[serde(default)]
97    pub model_parameters: HashMap<String, serde_json::Value>,
98    /// Token usage broken out by category.
99    pub usage: TokenUsage,
100    /// USD cost, structured. None when the model is not in the price table.
101    pub cost: Option<CostDetails>,
102    /// Time-to-first-token, when the provider reports it.
103    pub completion_start_time: Option<SystemTime>,
104    /// Provider's stop reason ("stop", "length", "tool_calls", …).
105    pub finish_reason: Option<String>,
106    /// Set when the prompt came from a `PromptStore`.
107    pub prompt_name: Option<String>,
108    /// Set when the prompt came from a `PromptStore`.
109    pub prompt_version: Option<u32>,
110}
111
112/// Token counts broken out by category.
113#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
114pub struct TokenUsage {
115    /// Input tokens (uncached portion).
116    pub input: u32,
117    /// Output tokens.
118    pub output: u32,
119    /// Cached input tokens that were read.
120    pub cache_read: u32,
121    /// Cached input tokens written this call.
122    pub cache_write: u32,
123}
124
125impl TokenUsage {
126    /// Sum of all token categories.
127    pub fn total(&self) -> u32 {
128        self.input + self.output + self.cache_read + self.cache_write
129    }
130}
131
132/// USD cost, broken out by token category.
133#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
134pub struct CostDetails {
135    /// Input-token cost.
136    pub input: f64,
137    /// Output-token cost.
138    pub output: f64,
139    /// Cache-read cost.
140    pub cache_read: f64,
141    /// Cache-write cost.
142    pub cache_write: f64,
143    /// Convenience: input + output + cache_read + cache_write.
144    pub total: f64,
145}
146
147/// Numeric, categorical, or boolean evaluation score.
148#[derive(Debug, Clone, Serialize, Deserialize)]
149#[serde(untagged)]
150pub enum ScoreValue {
151    /// Continuous numeric score.
152    Numeric(f64),
153    /// Categorical label (1–500 chars per Langfuse contract).
154    Categorical(String),
155    /// Boolean — serialized as 1.0 / 0.0 by Langfuse.
156    Boolean(bool),
157}
158
159/// Evaluation score attached to a run.
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct ScoreRecord {
162    /// The observation this score applies to.
163    pub run_id: Uuid,
164    /// Optional explicit trace pointer (for cross-trace scores).
165    pub trace_id: Option<Uuid>,
166    /// Optional session pointer (session-level scores).
167    pub session_id: Option<String>,
168    /// Score name (e.g. "novelty", "factuality").
169    pub name: String,
170    /// The score value.
171    pub value: ScoreValue,
172    /// Optional human comment.
173    pub comment: Option<String>,
174}
175
176/// In-flight span being assembled between `on_*_start` and `on_*_end`.
177/// Pure data; no behavior — `TracingHandler` drives it.
178#[derive(Debug, Clone)]
179pub struct SpanBuilder {
180    /// The span being assembled. `ended_at`, `output`, `level`, `status_message`,
181    /// and (for Generation) `generation` are filled at close time.
182    pub span: Span,
183}
184
185impl SpanBuilder {
186    /// Open a new span at `now`.
187    pub fn open(
188        run_id: Uuid,
189        parent_run_id: Option<Uuid>,
190        trace_id: Uuid,
191        kind: SpanKind,
192        name: impl Into<String>,
193        input: Option<serde_json::Value>,
194        now: SystemTime,
195    ) -> Self {
196        Self {
197            span: Span {
198                run_id,
199                parent_run_id,
200                trace_id,
201                kind,
202                name: name.into(),
203                started_at: now,
204                ended_at: None,
205                level: ObservationLevel::Default,
206                status_message: None,
207                input,
208                output: None,
209                session_id: None,
210                user_id: None,
211                tags: Vec::new(),
212                metadata: HashMap::new(),
213                generation: None,
214            },
215        }
216    }
217
218    /// Mark the span ended successfully and stamp the output.
219    pub fn finish_ok(mut self, output: Option<serde_json::Value>, now: SystemTime) -> Span {
220        self.span.ended_at = Some(now);
221        self.span.output = output;
222        self.span
223    }
224
225    /// Mark the span ended with an error.
226    pub fn finish_error(mut self, message: impl Into<String>, now: SystemTime) -> Span {
227        self.span.ended_at = Some(now);
228        self.span.level = ObservationLevel::Error;
229        self.span.status_message = Some(message.into());
230        self.span
231    }
232
233    /// Attach a populated `Generation` payload (used on `on_llm_end`).
234    pub fn with_generation(mut self, gen: Generation) -> Self {
235        self.span.kind = SpanKind::Generation;
236        self.span.generation = Some(gen);
237        self
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn span_kind_serializes_uppercase() {
247        let s = serde_json::to_string(&SpanKind::Generation).unwrap();
248        assert_eq!(s, "\"GENERATION\"");
249    }
250
251    #[test]
252    fn observation_level_default_is_default() {
253        assert_eq!(ObservationLevel::default(), ObservationLevel::Default);
254    }
255
256    #[test]
257    fn token_usage_total_sums_categories() {
258        let u = TokenUsage {
259            input: 10,
260            output: 20,
261            cache_read: 5,
262            cache_write: 3,
263        };
264        assert_eq!(u.total(), 38);
265    }
266
267    #[test]
268    fn score_value_numeric_serializes_as_number() {
269        let v = ScoreValue::Numeric(0.9);
270        assert_eq!(serde_json::to_string(&v).unwrap(), "0.9");
271    }
272
273    #[test]
274    fn score_value_categorical_serializes_as_string() {
275        let v = ScoreValue::Categorical("good".into());
276        assert_eq!(serde_json::to_string(&v).unwrap(), "\"good\"");
277    }
278
279    #[test]
280    fn span_builder_opens_with_default_level() {
281        let id = Uuid::new_v4();
282        let now = SystemTime::now();
283        let b = SpanBuilder::open(id, None, id, SpanKind::Chain, "x", None, now);
284        assert_eq!(b.span.level, ObservationLevel::Default);
285        assert!(b.span.ended_at.is_none());
286    }
287
288    #[test]
289    fn span_builder_finish_ok_sets_end_and_output() {
290        let id = Uuid::new_v4();
291        let now = SystemTime::now();
292        let b = SpanBuilder::open(id, None, id, SpanKind::Chain, "x", None, now);
293        let span = b.finish_ok(Some(serde_json::json!({"k": "v"})), now);
294        assert!(span.ended_at.is_some());
295        assert_eq!(span.level, ObservationLevel::Default);
296        assert_eq!(span.output, Some(serde_json::json!({"k": "v"})));
297    }
298
299    #[test]
300    fn span_builder_finish_error_sets_level_and_message() {
301        let id = Uuid::new_v4();
302        let now = SystemTime::now();
303        let b = SpanBuilder::open(id, None, id, SpanKind::Tool, "t", None, now);
304        let span = b.finish_error("boom", now);
305        assert_eq!(span.level, ObservationLevel::Error);
306        assert_eq!(span.status_message.as_deref(), Some("boom"));
307    }
308
309    #[test]
310    fn span_builder_with_generation_flips_kind() {
311        let id = Uuid::new_v4();
312        let b = SpanBuilder::open(id, None, id, SpanKind::Span, "g", None, SystemTime::now())
313            .with_generation(Generation {
314                model: "gpt-4o".into(),
315                provider: "openai".into(),
316                ..Default::default()
317            });
318        assert_eq!(b.span.kind, SpanKind::Generation);
319        assert!(b.span.generation.is_some());
320    }
321}
cognis_trace/span.rs

cognis_trace/
span.rs
