Skip to main content

rain_engine_core/
blob.rs

1use crate::{AttachmentRef, BlobDescriptor, MultimodalPayload};
2use async_trait::async_trait;
3use std::collections::HashMap;
4use std::sync::Arc;
5use thiserror::Error;
6use tokio::sync::RwLock;
7use uuid::Uuid;
8
9#[derive(Debug, Error, Clone, PartialEq, Eq)]
10#[error("{message}")]
11pub struct BlobStoreError {
12    pub message: String,
13}
14
15impl BlobStoreError {
16    pub fn new(message: impl Into<String>) -> Self {
17        Self {
18            message: message.into(),
19        }
20    }
21}
22
23#[async_trait]
24pub trait BlobStore: Send + Sync {
25    async fn put(
26        &self,
27        attachment_id: String,
28        payload: MultimodalPayload,
29    ) -> Result<AttachmentRef, BlobStoreError>;
30
31    async fn get(&self, descriptor: &BlobDescriptor) -> Result<MultimodalPayload, BlobStoreError>;
32}
33
34#[derive(Debug, Default, Clone)]
35pub struct InMemoryBlobStore {
36    values: Arc<RwLock<HashMap<String, MultimodalPayload>>>,
37}
38
39impl InMemoryBlobStore {
40    pub fn new() -> Self {
41        Self::default()
42    }
43}
44
45#[async_trait]
46impl BlobStore for InMemoryBlobStore {
47    async fn put(
48        &self,
49        attachment_id: String,
50        payload: MultimodalPayload,
51    ) -> Result<AttachmentRef, BlobStoreError> {
52        let uri = format!("memory://{}", Uuid::new_v4());
53        self.values
54            .write()
55            .await
56            .insert(uri.clone(), payload.clone());
57        Ok(AttachmentRef::blob(
58            attachment_id,
59            payload.mime_type,
60            payload.file_name,
61            BlobDescriptor {
62                uri,
63                size_bytes: payload.data.len(),
64            },
65        ))
66    }
67
68    async fn get(&self, descriptor: &BlobDescriptor) -> Result<MultimodalPayload, BlobStoreError> {
69        self.values
70            .read()
71            .await
72            .get(&descriptor.uri)
73            .cloned()
74            .ok_or_else(|| BlobStoreError::new("blob not found"))
75    }
76}