Skip to main content

litellm_rs/storage/files/
storage.rs

1//! FileStorage enum implementation with dispatch methods
2
3use crate::config::models::file_storage::FileStorageConfig;
4use crate::utils::error::gateway_error::{GatewayError, Result};
5use tracing::info;
6
7use super::local::LocalStorage;
8use super::s3::S3Storage;
9use super::types::{FileMetadata, FileOwnerScope, FileStorage, StoredFileMetadata};
10
11impl FileStorage {
12    /// Create a new file storage instance
13    pub async fn new(config: &FileStorageConfig) -> Result<Self> {
14        info!("Initializing file storage: {}", config.storage_type);
15
16        match config.storage_type.as_str() {
17            "local" => {
18                let default_path;
19                let path = match config.local_path.as_ref() {
20                    Some(p) => p.as_str(),
21                    None => {
22                        default_path = super::default_data_path();
23                        default_path.to_str().unwrap_or("/tmp/litellm-rs/data")
24                    }
25                };
26                Ok(FileStorage::Local(LocalStorage::new(path).await?))
27            }
28            "s3" => {
29                let s3_config = config.s3.as_ref().ok_or_else(|| {
30                    GatewayError::Config("S3 configuration not specified".to_string())
31                })?;
32                Ok(FileStorage::S3(S3Storage::new(s3_config).await?))
33            }
34            _ => Err(GatewayError::Config(format!(
35                "Unsupported storage type: {}",
36                config.storage_type
37            ))),
38        }
39    }
40
41    /// Store a file and return its ID
42    pub async fn store(&self, filename: &str, content: &[u8]) -> Result<String> {
43        self.store_with_purpose(filename, content, None).await
44    }
45
46    /// Store a file with optional OpenAI purpose metadata and return its ID
47    pub async fn store_with_purpose(
48        &self,
49        filename: &str,
50        content: &[u8],
51        purpose: Option<&str>,
52    ) -> Result<String> {
53        match self {
54            FileStorage::Local(storage) => {
55                storage.store_with_purpose(filename, content, purpose).await
56            }
57            FileStorage::S3(storage) => {
58                storage.store_with_purpose(filename, content, purpose).await
59            }
60        }
61    }
62
63    /// Store a file with a non-optional, canonical owner scope.
64    pub(crate) async fn store_owned_with_purpose(
65        &self,
66        filename: &str,
67        content: &[u8],
68        purpose: Option<&str>,
69        owner: FileOwnerScope,
70    ) -> Result<String> {
71        match self {
72            FileStorage::Local(storage) => {
73                storage
74                    .store_owned_with_purpose(filename, content, purpose, owner)
75                    .await
76            }
77            FileStorage::S3(storage) => {
78                storage
79                    .store_owned_with_purpose(filename, content, purpose, owner)
80                    .await
81            }
82        }
83    }
84
85    /// Retrieve file content by ID
86    pub async fn get(&self, file_id: &str) -> Result<Vec<u8>> {
87        match self {
88            FileStorage::Local(storage) => storage.get(file_id).await,
89            FileStorage::S3(storage) => storage.get(file_id).await,
90        }
91    }
92
93    /// Delete a file by ID
94    pub async fn delete(&self, file_id: &str) -> Result<()> {
95        match self {
96            FileStorage::Local(storage) => storage.delete(file_id).await,
97            FileStorage::S3(storage) => storage.delete(file_id).await,
98        }
99    }
100
101    /// Check if a file exists
102    pub async fn exists(&self, file_id: &str) -> Result<bool> {
103        match self {
104            FileStorage::Local(storage) => storage.exists(file_id).await,
105            FileStorage::S3(storage) => storage.exists(file_id).await,
106        }
107    }
108
109    /// Get file metadata
110    pub async fn metadata(&self, file_id: &str) -> Result<FileMetadata> {
111        match self {
112            FileStorage::Local(storage) => storage.metadata(file_id).await,
113            FileStorage::S3(storage) => storage.metadata(file_id).await,
114        }
115    }
116
117    /// Get public metadata plus its internal owner presence state.
118    pub(crate) async fn metadata_with_owner(&self, file_id: &str) -> Result<StoredFileMetadata> {
119        match self {
120            FileStorage::Local(storage) => storage.metadata_with_owner(file_id).await,
121            FileStorage::S3(storage) => storage.metadata_with_owner(file_id).await,
122        }
123    }
124
125    /// List files with pagination
126    pub async fn list(&self, prefix: Option<&str>, limit: Option<usize>) -> Result<Vec<String>> {
127        match self {
128            FileStorage::Local(storage) => storage.list(prefix, limit).await,
129            FileStorage::S3(storage) => storage.list(prefix, limit).await,
130        }
131    }
132
133    /// Health check
134    pub async fn health_check(&self) -> Result<()> {
135        match self {
136            FileStorage::Local(storage) => storage.health_check().await,
137            FileStorage::S3(storage) => storage.health_check().await,
138        }
139    }
140
141    /// Close storage connections
142    pub async fn close(&self) -> Result<()> {
143        match self {
144            FileStorage::Local(storage) => storage.close().await,
145            FileStorage::S3(storage) => storage.close().await,
146        }
147    }
148}