Skip to main content

ferrin_core/telemetry/
dispatcher.rs

1//! Fans telemetry events out to the configured integrations and to
2//! `tracing`.
3
4use std::sync::Arc;
5
6use ferrin_spec::BoxFuture;
7use ferrin_tool::ToolError;
8
9use super::AbortEvent;
10use super::EmbedEndEvent;
11use super::EmbedStartEvent;
12use super::EndEvent;
13use super::ErrorEvent;
14use super::ModelCallContext;
15use super::ModelCallEndEvent;
16use super::ModelCallOutcome;
17use super::ModelCallStartEvent;
18use super::RerankEndEvent;
19use super::RerankStartEvent;
20use super::StartEvent;
21use super::StepEndEvent;
22use super::StepStartEvent;
23use super::Telemetry;
24use super::TelemetryOptions;
25use super::ToolExecutionContext;
26use super::ToolExecutionEndEvent;
27use super::ToolExecutionStartEvent;
28use super::ToolOutcome;
29use crate::error::Error;
30
31/// Dispatches events to every integration of a [`TelemetryOptions`].
32#[derive(Clone, Debug)]
33pub(crate) struct TelemetryDispatcher {
34    options: Arc<TelemetryOptions>,
35}
36
37macro_rules! dispatch {
38    ($name:ident, $event:ty) => {
39        pub(crate) fn $name(&self, event: &$event) {
40            if !self.options.enabled {
41                return;
42            }
43            for integration in &self.options.integrations {
44                integration.$name(event);
45            }
46        }
47    };
48}
49
50impl TelemetryDispatcher {
51    pub(crate) fn new(options: TelemetryOptions) -> Self {
52        Self {
53            options: Arc::new(options),
54        }
55    }
56
57    pub(crate) fn record_inputs(&self) -> bool {
58        self.options.enabled && self.options.record_inputs
59    }
60
61    pub(crate) fn record_outputs(&self) -> bool {
62        self.options.enabled && self.options.record_outputs
63    }
64
65    dispatch!(on_start, StartEvent);
66    dispatch!(on_step_start, StepStartEvent);
67    dispatch!(on_language_model_call_start, ModelCallStartEvent);
68    dispatch!(on_language_model_call_end, ModelCallEndEvent);
69    dispatch!(on_tool_execution_start, ToolExecutionStartEvent);
70    dispatch!(on_tool_execution_end, ToolExecutionEndEvent);
71    dispatch!(on_step_end, StepEndEvent);
72    dispatch!(on_end, EndEvent);
73    dispatch!(on_abort, AbortEvent);
74
75    pub(crate) fn on_error(&self, event: &ErrorEvent<'_>) {
76        if !self.options.enabled {
77            return;
78        }
79        for integration in &self.options.integrations {
80            integration.on_error(event);
81        }
82    }
83
84    /// Wraps `call` with every integration's `execute_language_model_call`;
85    /// the first integration becomes the outermost wrapper.
86    pub(crate) fn execute_language_model_call<'a>(
87        &'a self,
88        ctx: &'a ModelCallContext,
89        call: BoxFuture<'a, Result<ModelCallOutcome, Error>>,
90    ) -> BoxFuture<'a, Result<ModelCallOutcome, Error>> {
91        if !self.options.enabled {
92            return call;
93        }
94        self.options
95            .integrations
96            .iter()
97            .rev()
98            .fold(call, |inner, integration| {
99                integration.execute_language_model_call(ctx, inner)
100            })
101    }
102
103    /// Wraps `call` with every integration's `execute_tool`.
104    pub(crate) fn execute_tool<'a>(
105        &'a self,
106        ctx: &'a ToolExecutionContext,
107        call: BoxFuture<'a, Result<ToolOutcome, ToolError>>,
108    ) -> BoxFuture<'a, Result<ToolOutcome, ToolError>> {
109        if !self.options.enabled {
110            return call;
111        }
112        self.options
113            .integrations
114            .iter()
115            .rev()
116            .fold(call, |inner, integration| {
117                integration.execute_tool(ctx, inner)
118            })
119    }
120}
121
122impl TelemetryDispatcher {
123    dispatch!(on_embed_start, EmbedStartEvent);
124    dispatch!(on_embed_end, EmbedEndEvent);
125    dispatch!(on_rerank_start, RerankStartEvent);
126    dispatch!(on_rerank_end, RerankEndEvent);
127}
128
129impl Default for TelemetryDispatcher {
130    fn default() -> Self {
131        Self::new(TelemetryOptions::default())
132    }
133}
134
135impl dyn Telemetry {
136    /// Convenience for tests: returns `true` when `self` is the same object.
137    #[must_use]
138    pub fn ptr_eq(this: &Arc<Self>, other: &Arc<Self>) -> bool {
139        Arc::ptr_eq(this, other)
140    }
141}