rig_memory/
blobs.rs

1//! Images kept as blobs: a saved message names each image's blob,
2//! `blob:<sha256>.<ext>`, in place of its data, and the store keeps the
3//! blob as the file `blobs/<sha256>.<ext>` ([`BLOBS`]). A message is saved
4//! in a push or a splice.
5
6use std::borrow::Cow;
7use std::io;
8use std::sync::Arc;
9
10use base64::Engine;
11use base64::engine::general_purpose::STANDARD;
12use rig_core::completion::Message;
13use rig_core::conversation::Change;
14use rig_core::message::DocumentSourceKind::{Base64, Raw, Unknown, Url};
15use rig_core::message::{Image, ImageMediaType};
16use sha2::{Digest, Sha256};
17
18use crate::Files;
19
20/// How a saved message names an image stored as a blob, in place of its
21/// data: `blob:<sha256>.<ext>`.
22const BLOB: &str = "blob:";
23
24/// Where a store keeps the blobs its messages name: files under `blobs/`,
25/// each named by a hash of its content.
26pub const BLOBS: &str = "blobs/";
27
28/// Stores the blob `name` in `files`: a name is a hash of the content, so
29/// writing it again changes nothing.
30pub(crate) fn put_blob(files: &(impl Files + ?Sized), name: &str, bytes: &[u8]) -> io::Result<()> {
31    files.write_file(&format!("{BLOBS}{name}"), bytes)
32}
33
34/// The blob `name` of `files`; [`io::ErrorKind::NotFound`] when there is
35/// none.
36pub(crate) fn read_blob(files: &(impl Files + ?Sized), name: &str) -> io::Result<Vec<u8>> {
37    files.read_file(&format!("{BLOBS}{name}"))
38}
39
40/// The names of the blobs `files` keeps.
41pub(crate) fn blob_names(files: &(impl Files + ?Sized)) -> io::Result<Vec<String>> {
42    let names = files.files()?.into_iter();
43    Ok(names
44        .filter_map(|name| Some(name.strip_prefix(BLOBS)?.to_owned()))
45        .collect())
46}
47
48fn inline(image: &Image) -> bool {
49    matches!(image.data, Base64(_) | Raw(_))
50}
51
52/// `message` as it is saved: the data of each image in it stored in
53/// `store` once per content, and named there instead. Borrowed when it has
54/// no image data to store; data that is not valid base64 stays inline.
55///
56/// # Errors
57///
58/// When the store cannot store a blob.
59pub fn store_images<'m>(message: &'m Message, store: &dyn Files) -> io::Result<Cow<'m, Message>> {
60    store_images_with(message, &mut |name, bytes| put_blob(store, &name, &bytes))
61}
62
63/// Where a save puts each blob it stores, by name.
64pub(crate) type PutBlob<'a> = dyn FnMut(String, Vec<u8>) -> io::Result<()> + 'a;
65
66/// [`store_images`], each blob given to `put`.
67pub(crate) fn store_images_with<'m>(
68    message: &'m Message,
69    put: &mut PutBlob<'_>,
70) -> io::Result<Cow<'m, Message>> {
71    if !message.images().any(inline) {
72        return Ok(Cow::Borrowed(message));
73    }
74    let mut message = message.clone();
75    for image in message.images_mut() {
76        let bytes = match &image.data {
77            Base64(data) => STANDARD.decode(data).ok(),
78            Raw(bytes) => Some(bytes.clone()),
79            _ => None,
80        };
81        if let Some(bytes) = bytes {
82            let extension = image.media_type.as_ref();
83            let extension = extension.map_or("bin", ImageMediaType::extension);
84            let name = format!("{:x}.{extension}", Sha256::digest(&bytes));
85            image.data = Url(format!("{BLOB}{name}"));
86            put(name, bytes)?;
87        }
88    }
89    Ok(Cow::Owned(message))
90}
91
92/// Puts the data of each image `message` names in `store` back in place.
93/// An image whose blob cannot be read has no data any more ([`Unknown`]),
94/// which a request sends as a placeholder; the first such failure is
95/// returned.
96///
97/// # Errors
98///
99/// The first blob that could not be read.
100pub fn load_images(message: &mut Message, store: &dyn Files) -> io::Result<()> {
101    load_images_with(message, &|name| read_blob(store, name))
102}
103
104/// [`load_images`] from the blobs `blob` reads.
105pub(crate) fn load_images_with(
106    message: &mut Message,
107    blob: &dyn Fn(&str) -> io::Result<Vec<u8>>,
108) -> io::Result<()> {
109    let mut gone = Ok(());
110    for image in message.images_mut() {
111        if let Url(url) = &image.data
112            && let Some(name) = url.strip_prefix(BLOB)
113        {
114            image.data = match blob(name) {
115                Ok(bytes) => Base64(STANDARD.encode(bytes)),
116                Err(failure) => {
117                    gone = gone.and(Err(failure));
118                    Unknown
119                }
120            };
121        }
122    }
123    gone
124}
125
126/// Adds the name of every blob `bytes` names (`blob:<name>`) to `named`.
127pub(crate) fn named_blobs(bytes: &[u8], named: &mut std::collections::BTreeSet<String>) {
128    let name_byte = |byte: &u8| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-');
129    let mut rest = bytes;
130    while let Some(at) = memchr::memmem::find(rest, BLOB.as_bytes()) {
131        let after = rest.get(at + BLOB.len()..).unwrap_or_default();
132        let len = after.iter().take_while(|byte| name_byte(byte)).count();
133        if let Some(name) = after
134            .get(..len)
135            .and_then(|name| std::str::from_utf8(name).ok())
136            && !name.is_empty()
137        {
138            named.insert(name.to_owned());
139        }
140        rest = after.get(len..).unwrap_or_default();
141    }
142}
143
144/// Calls `f` on each message `change` saves: the entries of a push or a
145/// splice.
146pub(crate) fn messages_mut(
147    change: &mut Change,
148    f: &mut dyn FnMut(&mut Message) -> io::Result<()>,
149) -> io::Result<()> {
150    match change {
151        Change::Push { entries } | Change::Splice { with: entries, .. } => entries
152            .iter_mut()
153            .try_for_each(|entry| f(Arc::make_mut(&mut entry.message))),
154        _ => Ok(()),
155    }
156}
157
158/// `change` as it is saved, its messages' images stored as blobs
159/// ([`store_images`], [`messages_mut`]); `None` when it has none to store.
160/// The whole value of a [`Change::Reset`] keeps its images inline.
161pub(crate) fn store_change(change: &Change, put: &mut PutBlob<'_>) -> io::Result<Option<Change>> {
162    let has_inline = match change {
163        Change::Push { entries } | Change::Splice { with: entries, .. } => entries
164            .iter()
165            .any(|entry| entry.message.images().any(inline)),
166        _ => false,
167    };
168    if !has_inline {
169        return Ok(None);
170    }
171    let mut change = change.clone();
172    messages_mut(&mut change, &mut |message| {
173        if let Cow::Owned(stored) = store_images_with(message, put)? {
174            *message = stored;
175        }
176        Ok(())
177    })?;
178    Ok(Some(change))
179}
rig_memory/blobs.rs

rig_memory/
blobs.rs
