gemini_adk_rs/
tape.rs

1//! Record model and resolver outputs once; replay them without the network.
2//!
3//! [`replay_session`](crate::live::replay::replay_session) re-drives a
4//! recorded Live session through the real control plane, but anything that
5//! calls a model or a remote service out of band still runs for real:
6//!
7//! - an LLM extractor;
8//! - an async resolver behind a slot;
9//! - a background agent.
10//!
11//! A replay would then cost money, need credentials, and could answer
12//! differently. A [`Tape`] closes that gap. While recording, every call's
13//! input and output are appended to it. While replaying, calls are answered
14//! from it, in the order they were recorded, and nothing leaves the process.
15//!
16//! - [`TapedLlm`] wraps any [`BaseLlm`]. [`TapedLlm::recording`] passes calls
17//!   through and records them; [`TapedLlm::replaying`] needs no inner model
18//!   and no credentials.
19//! - [`taped_resolver`] wraps a resolver's `fetch` the same way, for
20//!   `Extract::field_resolve` and `Conversation::resolve_slot`.
21//!
22//! Calls are matched on their canonical JSON input: the request, or the
23//! resolver's name and arguments. The same input recorded twice replays its
24//! two outputs in order. A replay that asks for something the tape does not
25//! hold fails loudly rather than calling out.
26//!
27//! ```
28//! # async fn demo() -> Result<(), Box<dyn std::error::Error>> {
29//! use std::sync::Arc;
30//! use gemini_adk_rs::llm::{BaseLlm, LlmRequest, LlmResponse, MockLlm};
31//! use gemini_adk_rs::tape::{MemoryTape, TapedLlm};
32//!
33//! let tape = Arc::new(MemoryTape::new());
34//!
35//! // Record against a real (here: scripted) model.
36//! let live = TapedLlm::recording(MockLlm::script([LlmResponse::from_text("Paris")]), tape.clone());
37//! live.generate(LlmRequest::from_text("Capital of France?")).await?;
38//!
39//! // Replay with no model at all.
40//! let offline = TapedLlm::replaying("gemini-2.5-flash", tape);
41//! let reply = offline.generate(LlmRequest::from_text("Capital of France?")).await?;
42//! assert_eq!(reply.text(), "Paris");
43//! # Ok(())
44//! # }
45//! ```
46
47use std::collections::{HashMap, VecDeque};
48use std::fs::{File, OpenOptions};
49use std::future::Future;
50use std::io::{BufRead, BufReader, BufWriter, Write};
51use std::path::{Path, PathBuf};
52use std::sync::Arc;
53
54use async_trait::async_trait;
55use parking_lot::Mutex;
56use serde::{Deserialize, Serialize};
57use serde_json::Value;
58
59use crate::llm::{BaseLlm, LlmError, LlmRequest, LlmResponse, ModelCapabilities};
60
61/// The kind of call a [`TapeEntry`] records.
62pub const LLM_CALL: &str = "llm";
63
64/// One recorded call: what went in and what came out.
65#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
66pub struct TapeEntry {
67    /// What kind of call this was: [`LLM_CALL`], or `resolver:{name}`.
68    pub kind: String,
69    /// The call's input as canonical JSON (object keys sorted).
70    pub input: String,
71    /// The call's output, or the error message it failed with.
72    pub output: Result<Value, String>,
73}
74
75/// Where recorded calls are kept.
76pub trait Tape: Send + Sync {
77    /// Append a call.
78    fn record(&self, entry: TapeEntry);
79
80    /// Take the next unreplayed output recorded for `kind` and `input`, in
81    /// recording order. `None` when the tape holds no (more) such call.
82    fn next(&self, kind: &str, input: &str) -> Option<Result<Value, String>>;
83}
84
85/// Recorded calls, queued per `(kind, input)` for replay.
86#[derive(Debug, Default)]
87struct Queues {
88    entries: Vec<TapeEntry>,
89    pending: HashMap<(String, String), VecDeque<Result<Value, String>>>,
90}
91
92impl Queues {
93    fn push(&mut self, entry: TapeEntry) {
94        self.pending
95            .entry((entry.kind.clone(), entry.input.clone()))
96            .or_default()
97            .push_back(entry.output.clone());
98        self.entries.push(entry);
99    }
100
101    fn next(&mut self, kind: &str, input: &str) -> Option<Result<Value, String>> {
102        self.pending
103            .get_mut(&(kind.to_string(), input.to_string()))?
104            .pop_front()
105    }
106}
107
108/// A tape held in memory.
109#[derive(Debug, Default)]
110pub struct MemoryTape {
111    queues: Mutex<Queues>,
112}
113
114impl MemoryTape {
115    /// An empty tape.
116    pub fn new() -> Self {
117        Self::default()
118    }
119
120    /// A tape holding `entries`, ready to replay.
121    pub fn from_entries(entries: impl IntoIterator<Item = TapeEntry>) -> Self {
122        let tape = Self::new();
123        for entry in entries {
124            tape.record(entry);
125        }
126        tape
127    }
128
129    /// Every call recorded so far, in order.
130    pub fn entries(&self) -> Vec<TapeEntry> {
131        self.queues.lock().entries.clone()
132    }
133}
134
135impl Tape for MemoryTape {
136    fn record(&self, entry: TapeEntry) {
137        self.queues.lock().push(entry);
138    }
139
140    fn next(&self, kind: &str, input: &str) -> Option<Result<Value, String>> {
141        self.queues.lock().next(kind, input)
142    }
143}
144
145/// A tape kept in a JSONL file, one [`TapeEntry`] per line.
146///
147/// [`FileTape::create`] starts an empty file to record into.
148/// [`FileTape::open`] loads an existing one for replay. Any calls recorded
149/// after that are appended to the same file.
150#[derive(Debug)]
151pub struct FileTape {
152    path: PathBuf,
153    queues: Mutex<Queues>,
154    writer: Mutex<BufWriter<File>>,
155}
156
157impl FileTape {
158    /// Create (or truncate) `path` and record into it.
159    pub fn create(path: impl AsRef<Path>) -> std::io::Result<Self> {
160        let path = path.as_ref().to_path_buf();
161        let file = File::create(&path)?;
162        Ok(Self {
163            path,
164            queues: Mutex::new(Queues::default()),
165            writer: Mutex::new(BufWriter::new(file)),
166        })
167    }
168
169    /// Load the calls recorded in `path`, ready to replay.
170    pub fn open(path: impl AsRef<Path>) -> std::io::Result<Self> {
171        let path = path.as_ref().to_path_buf();
172        let mut queues = Queues::default();
173        for (n, line) in BufReader::new(File::open(&path)?).lines().enumerate() {
174            let line = line?;
175            if line.trim().is_empty() {
176                continue;
177            }
178            let entry: TapeEntry = serde_json::from_str(&line).map_err(|e| {
179                std::io::Error::new(
180                    std::io::ErrorKind::InvalidData,
181                    format!("{}:{}: {e}", path.display(), n + 1),
182                )
183            })?;
184            queues.push(entry);
185        }
186        let file = OpenOptions::new().append(true).open(&path)?;
187        Ok(Self {
188            path,
189            queues: Mutex::new(queues),
190            writer: Mutex::new(BufWriter::new(file)),
191        })
192    }
193
194    /// The file this tape reads and writes.
195    pub fn path(&self) -> &Path {
196        &self.path
197    }
198}
199
200impl Tape for FileTape {
201    fn record(&self, entry: TapeEntry) {
202        match serde_json::to_string(&entry) {
203            Ok(line) => {
204                let mut writer = self.writer.lock();
205                if let Err(e) = writeln!(writer, "{line}").and_then(|()| writer.flush()) {
206                    tracing::warn!(path = %self.path.display(), "tape write failed: {e}");
207                }
208            }
209            Err(e) => tracing::warn!("tape entry not serializable: {e}"),
210        }
211        self.queues.lock().push(entry);
212    }
213
214    fn next(&self, kind: &str, input: &str) -> Option<Result<Value, String>> {
215        self.queues.lock().next(kind, input)
216    }
217}
218
219/// Canonical JSON for `value`: object keys sorted, no whitespace.
220///
221/// Values that fail to serialize fall back to their `Debug` form, so a key is
222/// always produced.
223pub fn canonical_json(value: &impl Serialize) -> String {
224    serde_json::to_value(value)
225        .map(|v| v.to_string())
226        .unwrap_or_else(|e| format!("<unserializable: {e}>"))
227}
228
229enum Mode<L> {
230    Recording(L),
231    Replaying { model_id: String },
232}
233
234/// A [`BaseLlm`] that records its calls to a [`Tape`] or answers them from
235/// one. See the [module docs](self).
236pub struct TapedLlm<L = Arc<dyn BaseLlm>> {
237    mode: Mode<L>,
238    tape: Arc<dyn Tape>,
239}
240
241impl<L: BaseLlm> TapedLlm<L> {
242    /// Pass every call through to `inner` and record it on `tape`.
243    pub fn recording(inner: L, tape: Arc<dyn Tape>) -> Self {
244        Self {
245            mode: Mode::Recording(inner),
246            tape,
247        }
248    }
249}
250
251impl TapedLlm {
252    /// Answer every call from `tape`. No model, network or credential is
253    /// used; a call the tape does not hold fails with [`LlmError::Config`].
254    pub fn replaying(model_id: impl Into<String>, tape: Arc<dyn Tape>) -> Self {
255        Self {
256            mode: Mode::Replaying {
257                model_id: model_id.into(),
258            },
259            tape,
260        }
261    }
262}
263
264#[async_trait]
265impl<L: BaseLlm> BaseLlm for TapedLlm<L> {
266    fn model_id(&self) -> &str {
267        match &self.mode {
268            Mode::Recording(inner) => inner.model_id(),
269            Mode::Replaying { model_id } => model_id,
270        }
271    }
272
273    fn capabilities(&self) -> ModelCapabilities {
274        match &self.mode {
275            Mode::Recording(inner) => inner.capabilities(),
276            Mode::Replaying { model_id } => ModelCapabilities::infer_from_id(model_id),
277        }
278    }
279
280    async fn generate(&self, request: LlmRequest) -> Result<LlmResponse, LlmError> {
281        let input = canonical_json(&request);
282        match &self.mode {
283            Mode::Recording(inner) => {
284                let result = inner.generate(request).await;
285                let output = match &result {
286                    Ok(response) => {
287                        serde_json::to_value(response).map_err(|e| format!("unserializable: {e}"))
288                    }
289                    Err(e) => Err(e.to_string()),
290                };
291                self.tape.record(TapeEntry {
292                    kind: LLM_CALL.into(),
293                    input,
294                    output,
295                });
296                result
297            }
298            Mode::Replaying { .. } => match self.tape.next(LLM_CALL, &input) {
299                Some(Ok(value)) => serde_json::from_value(value)
300                    .map_err(|e| LlmError::Other(format!("taped response unreadable: {e}"))),
301                Some(Err(message)) => Err(LlmError::Other(message)),
302                None => Err(LlmError::Config(format!(
303                    "the replay tape holds no (more) recorded call for this request: {}",
304                    truncate(&input, 200)
305                ))),
306            },
307        }
308    }
309
310    async fn warm_up(&self) -> Result<(), LlmError> {
311        match &self.mode {
312            Mode::Recording(inner) => inner.warm_up().await,
313            Mode::Replaying { .. } => Ok(()),
314        }
315    }
316}
317
318/// Whether a [`taped_resolver`] records calls or answers them from the tape.
319#[derive(Debug, Clone, Copy, PartialEq, Eq)]
320pub enum TapeMode {
321    /// Call through and record.
322    Record,
323    /// Answer from the tape; never call through.
324    Replay,
325}
326
327/// Wrap a resolver's `fetch` so its calls are recorded on, or answered from,
328/// `tape`.
329///
330/// `name` scopes the recording (use the field or slot name), so two
331/// resolvers given the same arguments do not answer for each other. The
332/// result has the shape `Extract::field_resolve` and
333/// `Conversation::resolve_slot` take.
334pub fn taped_resolver<F, Fut>(
335    tape: Arc<dyn Tape>,
336    mode: TapeMode,
337    name: impl Into<String>,
338    fetch: F,
339) -> impl Fn(Value) -> std::pin::Pin<Box<dyn Future<Output = Result<Value, String>> + Send>>
340+ Send
341+ Sync
342+ 'static
343where
344    F: Fn(Value) -> Fut + Send + Sync + 'static,
345    Fut: Future<Output = Result<Value, String>> + Send + 'static,
346{
347    let kind = format!("resolver:{}", name.into());
348    let fetch = Arc::new(fetch);
349    move |args: Value| {
350        let tape = tape.clone();
351        let kind = kind.clone();
352        let fetch = fetch.clone();
353        Box::pin(async move {
354            let input = canonical_json(&args);
355            match mode {
356                TapeMode::Record => {
357                    let output = fetch(args).await;
358                    tape.record(TapeEntry {
359                        kind,
360                        input,
361                        output: output.clone(),
362                    });
363                    output
364                }
365                TapeMode::Replay => tape.next(&kind, &input).unwrap_or_else(|| {
366                    Err(format!(
367                        "the replay tape holds no (more) recorded {kind} call for {}",
368                        truncate(&input, 200)
369                    ))
370                }),
371            }
372        })
373    }
374}
375
376fn truncate(s: &str, max: usize) -> &str {
377    match s.char_indices().nth(max) {
378        Some((i, _)) => &s[..i],
379        None => s,
380    }
381}
382
383#[cfg(test)]
384mod tests {
385    use super::*;
386    use crate::llm::MockLlm;
387    use serde_json::json;
388
389    #[tokio::test]
390    async fn a_replay_answers_in_recorded_order_without_the_model() {
391        let tape = Arc::new(MemoryTape::new());
392        let live = TapedLlm::recording(
393            MockLlm::script([LlmResponse::from_text("one"), LlmResponse::from_text("two")]),
394            tape.clone(),
395        );
396        let ask = || LlmRequest::from_text("again?");
397        assert_eq!(live.generate(ask()).await.unwrap().text(), "one");
398        assert_eq!(live.generate(ask()).await.unwrap().text(), "two");
399
400        let offline = TapedLlm::replaying("m", Arc::new(MemoryTape::from_entries(tape.entries())));
401        assert_eq!(offline.generate(ask()).await.unwrap().text(), "one");
402        assert_eq!(offline.generate(ask()).await.unwrap().text(), "two");
403        let exhausted = offline.generate(ask()).await.unwrap_err();
404        assert!(matches!(exhausted, LlmError::Config(_)), "{exhausted}");
405    }
406
407    #[tokio::test]
408    async fn a_replay_refuses_a_request_it_never_saw() {
409        let offline = TapedLlm::replaying("m", Arc::new(MemoryTape::new()));
410        let err = offline
411            .generate(LlmRequest::from_text("unseen"))
412            .await
413            .unwrap_err();
414        assert!(err.to_string().contains("no (more) recorded call"), "{err}");
415    }
416
417    #[tokio::test]
418    async fn a_recorded_failure_replays_as_a_failure() {
419        let tape = Arc::new(MemoryTape::new());
420        let failing = TapedLlm::recording(
421            MockLlm::from_fn(|_| Err(LlmError::RateLimited)),
422            tape.clone(),
423        );
424        assert!(failing.generate(LlmRequest::from_text("x")).await.is_err());
425
426        let offline = TapedLlm::replaying("m", tape);
427        let err = offline
428            .generate(LlmRequest::from_text("x"))
429            .await
430            .unwrap_err();
431        assert_eq!(err.to_string(), "Rate limited");
432    }
433
434    #[tokio::test]
435    async fn a_resolver_replays_per_name_and_arguments() {
436        let tape: Arc<dyn Tape> = Arc::new(MemoryTape::new());
437        let record = taped_resolver(
438            tape.clone(),
439            TapeMode::Record,
440            "balance",
441            |args| async move { Ok(json!({ "for": args["account"], "cents": 1200 })) },
442        );
443        let recorded = record(json!({ "account": "A1" })).await.unwrap();
444
445        let replay = taped_resolver(tape.clone(), TapeMode::Replay, "balance", |_| async {
446            Err::<Value, _>("must not be called".to_string())
447        });
448        assert_eq!(replay(json!({ "account": "A1" })).await.unwrap(), recorded);
449        assert!(replay(json!({ "account": "B2" })).await.is_err());
450
451        let other = taped_resolver(tape, TapeMode::Replay, "limit", |_| async {
452            Err::<Value, _>("must not be called".to_string())
453        });
454        assert!(other(json!({ "account": "A1" })).await.is_err());
455    }
456
457    #[test]
458    fn canonical_json_sorts_keys() {
459        assert_eq!(
460            canonical_json(&json!({ "b": 1, "a": { "d": 2, "c": 3 } })),
461            r#"{"a":{"c":3,"d":2},"b":1}"#
462        );
463    }
464
465    #[test]
466    fn a_file_tape_round_trips() {
467        let path = std::env::temp_dir().join(format!(
468            "gemini-adk-tape-{}-{}.jsonl",
469            std::process::id(),
470            std::time::SystemTime::now()
471                .duration_since(std::time::UNIX_EPOCH)
472                .unwrap()
473                .as_nanos()
474        ));
475        {
476            let tape = FileTape::create(&path).unwrap();
477            tape.record(TapeEntry {
478                kind: "resolver:x".into(),
479                input: "{}".into(),
480                output: Ok(json!(7)),
481            });
482        }
483        let tape = FileTape::open(&path).unwrap();
484        assert_eq!(tape.next("resolver:x", "{}"), Some(Ok(json!(7))));
485        assert_eq!(tape.next("resolver:x", "{}"), None);
486        std::fs::remove_file(path).unwrap();
487    }
488}
gemini_adk_rs/tape.rs

gemini_adk_rs/
tape.rs
