rig_core/client/
gemini_caching.rs

1//! The [`Caching`] transport and the cache book's I/O: the part of Gemini's
2//! automatic caching that sends requests. What to read, create and retire is
3//! decided by [`CacheBook`] in [`crate::providers::gemini::caching`], from
4//! bytes alone.
5
6use futures::StreamExt;
7use serde::Deserialize;
8
9use crate::driver::{Exchange, Model, Opened, Opening, Transport};
10use crate::error::ProviderError;
11use crate::providers::gemini::cached_content::CachedContents;
12use crate::providers::gemini::caching::{
13    CacheBook, Create, Lease, Parsed, cache_body, digests, is_user_text, parse, short_digest,
14    stripped,
15};
16use crate::providers::gemini::completion::{GenerateContent, ThoughtReplay};
17use crate::providers::internal::wire::classify_untyped_line;
18use crate::wire::{Body, Encoded, Framing, Mode, WireEvent, WireFrame};
19
20impl CacheBook {
21    /// Keep each lease only if Google still has a cache under its name whose
22    /// display name ends with the lease's digest. Returns how many survived.
23    pub async fn prove<T>(&self, caches: &Model<CachedContents, T>) -> usize
24    where
25        T: Transport<CachedContents>,
26    {
27        let mut survived = 0;
28        for lease in self.leases() {
29            let alive = match caches.get(&lease.name).await {
30                Ok(resource) => resource
31                    .display_name
32                    .as_deref()
33                    .is_some_and(|name| name.ends_with(short_digest(&lease.digest))),
34                Err(_) => false,
35            };
36            if alive {
37                survived += 1;
38            } else {
39                self.lost(&lease);
40            }
41        }
42        survived
43    }
44
45    /// Delete every cache the book still holds. A cache already gone (403)
46    /// counts as deleted.
47    pub async fn close<T>(&self, caches: &Model<CachedContents, T>)
48    where
49        T: Transport<CachedContents>,
50    {
51        for lease in self.leases() {
52            let result = caches.delete(&lease.name).await;
53            let gone = matches!(result, Ok(()) | Err(ProviderError::CacheExpired { .. }));
54            if !gone {
55                tracing::warn!(target: "gemini.cache", name = %lease.name, "cache delete failed");
56            }
57            self.retired(&lease);
58        }
59    }
60}
61
62/// A transport that caches GenerateContent requests through `inner`, as
63/// its [`CacheBook`] decides. Build it with [`Model::caching`].
64#[derive(Clone)]
65pub struct Caching<T> {
66    inner: T,
67    config: crate::providers::gemini::GeminiConfig,
68    book: CacheBook,
69    thought_replay: ThoughtReplay,
70}
71
72impl<T> std::fmt::Debug for Caching<T> {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("Caching")
75            .field("book", &self.book)
76            .field("thought_replay", &self.thought_replay)
77            .finish_non_exhaustive()
78    }
79}
80
81impl<T> Model<GenerateContent, T> {
82    /// This model, reading and creating explicit caches as `book` decides.
83    /// See [`crate::providers::gemini::caching`].
84    pub fn caching(self, book: &CacheBook) -> Model<GenerateContent, Caching<T>> {
85        let caching = Caching {
86            inner: self.transport,
87            config: self.wire.provider.clone(),
88            book: book.clone(),
89            thought_replay: self.wire.thought_replay,
90        };
91        Model::new(self.wire, caching)
92    }
93}
94
95fn model_of(path: &str) -> Option<String> {
96    let rest = path.split("/models/").nth(1)?;
97    let (model, verb) = rest.split_once(':')?;
98    let verb = verb.split('?').next()?;
99    matches!(verb, "generateContent" | "streamGenerateContent").then(|| model.to_owned())
100}
101
102/// `payload` with `bytes` as its body.
103fn with_body(payload: &Encoded, bytes: Vec<u8>) -> Result<Encoded, ProviderError> {
104    let mut builder = http::Request::builder()
105        .method(payload.request.method().clone())
106        .uri(payload.request.uri().clone());
107    for (name, value) in payload.request.headers() {
108        builder = builder.header(name, value);
109    }
110    let request = builder
111        .body(Body::Bytes(bytes))
112        .map_err(|error| ProviderError::request(error.to_string()))?;
113    Ok(Encoded {
114        request,
115        framing: payload.framing,
116        request_id_header: payload.request_id_header,
117        relaxed_content_type: payload.relaxed_content_type,
118        route: payload.route,
119        project: payload.project,
120        analysis_only: payload.analysis_only,
121    })
122}
123
124/// A `cachedContents` reply: its status (when it failed) and its body.
125struct Reply {
126    status: Option<u16>,
127    body: String,
128    error: Option<String>,
129}
130
131impl<T> Caching<T>
132where
133    T: Transport<GenerateContent>,
134{
135    /// Send one `cachedContents` request through the inner transport.
136    async fn resource(&self, method: http::Method, path: &str, body: Option<Vec<u8>>) -> Reply {
137        let request = http::Request::builder()
138            .method(method)
139            .uri(self.config.uri(path))
140            .header("Content-Type", "application/json")
141            .body(Body::Bytes(body.unwrap_or_default()));
142        let request = match request {
143            Ok(request) => request,
144            Err(error) => {
145                return Reply {
146                    status: None,
147                    body: String::new(),
148                    error: Some(error.to_string()),
149                };
150            }
151        };
152        let exchange = Exchange {
153            mode: Mode::Unary,
154            observation: None,
155            extensions: http::Extensions::new(),
156        };
157        let opened = match self
158            .inner
159            .send(Encoded::new(request, Framing::Whole), exchange)
160            .await
161        {
162            Ok(opened) => opened,
163            Err(error) => {
164                return Reply {
165                    status: error.provider_response_status().map(|s| s.as_u16()),
166                    body: String::new(),
167                    error: Some(error.to_string()),
168                };
169            }
170        };
171        let mut frames = opened.frames;
172        let mut body = String::new();
173        while let Some(frame) = frames.next().await {
174            match frame {
175                Ok(frame) => body.push_str(&frame.as_str()),
176                Err(error) => {
177                    return Reply {
178                        status: error
179                            .provider_response_status()
180                            .map(|status| status.as_u16()),
181                        body: error
182                            .provider_response_body()
183                            .unwrap_or_default()
184                            .to_owned(),
185                        error: Some(error.to_string()),
186                    };
187                }
188            }
189        }
190        Reply {
191            status: None,
192            body,
193            error: None,
194        }
195    }
196
197    async fn delete(&self, lease: &Lease) {
198        let reply = self
199            .resource(
200                http::Method::DELETE,
201                &format!("/v1beta/{}", lease.name),
202                None,
203            )
204            .await;
205        if reply.error.is_some() && reply.status != Some(403) && reply.status != Some(404) {
206            let status = reply.status.unwrap_or_default();
207            tracing::warn!(target: "gemini.cache", name = %lease.name, status, "cache delete failed");
208        }
209        self.book.retired(lease);
210    }
211
212    async fn extend(&self, lease: &Lease, ttl_secs: u64) {
213        let body = format!("{{\"ttl\":\"{ttl_secs}s\"}}").into_bytes();
214        let reply = self
215            .resource(
216                http::Method::PATCH,
217                &format!("/v1beta/{}?updateMask=ttl", lease.name),
218                Some(body),
219            )
220            .await;
221        if reply.error.is_none() {
222            self.book.extended(lease, ttl_secs);
223        }
224    }
225
226    /// Create the cache `create` describes, unless one already exists for
227    /// its digest. Returns the lease to read.
228    async fn create(
229        &self,
230        model: &str,
231        parsed: &Parsed,
232        create: &Create,
233        line: &str,
234        coverable: u64,
235    ) -> Option<Lease> {
236        let _single = self.book.create.lock().await;
237        if let Some(existing) = self.book.book().leases.get(&create.digest).cloned() {
238            return Some(existing);
239        }
240        let display_name = format!(
241            "{}{}",
242            self.book.display_prefix,
243            short_digest(&create.digest)
244        );
245        let body = cache_body(model, parsed, create.covers, &display_name, create.ttl_secs)?;
246        let reply = self
247            .resource(http::Method::POST, "/v1beta/cachedContents", Some(body))
248            .await;
249        if reply.error.is_some() {
250            #[derive(Deserialize)]
251            struct Envelope {
252                error: ErrorBody,
253            }
254            #[derive(Deserialize)]
255            struct ErrorBody {
256                message: String,
257            }
258            let message = known(classify_untyped_line::<Envelope>(reply.body.as_bytes()))
259                .map(|envelope| envelope.error.message)
260                .or(reply.error)
261                .unwrap_or_default();
262            self.book
263                .create_failed(line, coverable, reply.status, message);
264            return None;
265        }
266        #[derive(Deserialize)]
267        #[serde(rename_all = "camelCase")]
268        struct Resource {
269            name: String,
270            #[serde(default)]
271            usage_metadata: Option<Usage>,
272        }
273        #[derive(Deserialize)]
274        #[serde(rename_all = "camelCase")]
275        struct Usage {
276            #[serde(default)]
277            total_token_count: u64,
278        }
279        let resource = known(classify_untyped_line::<Resource>(reply.body.as_bytes()))?;
280        let tokens = resource
281            .usage_metadata
282            .map_or(0, |usage| usage.total_token_count);
283        Some(
284            self.book
285                .created(line, create, resource.name, tokens, model),
286        )
287    }
288}
289
290/// The decoded value of a classified payload, when it decoded.
291fn known<T>(event: WireEvent<T>) -> Option<T> {
292    match event {
293        WireEvent::Known(value) => Some(value),
294        _ => None,
295    }
296}
297
298/// The final `cachedContentTokenCount` a reply frame reports, when the frame
299/// ends the reply.
300fn final_cached(frame: &WireFrame) -> Option<u64> {
301    #[derive(Deserialize)]
302    #[serde(rename_all = "camelCase")]
303    struct Usage {
304        #[serde(default)]
305        cached_content_token_count: u64,
306    }
307    #[derive(Deserialize)]
308    #[serde(rename_all = "camelCase")]
309    struct Candidate {
310        finish_reason: Option<String>,
311    }
312    #[derive(Deserialize)]
313    #[serde(rename_all = "camelCase")]
314    struct Reply {
315        usage_metadata: Option<Usage>,
316        #[serde(default)]
317        candidates: Vec<Candidate>,
318    }
319    let reply = known(classify_untyped_line::<Reply>(frame.as_str().as_bytes()))?;
320    let finished = reply
321        .candidates
322        .iter()
323        .any(|candidate| candidate.finish_reason.is_some());
324    finished.then_some(reply.usage_metadata?.cached_content_token_count)
325}
326
327impl<T> Transport<GenerateContent> for Caching<T>
328where
329    T: Transport<GenerateContent>,
330{
331    fn send(&self, payload: Encoded, exchange: Exchange) -> Opening<WireFrame> {
332        let this = self.clone();
333        Opening::new(async move {
334            let path = payload.request.uri().path().to_owned();
335            let parsed = match payload.request.body() {
336                Body::Bytes(bytes) => parse(bytes),
337                Body::Multipart(_) => None,
338            };
339            let (Some(model), Some(parsed)) = (model_of(&path), parsed) else {
340                return this.inner.send(payload, exchange).await;
341            };
342            if parsed.has_cached_content {
343                return this.inner.send(payload, exchange).await;
344            }
345
346            let digested = digests(&model, &parsed);
347            let roll_allowed = match this.thought_replay {
348                ThoughtReplay::All => true,
349                ThoughtReplay::CurrentTurn => {
350                    parsed.contents.last().is_some_and(|c| is_user_text(c))
351                }
352            };
353            let plan = this.book.plan(&model, &digested, roll_allowed);
354            for idle in &plan.retire {
355                this.delete(idle).await;
356            }
357            let mut read = plan.read.clone();
358            let mut line = plan.line.clone();
359            if let Some(create) = &plan.create
360                && let Some(lease) = this
361                    .create(&model, &parsed, create, &plan.line, plan.coverable)
362                    .await
363            {
364                if let Some(replaced) = &create.replaces
365                    && replaced.name != lease.name
366                {
367                    this.delete(replaced).await;
368                }
369                if lease.covers > 0 {
370                    line = lease.digest.clone();
371                }
372                read = Some(lease);
373            }
374            if let Some((lease, ttl_secs)) = &plan.extend
375                && read.as_ref().is_some_and(|read| read.name == lease.name)
376            {
377                this.extend(lease, *ttl_secs).await;
378            }
379
380            // A request that reads a cache keeps the encoded one to resend
381            // inline; any other goes out as encoded.
382            let (sent, inline) = match read.as_ref().and_then(|lease| stripped(&parsed, lease)) {
383                Some(body) => (with_body(&payload, body)?, Some(payload)),
384                None => {
385                    read = None;
386                    (payload, None)
387                }
388            };
389            let mut opened: Opened<WireFrame> = this.inner.send(sent, exchange.clone()).await?;
390            if let Some(lease) = read.clone() {
391                let first = opened.frames.next().await;
392                let forbidden = matches!(
393                    &first,
394                    Some(Err(error)) if error.provider_response_status() == Some(http::StatusCode::FORBIDDEN)
395                );
396                if let (true, Some(inline)) = (forbidden, inline) {
397                    // Deleted elsewhere or expired early: forget it and send inline, once.
398                    this.book.lost(&lease);
399                    read = None;
400                    opened = this.inner.send(inline, exchange).await?;
401                } else {
402                    this.book.touched(&lease);
403                    let rest =
404                        std::mem::replace(&mut opened.frames, Box::pin(futures::stream::empty()));
405                    opened.frames = Box::pin(futures::stream::iter(first).chain(rest));
406                }
407            }
408
409            let book = this.book.clone();
410            let read_name = read.map(|lease| lease.name);
411            let coverable = plan.coverable;
412            Ok(opened.map_frames(move |frames| {
413                frames.inspect(move |frame| {
414                    if let Ok(frame) = frame
415                        && let Some(cached) = final_cached(frame)
416                    {
417                        book.observe(&line, read_name.as_deref(), coverable, cached);
418                    }
419                })
420            }))
421        })
422    }
423}
rig_core/client/gemini_caching.rs

rig_core/client/
gemini_caching.rs
