1use std::collections::HashMap;
7
8use async_trait::async_trait;
9use serde_json::Value;
10use uuid::Uuid;
11
12use crate::messages::BaseMessage;
13use crate::outputs::LLMResult;
14use crate::tracers::core::{TracerCore, TracerError};
15use crate::tracers::schemas::Run;
16
17pub trait BaseTracer: TracerCore {
22 fn persist_run_impl(&mut self, run: &Run);
24
25 fn start_trace_impl(&mut self, run: &mut Run) {
27 self.start_trace(run);
28 self.on_run_create(run);
29 }
30
31 fn end_trace_impl(&mut self, run: &Run) {
33 if run.parent_run_id.is_none() {
34 self.persist_run_impl(run);
35 }
36 self.end_trace(run);
37 self.on_run_update(run);
38 }
39
40 #[allow(clippy::too_many_arguments)]
42 fn handle_chat_model_start(
43 &mut self,
44 serialized: HashMap<String, Value>,
45 messages: &[Vec<BaseMessage>],
46 run_id: Uuid,
47 parent_run_id: Option<Uuid>,
48 tags: Option<Vec<String>>,
49 metadata: Option<HashMap<String, Value>>,
50 name: Option<String>,
51 extra: HashMap<String, Value>,
52 ) -> Result<Run, TracerError> {
53 let mut chat_model_run = self.create_chat_model_run(
54 serialized,
55 messages,
56 run_id,
57 parent_run_id,
58 tags,
59 metadata,
60 name,
61 extra,
62 )?;
63 self.start_trace_impl(&mut chat_model_run);
64 self.on_chat_model_start(&chat_model_run);
65 Ok(chat_model_run)
66 }
67
68 #[allow(clippy::too_many_arguments)]
70 fn handle_llm_start(
71 &mut self,
72 serialized: HashMap<String, Value>,
73 prompts: &[String],
74 run_id: Uuid,
75 parent_run_id: Option<Uuid>,
76 tags: Option<Vec<String>>,
77 metadata: Option<HashMap<String, Value>>,
78 name: Option<String>,
79 extra: HashMap<String, Value>,
80 ) -> Run {
81 let mut llm_run = self.create_llm_run(
82 serialized,
83 prompts,
84 run_id,
85 parent_run_id,
86 tags,
87 metadata,
88 name,
89 extra,
90 );
91 self.start_trace_impl(&mut llm_run);
92 self.on_llm_start(&llm_run);
93 llm_run
94 }
95
96 fn handle_llm_new_token(
98 &mut self,
99 token: &str,
100 run_id: Uuid,
101 chunk: Option<&dyn std::any::Any>,
102 parent_run_id: Option<Uuid>,
103 ) -> Result<Run, TracerError> {
104 let llm_run = self.llm_run_with_token_event(token, run_id, chunk, parent_run_id)?;
105 self.on_llm_new_token(&llm_run, token, chunk);
106 Ok(llm_run)
107 }
108
109 fn handle_retry(
111 &mut self,
112 retry_state: &HashMap<String, Value>,
113 run_id: Uuid,
114 ) -> Result<Run, TracerError> {
115 self.llm_run_with_retry_event(retry_state, run_id)
116 }
117
118 fn handle_llm_end(&mut self, response: &LLMResult, run_id: Uuid) -> Result<Run, TracerError> {
120 let llm_run = self.complete_llm_run(response, run_id)?;
121 self.end_trace_impl(&llm_run);
122 self.on_llm_end(&llm_run);
123 Ok(llm_run)
124 }
125
126 fn handle_llm_error(
128 &mut self,
129 error: &dyn std::error::Error,
130 run_id: Uuid,
131 response: Option<&LLMResult>,
132 ) -> Result<Run, TracerError> {
133 let llm_run = self.errored_llm_run(error, run_id, response)?;
134 self.end_trace_impl(&llm_run);
135 self.on_llm_error(&llm_run);
136 Ok(llm_run)
137 }
138
139 #[allow(clippy::too_many_arguments)]
141 fn handle_chain_start(
142 &mut self,
143 serialized: HashMap<String, Value>,
144 inputs: HashMap<String, Value>,
145 run_id: Uuid,
146 parent_run_id: Option<Uuid>,
147 tags: Option<Vec<String>>,
148 metadata: Option<HashMap<String, Value>>,
149 run_type: Option<String>,
150 name: Option<String>,
151 extra: HashMap<String, Value>,
152 ) -> Run {
153 let mut chain_run = self.create_chain_run(
154 serialized,
155 inputs,
156 run_id,
157 parent_run_id,
158 tags,
159 metadata,
160 run_type,
161 name,
162 extra,
163 );
164 self.start_trace_impl(&mut chain_run);
165 self.on_chain_start(&chain_run);
166 chain_run
167 }
168
169 fn handle_chain_end(
171 &mut self,
172 outputs: HashMap<String, Value>,
173 run_id: Uuid,
174 inputs: Option<HashMap<String, Value>>,
175 ) -> Result<Run, TracerError> {
176 let chain_run = self.complete_chain_run(outputs, run_id, inputs)?;
177 self.end_trace_impl(&chain_run);
178 self.on_chain_end(&chain_run);
179 Ok(chain_run)
180 }
181
182 fn handle_chain_error(
184 &mut self,
185 error: &dyn std::error::Error,
186 run_id: Uuid,
187 inputs: Option<HashMap<String, Value>>,
188 ) -> Result<Run, TracerError> {
189 let chain_run = self.errored_chain_run(error, run_id, inputs)?;
190 self.end_trace_impl(&chain_run);
191 self.on_chain_error(&chain_run);
192 Ok(chain_run)
193 }
194
195 #[allow(clippy::too_many_arguments)]
197 fn handle_tool_start(
198 &mut self,
199 serialized: HashMap<String, Value>,
200 input_str: &str,
201 run_id: Uuid,
202 parent_run_id: Option<Uuid>,
203 tags: Option<Vec<String>>,
204 metadata: Option<HashMap<String, Value>>,
205 name: Option<String>,
206 inputs: Option<HashMap<String, Value>>,
207 extra: HashMap<String, Value>,
208 ) -> Run {
209 let mut tool_run = self.create_tool_run(
210 serialized,
211 input_str,
212 run_id,
213 parent_run_id,
214 tags,
215 metadata,
216 name,
217 inputs,
218 extra,
219 );
220 self.start_trace_impl(&mut tool_run);
221 self.on_tool_start(&tool_run);
222 tool_run
223 }
224
225 fn handle_tool_end(&mut self, output: Value, run_id: Uuid) -> Result<Run, TracerError> {
227 let tool_run = self.complete_tool_run(output, run_id)?;
228 self.end_trace_impl(&tool_run);
229 self.on_tool_end(&tool_run);
230 Ok(tool_run)
231 }
232
233 fn handle_tool_error(
235 &mut self,
236 error: &dyn std::error::Error,
237 run_id: Uuid,
238 ) -> Result<Run, TracerError> {
239 let tool_run = self.errored_tool_run(error, run_id)?;
240 self.end_trace_impl(&tool_run);
241 self.on_tool_error(&tool_run);
242 Ok(tool_run)
243 }
244
245 #[allow(clippy::too_many_arguments)]
247 fn handle_retriever_start(
248 &mut self,
249 serialized: HashMap<String, Value>,
250 query: &str,
251 run_id: Uuid,
252 parent_run_id: Option<Uuid>,
253 tags: Option<Vec<String>>,
254 metadata: Option<HashMap<String, Value>>,
255 name: Option<String>,
256 extra: HashMap<String, Value>,
257 ) -> Run {
258 let mut retrieval_run = self.create_retrieval_run(
259 serialized,
260 query,
261 run_id,
262 parent_run_id,
263 tags,
264 metadata,
265 name,
266 extra,
267 );
268 self.start_trace_impl(&mut retrieval_run);
269 self.on_retriever_start(&retrieval_run);
270 retrieval_run
271 }
272
273 fn handle_retriever_end(
275 &mut self,
276 documents: Vec<Value>,
277 run_id: Uuid,
278 ) -> Result<Run, TracerError> {
279 let retrieval_run = self.complete_retrieval_run(documents, run_id)?;
280 self.end_trace_impl(&retrieval_run);
281 self.on_retriever_end(&retrieval_run);
282 Ok(retrieval_run)
283 }
284
285 fn handle_retriever_error(
287 &mut self,
288 error: &dyn std::error::Error,
289 run_id: Uuid,
290 ) -> Result<Run, TracerError> {
291 let retrieval_run = self.errored_retrieval_run(error, run_id)?;
292 self.end_trace_impl(&retrieval_run);
293 self.on_retriever_error(&retrieval_run);
294 Ok(retrieval_run)
295 }
296}
297
298#[async_trait]
302pub trait AsyncBaseTracer: TracerCore + Send + Sync {
303 async fn persist_run_async(&mut self, run: &Run);
305
306 async fn start_trace_async(&mut self, run: &mut Run) {
308 self.start_trace(run);
309 self.on_run_create_async(run).await;
310 }
311
312 async fn end_trace_async(&mut self, run: &Run) {
314 if run.parent_run_id.is_none() {
315 self.persist_run_async(run).await;
316 }
317 self.end_trace(run);
318 self.on_run_update_async(run).await;
319 }
320
321 #[allow(clippy::too_many_arguments)]
323 async fn handle_chat_model_start_async(
324 &mut self,
325 serialized: HashMap<String, Value>,
326 messages: &[Vec<BaseMessage>],
327 run_id: Uuid,
328 parent_run_id: Option<Uuid>,
329 tags: Option<Vec<String>>,
330 metadata: Option<HashMap<String, Value>>,
331 name: Option<String>,
332 extra: HashMap<String, Value>,
333 ) -> Result<Run, TracerError> {
334 let mut chat_model_run = self.create_chat_model_run(
335 serialized,
336 messages,
337 run_id,
338 parent_run_id,
339 tags,
340 metadata,
341 name,
342 extra,
343 )?;
344
345 self.start_trace_async(&mut chat_model_run).await;
347 self.on_chat_model_start_async(&chat_model_run).await;
348
349 Ok(chat_model_run)
350 }
351
352 #[allow(clippy::too_many_arguments)]
354 async fn handle_llm_start_async(
355 &mut self,
356 serialized: HashMap<String, Value>,
357 prompts: &[String],
358 run_id: Uuid,
359 parent_run_id: Option<Uuid>,
360 tags: Option<Vec<String>>,
361 metadata: Option<HashMap<String, Value>>,
362 name: Option<String>,
363 extra: HashMap<String, Value>,
364 ) -> Run {
365 let mut llm_run = self.create_llm_run(
366 serialized,
367 prompts,
368 run_id,
369 parent_run_id,
370 tags,
371 metadata,
372 name,
373 extra,
374 );
375
376 self.start_trace_async(&mut llm_run).await;
377 self.on_llm_start_async(&llm_run).await;
378
379 llm_run
380 }
381
382 async fn handle_llm_new_token_async(
384 &mut self,
385 token: &str,
386 run_id: Uuid,
387 chunk: Option<&(dyn std::any::Any + Send + Sync)>,
388 parent_run_id: Option<Uuid>,
389 ) -> Result<Run, TracerError> {
390 let llm_run = self.llm_run_with_token_event(
391 token,
392 run_id,
393 chunk.map(|c| c as &dyn std::any::Any),
394 parent_run_id,
395 )?;
396 self.on_llm_new_token_async(&llm_run, token, chunk).await;
397 Ok(llm_run)
398 }
399
400 async fn handle_retry_async(
402 &mut self,
403 retry_state: &HashMap<String, Value>,
404 run_id: Uuid,
405 ) -> Result<Run, TracerError> {
406 self.llm_run_with_retry_event(retry_state, run_id)
407 }
408
409 async fn handle_llm_end_async(
411 &mut self,
412 response: &LLMResult,
413 run_id: Uuid,
414 ) -> Result<Run, TracerError> {
415 let llm_run = self.complete_llm_run(response, run_id)?;
416
417 self.on_llm_end_async(&llm_run).await;
419 self.end_trace_async(&llm_run).await;
420
421 Ok(llm_run)
422 }
423
424 async fn handle_llm_error_async(
426 &mut self,
427 error: &(dyn std::error::Error + Send + Sync),
428 run_id: Uuid,
429 response: Option<&LLMResult>,
430 ) -> Result<Run, TracerError> {
431 let llm_run = self.errored_llm_run(error, run_id, response)?;
432
433 self.on_llm_error_async(&llm_run).await;
434 self.end_trace_async(&llm_run).await;
435
436 Ok(llm_run)
437 }
438
439 #[allow(clippy::too_many_arguments)]
441 async fn handle_chain_start_async(
442 &mut self,
443 serialized: HashMap<String, Value>,
444 inputs: HashMap<String, Value>,
445 run_id: Uuid,
446 parent_run_id: Option<Uuid>,
447 tags: Option<Vec<String>>,
448 metadata: Option<HashMap<String, Value>>,
449 run_type: Option<String>,
450 name: Option<String>,
451 extra: HashMap<String, Value>,
452 ) -> Run {
453 let mut chain_run = self.create_chain_run(
454 serialized,
455 inputs,
456 run_id,
457 parent_run_id,
458 tags,
459 metadata,
460 run_type,
461 name,
462 extra,
463 );
464
465 self.start_trace_async(&mut chain_run).await;
466 self.on_chain_start_async(&chain_run).await;
467
468 chain_run
469 }
470
471 async fn handle_chain_end_async(
473 &mut self,
474 outputs: HashMap<String, Value>,
475 run_id: Uuid,
476 inputs: Option<HashMap<String, Value>>,
477 ) -> Result<Run, TracerError> {
478 let chain_run = self.complete_chain_run(outputs, run_id, inputs)?;
479
480 self.end_trace_async(&chain_run).await;
481 self.on_chain_end_async(&chain_run).await;
482
483 Ok(chain_run)
484 }
485
486 async fn handle_chain_error_async(
488 &mut self,
489 error: &(dyn std::error::Error + Send + Sync),
490 run_id: Uuid,
491 inputs: Option<HashMap<String, Value>>,
492 ) -> Result<Run, TracerError> {
493 let chain_run = self.errored_chain_run(error, run_id, inputs)?;
494
495 self.end_trace_async(&chain_run).await;
496 self.on_chain_error_async(&chain_run).await;
497
498 Ok(chain_run)
499 }
500
501 #[allow(clippy::too_many_arguments)]
503 async fn handle_tool_start_async(
504 &mut self,
505 serialized: HashMap<String, Value>,
506 input_str: &str,
507 run_id: Uuid,
508 parent_run_id: Option<Uuid>,
509 tags: Option<Vec<String>>,
510 metadata: Option<HashMap<String, Value>>,
511 name: Option<String>,
512 inputs: Option<HashMap<String, Value>>,
513 extra: HashMap<String, Value>,
514 ) -> Run {
515 let mut tool_run = self.create_tool_run(
516 serialized,
517 input_str,
518 run_id,
519 parent_run_id,
520 tags,
521 metadata,
522 name,
523 inputs,
524 extra,
525 );
526
527 self.start_trace_async(&mut tool_run).await;
528 self.on_tool_start_async(&tool_run).await;
529
530 tool_run
531 }
532
533 async fn handle_tool_end_async(
535 &mut self,
536 output: Value,
537 run_id: Uuid,
538 ) -> Result<Run, TracerError> {
539 let tool_run = self.complete_tool_run(output, run_id)?;
540
541 self.end_trace_async(&tool_run).await;
542 self.on_tool_end_async(&tool_run).await;
543
544 Ok(tool_run)
545 }
546
547 async fn handle_tool_error_async(
549 &mut self,
550 error: &(dyn std::error::Error + Send + Sync),
551 run_id: Uuid,
552 ) -> Result<Run, TracerError> {
553 let tool_run = self.errored_tool_run(error, run_id)?;
554
555 self.end_trace_async(&tool_run).await;
556 self.on_tool_error_async(&tool_run).await;
557
558 Ok(tool_run)
559 }
560
561 #[allow(clippy::too_many_arguments)]
563 async fn handle_retriever_start_async(
564 &mut self,
565 serialized: HashMap<String, Value>,
566 query: &str,
567 run_id: Uuid,
568 parent_run_id: Option<Uuid>,
569 tags: Option<Vec<String>>,
570 metadata: Option<HashMap<String, Value>>,
571 name: Option<String>,
572 extra: HashMap<String, Value>,
573 ) -> Run {
574 let mut retriever_run = self.create_retrieval_run(
575 serialized,
576 query,
577 run_id,
578 parent_run_id,
579 tags,
580 metadata,
581 name,
582 extra,
583 );
584
585 self.start_trace_async(&mut retriever_run).await;
586 self.on_retriever_start_async(&retriever_run).await;
587
588 retriever_run
589 }
590
591 async fn handle_retriever_end_async(
593 &mut self,
594 documents: Vec<Value>,
595 run_id: Uuid,
596 ) -> Result<Run, TracerError> {
597 let retrieval_run = self.complete_retrieval_run(documents, run_id)?;
598
599 self.end_trace_async(&retrieval_run).await;
600 self.on_retriever_end_async(&retrieval_run).await;
601
602 Ok(retrieval_run)
603 }
604
605 async fn handle_retriever_error_async(
607 &mut self,
608 error: &(dyn std::error::Error + Send + Sync),
609 run_id: Uuid,
610 ) -> Result<Run, TracerError> {
611 let retrieval_run = self.errored_retrieval_run(error, run_id)?;
612
613 self.end_trace_async(&retrieval_run).await;
614 self.on_retriever_error_async(&retrieval_run).await;
615
616 Ok(retrieval_run)
617 }
618
619 async fn on_run_create_async(&mut self, _run: &Run) {}
623
624 async fn on_run_update_async(&mut self, _run: &Run) {}
626
627 async fn on_llm_start_async(&mut self, _run: &Run) {}
629
630 async fn on_llm_new_token_async(
632 &mut self,
633 _run: &Run,
634 _token: &str,
635 _chunk: Option<&(dyn std::any::Any + Send + Sync)>,
636 ) {
637 }
638
639 async fn on_llm_end_async(&mut self, _run: &Run) {}
641
642 async fn on_llm_error_async(&mut self, _run: &Run) {}
644
645 async fn on_chain_start_async(&mut self, _run: &Run) {}
647
648 async fn on_chain_end_async(&mut self, _run: &Run) {}
650
651 async fn on_chain_error_async(&mut self, _run: &Run) {}
653
654 async fn on_tool_start_async(&mut self, _run: &Run) {}
656
657 async fn on_tool_end_async(&mut self, _run: &Run) {}
659
660 async fn on_tool_error_async(&mut self, _run: &Run) {}
662
663 async fn on_chat_model_start_async(&mut self, _run: &Run) {}
665
666 async fn on_retriever_start_async(&mut self, _run: &Run) {}
668
669 async fn on_retriever_end_async(&mut self, _run: &Run) {}
671
672 async fn on_retriever_error_async(&mut self, _run: &Run) {}
674}
675
676#[cfg(test)]
677mod tests {
678 use super::*;
679 use crate::tracers::core::TracerCoreConfig;
680
681 #[derive(Debug)]
682 struct TestBaseTracer {
683 config: TracerCoreConfig,
684 run_map: HashMap<String, Run>,
685 order_map: HashMap<Uuid, (Uuid, String)>,
686 persisted_runs: Vec<Run>,
687 }
688
689 impl TestBaseTracer {
690 fn new() -> Self {
691 Self {
692 config: TracerCoreConfig::default(),
693 run_map: HashMap::new(),
694 order_map: HashMap::new(),
695 persisted_runs: Vec::new(),
696 }
697 }
698 }
699
700 impl TracerCore for TestBaseTracer {
701 fn config(&self) -> &TracerCoreConfig {
702 &self.config
703 }
704
705 fn config_mut(&mut self) -> &mut TracerCoreConfig {
706 &mut self.config
707 }
708
709 fn run_map(&self) -> &HashMap<String, Run> {
710 &self.run_map
711 }
712
713 fn run_map_mut(&mut self) -> &mut HashMap<String, Run> {
714 &mut self.run_map
715 }
716
717 fn order_map(&self) -> &HashMap<Uuid, (Uuid, String)> {
718 &self.order_map
719 }
720
721 fn order_map_mut(&mut self) -> &mut HashMap<Uuid, (Uuid, String)> {
722 &mut self.order_map
723 }
724
725 fn persist_run(&mut self, _run: &Run) {}
726 }
727
728 impl BaseTracer for TestBaseTracer {
729 fn persist_run_impl(&mut self, run: &Run) {
730 self.persisted_runs.push(run.clone());
731 }
732 }
733
734 #[test]
735 fn test_base_tracer_chain_lifecycle() {
736 let mut tracer = TestBaseTracer::new();
737
738 let run = tracer.handle_chain_start(
740 HashMap::new(),
741 HashMap::new(),
742 Uuid::new_v4(),
743 None,
744 None,
745 None,
746 None,
747 Some("test_chain".to_string()),
748 HashMap::new(),
749 );
750
751 assert_eq!(run.name, "test_chain");
752 assert_eq!(run.run_type, "chain");
753 assert!(tracer.run_map.contains_key(&run.id.to_string()));
754
755 let run_id = run.id;
757 let result = tracer.handle_chain_end(
758 [("output".to_string(), Value::String("result".to_string()))]
759 .into_iter()
760 .collect(),
761 run_id,
762 None,
763 );
764
765 assert!(result.is_ok());
766 let completed_run = result.unwrap();
767 assert!(completed_run.end_time.is_some());
768 assert!(!tracer.run_map.contains_key(&run_id.to_string()));
769 assert_eq!(tracer.persisted_runs.len(), 1);
770 }
771
772 #[test]
773 fn test_base_tracer_tool_lifecycle() {
774 let mut tracer = TestBaseTracer::new();
775
776 let run = tracer.handle_tool_start(
778 HashMap::new(),
779 "test input",
780 Uuid::new_v4(),
781 None,
782 None,
783 None,
784 Some("test_tool".to_string()),
785 None,
786 HashMap::new(),
787 );
788
789 assert_eq!(run.name, "test_tool");
790 assert_eq!(run.run_type, "tool");
791
792 let result = tracer.handle_tool_end(Value::String("output".to_string()), run.id);
794
795 assert!(result.is_ok());
796 }
797
798 #[test]
799 fn test_base_tracer_error_handling() {
800 let mut tracer = TestBaseTracer::new();
801
802 let run = tracer.handle_chain_start(
804 HashMap::new(),
805 HashMap::new(),
806 Uuid::new_v4(),
807 None,
808 None,
809 None,
810 None,
811 None,
812 HashMap::new(),
813 );
814
815 let error = std::io::Error::other("test error");
817 let result = tracer.handle_chain_error(&error, run.id, None);
818
819 assert!(result.is_ok());
820 let errored_run = result.unwrap();
821 assert!(errored_run.error.is_some());
822 assert!(errored_run.end_time.is_some());
823 }
824}