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}