Skip to main content

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