Skip to main content

langsmith_rust/tracing/
decorator.rs

1use crate::error::Result;
2use crate::models::run::RunType;
3use crate::tracing::tracer::Tracer;
4use crate::utils::serialization::{ensure_inputs_object, ensure_outputs_object};
5use serde::Serialize;
6use std::future::Future;
7
8/// Helper function to trace a node execution
9/// 
10/// This function wraps a node execution with tracing:
11/// 1. Creates a tracer with serialized inputs
12/// 2. Posts the run to LangSmith (start_time, inputs)
13/// 3. Executes the function
14/// 4. Updates the run with outputs and end_time
15/// 5. Handles errors appropriately
16pub async fn trace_node<F, Fut, I, O>(
17    name: &str,
18    run_type: RunType,
19    inputs: I,
20    f: F,
21) -> Result<O>
22where
23    F: FnOnce(I) -> Fut,
24    Fut: Future<Output = Result<O>>,
25    I: Serialize,
26    O: Serialize,
27{
28    // Check if tracing is enabled
29    if !crate::config::Config::is_tracing_enabled() {
30        return f(inputs).await;
31    }
32
33    // 1. Serialize inputs - ensure it's always an object
34    let inputs_value = ensure_inputs_object(&inputs)
35        .map_err(|e| crate::error::LangSmithError::Serialization(e))?;
36
37    // 2. Create tracer
38    let mut tracer = Tracer::new(name, run_type, inputs_value);
39
40    // 3. POST /runs - save initial run (start_time, inputs)
41    if let Err(e) = tracer.post().await {
42        // Log error but don't fail the node execution
43        eprintln!("LangSmith tracing error (post): {}", e);
44    }
45
46    // 4. Execute the function
47    match f(inputs).await {
48        Ok(output) => {
49            // 5. Serialize outputs - ensure it's always an object
50            let output_value = ensure_outputs_object(&output)
51                .map_err(|e| crate::error::LangSmithError::Serialization(e))?;
52
53            // 6. Mark run as finished and PATCH /runs/{run_id} - save outputs and end_time
54            tracer.end(output_value);
55            if let Err(e) = tracer.patch().await {
56                // Log error but don't fail the node execution
57                eprintln!("LangSmith tracing error (patch): {}", e);
58            }
59
60            Ok(output)
61        }
62        Err(e) => {
63            // In case of error, mark run with error
64            tracer.set_error(&e.to_string());
65            if let Err(trace_err) = tracer.patch().await {
66                eprintln!("LangSmith tracing error (patch): {}", trace_err);
67            }
68            Err(e)
69        }
70    }
71}
72
73/// Synchronous version of trace_node
74pub fn trace_node_sync<F, I, O>(
75    name: &str,
76    run_type: RunType,
77    inputs: I,
78    f: F,
79) -> Result<O>
80where
81    F: FnOnce(I) -> Result<O>,
82    I: Serialize,
83    O: Serialize,
84{
85    // Check if tracing is enabled
86    if !crate::config::Config::is_tracing_enabled() {
87        return f(inputs);
88    }
89
90    // 1. Serialize inputs - ensure it's always an object
91    let inputs_value = ensure_inputs_object(&inputs)
92        .map_err(|e| crate::error::LangSmithError::Serialization(e))?;
93
94    // 2. Create tracer
95    let mut tracer = Tracer::new(name, run_type, inputs_value);
96
97    // 3. POST /runs - save initial run (start_time, inputs)
98    // For sync version, we need to use tokio runtime
99    let rt = tokio::runtime::Runtime::new().unwrap();
100    if let Err(e) = rt.block_on(tracer.post()) {
101        eprintln!("LangSmith tracing error (post): {}", e);
102    }
103
104    // 4. Execute the function
105    match f(inputs) {
106        Ok(output) => {
107            // 5. Serialize outputs - ensure it's always an object
108            let output_value = ensure_outputs_object(&output)
109                .map_err(|e| crate::error::LangSmithError::Serialization(e))?;
110
111            // 6. Mark run as finished and PATCH /runs/{run_id} - save outputs and end_time
112            tracer.end(output_value);
113            if let Err(e) = rt.block_on(tracer.patch()) {
114                eprintln!("LangSmith tracing error (patch): {}", e);
115            }
116
117            Ok(output)
118        }
119        Err(e) => {
120            // In case of error, mark run with error
121            tracer.set_error(&e.to_string());
122            if let Err(trace_err) = rt.block_on(tracer.patch()) {
123                eprintln!("LangSmith tracing error (patch): {}", trace_err);
124            }
125            Err(e)
126        }
127    }
128}
129