agentai_core/
storage.rs

1//! Storage for media: images, audio and other files, kept under an ID.
2//!
3//! [`MediaSave`] stores a [`Media`] and returns its ID, [`MediaRead`] returns it by that ID.
4//! A type with both is a [`MediaStorage`]. Implement them for your blob store, database or
5//! directory; [`MemoryStorage`] keeps the media in memory.
6
7use crate::message::Media;
8use async_trait::async_trait;
9use std::collections::HashMap;
10use std::sync::{Arc, Mutex};
11use thiserror::Error;
12
13/// Errors of a media storage.
14#[derive(Debug, Error)]
15#[non_exhaustive]
16pub enum StorageError {
17    /// Nothing is stored under this ID.
18    #[error("no media with id '{0}'")]
19    NotFound(String),
20    /// Any other failure, for example of the disk or the network.
21    #[error(transparent)]
22    Other(Box<dyn std::error::Error + Send + Sync>),
23}
24
25impl StorageError {
26    /// Creates a [`StorageError::Other`] from any error.
27    pub fn other(err: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> Self {
28        Self::Other(err.into())
29    }
30}
31
32/// Reads media by ID.
33#[async_trait]
34pub trait MediaRead: Send + Sync + 'static {
35    /// Returns the media stored under `id`, or [`StorageError::NotFound`].
36    async fn read(&self, id: &str) -> Result<Media, StorageError>;
37}
38
39/// Saves media and chooses its ID.
40#[async_trait]
41pub trait MediaSave: Send + Sync + 'static {
42    /// Stores `media` and returns the ID to read it with. The storage chooses the ID; it may
43    /// use [`Media::name`].
44    async fn save(&self, media: &Media) -> Result<String, StorageError>;
45}
46
47/// A storage that can save and read media. Implemented for every type with both traits.
48pub trait MediaStorage: MediaRead + MediaSave {}
49
50impl<T: MediaRead + MediaSave> MediaStorage for T {}
51
52#[async_trait]
53impl<T: MediaRead + ?Sized> MediaRead for Arc<T> {
54    async fn read(&self, id: &str) -> Result<Media, StorageError> {
55        (**self).read(id).await
56    }
57}
58
59#[async_trait]
60impl<T: MediaSave + ?Sized> MediaSave for Arc<T> {
61    async fn save(&self, media: &Media) -> Result<String, StorageError> {
62        (**self).save(media).await
63    }
64}
65
66/// [`MediaStorage`] in memory. The ID is the media's name with a number in front
67/// (`1-cat.png`), or `media-1` for media without a name.
68///
69/// Clones share the same storage.
70#[derive(Debug, Clone, Default)]
71pub struct MemoryStorage {
72    inner: Arc<Mutex<MemoryInner>>,
73}
74
75#[derive(Debug, Default)]
76struct MemoryInner {
77    count: usize,
78    media: HashMap<String, Media>,
79}
80
81impl MemoryStorage {
82    /// Creates an empty storage.
83    pub fn new() -> Self {
84        Self::default()
85    }
86
87    /// Returns the IDs of the stored media.
88    pub fn ids(&self) -> Vec<String> {
89        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
90        inner.media.keys().cloned().collect()
91    }
92}
93
94#[async_trait]
95impl MediaRead for MemoryStorage {
96    async fn read(&self, id: &str) -> Result<Media, StorageError> {
97        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
98        inner
99            .media
100            .get(id)
101            .cloned()
102            .ok_or_else(|| StorageError::NotFound(id.to_string()))
103    }
104}
105
106#[async_trait]
107impl MediaSave for MemoryStorage {
108    async fn save(&self, media: &Media) -> Result<String, StorageError> {
109        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
110        inner.count += 1;
111        let id = match &media.name {
112            Some(name) => format!("{}-{name}", inner.count),
113            None => format!("media-{}", inner.count),
114        };
115        inner.media.insert(id.clone(), media.clone());
116        Ok(id)
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[tokio::test]
125    async fn memory_storage_reads_back_what_was_saved_under_distinct_ids() {
126        let storage = MemoryStorage::new();
127        let cat = Media::from_bytes("image/png", b"cat").with_name("a.png");
128        let dog = Media::from_bytes("image/png", b"dog").with_name("a.png");
129
130        let cat_id = storage.save(&cat).await.unwrap();
131        let dog_id = storage.save(&dog).await.unwrap();
132        let unnamed = storage
133            .save(&Media::from_bytes("image/png", b"x"))
134            .await
135            .unwrap();
136
137        assert_ne!(cat_id, dog_id);
138        assert_eq!(unnamed, "media-3");
139        assert_eq!(storage.read(&cat_id).await.unwrap(), cat);
140        assert_eq!(storage.read(&dog_id).await.unwrap(), dog);
141        assert!(matches!(
142            storage.read("missing").await,
143            Err(StorageError::NotFound(id)) if id == "missing"
144        ));
145    }
146}
agentai_core/storage.rs

agentai_core/
storage.rs
