Skip to main content

langsmith_rust/tracing/
graph.rs

1use crate::error::Result;
2use crate::models::run::RunType;
3use crate::tracing::scope::RunScope;
4use crate::tracing::tracer::Tracer;
5use serde_json::Value;
6
7/// Opinionated tracing helpers to build a Graph-style hierarchy in LangSmith:
8/// - Root run named `Graph` (RunType::Chain)
9/// - Step runs under root (e.g. `chatbot`, `should_continue`, `tools`)
10/// - Nested runs under steps (e.g. `ChatOpenAI`, tool runs)
11///
12/// This is intentionally generic: inputs/outputs are `serde_json::Value` so any app can
13/// provide LangSmith-compatible payloads such as `{ \"messages\": [...] }`.
14pub struct GraphTrace {
15    root: RunScope,
16}
17
18impl GraphTrace {
19    /// Starts the root Graph run (name: `Graph`, type: Chain) and POSTs it.
20    pub async fn start_root(inputs: Value, thread_id: Option<String>) -> Result<Self> {
21        let mut root = RunScope::root_value("Graph", RunType::Chain, inputs);
22        if let Some(tid) = thread_id {
23            root = root.with_thread_id(tid);
24        }
25        root.post_start().await?;
26        Ok(Self { root })
27    }
28
29    pub fn root_scope(&self) -> &RunScope {
30        &self.root
31    }
32
33    pub fn root_tracer(&self) -> &Tracer {
34        self.root.tracer()
35    }
36
37    /// Starts a top-level step/node under the root run (POSTs it) and returns the scope.
38    /// Use this for nodes like "chatbot", "tools", etc. that may contain nested runs.
39    pub async fn start_node_iteration(&self, node_name: &str, inputs: Value) -> Result<RunScope> {
40        let mut step = self.root.child_value(node_name, RunType::Chain, inputs);
41        step.post_start().await?;
42        Ok(step)
43    }
44
45    /// Traces an LLM call within a parent node (e.g., within "chatbot").
46    /// `llm_name` is typically "ChatOpenAI" or similar.
47    /// `model_name` is optional (e.g., "gpt-4o-mini") and will be added to inputs if provided.
48    pub async fn trace_llm_call(
49        &self,
50        parent_node: &RunScope,
51        llm_name: &str,
52        inputs: Value,
53        outputs: Value,
54        model_name: Option<&str>,
55    ) -> Result<()> {
56        let mut llm_inputs = inputs;
57        if let Some(model) = model_name {
58            if let Some(obj) = llm_inputs.as_object_mut() {
59                obj.insert("model".to_string(), serde_json::json!(model));
60            }
61        }
62        let mut llm = parent_node.child_value(llm_name, RunType::Llm, llm_inputs);
63        llm.post_start().await?;
64        llm.end_ok(outputs).await
65    }
66
67    /// Traces a routing/decision step (e.g., "should_continue").
68    pub async fn trace_decision(
69        &self,
70        parent_node: &RunScope,
71        decision_name: &str,
72        inputs: Value,
73        outputs: Value,
74    ) -> Result<()> {
75        let mut decision = parent_node.child_value(decision_name, RunType::Chain, inputs);
76        decision.post_start().await?;
77        decision.end_ok(outputs).await
78    }
79
80    /// Traces a tool call within a parent node (e.g., within "tools").
81    /// `tool_name` should be the actual tool name (e.g., "calculator").
82    pub async fn trace_tool_call(
83        &self,
84        parent_node: &RunScope,
85        tool_name: &str,
86        inputs: Value,
87        outputs: Value,
88    ) -> Result<()> {
89        // Format tool name as "tool/{name}" to match LangGraph Python convention
90        let formatted_name = format!("tool/{}", tool_name);
91        let mut tool = parent_node.child_value(&formatted_name, RunType::Tool, inputs);
92        tool.post_start().await?;
93        tool.end_ok(outputs).await
94    }
95
96    /// Ends the root run with the provided outputs (PATCH). Consumes self.
97    pub async fn end_root(self, outputs: Value) -> Result<()> {
98        self.root.end_ok(outputs).await
99    }
100}
101
102