rig_tap/
query.rs

1//! In-process query helpers for captured observability events.
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use serde::{Deserialize, Serialize};
6
7use crate::event::ObservabilityEvent;
8
9/// Predicate used by [`EventQuery`] to select observability events.
10///
11/// Every configured field must match. Tick bounds are inclusive.
12#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct EventFilter {
14    /// Conversation identifier to match.
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub conversation_id: Option<String>,
17    /// Wire event kind to match, such as `"tool.completed"`.
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub kind: Option<String>,
20    /// Inclusive lower tick bound.
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub min_tick: Option<u64>,
23    /// Inclusive upper tick bound.
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub max_tick: Option<u64>,
26    /// Compose kernel identifier to match.
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub kernel_id: Option<String>,
29    /// Tool name or retry target to match.
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub tool_name: Option<String>,
32    /// Tool call correlation identifier to match.
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub call_id: Option<String>,
35    /// Compose skill identifier to match.
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub skill_id: Option<String>,
38    /// Prompt model identifier to match.
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub model: Option<String>,
41}
42
43impl EventFilter {
44    /// Build an empty filter that matches every event.
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Match a conversation identifier.
50    pub fn conversation_id(mut self, conversation_id: impl Into<String>) -> Self {
51        self.conversation_id = Some(conversation_id.into());
52        self
53    }
54
55    /// Match a wire event kind, such as `"prompt.started"`.
56    pub fn kind(mut self, kind: impl Into<String>) -> Self {
57        self.kind = Some(kind.into());
58        self
59    }
60
61    /// Match events at or after `tick`.
62    pub fn min_tick(mut self, tick: u64) -> Self {
63        self.min_tick = Some(tick);
64        self
65    }
66
67    /// Match events at or before `tick`.
68    pub fn max_tick(mut self, tick: u64) -> Self {
69        self.max_tick = Some(tick);
70        self
71    }
72
73    /// Match a compose kernel identifier.
74    pub fn kernel_id(mut self, kernel_id: impl Into<String>) -> Self {
75        self.kernel_id = Some(kernel_id.into());
76        self
77    }
78
79    /// Match a tool name or retry target.
80    pub fn tool_name(mut self, tool_name: impl Into<String>) -> Self {
81        self.tool_name = Some(tool_name.into());
82        self
83    }
84
85    /// Match a tool call correlation identifier.
86    pub fn call_id(mut self, call_id: impl Into<String>) -> Self {
87        self.call_id = Some(call_id.into());
88        self
89    }
90
91    /// Match a compose skill identifier.
92    pub fn skill_id(mut self, skill_id: impl Into<String>) -> Self {
93        self.skill_id = Some(skill_id.into());
94        self
95    }
96
97    /// Match a prompt model identifier.
98    pub fn model(mut self, model: impl Into<String>) -> Self {
99        self.model = Some(model.into());
100        self
101    }
102
103    /// Return `true` if `event` satisfies every configured predicate.
104    pub fn matches(&self, event: &ObservabilityEvent) -> bool {
105        if self
106            .conversation_id
107            .as_ref()
108            .is_some_and(|expected| expected != &event.conversation_id)
109        {
110            return false;
111        }
112        if self
113            .kind
114            .as_ref()
115            .is_some_and(|expected| expected != event.kind.discriminant())
116        {
117            return false;
118        }
119        if self.min_tick.is_some_and(|min_tick| event.tick < min_tick) {
120            return false;
121        }
122        if self.max_tick.is_some_and(|max_tick| event.tick > max_tick) {
123            return false;
124        }
125
126        let fields = event.kind.scalar_fields();
127        if self
128            .kernel_id
129            .as_ref()
130            .is_some_and(|expected| expected != fields.kernel_id)
131        {
132            return false;
133        }
134        if self
135            .tool_name
136            .as_ref()
137            .is_some_and(|expected| expected != fields.tool_name)
138        {
139            return false;
140        }
141        if self
142            .call_id
143            .as_ref()
144            .is_some_and(|expected| expected != fields.call_id)
145        {
146            return false;
147        }
148        if self
149            .skill_id
150            .as_ref()
151            .is_some_and(|expected| expected != fields.skill_id)
152        {
153            return false;
154        }
155        if self
156            .model
157            .as_ref()
158            .is_some_and(|expected| expected != fields.model)
159        {
160            return false;
161        }
162
163        true
164    }
165}
166
167/// Immutable query view over a snapshot of [`ObservabilityEvent`] values.
168///
169/// `EventQuery` is intentionally in-process and host-owned. It is useful for
170/// tests, demos, and small local dashboards; production exporters should keep
171/// using `tracing` sinks and external stores.
172#[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
173pub struct EventQuery {
174    events: Vec<ObservabilityEvent>,
175}
176
177impl EventQuery {
178    /// Build a query view over `events` in their existing order.
179    pub fn new(events: Vec<ObservabilityEvent>) -> Self {
180        Self { events }
181    }
182
183    /// Return all events in snapshot order.
184    pub fn all(&self) -> &[ObservabilityEvent] {
185        &self.events
186    }
187
188    /// Return the number of events in this snapshot.
189    pub fn len(&self) -> usize {
190        self.events.len()
191    }
192
193    /// Return `true` when this snapshot has no events.
194    pub fn is_empty(&self) -> bool {
195        self.events.is_empty()
196    }
197
198    /// Return all events matching `filter` in snapshot order.
199    pub fn filter(&self, filter: &EventFilter) -> Vec<ObservabilityEvent> {
200        self.events
201            .iter()
202            .filter(|event| filter.matches(event))
203            .cloned()
204            .collect()
205    }
206
207    /// Return up to `limit` most recent events in ascending snapshot order.
208    pub fn latest(&self, limit: usize) -> Vec<ObservabilityEvent> {
209        let mut events = self
210            .events
211            .iter()
212            .rev()
213            .take(limit)
214            .cloned()
215            .collect::<Vec<_>>();
216        events.reverse();
217        events
218    }
219
220    /// Count events by wire event kind.
221    pub fn count_by_kind(&self) -> BTreeMap<String, usize> {
222        let mut counts = BTreeMap::new();
223        for event in &self.events {
224            let count = counts
225                .entry(event.kind.discriminant().to_string())
226                .or_insert(0);
227            *count += 1;
228        }
229        counts
230    }
231
232    /// Return conversation identifiers present in this snapshot.
233    pub fn conversations(&self) -> Vec<String> {
234        self.events
235            .iter()
236            .map(|event| event.conversation_id.clone())
237            .collect::<BTreeSet<_>>()
238            .into_iter()
239            .collect()
240    }
241}
242
243impl From<Vec<ObservabilityEvent>> for EventQuery {
244    fn from(events: Vec<ObservabilityEvent>) -> Self {
245        Self::new(events)
246    }
247}
248
249#[cfg(test)]
250#[allow(
251    clippy::unwrap_used,
252    clippy::panic,
253    clippy::indexing_slicing,
254    clippy::expect_used
255)]
256mod tests {
257    use super::*;
258    use crate::event::{EventKind, SCHEMA_VERSION};
259
260    fn event(tick: u64, conversation_id: &str, kind: EventKind) -> ObservabilityEvent {
261        ObservabilityEvent {
262            version: SCHEMA_VERSION,
263            occurred_at_millis: 1_715_000_000_000 + tick,
264            tick,
265            conversation_id: conversation_id.into(),
266            span_id: None,
267            kind,
268        }
269    }
270
271    #[test]
272    fn filter_matches_conversation_kind_and_tick_window() {
273        let query = EventQuery::new(vec![
274            event(
275                1,
276                "a",
277                EventKind::PromptStarted {
278                    model: "m".into(),
279                    messages_in: 1,
280                },
281            ),
282            event(
283                2,
284                "a",
285                EventKind::ToolCompleted {
286                    tool_name: "search".into(),
287                    provider_call_id: None,
288                    call_id: "call-1".into(),
289                    result: "ok".into(),
290                    truncated: false,
291                    duration_ms: None,
292                },
293            ),
294            event(
295                3,
296                "b",
297                EventKind::ToolCompleted {
298                    tool_name: "search".into(),
299                    provider_call_id: None,
300                    call_id: "call-2".into(),
301                    result: "ok".into(),
302                    truncated: false,
303                    duration_ms: None,
304                },
305            ),
306        ]);
307
308        let matches = query.filter(
309            &EventFilter::new()
310                .conversation_id("a")
311                .kind("tool.completed")
312                .min_tick(2)
313                .max_tick(3),
314        );
315
316        assert_eq!(matches.len(), 1);
317        assert_eq!(matches[0].tick, 2);
318    }
319
320    #[test]
321    fn filter_matches_scalar_fields() {
322        let query = EventQuery::new(vec![event(
323            1,
324            "thread",
325            EventKind::ComposeSkillResolved {
326                kernel_id: "kernel".into(),
327                skill_id: "retrieval".into(),
328                applies: true,
329                delta: Some(0.2),
330                confidence: Some(0.8),
331            },
332        )]);
333
334        let matches = query.filter(&EventFilter::new().kernel_id("kernel").skill_id("retrieval"));
335
336        assert_eq!(matches.len(), 1);
337        assert!(
338            query
339                .filter(&EventFilter::new().tool_name("search"))
340                .is_empty()
341        );
342    }
343
344    #[test]
345    fn query_summarizes_conversations_and_kinds() {
346        let query = EventQuery::new(vec![
347            event(
348                1,
349                "b",
350                EventKind::PromptStarted {
351                    model: "m".into(),
352                    messages_in: 1,
353                },
354            ),
355            event(
356                2,
357                "a",
358                EventKind::PromptStarted {
359                    model: "m".into(),
360                    messages_in: 2,
361                },
362            ),
363            event(
364                3,
365                "b",
366                EventKind::ContextSampled {
367                    message_count: 1,
368                    byte_size: 2,
369                    token_estimate: None,
370                },
371            ),
372        ]);
373
374        assert_eq!(query.conversations(), vec!["a", "b"]);
375        assert_eq!(query.count_by_kind().get("prompt.started"), Some(&2));
376        assert_eq!(
377            query
378                .latest(2)
379                .iter()
380                .map(|event| event.tick)
381                .collect::<Vec<_>>(),
382            vec![2, 3]
383        );
384    }
385}
rig_tap/query.rs

rig_tap/
query.rs
