rig_tap/
observed_memory.rs

1//! [`ObservedMemory`]: a [`rig::memory::ConversationMemory`] decorator that
2//! emits a [`EventKind::ContextSampled`](crate::EventKind::ContextSampled)
3//! event on every [`load`](rig::memory::ConversationMemory::load).
4//!
5//! Use it to feed the active-context-size curve of a funnel-style observability
6//! UI: every time the agent loads history before a turn, the decorator measures
7//! the size of the loaded conversation and emits a sample.
8
9use rig::completion::Message;
10use rig::memory::{ConversationMemory, MemoryError};
11use rig::wasm_compat::WasmBoxedFuture;
12
13use crate::emit::emit_kind;
14use crate::event::EventKind;
15
16/// Wraps any [`ConversationMemory`] and emits a `context.sampled` event on
17/// every `load`. `append` and `clear` pass through unchanged.
18///
19/// # Example
20///
21/// ```no_run
22/// use rig::memory::InMemoryConversationMemory;
23/// use rig_tap::ObservedMemory;
24///
25/// let inner = InMemoryConversationMemory::new();
26/// let observed = ObservedMemory::new(inner);
27/// // pass `observed` to `agent.memory(observed)`
28/// ```
29pub struct ObservedMemory<M> {
30    inner: M,
31}
32
33impl<M> ObservedMemory<M> {
34    /// Wrap `inner` so its loads emit `context.sampled` events.
35    pub fn new(inner: M) -> Self {
36        Self { inner }
37    }
38
39    /// Return a reference to the wrapped memory.
40    pub fn inner(&self) -> &M {
41        &self.inner
42    }
43
44    /// Consume the decorator and return the wrapped memory.
45    pub fn into_inner(self) -> M {
46        self.inner
47    }
48}
49
50impl<M> std::fmt::Debug for ObservedMemory<M>
51where
52    M: std::fmt::Debug,
53{
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        f.debug_struct("ObservedMemory")
56            .field("inner", &self.inner)
57            .finish()
58    }
59}
60
61impl<M> Clone for ObservedMemory<M>
62where
63    M: Clone,
64{
65    fn clone(&self) -> Self {
66        Self {
67            inner: self.inner.clone(),
68        }
69    }
70}
71
72impl<M> ConversationMemory for ObservedMemory<M>
73where
74    M: ConversationMemory,
75{
76    fn load<'a>(
77        &'a self,
78        conversation_id: &'a str,
79    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
80        Box::pin(async move {
81            let messages = self.inner.load(conversation_id).await?;
82            let message_count = messages.len();
83            let byte_size = approx_json_size(&messages);
84            emit_kind(
85                conversation_id,
86                EventKind::ContextSampled {
87                    message_count,
88                    byte_size,
89                    token_estimate: None,
90                },
91            );
92            Ok(messages)
93        })
94    }
95
96    fn append<'a>(
97        &'a self,
98        conversation_id: &'a str,
99        messages: Vec<Message>,
100    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
101        self.inner.append(conversation_id, messages)
102    }
103
104    fn clear<'a>(
105        &'a self,
106        conversation_id: &'a str,
107    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
108        self.inner.clear(conversation_id)
109    }
110}
111
112/// Counts bytes that a streaming JSON serializer would emit, without ever
113/// materializing the JSON. Implements [`std::io::Write`] so we can pass it
114/// to [`serde_json::to_writer`].
115#[derive(Default)]
116struct CountingWriter {
117    bytes: usize,
118}
119
120impl std::io::Write for CountingWriter {
121    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
122        self.bytes = self.bytes.saturating_add(buf.len());
123        Ok(buf.len())
124    }
125    fn flush(&mut self) -> std::io::Result<()> {
126        Ok(())
127    }
128}
129
130/// Approximate JSON byte size of `messages`. Cheap: serializes through a
131/// counter-only writer, no allocations of the encoded payload. Returns `0`
132/// on the (unreachable in practice) case where `Message` serialization
133/// fails.
134fn approx_json_size(messages: &[Message]) -> usize {
135    let mut writer = CountingWriter::default();
136    if serde_json::to_writer(&mut writer, messages).is_err() {
137        return 0;
138    }
139    writer.bytes
140}
rig_tap/observed_memory.rs

rig_tap/
observed_memory.rs
