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| ev_start.lock().unwrap().push(format!("start:{}", r.name)))
211            .with_run_end(move |r| ev_end.lock().unwrap().push(format!("end:{}", r.name)));
212
213        let h = run("r1");
214        handler.on_run_start(&h).await;
215        handler.on_run_end(&h).await;
216
217        let got = events.lock().unwrap().clone();
218        assert_eq!(got, vec!["start:r1".to_string(), "end:r1".to_string()]);
219    }
220
221    #[tokio::test]
222    async fn test_run_error_hook() {
223        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
224        let events2 = Arc::clone(&events);
225
226        let handler = GenericHandler::new()
227            .with_run_error(move |_r, e| events2.lock().unwrap().push(format!("error:{e}")));
228
229        handler.on_run_error(&run("r1"), "boom").await;
230
231        let got = events.lock().unwrap().clone();
232        assert_eq!(got, vec!["error:boom".to_string()]);
233    }
234
235    #[tokio::test]
236    async fn test_llm_hooks_delegate_to_run_when_unset() {
237        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
238        let events2 = Arc::clone(&events);
239
240        // Only the generic layer is set: typed hooks must forward to it.
241        let handler = GenericHandler::new()
242            .with_run_start(move |r| events2.lock().unwrap().push(format!("start:{}", r.name)));
243
244        let h = run("llm1");
245        handler.on_llm_start(&h, &[]).await;
246
247        let got = events.lock().unwrap().clone();
248        assert_eq!(got, vec!["start:llm1".to_string()]);
249    }
250
251    #[tokio::test]
252    async fn test_llm_token_and_thinking_hooks() {
253        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
254        let ev_tokens = Arc::clone(&events);
255        let ev_thinking = Arc::clone(&events);
256
257        let handler = GenericHandler::new()
258            .with_llm_new_token(move |_r, t| ev_tokens.lock().unwrap().push(format!("tok:{t}")))
259            .with_llm_thinking(move |_r, t| ev_thinking.lock().unwrap().push(format!("think:{t}")));
260
261        let h = run("llm1");
262        handler.on_llm_new_token(&h, "hello").await;
263        handler.on_llm_thinking(&h, "hmm").await;
264
265        let got = events.lock().unwrap().clone();
266        assert_eq!(got, vec!["tok:hello".to_string(), "think:hmm".to_string()]);
267    }
268
269    #[tokio::test]
270    async fn test_tool_and_retriever_hooks() {
271        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
272        let ev_tool = Arc::clone(&events);
273        let ev_retriever = Arc::clone(&events);
274
275        let handler = GenericHandler::new()
276            .with_tool_start(move |_r, name, input| {
277                ev_tool.lock().unwrap().push(format!("tool:{name}:{input}"))
278            })
279            .with_retriever_start(move |_r, q| {
280                ev_retriever.lock().unwrap().push(format!("retriever:{q}"))
281            });
282
283        let h = run("t1");
284        handler.on_tool_start(&h, "search", "rust").await;
285        handler.on_retriever_start(&h, "doc").await;
286
287        let got = events.lock().unwrap().clone();
288        assert_eq!(
289            got,
290            vec!["tool:search:rust".to_string(), "retriever:doc".to_string()]
291        );
292    }
293}
lc_callbacks/handlers/generic_handler.rs

lc_callbacks/handlers/
generic_handler.rs
