1use std::fmt;
4use std::future::Future;
5use std::panic::AssertUnwindSafe;
6use std::panic::catch_unwind;
7use std::sync::Arc;
8
9use ferrin_spec::BoxFuture;
10use futures_util::FutureExt;
11use futures_util::future::join_all;
12
13use crate::generate_text::StepResult;
14use crate::stream_text::StreamEvent;
15use crate::telemetry::AbortEvent;
16use crate::telemetry::EndEvent;
17use crate::telemetry::ModelCallEndEvent;
18use crate::telemetry::ModelCallStartEvent;
19use crate::telemetry::StartEvent;
20use crate::telemetry::StepStartEvent;
21use crate::telemetry::ToolExecutionEndEvent;
22use crate::telemetry::ToolExecutionStartEvent;
23
24pub trait HookFn<E>: Send + Sync + 'static {
32 fn call(&self, event: Arc<E>) -> BoxFuture<'static, ()>;
34}
35
36impl<E, F, Fut> HookFn<E> for F
37where
38 F: Fn(Arc<E>) -> Fut + Send + Sync + 'static,
39 Fut: Future<Output = ()> + Send + 'static,
40{
41 fn call(&self, event: Arc<E>) -> BoxFuture<'static, ()> {
42 Box::pin(self(event))
43 }
44}
45
46pub type HookList<E> = Vec<Arc<dyn HookFn<E>>>;
48
49#[derive(Clone, Default)]
51pub struct Hooks {
52 pub on_start: HookList<StartEvent>,
54 pub on_step_start: HookList<StepStartEvent>,
56 pub on_language_model_call_start: HookList<ModelCallStartEvent>,
58 pub on_language_model_call_end: HookList<ModelCallEndEvent>,
60 pub on_tool_execution_start: HookList<ToolExecutionStartEvent>,
62 pub on_tool_execution_end: HookList<ToolExecutionEndEvent>,
64 pub on_step_end: HookList<StepResult>,
66 pub on_end: HookList<EndEvent>,
68 pub on_chunk: HookList<StreamEvent>,
70 pub on_abort: HookList<AbortEvent>,
72}
73
74impl Hooks {
75 #[must_use]
79 pub fn merged(mut self, other: Hooks) -> Hooks {
80 self.on_start.extend(other.on_start);
81 self.on_step_start.extend(other.on_step_start);
82 self.on_language_model_call_start
83 .extend(other.on_language_model_call_start);
84 self.on_language_model_call_end
85 .extend(other.on_language_model_call_end);
86 self.on_tool_execution_start
87 .extend(other.on_tool_execution_start);
88 self.on_tool_execution_end
89 .extend(other.on_tool_execution_end);
90 self.on_step_end.extend(other.on_step_end);
91 self.on_end.extend(other.on_end);
92 self.on_chunk.extend(other.on_chunk);
93 self.on_abort.extend(other.on_abort);
94 self
95 }
96
97 pub async fn emit<E: 'static>(list: &[Arc<dyn HookFn<E>>], event: Arc<E>) {
103 let futures = list.iter().filter_map(|hook| {
104 catch_unwind(AssertUnwindSafe(|| hook.call(Arc::clone(&event))))
105 .ok()
106 .map(|future| AssertUnwindSafe(future).catch_unwind())
107 });
108 let _ = join_all(futures).await;
109 }
110}
111
112impl fmt::Debug for Hooks {
113 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114 f.debug_struct("Hooks")
115 .field("on_start", &self.on_start.len())
116 .field("on_step_start", &self.on_step_start.len())
117 .field(
118 "on_language_model_call_start",
119 &self.on_language_model_call_start.len(),
120 )
121 .field(
122 "on_language_model_call_end",
123 &self.on_language_model_call_end.len(),
124 )
125 .field(
126 "on_tool_execution_start",
127 &self.on_tool_execution_start.len(),
128 )
129 .field("on_tool_execution_end", &self.on_tool_execution_end.len())
130 .field("on_step_end", &self.on_step_end.len())
131 .field("on_end", &self.on_end.len())
132 .field("on_chunk", &self.on_chunk.len())
133 .field("on_abort", &self.on_abort.len())
134 .finish()
135 }
136}