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};
13
14/// Caller-supplied resolver for the conversation ID stamped on emitted
15/// events. Consulted on every emission; when it returns `Some(id)`, that
16/// value wins over [`TelemetryHookConfig::conversation_id`].
17///
18/// Use this when the host runtime threads a per-request conversation ID
19/// through e.g. a task-local, a `tracing::Span` field, or a request-scoped
20/// context object. The Rig `PromptHook` signature does not currently
21/// propagate a conversation ID; this is the escape hatch.
22pub type ConversationIdResolver = Arc<dyn Fn() -> Option<String> + Send + Sync>;
23
24/// Caller-supplied resolver that pulls the *actual* model identifier out
25/// of a provider response. Useful for routed providers (OpenRouter,
26/// Bedrock model-routing, vendor multi-model endpoints) where the model
27/// recorded at hook construction is a logical alias and the response's
28/// raw payload carries the concrete model that served the request.
29///
30/// When set and the resolver returns `Some(model)`, that value is stamped
31/// on `prompt.completed` instead of [`TelemetryHookConfig::model`].
32pub type ModelResolver<R> = Arc<dyn Fn(&CompletionResponse<R>) -> Option<String> + Send + Sync>;
33
34/// Conversation identifier to stamp on emitted events when the agent runtime
35/// does not surface one to the hook. The Rig `PromptHook` signature does not
36/// currently propagate the conversation ID, so the hook stamps events with a
37/// constant chosen by the caller (typically `"default"` for single-thread
38/// agents, or a unique value per agent instance for multi-thread setups).
39///
40/// For per-request resolution see
41/// [`TelemetryHook::with_conversation_id_resolver`].
42#[derive(Debug, Clone)]
43pub struct TelemetryHookConfig {
44 /// Default model label (e.g. `"gpt-4o"`) recorded on `prompt.*` events.
45 /// For routed providers, prefer [`TelemetryHook::with_model_resolver`]
46 /// to extract the model name from the actual response payload.
47 pub model: String,
48 /// Default conversation ID stamped on every emitted event when no
49 /// per-request resolver is registered or the resolver returns `None`.
50 pub conversation_id: String,
51 /// Maximum byte length of inline `args_json` / `result` payloads before
52 /// truncation. Defaults to [`PAYLOAD_TRUNCATE_BYTES`].
53 pub payload_truncate_bytes: usize,
54}
55
56impl TelemetryHookConfig {
57 /// Build a config with the given model label and conversation ID, using
58 /// the default truncation threshold.
59 pub fn new(model: impl Into<String>, conversation_id: impl Into<String>) -> Self {
60 Self {
61 model: model.into(),
62 conversation_id: conversation_id.into(),
63 payload_truncate_bytes: PAYLOAD_TRUNCATE_BYTES,
64 }
65 }
66}
67
68/// Per-request hook that emits structured observability events from the five
69/// [`PromptHook`] lifecycle methods.
70///
71/// `M` is the [`CompletionModel`] used by the agent. The hook is generic so a
72/// single `rig-tap` build can attach to OpenAI, Anthropic, Ollama, etc.
73///
74/// # Example
75///
76/// ```no_run
77/// use rig_tap::{TelemetryHook, TelemetryHookConfig};
78///
79/// # fn make_hook<M: rig::completion::CompletionModel>() -> TelemetryHook<M> {
80/// TelemetryHook::new(TelemetryHookConfig::new("gpt-4o", "thread-1"))
81/// # }
82/// ```
83pub struct TelemetryHook<M: CompletionModel> {
84 config: TelemetryHookConfig,
85 conversation_id_resolver: Option<ConversationIdResolver>,
86 model_resolver: Option<ModelResolver<M::Response>>,
87 _model: PhantomData<fn() -> M>,
88}
89
90impl<M: CompletionModel> TelemetryHook<M> {
91 /// Build a hook from `config`.
92 pub fn new(config: TelemetryHookConfig) -> Self {
93 Self {
94 config,
95 conversation_id_resolver: None,
96 model_resolver: None,
97 _model: PhantomData,
98 }
99 }
100
101 /// Convenience: build a hook stamping events with `model` and
102 /// `conversation_id` and default truncation.
103 pub fn with_defaults(model: impl Into<String>, conversation_id: impl Into<String>) -> Self {
104 Self::new(TelemetryHookConfig::new(model, conversation_id))
105 }
106
107 /// Register a per-request resolver for the conversation ID. The
108 /// resolver is consulted on every emission; if it returns `Some(id)`,
109 /// that value is stamped on the event instead of
110 /// [`TelemetryHookConfig::conversation_id`].
111 ///
112 /// Typical wiring: a `tokio::task_local!` (or equivalent) set by the
113 /// host on every request, read by the closure.
114 #[must_use]
115 pub fn with_conversation_id_resolver<F>(mut self, resolver: F) -> Self
116 where
117 F: Fn() -> Option<String> + Send + Sync + 'static,
118 {
119 self.conversation_id_resolver = Some(Arc::new(resolver));
120 self
121 }
122
123 /// Register a resolver that extracts the concrete model identifier
124 /// from each [`CompletionResponse`]. When the resolver returns
125 /// `Some(model)`, that value is stamped on `prompt.completed`
126 /// instead of [`TelemetryHookConfig::model`].
127 ///
128 /// Use this with routed providers (OpenRouter, Bedrock routing,
129 /// vendor multi-model endpoints) where the configured model name is
130 /// a logical alias and the response payload carries the actual
131 /// model that served the request.
132 #[must_use]
133 pub fn with_model_resolver<F>(mut self, resolver: F) -> Self
134 where
135 F: Fn(&CompletionResponse<M::Response>) -> Option<String> + Send + Sync + 'static,
136 {
137 self.model_resolver = Some(Arc::new(resolver));
138 self
139 }
140
141 fn resolved_conversation_id(&self) -> String {
142 self.conversation_id_resolver
143 .as_ref()
144 .and_then(|f| f())
145 .unwrap_or_else(|| self.config.conversation_id.clone())
146 }
147
148 fn resolved_model(&self, response: &CompletionResponse<M::Response>) -> String {
149 self.model_resolver
150 .as_ref()
151 .and_then(|f| f(response))
152 .unwrap_or_else(|| self.config.model.clone())
153 }
154}
155
156impl<M: CompletionModel> Clone for TelemetryHook<M> {
157 fn clone(&self) -> Self {
158 Self {
159 config: self.config.clone(),
160 conversation_id_resolver: self.conversation_id_resolver.clone(),
161 model_resolver: self.model_resolver.clone(),
162 _model: PhantomData,
163 }
164 }
165}
166
167impl<M: CompletionModel> std::fmt::Debug for TelemetryHook<M> {
168 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
169 f.debug_struct("TelemetryHook")
170 .field("config", &self.config)
171 .field(
172 "conversation_id_resolver",
173 &self.conversation_id_resolver.as_ref().map(|_| "<fn>"),
174 )
175 .field(
176 "model_resolver",
177 &self.model_resolver.as_ref().map(|_| "<fn>"),
178 )
179 .finish_non_exhaustive()
180 }
181}
182
183impl<M> PromptHook<M> for TelemetryHook<M>
184where
185 M: CompletionModel,
186{
187 async fn on_completion_call(&self, _prompt: &Message, history: &[Message]) -> HookAction {
188 // `messages_in` counts the prompt + prior history that will be sent
189 // to the provider.
190 let messages_in = history.len().saturating_add(1);
191 emit_kind(
192 self.resolved_conversation_id(),
193 EventKind::PromptStarted {
194 model: self.config.model.clone(),
195 messages_in,
196 },
197 );
198 HookAction::cont()
199 }
200
201 async fn on_completion_response(
202 &self,
203 _prompt: &Message,
204 response: &CompletionResponse<M::Response>,
205 ) -> HookAction {
206 let usage = response.usage;
207 emit_kind(
208 self.resolved_conversation_id(),
209 EventKind::PromptCompleted {
210 model: self.resolved_model(response),
211 tokens_in: positive(usage.input_tokens),
212 tokens_out: positive(usage.output_tokens),
213 response_id: response.message_id.clone(),
214 },
215 );
216 HookAction::cont()
217 }
218
219 async fn on_tool_call(
220 &self,
221 tool_name: &str,
222 tool_call_id: Option<String>,
223 internal_call_id: &str,
224 args: &str,
225 ) -> ToolCallHookAction {
226 let (args_json, truncated) = truncate_utf8(args, self.config.payload_truncate_bytes);
227 emit_kind(
228 self.resolved_conversation_id(),
229 EventKind::ToolInvoked {
230 tool_name: tool_name.to_string(),
231 provider_call_id: tool_call_id,
232 call_id: internal_call_id.to_string(),
233 args_json,
234 truncated,
235 },
236 );
237 ToolCallHookAction::cont()
238 }
239
240 async fn on_tool_result(
241 &self,
242 tool_name: &str,
243 tool_call_id: Option<String>,
244 internal_call_id: &str,
245 _args: &str,
246 result: &str,
247 ) -> HookAction {
248 let (result, truncated) = truncate_utf8(result, self.config.payload_truncate_bytes);
249 emit_kind(
250 self.resolved_conversation_id(),
251 EventKind::ToolCompleted {
252 tool_name: tool_name.to_string(),
253 provider_call_id: tool_call_id,
254 call_id: internal_call_id.to_string(),
255 result,
256 truncated,
257 },
258 );
259 HookAction::cont()
260 }
261}
262
263fn positive(value: u64) -> Option<u64> {
264 if value == 0 { None } else { Some(value) }
265}