Skip to main content

ferrin_core/telemetry/
dispatcher.rs

1//! Fans telemetry events out to the configured integrations and to
2//! `tracing`.
3
4mod modalities;
5
6use std::panic::AssertUnwindSafe;
7use std::panic::catch_unwind;
8use std::sync::Arc;
9
10use ferrin_spec::BoxFuture;
11use ferrin_tool::ToolError;
12use futures_util::FutureExt;
13use futures_util::future::join_all;
14
15use super::AbortEvent;
16use super::EmbedEndEvent;
17use super::EmbedStartEvent;
18use super::EndEvent;
19use super::ErrorEvent;
20use super::ModelCallContext;
21use super::ModelCallEndEvent;
22use super::ModelCallOutcome;
23use super::ModelCallStartEvent;
24use super::RerankEndEvent;
25use super::RerankStartEvent;
26use super::StartEvent;
27use super::StepEndEvent;
28use super::StepStartEvent;
29use super::Telemetry;
30use super::TelemetryOptions;
31use super::ToolExecutionContext;
32use super::ToolExecutionEndEvent;
33use super::ToolExecutionStartEvent;
34use super::ToolOutcome;
35use crate::error::Error;
36
37/// Dispatches events to every integration of a [`TelemetryOptions`].
38#[derive(Clone, Debug)]
39pub(crate) struct TelemetryDispatcher {
40    options: Arc<TelemetryOptions>,
41}
42
43macro_rules! dispatch {
44    ($name:ident, $event:ty) => {
45        pub(crate) async fn $name(&self, event: &$event) {
46            if !self.options.enabled {
47                return;
48            }
49            self.dispatch(|integration| integration.$name(event)).await;
50        }
51    };
52}
53
54macro_rules! dispatch_with_context {
55    ($name:ident, $event:ty) => {
56        pub(crate) async fn $name(&self, event: &$event) {
57            if !self.options.enabled {
58                return;
59            }
60            let mut recorded = event.clone();
61            if !self.options.include_runtime_context {
62                recorded.runtime_context = None;
63            }
64            self.dispatch(|integration| integration.$name(&recorded))
65                .await;
66        }
67    };
68}
69
70impl TelemetryDispatcher {
71    async fn dispatch<'a>(
72        &'a self,
73        callback: impl Fn(&'a dyn Telemetry) -> BoxFuture<'a, ()> + Send + Sync,
74    ) {
75        let futures = self.options.integrations.iter().filter_map(|integration| {
76            catch_unwind(AssertUnwindSafe(|| callback(integration.as_ref())))
77                .ok()
78                .map(|future| AssertUnwindSafe(future).catch_unwind())
79        });
80        let _ = join_all(futures).await;
81    }
82
83    pub(crate) fn new(options: TelemetryOptions) -> Self {
84        Self {
85            options: Arc::new(options),
86        }
87    }
88
89    pub(crate) fn record_inputs(&self) -> bool {
90        self.options.enabled && self.options.record_inputs
91    }
92
93    pub(crate) fn record_outputs(&self) -> bool {
94        self.options.enabled && self.options.record_outputs
95    }
96
97    dispatch_with_context!(on_start, StartEvent);
98    dispatch_with_context!(on_step_start, StepStartEvent);
99    dispatch_with_context!(on_language_model_call_start, ModelCallStartEvent);
100
101    dispatch_with_context!(on_tool_execution_start, ToolExecutionStartEvent);
102    pub(crate) async fn on_tool_execution_end(&self, event: &ToolExecutionEndEvent) {
103        if !self.options.enabled {
104            return;
105        }
106        let mut recorded = event.clone();
107        if !self.options.include_runtime_context {
108            recorded.runtime_context = None;
109        }
110        if !self.record_outputs() {
111            recorded.output = None;
112            recorded.error = recorded
113                .error
114                .map(|_| crate::generate_text::ToolErrorInfo::text(super::redact::REDACTED));
115        }
116        self.dispatch(|integration| integration.on_tool_execution_end(&recorded))
117            .await;
118    }
119
120    dispatch!(on_abort, AbortEvent);
121
122    pub(crate) async fn on_language_model_call_end(&self, event: &ModelCallEndEvent) {
123        if !self.options.enabled {
124            return;
125        }
126        let mut recorded = event.clone();
127        if !self.options.include_runtime_context {
128            recorded.runtime_context = None;
129        }
130        if !(self.record_inputs() && self.record_outputs()) {
131            recorded.warnings = super::redact::warnings(&recorded.warnings);
132        }
133        if !self.record_outputs() {
134            recorded.content = None;
135            recorded.response.body = None;
136        }
137        self.dispatch(|integration| integration.on_language_model_call_end(&recorded))
138            .await;
139    }
140
141    pub(crate) async fn on_step_end(&self, event: &StepEndEvent) {
142        if !self.options.enabled {
143            return;
144        }
145        let recorded = StepEndEvent {
146            call_id: event.call_id.clone(),
147            step: Arc::new(self.recorded_step(&event.step)),
148        };
149        self.dispatch(|integration| integration.on_step_end(&recorded))
150            .await;
151    }
152
153    pub(crate) async fn on_end(&self, event: &EndEvent) {
154        if !self.options.enabled {
155            return;
156        }
157        let recorded = EndEvent {
158            runtime_context: self
159                .options
160                .include_runtime_context
161                .then(|| event.runtime_context.clone())
162                .flatten(),
163            call_id: event.call_id.clone(),
164            steps: event
165                .steps
166                .iter()
167                .map(|step| self.recorded_step(step))
168                .collect(),
169            total_usage: event.total_usage.clone(),
170            output_recorded: self
171                .record_outputs()
172                .then(|| event.output_recorded.clone())
173                .flatten(),
174        };
175        self.dispatch(|integration| integration.on_end(&recorded))
176            .await;
177    }
178
179    /// Filters the telemetry copy without changing application hooks or results.
180    fn recorded_step(
181        &self,
182        step: &crate::generate_text::StepResult,
183    ) -> crate::generate_text::StepResult {
184        let mut recorded = step.clone();
185        if !self.options.include_runtime_context {
186            recorded.runtime_context = None;
187        }
188        if !self.options.include_tools_context {
189            recorded.tools_context = None;
190        }
191        if !(self.record_inputs() && self.record_outputs()) {
192            recorded.warnings = super::redact::warnings(&recorded.warnings);
193        }
194        if !self.record_inputs() {
195            recorded.request.body = None;
196            recorded.request.messages = None;
197        }
198        if !self.record_outputs() {
199            recorded.content.clear();
200            recorded.response.body = None;
201            recorded.response.messages.clear();
202            recorded.provider_metadata = None;
203        }
204        recorded
205    }
206
207    pub(crate) async fn on_error(&self, event: &ErrorEvent<'_>) {
208        if !self.options.enabled {
209            return;
210        }
211        let error = (!(self.record_inputs() && self.record_outputs()))
212            .then(|| super::redact::redact_error(event.error, &self.options));
213        let recorded = ErrorEvent {
214            call_id: event.call_id,
215            error: error.as_ref().unwrap_or(event.error),
216            phase: event.phase,
217        };
218        self.dispatch(|integration| integration.on_error(&recorded))
219            .await;
220    }
221
222    /// Wraps `call` with every integration's `execute_language_model_call`;
223    /// the last integration becomes the outermost wrapper.
224    pub(crate) fn execute_language_model_call<'a>(
225        &'a self,
226        ctx: &'a ModelCallContext,
227        call: BoxFuture<'a, Result<ModelCallOutcome, Error>>,
228    ) -> BoxFuture<'a, Result<ModelCallOutcome, Error>> {
229        if !self.options.enabled {
230            return call;
231        }
232        self.options
233            .integrations
234            .iter()
235            .fold(call, |inner, integration| {
236                integration.execute_language_model_call(ctx, inner)
237            })
238    }
239
240    /// Wraps `call` with every integration's `execute_tool`.
241    pub(crate) fn execute_tool<'a>(
242        &'a self,
243        ctx: &'a ToolExecutionContext,
244        call: BoxFuture<'a, Result<ToolOutcome, ToolError>>,
245    ) -> BoxFuture<'a, Result<ToolOutcome, ToolError>> {
246        if !self.options.enabled {
247            return call;
248        }
249        self.options
250            .integrations
251            .iter()
252            .fold(call, |inner, integration| {
253                integration.execute_tool(ctx, inner)
254            })
255    }
256}
257
258impl TelemetryDispatcher {
259    dispatch!(on_embed_start, EmbedStartEvent);
260    dispatch!(on_embed_end, EmbedEndEvent);
261    dispatch!(on_rerank_start, RerankStartEvent);
262    dispatch!(on_rerank_end, RerankEndEvent);
263}
264
265impl Default for TelemetryDispatcher {
266    fn default() -> Self {
267        Self::new(TelemetryOptions::default())
268    }
269}
270
271impl dyn Telemetry {
272    /// Convenience for tests: returns `true` when `self` is the same object.
273    #[must_use]
274    pub fn ptr_eq(this: &Arc<Self>, other: &Arc<Self>) -> bool {
275        Arc::ptr_eq(this, other)
276    }
277}