lc_callbacks/handlers/
generic_handler.rs

1// lc-callbacks/src/handlers/generic_handler.rs
2//! Closure-based callback handler
3//!
4//! Lets you wire up a [`CallbackHandler`] from plain closures without defining
5//! a new struct. Every hook is optional: hooks you don't set fall back to the
6//! generic lifecycle layer (`on_run_start` / `on_run_end` / `on_run_error`),
7//! exactly like the [`CallbackHandler`] trait defaults — so you can observe
8//! only what you care about.
9//!
10//! # Example
11//!
12//! ```ignore
13//! use lc_callbacks::{CallbackManager, GenericHandler};
14//! use std::sync::Arc;
15//!
16//! let manager = CallbackManager::new().add_handler(Arc::new(
17//!     GenericHandler::new()
18//!         .with_llm_start(|run, _msgs| println!("llm start: {}", run.name))
19//!         .with_llm_new_token(|_run, token| print!("{}", token)),
20//! ));
21//! ```
22
23use async_trait::async_trait;
24use std::sync::Arc;
25
26use crate::{CallbackHandler, RunTree};
27use lc_schema::Message;
28
29type RunStartFn = Arc<dyn Fn(&RunTree) + Send + Sync>;
30type RunEndFn = Arc<dyn Fn(&RunTree) + Send + Sync>;
31type RunErrorFn = Arc<dyn Fn(&RunTree, &str) + Send + Sync>;
32type LlmStartFn = Arc<dyn Fn(&RunTree, &[Message]) + Send + Sync>;
33type LlmEndFn = Arc<dyn Fn(&RunTree, &str) + Send + Sync>;
34type TokenFn = Arc<dyn Fn(&RunTree, &str) + Send + Sync>;
35type ToolStartFn = Arc<dyn Fn(&RunTree, &str, &str) + Send + Sync>;
36type RetrieverStartFn = Arc<dyn Fn(&RunTree, &str) + Send + Sync>;
37
38/// Closure-based [`CallbackHandler`] for ad-hoc observability hooks.
39///
40/// Unset hooks delegate to the generic lifecycle layer, matching the
41/// [`CallbackHandler`] trait defaults (e.g. an unset `on_llm_start` forwards
42/// to `on_run_start`).
43#[derive(Default)]
44pub struct GenericHandler {
45    on_run_start: Option<RunStartFn>,
46    on_run_end: Option<RunEndFn>,
47    on_run_error: Option<RunErrorFn>,
48    on_llm_start: Option<LlmStartFn>,
49    on_llm_end: Option<LlmEndFn>,
50    on_llm_new_token: Option<TokenFn>,
51    on_llm_thinking: Option<TokenFn>,
52    on_tool_start: Option<ToolStartFn>,
53    on_retriever_start: Option<RetrieverStartFn>,
54}
55
56impl GenericHandler {
57    /// Create a handler with no hooks set.
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    /// Set the generic run-start hook.
63    pub fn with_run_start(mut self, f: impl Fn(&RunTree) + Send + Sync + 'static) -> Self {
64        self.on_run_start = Some(Arc::new(f));
65        self
66    }
67
68    /// Set the generic run-end hook.
69    pub fn with_run_end(mut self, f: impl Fn(&RunTree) + Send + Sync + 'static) -> Self {
70        self.on_run_end = Some(Arc::new(f));
71        self
72    }
73
74    /// Set the generic run-error hook.
75    pub fn with_run_error(mut self, f: impl Fn(&RunTree, &str) + Send + Sync + 'static) -> Self {
76        self.on_run_error = Some(Arc::new(f));
77        self
78    }
79
80    /// Set the LLM start hook (falls back to the run-start hook when unset).
81    pub fn with_llm_start(
82        mut self,
83        f: impl Fn(&RunTree, &[Message]) + Send + Sync + 'static,
84    ) -> Self {
85        self.on_llm_start = Some(Arc::new(f));
86        self
87    }
88
89    /// Set the LLM end hook (falls back to the run-end hook when unset).
90    pub fn with_llm_end(mut self, f: impl Fn(&RunTree, &str) + Send + Sync + 'static) -> Self {
91        self.on_llm_end = Some(Arc::new(f));
92        self
93    }
94
95    /// Set the streaming-token hook (default: no-op).
96    pub fn with_llm_new_token(
97        mut self,
98        f: impl Fn(&RunTree, &str) + Send + Sync + 'static,
99    ) -> Self {
100        self.on_llm_new_token = Some(Arc::new(f));
101        self
102    }
103
104    /// Set the extended-thinking token hook (default: no-op).
105    pub fn with_llm_thinking(mut self, f: impl Fn(&RunTree, &str) + Send + Sync + 'static) -> Self {
106        self.on_llm_thinking = Some(Arc::new(f));
107        self
108    }
109
110    /// Set the tool start hook (falls back to the run-start hook when unset).
111    pub fn with_tool_start(
112        mut self,
113        f: impl Fn(&RunTree, &str, &str) + Send + Sync + 'static,
114    ) -> Self {
115        self.on_tool_start = Some(Arc::new(f));
116        self
117    }
118
119    /// Set the retriever start hook (falls back to the run-start hook when unset).
120    pub fn with_retriever_start(
121        mut self,
122        f: impl Fn(&RunTree, &str) + Send + Sync + 'static,
123    ) -> Self {
124        self.on_retriever_start = Some(Arc::new(f));
125        self
126    }
127}
128
129#[async_trait]
130impl CallbackHandler for GenericHandler {
131    async fn on_run_start(&self, run: &RunTree) {
132        if let Some(f) = &self.on_run_start {
133            f(run);
134        }
135    }
136
137    async fn on_run_end(&self, run: &RunTree) {
138        if let Some(f) = &self.on_run_end {
139            f(run);
140        }
141    }
142
143    async fn on_run_error(&self, run: &RunTree, error: &str) {
144        if let Some(f) = &self.on_run_error {
145            f(run, error);
146        }
147    }
148
149    async fn on_llm_start(&self, run: &RunTree, messages: &[Message]) {
150        if let Some(f) = &self.on_llm_start {
151            f(run, messages);
152        } else {
153            self.on_run_start(run).await;
154        }
155    }
156
157    async fn on_llm_end(&self, run: &RunTree, response: &str) {
158        if let Some(f) = &self.on_llm_end {
159            f(run, response);
160        } else {
161            self.on_run_end(run).await;
162        }
163    }
164
165    async fn on_llm_new_token(&self, run: &RunTree, token: &str) {
166        if let Some(f) = &self.on_llm_new_token {
167            f(run, token);
168        }
169    }
170
171    async fn on_llm_thinking(&self, run: &RunTree, thinking: &str) {
172        if let Some(f) = &self.on_llm_thinking {
173            f(run, thinking);
174        }
175    }
176
177    async fn on_tool_start(&self, run: &RunTree, tool_name: &str, input: &str) {
178        if let Some(f) = &self.on_tool_start {
179            f(run, tool_name, input);
180        } else {
181            self.on_run_start(run).await;
182        }
183    }
184
185    async fn on_retriever_start(&self, run: &RunTree, query: &str) {
186        if let Some(f) = &self.on_retriever_start {
187            f(run, query);
188        } else {
189            self.on_run_start(run).await;
190        }
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use std::sync::Mutex;
198
199    fn run(name: &str) -> RunTree {
200        RunTree::new(name, crate::RunType::Chain, serde_json::json!({}))
201    }
202
203    #[tokio::test]
204    async fn test_run_hooks_fire() {
205        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
206        let ev_start = Arc::clone(&events);
207        let ev_end = Arc::clone(&events);
208
209        let handler = GenericHandler::new()
210            .with_run_start(move |r| {
211                ev_start
212                    .lock()
213                    .unwrap_or_else(|e| e.into_inner())
214                    .push(format!("start:{}", r.name))
215            })
216            .with_run_end(move |r| {
217                ev_end
218                    .lock()
219                    .unwrap_or_else(|e| e.into_inner())
220                    .push(format!("end:{}", r.name))
221            });
222
223        let h = run("r1");
224        handler.on_run_start(&h).await;
225        handler.on_run_end(&h).await;
226
227        let got = events.lock().unwrap_or_else(|e| e.into_inner()).clone();
228        assert_eq!(got, vec!["start:r1".to_string(), "end:r1".to_string()]);
229    }
230
231    #[tokio::test]
232    async fn test_run_error_hook() {
233        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
234        let events2 = Arc::clone(&events);
235
236        let handler = GenericHandler::new().with_run_error(move |_r, e| {
237            events2
238                .lock()
239                .unwrap_or_else(|e| e.into_inner())
240                .push(format!("error:{e}"))
241        });
242
243        handler.on_run_error(&run("r1"), "boom").await;
244
245        let got = events.lock().unwrap_or_else(|e| e.into_inner()).clone();
246        assert_eq!(got, vec!["error:boom".to_string()]);
247    }
248
249    #[tokio::test]
250    async fn test_llm_hooks_delegate_to_run_when_unset() {
251        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
252        let events2 = Arc::clone(&events);
253
254        // Only the generic layer is set: typed hooks must forward to it.
255        let handler = GenericHandler::new().with_run_start(move |r| {
256            events2
257                .lock()
258                .unwrap_or_else(|e| e.into_inner())
259                .push(format!("start:{}", r.name))
260        });
261
262        let h = run("llm1");
263        handler.on_llm_start(&h, &[]).await;
264
265        let got = events.lock().unwrap_or_else(|e| e.into_inner()).clone();
266        assert_eq!(got, vec!["start:llm1".to_string()]);
267    }
268
269    #[tokio::test]
270    async fn test_llm_token_and_thinking_hooks() {
271        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
272        let ev_tokens = Arc::clone(&events);
273        let ev_thinking = Arc::clone(&events);
274
275        let handler = GenericHandler::new()
276            .with_llm_new_token(move |_r, t| {
277                ev_tokens
278                    .lock()
279                    .unwrap_or_else(|e| e.into_inner())
280                    .push(format!("tok:{t}"))
281            })
282            .with_llm_thinking(move |_r, t| {
283                ev_thinking
284                    .lock()
285                    .unwrap_or_else(|e| e.into_inner())
286                    .push(format!("think:{t}"))
287            });
288
289        let h = run("llm1");
290        handler.on_llm_new_token(&h, "hello").await;
291        handler.on_llm_thinking(&h, "hmm").await;
292
293        let got = events.lock().unwrap_or_else(|e| e.into_inner()).clone();
294        assert_eq!(got, vec!["tok:hello".to_string(), "think:hmm".to_string()]);
295    }
296
297    #[tokio::test]
298    async fn test_tool_and_retriever_hooks() {
299        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
300        let ev_tool = Arc::clone(&events);
301        let ev_retriever = Arc::clone(&events);
302
303        let handler = GenericHandler::new()
304            .with_tool_start(move |_r, name, input| {
305                ev_tool
306                    .lock()
307                    .unwrap_or_else(|e| e.into_inner())
308                    .push(format!("tool:{name}:{input}"))
309            })
310            .with_retriever_start(move |_r, q| {
311                ev_retriever
312                    .lock()
313                    .unwrap_or_else(|e| e.into_inner())
314                    .push(format!("retriever:{q}"))
315            });
316
317        let h = run("t1");
318        handler.on_tool_start(&h, "search", "rust").await;
319        handler.on_retriever_start(&h, "doc").await;
320
321        let got = events.lock().unwrap_or_else(|e| e.into_inner()).clone();
322        assert_eq!(
323            got,
324            vec!["tool:search:rust".to_string(), "retriever:doc".to_string()]
325        );
326    }
327}
lc_callbacks/handlers/generic_handler.rs

lc_callbacks/handlers/
generic_handler.rs
