lc-callbacks 0.20.2

// lc-callbacks/src/run_tree.rs
//! Run tree data structure for tracing

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;

use super::RunType;

/// Run tree node for tracing
///
/// Each run is a node in a tree, with optional parent and children.
/// The entire trace forms a tree structure, with the root being the top-level call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RunTree {
    /// Unique run ID (UUID v7 with timestamp)
    pub id: Uuid,

    /// Run name
    pub name: String,

    /// Run type
    pub run_type: RunType,

    /// Input data
    pub inputs: serde_json::Value,

    /// Output data (set when run ends)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub outputs: Option<serde_json::Value>,

    /// Error message (if run failed)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,

    /// Parent run ID
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parent_run_id: Option<Uuid>,

    /// Trace ID (ID of the root run)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub trace_id: Option<Uuid>,

    /// Start time
    pub start_time: DateTime<Utc>,

    /// End time
    #[serde(skip_serializing_if = "Option::is_none")]
    pub end_time: Option<DateTime<Utc>>,

    /// Metadata
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,

    /// Tags
    #[serde(default)]
    pub tags: Vec<String>,

    /// Project name (LangSmith project)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub project_name: Option<String>,

    /// Serialized representation of the component
    #[serde(default)]
    pub serialized: serde_json::Value,
}

impl RunTree {
    /// Create a new run
    pub fn new(name: impl Into<String>, run_type: RunType, inputs: serde_json::Value) -> Self {
        Self {
            id: Uuid::now_v7(),
            name: name.into(),
            run_type,
            inputs,
            outputs: None,
            error: None,
            parent_run_id: None,
            trace_id: None,
            start_time: Utc::now(),
            end_time: None,
            metadata: HashMap::new(),
            tags: Vec::new(),
            project_name: None,
            serialized: serde_json::Value::Null,
        }
    }

    /// Create a child run from this run
    pub fn create_child(
        &self,
        name: impl Into<String>,
        run_type: RunType,
        inputs: serde_json::Value,
    ) -> Self {
        let mut child = Self::new(name, run_type, inputs);
        child.parent_run_id = Some(self.id);
        child.trace_id = self.trace_id.or(Some(self.id));
        child.project_name = self.project_name.clone();
        child
    }

    /// End the run with outputs
    pub fn end(&mut self, outputs: serde_json::Value) {
        self.outputs = Some(outputs);
        self.end_time = Some(Utc::now());
    }

    /// End the run with an error
    pub fn end_with_error(&mut self, error: impl Into<String>) {
        self.error = Some(error.into());
        self.end_time = Some(Utc::now());
    }

    /// Add a tag
    pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
        self.tags.push(tag.into());
        self
    }

    /// Add metadata
    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
        self.metadata.insert(key.into(), value);
        self
    }

    /// Set project name
    pub fn with_project(mut self, project: impl Into<String>) -> Self {
        self.project_name = Some(project.into());
        self
    }

    /// Calculate run duration in milliseconds
    pub fn duration_ms(&self) -> Option<i64> {
        self.end_time
            .map(|end| (end - self.start_time).num_milliseconds())
    }
}

/// LangSmith run 的 extra 载荷:自定义 metadata 放 `extra.metadata`。
/// LangSmith v2 忽略 run 顶层 `metadata`,只认 `extra.metadata`。
#[derive(Debug, Default, Serialize)]
pub struct RunExtra {
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl RunExtra {
    pub fn is_empty(&self) -> bool {
        self.metadata.is_empty()
    }
}

/// Simplified run structure for API requests
#[derive(Debug, Serialize)]
pub struct RunCreate {
    /// Run ID.
    pub id: String,
    /// Run name.
    pub name: String,
    /// Run type as a string.
    pub run_type: String,
    /// Inputs for the run.
    pub inputs: serde_json::Value,
    /// Outputs of the run.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub outputs: Option<serde_json::Value>,
    /// Error message, if the run failed.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// ID of the parent run, if any.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parent_run_id: Option<String>,
    /// Start time as an RFC 3339 timestamp.
    pub start_time: String,
    /// End time as an RFC 3339 timestamp, if the run finished.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub end_time: Option<String>,
    /// Session/project name.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session_name: Option<String>,
    /// Tags attached to the run.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tags: Vec<String>,
    /// 额外载荷:自定义 metadata 放 `extra.metadata`(顶层 `metadata` 会被 LangSmith v2 忽略)。
    #[serde(default, skip_serializing_if = "RunExtra::is_empty")]
    pub extra: RunExtra,
}

impl From<&RunTree> for RunCreate {
    fn from(run: &RunTree) -> Self {
        Self {
            id: run.id.to_string(),
            name: run.name.clone(),
            run_type: run.run_type.as_str().to_string(),
            inputs: run.inputs.clone(),
            outputs: run.outputs.clone(),
            error: run.error.clone(),
            parent_run_id: run.parent_run_id.map(|id| id.to_string()),
            start_time: run.start_time.to_rfc3339(),
            end_time: run.end_time.map(|t| t.to_rfc3339()),
            session_name: run.project_name.clone(),
            tags: run.tags.clone(),
            extra: RunExtra { metadata: run.metadata.clone() },
        }
    }
}

/// Run update structure for PATCH requests
#[derive(Debug, Serialize)]
pub struct RunUpdate {
    /// Updated outputs of the run.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub outputs: Option<serde_json::Value>,
    /// Updated error message.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// Updated end time as an RFC 3339 timestamp.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub end_time: Option<String>,
}

impl From<&RunTree> for RunUpdate {
    fn from(run: &RunTree) -> Self {
        Self {
            outputs: run.outputs.clone(),
            error: run.error.clone(),
            end_time: run.end_time.map(|t| t.to_rfc3339()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn run_create_puts_metadata_into_extra() {
        let run = RunTree::new("t", RunType::Chain, serde_json::json!({}))
            .with_metadata("obs", serde_json::json!({"total_tokens": 451}));
        let body = serde_json::to_value(RunCreate::from(&run)).unwrap();
        // LangSmith v2 忽略顶层 metadata,只认 extra.metadata
        assert!(body.get("metadata").is_none(), "顶层 metadata 不应再出现");
        assert_eq!(body["extra"]["metadata"]["obs"]["total_tokens"], 451);
    }
}