systemprompt_ai/services/storage/

image_storage.rs

//! Storage for generated images behind the [`FileStorage`] backend.
//!
//! [`ImageStorage`] persists decoded image bytes under a configured base path
//! relative to the storage root, optionally sharding by capture date, and
//! returns both the storage id and the public URL. [`StorageConfig`] holds
//! the base path, URL prefix, size cap, and date-organisation flag and
//! validates them before any write.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use crate::error::AiError;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64;
use chrono::{Datelike, Utc};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
use systemprompt_traits::{FileStorage, ImageStorageConfig, StoredFileId};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageConfig {
    pub base_path: PathBuf,
    pub url_prefix: String,
    #[serde(default = "default_max_file_size")]
    pub max_file_size_bytes: usize,
    #[serde(default = "default_organize_by_date")]
    pub organize_by_date: bool,
}

const fn default_max_file_size() -> usize {
    10 * 1024 * 1024
}

const fn default_organize_by_date() -> bool {
    true
}

impl StorageConfig {
    pub fn from_image_storage_config(config: ImageStorageConfig) -> Self {
        Self {
            base_path: config.base_path,
            url_prefix: config.url_prefix,
            max_file_size_bytes: default_max_file_size(),
            organize_by_date: true,
        }
    }
}

impl StorageConfig {
    pub const fn new(base_path: PathBuf, url_prefix: String) -> Self {
        Self {
            base_path,
            url_prefix,
            max_file_size_bytes: default_max_file_size(),
            organize_by_date: default_organize_by_date(),
        }
    }

    pub const fn validate(&self) -> Result<(), StorageConfigError> {
        if self.url_prefix.is_empty() {
            return Err(StorageConfigError::EmptyUrlPrefix);
        }

        if self.max_file_size_bytes == 0 {
            return Err(StorageConfigError::ZeroMaxFileSize);
        }

        Ok(())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum StorageConfigError {
    #[error("url_prefix cannot be empty")]
    EmptyUrlPrefix,

    #[error("max_file_size_bytes must be greater than 0")]
    ZeroMaxFileSize,
}

/// A decoded image written to storage: its key, public URL and the size of
/// the bytes actually stored (not the base64 transport length).
#[derive(Debug, Clone)]
pub struct StoredImage {
    pub id: StoredFileId,
    pub public_url: String,
    pub size_bytes: usize,
}

pub struct ImageStorage {
    config: StorageConfig,
    storage: Arc<dyn FileStorage>,
}

impl std::fmt::Debug for ImageStorage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ImageStorage")
            .field("config", &self.config)
            .finish_non_exhaustive()
    }
}

impl ImageStorage {
    pub fn new(config: StorageConfig, storage: Arc<dyn FileStorage>) -> Result<Self, AiError> {
        config.validate()?;

        Ok(Self { config, storage })
    }

    pub async fn save_base64_image(
        &self,
        base64_data: &str,
        mime_type: &str,
    ) -> Result<StoredImage, AiError> {
        let image_bytes = BASE64.decode(base64_data)?;

        let (id, public_url) = self.save_image_bytes(&image_bytes, mime_type).await?;
        Ok(StoredImage {
            id,
            public_url,
            size_bytes: image_bytes.len(),
        })
    }

    pub async fn save_image_bytes(
        &self,
        image_bytes: &[u8],
        mime_type: &str,
    ) -> Result<(StoredFileId, String), AiError> {
        if image_bytes.len() > self.config.max_file_size_bytes {
            return Err(AiError::ImageTooLarge {
                size: image_bytes.len(),
                max: self.config.max_file_size_bytes,
            });
        }

        let extension = Self::mime_type_to_extension(mime_type);
        let filename = format!(
            "{}_{}.{}",
            Uuid::new_v4(),
            Utc::now().timestamp(),
            extension
        );

        let partition = self.date_partition();
        let relative_path = self.config.base_path.join(&partition).join(&filename);
        let url_path = format!(
            "{}/{}{}",
            self.config.url_prefix,
            if partition.is_empty() {
                String::new()
            } else {
                format!("{partition}/")
            },
            filename
        );

        let stored_id = self
            .storage
            .store(&relative_path, image_bytes)
            .await
            .map_err(|source| AiError::Storage {
                context: format!("failed to write image file {}", relative_path.display()),
                source,
            })?;

        Ok((stored_id, url_path))
    }

    pub async fn delete_image(&self, id: &StoredFileId) -> Result<(), AiError> {
        self.storage
            .delete(id)
            .await
            .map_err(|source| AiError::Storage {
                context: format!("failed to delete file {id}"),
                source,
            })
    }

    pub async fn exists(&self, id: &StoredFileId) -> Result<bool, AiError> {
        self.storage
            .exists(id)
            .await
            .map_err(|source| AiError::Storage {
                context: format!("failed to stat file {id}"),
                source,
            })
    }

    pub fn get_full_path(&self, relative_path: &str) -> PathBuf {
        self.config.base_path.join(relative_path)
    }

    fn date_partition(&self) -> String {
        if !self.config.organize_by_date {
            return String::new();
        }
        let now = Utc::now();
        format!("{:04}/{:02}/{:02}", now.year(), now.month(), now.day())
    }

    fn mime_type_to_extension(mime_type: &str) -> String {
        match mime_type {
            "image/jpeg" | "image/jpg" => "jpg",
            "image/webp" => "webp",
            "image/gif" => "gif",
            _ => "png",
        }
        .to_owned()
    }
}