Skip to main content

agent_diva_manager/
file_service.rs

1use agent_diva_core::attachment::FileAttachment;
2use agent_diva_files::handle::FileMetadata;
3use agent_diva_files::FileManager;
4use std::sync::Arc;
5
6/// File service for handling file uploads and downloads
7///
8/// Uses a shared Arc<FileManager> to ensure data consistency
9/// across all components. Do NOT create multiple FileManager instances.
10pub struct FileService {
11    manager: Arc<FileManager>,
12}
13
14impl FileService {
15    /// Create a new file service with the given file manager
16    ///
17    /// # Arguments
18    /// * `manager` - Shared FileManager instance
19    ///
20    /// # Example
21    /// ```ignore
22    /// use agent_diva_files::{FileManager, FileConfig};
23    /// use agent_diva_manager::file_service::FileService;
24    /// use std::sync::Arc;
25    ///
26    /// let config = FileConfig::default();
27    /// let manager = Arc::new(FileManager::new(config).await?);
28    /// let file_service = FileService::new(manager);
29    /// ```
30    pub fn new(manager: Arc<FileManager>) -> Self {
31        Self { manager }
32    }
33
34    /// Upload a file and return a FileAttachment
35    ///
36    /// # Arguments
37    /// * `file_name` - Original file name
38    /// * `bytes` - File content
39    /// * `channel` - Source channel (e.g., "ui", "telegram")
40    /// * `message_id` - Optional message ID for association
41    pub async fn upload_file(
42        &self,
43        file_name: &str,
44        bytes: Vec<u8>,
45        channel: &str,
46        message_id: Option<&str>,
47    ) -> anyhow::Result<FileAttachment> {
48        let mime_type = mime_guess::from_path(file_name)
49            .first()
50            .map(|m| m.to_string());
51
52        let metadata = FileMetadata {
53            name: file_name.to_string(),
54            size: bytes.len() as u64,
55            mime_type,
56            source: Some(channel.to_string()),
57            created_at: chrono::Utc::now(),
58            last_accessed_at: None,
59            preview: None,
60        };
61
62        let handle = self.manager.store(&bytes, metadata).await?;
63
64        let attachment = FileAttachment::from_handle(handle, channel, message_id);
65
66        Ok(attachment)
67    }
68
69    /// Read file content by ID
70    ///
71    /// # Arguments
72    /// * `file_id` - File ID (SHA256 hash)
73    pub async fn read_file(&self, file_id: &str) -> anyhow::Result<Vec<u8>> {
74        let handle = self.manager.get(file_id).await?;
75        let content = self.manager.read(&handle).await?;
76
77        Ok(content)
78    }
79
80    /// Get the underlying FileManager
81    pub fn manager(&self) -> &FileManager {
82        &self.manager
83    }
84
85    /// Get a clone of the Arc<FileManager>
86    pub fn manager_arc(&self) -> Arc<FileManager> {
87        self.manager.clone()
88    }
89}
90
91impl Clone for FileService {
92    fn clone(&self) -> Self {
93        Self {
94            manager: self.manager.clone(),
95        }
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use agent_diva_files::FileConfig;
103    use tempfile::TempDir;
104
105    async fn create_test_service() -> (FileService, TempDir) {
106        let temp_dir = TempDir::new().unwrap();
107        let config = FileConfig::with_path(temp_dir.path().to_path_buf());
108        let manager = Arc::new(FileManager::new(config).await.unwrap());
109        let service = FileService::new(manager);
110        (service, temp_dir)
111    }
112
113    #[tokio::test]
114    async fn test_upload_and_read() {
115        let (service, _temp) = create_test_service().await;
116
117        let content = b"hello world";
118        let attachment = service
119            .upload_file("test.txt", content.to_vec(), "test", None)
120            .await
121            .unwrap();
122
123        assert!(attachment.file_id.starts_with("sha256:"));
124        assert_eq!(attachment.file_name, "test.txt");
125
126        // Read back
127        let read_content = service.read_file(&attachment.file_id).await.unwrap();
128        assert_eq!(read_content, content);
129    }
130
131    #[tokio::test]
132    async fn test_clone() {
133        let (service, _temp) = create_test_service().await;
134        let cloned = service.clone();
135
136        // Both should share the same manager
137        let content = b"test";
138        let attachment = service
139            .upload_file("test.txt", content.to_vec(), "test", None)
140            .await
141            .unwrap();
142
143        // Clone should be able to read
144        let read_content = cloned.read_file(&attachment.file_id).await.unwrap();
145        assert_eq!(read_content, content);
146    }
147}