trace_weft/
builder.rs

1use std::collections::HashMap;
2use trace_weft_core::{
3    BlobRef, CapturePolicy, CostEstimate, RunId, SpanId, SpanRecord, SpanStatus, TokenUsage,
4    TraceId, TraceWeftSpanKind,
5};
6use uuid::Uuid;
7
8pub struct SpanBuilder {
9    pub span: SpanRecord,
10}
11
12impl SpanBuilder {
13    pub fn new(kind: TraceWeftSpanKind, name: impl Into<String>) -> Self {
14        let now = std::time::SystemTime::now()
15            .duration_since(std::time::UNIX_EPOCH)
16            .unwrap_or_default()
17            .as_millis() as u64;
18
19        Self {
20            span: SpanRecord {
21                trace_id: TraceId(Uuid::now_v7()),
22                span_id: SpanId(Uuid::now_v7()),
23                parent_span_id: None,
24                run_id: RunId(Uuid::now_v7()),
25                session_id: None,
26                user_id_hash: None,
27                project_id: None,
28                span_kind: kind,
29                name: name.into(),
30                start_time: now,
31                end_time: None,
32                status: SpanStatus::InProgress,
33                status_message: None,
34                error_type: None,
35                error_message_redacted: None,
36                attributes: HashMap::new(),
37                otel_attributes: HashMap::new(),
38                openinference_attributes: HashMap::new(),
39                memory_state: None,
40                input_ref: None,
41                output_ref: None,
42                prompt_template_id: None,
43                prompt_version: None,
44                model_provider: None,
45                model_name: None,
46                tool_name: None,
47                tool_schema_hash: None,
48                retrieval_query_hash: None,
49                retrieved_document_refs: vec![],
50                token_usage: None,
51                cost_estimate: None,
52                latency_ms: None,
53                retry_count: None,
54                cache_hit: None,
55                redaction_policy: CapturePolicy::MetadataOnly,
56                schema_version: "1.0".to_string(),
57            },
58        }
59    }
60
61    pub fn provider(mut self, provider: impl Into<String>) -> Self {
62        self.span.model_provider = Some(provider.into());
63        self
64    }
65
66    pub fn model(mut self, model: impl Into<String>) -> Self {
67        self.span.model_name = Some(model.into());
68        self
69    }
70
71    pub fn prompt_version(mut self, version: impl Into<String>) -> Self {
72        self.span.prompt_version = Some(version.into());
73        self
74    }
75
76    pub fn tool_name(mut self, tool: impl Into<String>) -> Self {
77        self.span.tool_name = Some(tool.into());
78        self
79    }
80
81    pub fn input_ref(mut self, blob_ref: BlobRef) -> Self {
82        self.span.input_ref = Some(blob_ref);
83        self
84    }
85
86    pub fn output_ref(mut self, blob_ref: BlobRef) -> Self {
87        self.span.output_ref = Some(blob_ref);
88        self
89    }
90
91    pub fn token_usage(mut self, usage: TokenUsage) -> Self {
92        self.span.token_usage = Some(usage);
93        self
94    }
95
96    pub fn cost(mut self, cost: CostEstimate) -> Self {
97        self.span.cost_estimate = Some(cost);
98        self
99    }
100
101    pub fn cache_hit(mut self, hit: bool) -> Self {
102        self.span.cache_hit = Some(hit);
103        self
104    }
105
106    /// Record a retrieval query hash and the documents it returned.
107    pub fn retrieval(mut self, query_hash: impl Into<String>, doc_refs: Vec<BlobRef>) -> Self {
108        self.span.retrieval_query_hash = Some(query_hash.into());
109        self.span.retrieved_document_refs = doc_refs;
110        self
111    }
112
113    /// Insert a single free-form attribute.
114    pub fn attribute(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
115        self.span.attributes.insert(key.into(), value);
116        self
117    }
118
119    /// Merge a map of free-form attributes into the span.
120    pub fn attributes(mut self, attrs: HashMap<String, serde_json::Value>) -> Self {
121        self.span.attributes.extend(attrs);
122        self
123    }
124
125    pub fn with_parent(mut self, trace_id: TraceId, run_id: RunId, parent_id: SpanId) -> Self {
126        self.span.trace_id = trace_id;
127        self.span.run_id = run_id;
128        self.span.parent_span_id = Some(parent_id);
129        self
130    }
131
132    pub async fn wait_for_approval(mut self) -> Result<crate::hitl::HitlResponse, String> {
133        crate::context::link_to_ambient(&mut self.span);
134        let span_id = self.span.span_id.0.to_string();
135        self.span.status = SpanStatus::PendingApproval;
136
137        let rx = crate::hitl::register_approval(span_id);
138
139        // Record the span so the UI sees it is pending
140        crate::record_span(self.span.clone()).await;
141
142        // Wait for the approval/rejection from the UI/server
143        match rx.await {
144            Ok(response) => {
145                // we should end the span? actually this is a breakpoint.
146                // a breakpoint span is its own span. So we can just mark it done.
147                self.span.end_time = Some(
148                    std::time::SystemTime::now()
149                        .duration_since(std::time::UNIX_EPOCH)
150                        .unwrap_or_default()
151                        .as_millis() as u64,
152                );
153                self.span.latency_ms = Some(self.span.end_time.unwrap() - self.span.start_time);
154                self.span.status = SpanStatus::Ok;
155                crate::record_span(self.span).await;
156                Ok(response)
157            }
158            Err(_) => Err("Hitl approval channel closed unexpectedly".to_string()),
159        }
160    }
161
162    pub async fn run<F, Fut, T, E>(self, f: F) -> Result<T, E>
163    where
164        F: FnOnce() -> Fut,
165        Fut: std::future::Future<Output = Result<T, E>>,
166        E: std::fmt::Debug + std::fmt::Display + 'static,
167        T: serde::de::DeserializeOwned,
168    {
169        let mut span = self.span;
170        crate::context::link_to_ambient(&mut span);
171
172        // Mock / Replay interception
173        if let Some(mocked) = crate::replay::get_mocked_output(&span.name) {
174            span.end_time = Some(span.start_time);
175            span.latency_ms = Some(0);
176            span.status = SpanStatus::Ok;
177            span.attributes
178                .insert("replayed".to_string(), serde_json::json!(true));
179            crate::record_span(span.clone()).await;
180
181            if let Ok(value) = serde_json::from_value::<T>(mocked) {
182                return Ok(value);
183            }
184        }
185
186        // Install this span as the ambient parent for spans created inside `f`.
187        let ctx = crate::context::SpanContext::of(&span);
188        let result = crate::context::scope_current(ctx, f()).await;
189        span.end_time = Some(
190            std::time::SystemTime::now()
191                .duration_since(std::time::UNIX_EPOCH)
192                .unwrap_or_default()
193                .as_millis() as u64,
194        );
195        span.latency_ms = Some(span.end_time.unwrap() - span.start_time);
196
197        match &result {
198            Ok(_) => {
199                span.status = SpanStatus::Ok;
200            }
201            Err(e) => {
202                span.status = SpanStatus::Error;
203                span.error_type = Some(format!("{:?}", e)); // Naive type extraction
204                span.error_message_redacted = Some(e.to_string()); // Naive redaction
205            }
206        }
207
208        crate::record_span(span).await;
209
210        result
211    }
212
213    /// Like [`run`](Self::run) but for closures that don't return `Result`. The
214    /// span always completes with `Ok` status. Replay mocking (which is keyed on
215    /// deserializing a mocked value) applies only to `run`, not here.
216    pub async fn run_infallible<F, Fut, T>(self, f: F) -> T
217    where
218        F: FnOnce() -> Fut,
219        Fut: std::future::Future<Output = T>,
220    {
221        let mut span = self.span;
222        crate::context::link_to_ambient(&mut span);
223
224        let ctx = crate::context::SpanContext::of(&span);
225        let result = crate::context::scope_current(ctx, f()).await;
226
227        span.end_time = Some(
228            std::time::SystemTime::now()
229                .duration_since(std::time::UNIX_EPOCH)
230                .unwrap_or_default()
231                .as_millis() as u64,
232        );
233        span.latency_ms = Some(span.end_time.unwrap() - span.start_time);
234        span.status = SpanStatus::Ok;
235        crate::record_span(span).await;
236
237        result
238    }
239}
240
241pub fn llm_call(name: impl Into<String>) -> SpanBuilder {
242    SpanBuilder::new(TraceWeftSpanKind::LlmCall, name)
243}
244
245pub fn tool(name: impl Into<String>) -> SpanBuilder {
246    SpanBuilder::new(TraceWeftSpanKind::Tool, name)
247}
248
249pub fn agent(name: impl Into<String>) -> SpanBuilder {
250    SpanBuilder::new(TraceWeftSpanKind::Agent, name)
251}
trace_weft/builder.rs

trace_weft/
builder.rs
