rig_tap/hook.rs
1//! [`TelemetryHook`]: a [`rig::agent::PromptHook`] that emits
2//! `prompt.*` and `tool.*` [`ObservabilityEvent`](crate::ObservabilityEvent)s
3//! for every prompt and tool call.
4
5use std::marker::PhantomData;
6use std::sync::Arc;
7
8use rig::agent::{HookAction, PromptHook, ToolCallHookAction};
9use rig::completion::{CompletionModel, CompletionResponse, Message};
10
11use crate::emit::emit_kind;
12use crate::event::{EventKind, PAYLOAD_TRUNCATE_BYTES, truncate_utf8};
13use crate::sampling::{AlwaysSample, SamplingPolicy};
14
15/// Caller-supplied resolver for the conversation ID stamped on emitted
16/// events. Consulted on every emission; when it returns `Some(id)`, that
17/// value wins over [`TelemetryHookConfig::conversation_id`].
18///
19/// Use this when the host runtime threads a per-request conversation ID
20/// through e.g. a task-local, a `tracing::Span` field, or a request-scoped
21/// context object. The Rig `PromptHook` signature does not currently
22/// propagate a conversation ID; this is the escape hatch.
23pub type ConversationIdResolver = Arc<dyn Fn() -> Option<String> + Send + Sync>;
24
25/// Caller-supplied resolver that pulls the *actual* model identifier out
26/// of a provider response. Useful for routed providers (OpenRouter,
27/// Bedrock model-routing, vendor multi-model endpoints) where the model
28/// recorded at hook construction is a logical alias and the response's
29/// raw payload carries the concrete model that served the request.
30///
31/// When set and the resolver returns `Some(model)`, that value is stamped
32/// on `prompt.completed` instead of [`TelemetryHookConfig::model`].
33pub type ModelResolver<R> = Arc<dyn Fn(&CompletionResponse<R>) -> Option<String> + Send + Sync>;
34
35/// Caller-supplied resolver that returns the chain ancestor for the current
36/// turn (the `previous_response_id` argument sent to the provider) so it can
37/// be stamped on `prompt.completed`.
38///
39/// Useful for stateful endpoints — OpenAI Responses, future Anthropic and
40/// Google equivalents — where the host runtime tracks the chain itself
41/// (typically in a task-local or session object) and the provider response
42/// payload does not echo the value back. The Rig `PromptHook` signature
43/// does not currently propagate it, so this is the escape hatch.
44///
45/// When set and the resolver returns `Some(id)`, that value is stamped on
46/// `prompt.completed`; `None` leaves the field unset.
47pub type PreviousResponseIdResolver<R> =
48 Arc<dyn Fn(&CompletionResponse<R>) -> Option<String> + Send + Sync>;
49
50/// Conversation identifier to stamp on emitted events when the agent runtime
51/// does not surface one to the hook. The Rig `PromptHook` signature does not
52/// currently propagate the conversation ID, so the hook stamps events with a
53/// constant chosen by the caller (typically `"default"` for single-thread
54/// agents, or a unique value per agent instance for multi-thread setups).
55///
56/// For per-request resolution see
57/// [`TelemetryHook::with_conversation_id_resolver`].
58#[derive(Debug, Clone)]
59pub struct TelemetryHookConfig {
60 /// Default model label (e.g. `"gpt-4o"`) recorded on `prompt.*` events.
61 /// For routed providers, prefer [`TelemetryHook::with_model_resolver`]
62 /// to extract the model name from the actual response payload.
63 pub model: String,
64 /// Default conversation ID stamped on every emitted event when no
65 /// per-request resolver is registered or the resolver returns `None`.
66 pub conversation_id: String,
67 /// Maximum byte length of inline `args_json` / `result` payloads before
68 /// truncation. Defaults to [`PAYLOAD_TRUNCATE_BYTES`].
69 pub payload_truncate_bytes: usize,
70}
71
72impl TelemetryHookConfig {
73 /// Build a config with the given model label and conversation ID, using
74 /// the default truncation threshold.
75 pub fn new(model: impl Into<String>, conversation_id: impl Into<String>) -> Self {
76 Self {
77 model: model.into(),
78 conversation_id: conversation_id.into(),
79 payload_truncate_bytes: PAYLOAD_TRUNCATE_BYTES,
80 }
81 }
82}
83
84/// Per-request hook that emits structured observability events from the five
85/// [`PromptHook`] lifecycle methods.
86///
87/// `M` is the [`CompletionModel`] used by the agent. The hook is generic so a
88/// single `rig-tap` build can attach to OpenAI, Anthropic, Ollama, etc.
89///
90/// # Example
91///
92/// ```no_run
93/// use rig_tap::{TelemetryHook, TelemetryHookConfig};
94///
95/// # fn make_hook<M: rig::completion::CompletionModel>() -> TelemetryHook<M> {
96/// TelemetryHook::new(TelemetryHookConfig::new("gpt-4o", "thread-1"))
97/// # }
98/// ```
99pub struct TelemetryHook<M: CompletionModel> {
100 config: TelemetryHookConfig,
101 conversation_id_resolver: Option<ConversationIdResolver>,
102 model_resolver: Option<ModelResolver<M::Response>>,
103 previous_response_id_resolver: Option<PreviousResponseIdResolver<M::Response>>,
104 sampling: Arc<dyn SamplingPolicy>,
105 _model: PhantomData<fn() -> M>,
106}
107
108impl<M: CompletionModel> TelemetryHook<M> {
109 /// Build a hook from `config`.
110 pub fn new(config: TelemetryHookConfig) -> Self {
111 Self {
112 config,
113 conversation_id_resolver: None,
114 model_resolver: None,
115 previous_response_id_resolver: None,
116 sampling: Arc::new(AlwaysSample),
117 _model: PhantomData,
118 }
119 }
120
121 /// Convenience: build a hook stamping events with `model` and
122 /// `conversation_id` and default truncation.
123 pub fn with_defaults(model: impl Into<String>, conversation_id: impl Into<String>) -> Self {
124 Self::new(TelemetryHookConfig::new(model, conversation_id))
125 }
126
127 /// Register a per-request resolver for the conversation ID. The
128 /// resolver is consulted on every emission; if it returns `Some(id)`,
129 /// that value is stamped on the event instead of
130 /// [`TelemetryHookConfig::conversation_id`].
131 ///
132 /// Typical wiring: a `tokio::task_local!` (or equivalent) set by the
133 /// host on every request, read by the closure.
134 #[must_use]
135 pub fn with_conversation_id_resolver<F>(mut self, resolver: F) -> Self
136 where
137 F: Fn() -> Option<String> + Send + Sync + 'static,
138 {
139 self.conversation_id_resolver = Some(Arc::new(resolver));
140 self
141 }
142
143 /// Register a resolver that extracts the concrete model identifier
144 /// from each [`CompletionResponse`]. When the resolver returns
145 /// `Some(model)`, that value is stamped on `prompt.completed`
146 /// instead of [`TelemetryHookConfig::model`].
147 ///
148 /// Use this with routed providers (OpenRouter, Bedrock routing,
149 /// vendor multi-model endpoints) where the configured model name is
150 /// a logical alias and the response payload carries the actual
151 /// model that served the request.
152 #[must_use]
153 pub fn with_model_resolver<F>(mut self, resolver: F) -> Self
154 where
155 F: Fn(&CompletionResponse<M::Response>) -> Option<String> + Send + Sync + 'static,
156 {
157 self.model_resolver = Some(Arc::new(resolver));
158 self
159 }
160
161 /// Register a resolver that returns the chain ancestor
162 /// (`previous_response_id`) sent to the provider for the current turn.
163 /// When the resolver returns `Some(id)`, that value is stamped on
164 /// `prompt.completed`'s `previous_response_id` field.
165 ///
166 /// Use this with stateful endpoints — OpenAI Responses, future
167 /// Anthropic/Google equivalents — where the host runtime tracks the
168 /// chain (typically in a task-local or session object) and the
169 /// provider response payload does not echo the value back.
170 #[must_use]
171 pub fn with_previous_response_id_resolver<F>(mut self, resolver: F) -> Self
172 where
173 F: Fn(&CompletionResponse<M::Response>) -> Option<String> + Send + Sync + 'static,
174 {
175 self.previous_response_id_resolver = Some(Arc::new(resolver));
176 self
177 }
178
179 /// Install a [`SamplingPolicy`] that gates every `prompt.*` and
180 /// `tool.*` emission from this hook. The default policy is
181 /// [`AlwaysSample`](crate::AlwaysSample).
182 ///
183 /// Pairing: the hook passes the resolved conversation id as the
184 /// correlator for `prompt.*` events and the internal call id for
185 /// `tool.*` events. Policies that hash the correlator (such as
186 /// [`RatePolicy`](crate::RatePolicy)) therefore keep
187 /// `tool.invoked` / `tool.completed` pairs coherent
188 /// automatically.
189 ///
190 /// # Example
191 ///
192 /// ```no_run
193 /// use std::sync::Arc;
194 /// use rig_tap::{RatePolicy, TelemetryHook, TelemetryHookConfig};
195 ///
196 /// # fn make_hook<M: rig::completion::CompletionModel>() -> TelemetryHook<M> {
197 /// TelemetryHook::new(TelemetryHookConfig::new("gpt-4o", "thread-1"))
198 /// .with_sampling_policy(Arc::new(
199 /// RatePolicy::new()
200 /// .with_rate("tool.invoked", 0.1)
201 /// .with_rate("tool.completed", 0.1),
202 /// ))
203 /// # }
204 /// ```
205 #[must_use]
206 pub fn with_sampling_policy(mut self, policy: Arc<dyn SamplingPolicy>) -> Self {
207 self.sampling = policy;
208 self
209 }
210
211 fn resolved_conversation_id(&self) -> String {
212 self.conversation_id_resolver
213 .as_ref()
214 .and_then(|f| f())
215 .unwrap_or_else(|| self.config.conversation_id.clone())
216 }
217
218 fn resolved_model(&self, response: &CompletionResponse<M::Response>) -> String {
219 self.model_resolver
220 .as_ref()
221 .and_then(|f| f(response))
222 .unwrap_or_else(|| self.config.model.clone())
223 }
224
225 fn resolved_previous_response_id(
226 &self,
227 response: &CompletionResponse<M::Response>,
228 ) -> Option<String> {
229 self.previous_response_id_resolver
230 .as_ref()
231 .and_then(|f| f(response))
232 }
233}
234
235impl<M: CompletionModel> Clone for TelemetryHook<M> {
236 fn clone(&self) -> Self {
237 Self {
238 config: self.config.clone(),
239 conversation_id_resolver: self.conversation_id_resolver.clone(),
240 model_resolver: self.model_resolver.clone(),
241 previous_response_id_resolver: self.previous_response_id_resolver.clone(),
242 sampling: self.sampling.clone(),
243 _model: PhantomData,
244 }
245 }
246}
247
248impl<M: CompletionModel> std::fmt::Debug for TelemetryHook<M> {
249 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
250 f.debug_struct("TelemetryHook")
251 .field("config", &self.config)
252 .field(
253 "conversation_id_resolver",
254 &self.conversation_id_resolver.as_ref().map(|_| "<fn>"),
255 )
256 .field(
257 "model_resolver",
258 &self.model_resolver.as_ref().map(|_| "<fn>"),
259 )
260 .field(
261 "previous_response_id_resolver",
262 &self.previous_response_id_resolver.as_ref().map(|_| "<fn>"),
263 )
264 .field("sampling", &self.sampling)
265 .finish_non_exhaustive()
266 }
267}
268
269impl<M> PromptHook<M> for TelemetryHook<M>
270where
271 M: CompletionModel,
272{
273 async fn on_completion_call(&self, _prompt: &Message, history: &[Message]) -> HookAction {
274 // `messages_in` counts the prompt + prior history that will be sent
275 // to the provider.
276 let messages_in = history.len().saturating_add(1);
277 let conversation_id = self.resolved_conversation_id();
278 if self
279 .sampling
280 .should_sample("prompt.started", &conversation_id)
281 {
282 emit_kind(
283 conversation_id,
284 EventKind::PromptStarted {
285 model: self.config.model.clone(),
286 messages_in,
287 },
288 );
289 }
290 HookAction::cont()
291 }
292
293 async fn on_completion_response(
294 &self,
295 _prompt: &Message,
296 response: &CompletionResponse<M::Response>,
297 ) -> HookAction {
298 let usage = response.usage;
299 let conversation_id = self.resolved_conversation_id();
300 if self
301 .sampling
302 .should_sample("prompt.completed", &conversation_id)
303 {
304 emit_kind(
305 conversation_id,
306 EventKind::PromptCompleted {
307 model: self.resolved_model(response),
308 tokens_in: positive(usage.input_tokens),
309 tokens_out: positive(usage.output_tokens),
310 response_id: response.message_id.clone(),
311 previous_response_id: self.resolved_previous_response_id(response),
312 },
313 );
314 }
315 HookAction::cont()
316 }
317
318 async fn on_tool_call(
319 &self,
320 tool_name: &str,
321 tool_call_id: Option<String>,
322 internal_call_id: &str,
323 args: &str,
324 ) -> ToolCallHookAction {
325 let (args_json, truncated) = truncate_utf8(args, self.config.payload_truncate_bytes);
326 if self
327 .sampling
328 .should_sample("tool.invoked", internal_call_id)
329 {
330 emit_kind(
331 self.resolved_conversation_id(),
332 EventKind::ToolInvoked {
333 tool_name: tool_name.to_string(),
334 provider_call_id: tool_call_id,
335 call_id: internal_call_id.to_string(),
336 args_json,
337 truncated,
338 },
339 );
340 }
341 ToolCallHookAction::cont()
342 }
343
344 async fn on_tool_result(
345 &self,
346 tool_name: &str,
347 tool_call_id: Option<String>,
348 internal_call_id: &str,
349 _args: &str,
350 result: &str,
351 ) -> HookAction {
352 let (result, truncated) = truncate_utf8(result, self.config.payload_truncate_bytes);
353 if self
354 .sampling
355 .should_sample("tool.completed", internal_call_id)
356 {
357 emit_kind(
358 self.resolved_conversation_id(),
359 EventKind::ToolCompleted {
360 tool_name: tool_name.to_string(),
361 provider_call_id: tool_call_id,
362 call_id: internal_call_id.to_string(),
363 result,
364 truncated,
365 },
366 );
367 }
368 HookAction::cont()
369 }
370}
371
372fn positive(value: u64) -> Option<u64> {
373 if value == 0 { None } else { Some(value) }
374}