lc_callbacks/handlers/
langsmith_handler.rs

1// lc-callbacks/src/handlers/langsmith_handler.rs
2//! LangSmith callback handler
3
4use async_trait::async_trait;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8use crate::{CallbackHandler, LangSmithClient, LangSmithConfig, RunTree};
9use lc_schema::Message;
10
11/// LangSmith callback handler
12///
13/// Automatically sends trace data to LangSmith.
14pub struct LangSmithHandler {
15    client: Arc<LangSmithClient>,
16    active_runs: Arc<RwLock<Vec<RunTree>>>,
17    async_mode: bool,
18}
19
20impl LangSmithHandler {
21    pub fn new(config: LangSmithConfig) -> Self {
22        Self {
23            client: Arc::new(LangSmithClient::new(config)),
24            active_runs: Arc::new(RwLock::new(Vec::new())),
25            async_mode: false, // 默认同步模式，确保请求完成
26        }
27    }
28
29    pub fn from_env() -> Result<Self, String> {
30        let config = LangSmithConfig::from_env()?;
31        Ok(Self::new(config))
32    }
33
34    pub fn with_async_mode(mut self, async_mode: bool) -> Self {
35        self.async_mode = async_mode;
36        self
37    }
38
39    async fn push_run(&self, run: RunTree) {
40        self.active_runs.write().await.push(run);
41    }
42
43    async fn pop_run(&self) {
44        self.active_runs.write().await.pop();
45    }
46}
47
48#[async_trait]
49impl CallbackHandler for LangSmithHandler {
50    async fn on_run_start(&self, run: &RunTree) {
51        if !self.client.is_tracing_enabled() {
52            return;
53        }
54
55        if self.async_mode {
56            let run = run.clone();
57            let client = Arc::clone(&self.client);
58            tokio::spawn(async move {
59                if let Err(e) = client.create_run(&run).await {
60                    eprintln!("[LangSmith] create_run 失败: {}", e);
61                }
62            });
63        } else if let Err(e) = self.client.create_run(run).await {
64            eprintln!("[LangSmith] create_run 失败: {}", e);
65        }
66
67        self.push_run(run.clone()).await;
68    }
69
70    async fn on_run_end(&self, run: &RunTree) {
71        if !self.client.is_tracing_enabled() {
72            return;
73        }
74
75        if self.async_mode {
76            let run = run.clone();
77            let client = Arc::clone(&self.client);
78            tokio::spawn(async move {
79                if let Err(e) = client.update_run(&run).await {
80                    eprintln!("[LangSmith] update_run 失败: {}", e);
81                }
82            });
83        } else if let Err(e) = self.client.update_run(run).await {
84            eprintln!("[LangSmith] update_run 失败: {}", e);
85        }
86
87        self.pop_run().await;
88    }
89
90    async fn on_run_error(&self, run: &RunTree, error: &str) {
91        let mut run = run.clone();
92        run.end_with_error(error);
93        self.on_run_end(&run).await;
94    }
95
96    async fn on_llm_start(&self, run: &RunTree, _messages: &[Message]) {
97        self.on_run_start(run).await;
98    }
99
100    async fn on_llm_end(&self, run: &RunTree, _response: &str) {
101        self.on_run_end(run).await;
102    }
103
104    async fn on_llm_error(&self, run: &RunTree, error: &str) {
105        self.on_run_error(run, error).await;
106    }
107
108    async fn on_chain_start(&self, run: &RunTree, _inputs: &serde_json::Value) {
109        self.on_run_start(run).await;
110    }
111
112    async fn on_chain_end(&self, run: &RunTree, _outputs: &serde_json::Value) {
113        self.on_run_end(run).await;
114    }
115
116    async fn on_chain_error(&self, run: &RunTree, error: &str) {
117        self.on_run_error(run, error).await;
118    }
119
120    async fn on_tool_start(&self, run: &RunTree, _tool_name: &str, _input: &str) {
121        self.on_run_start(run).await;
122    }
123
124    async fn on_tool_end(&self, run: &RunTree, _output: &str) {
125        self.on_run_end(run).await;
126    }
127
128    async fn on_tool_error(&self, run: &RunTree, error: &str) {
129        self.on_run_error(run, error).await;
130    }
131
132    async fn on_retriever_start(&self, run: &RunTree, _query: &str) {
133        self.on_run_start(run).await;
134    }
135
136    async fn on_retriever_end(&self, run: &RunTree, _documents: &[serde_json::Value]) {
137        self.on_run_end(run).await;
138    }
139
140    async fn on_retriever_error(&self, run: &RunTree, error: &str) {
141        self.on_run_error(run, error).await;
142    }
143}
lc_callbacks/handlers/langsmith_handler.rs

lc_callbacks/handlers/
langsmith_handler.rs
