Skip to main content

agent_diva_core/
attachment.rs

1//! File attachment types for agent-diva
2//!
3//! This module provides the `FileAttachment` struct that represents
4//! a file managed by the agent-diva-files content-addressed storage system.
5//!
6//! ## Overview
7//!
8//! When a file is uploaded through any channel (Telegram, Discord, etc.),
9//! it gets stored in the content-addressed storage. The `FileAttachment`
10//! struct provides a unified view of such stored files across all channels.
11//!
12//! ## Usage
13//!
14//! ```ignore
15//! use agent_diva_core::attachment::FileAttachment;
16//! use agent_diva_files::{FileManager, FileConfig};
17//! use agent_diva_files::handle::FileMetadata;
18//! use std::path::PathBuf;
19//!
20//! // Create a FileHandle first (see agent-diva-files crate)
21//! let config = FileConfig::with_path(PathBuf::from("./data"));
22//! let manager = FileManager::new(config).await?;
23//! let metadata = FileMetadata {
24//!     name: "document.pdf".to_string(),
25//!     size: 1024,
26//!     mime_type: Some("application/pdf".to_string()),
27//!     source: Some("telegram".to_string()),
28//!     created_at: chrono::Utc::now(),
29//!     last_accessed_at: None,
30//!     preview: None,
31//! };
32//! let handle = manager.store(b"dummy content", metadata).await?;
33//! let attachment = FileAttachment::from_handle(handle, "telegram", Some("msg_123"));
34//! ```
35
36use agent_diva_files::handle::FileMetadata;
37use agent_diva_files::FileHandle;
38use chrono::{DateTime, Utc};
39use serde::{Deserialize, Serialize};
40
41/// Unified file attachment representation
42///
43/// This struct wraps a `FileHandle` from agent-diva-files and adds
44/// channel-specific metadata for tracking which channel and message
45/// the file was associated with.
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct FileAttachment {
48    /// Content-addressed file ID (SHA256 hash)
49    pub file_id: String,
50
51    /// Original filename as uploaded
52    pub filename: String,
53
54    /// File size in bytes
55    pub size: u64,
56
57    /// MIME type if known
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub mime_type: Option<String>,
60
61    /// Source channel (e.g., "telegram", "discord", "slack")
62    pub channel: String,
63
64    /// Associated message ID from the channel
65    #[serde(skip_serializing_if = "Option::is_none")]
66    pub message_id: Option<String>,
67
68    /// User who uploaded the file
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub uploaded_by: Option<String>,
71
72    /// When the file was stored
73    pub stored_at: DateTime<Utc>,
74
75    /// Reference count (number of channel associations)
76    pub ref_count: usize,
77}
78
79impl FileAttachment {
80    /// Create a `FileAttachment` from a `FileHandle` and channel info
81    ///
82    /// # Arguments
83    /// * `handle` - The FileHandle from agent-diva-files
84    /// * `channel` - The source channel identifier
85    /// * `message_id` - Optional message ID from the channel
86    ///
87    /// # Example
88    /// ```ignore
89    /// use agent_diva_core::attachment::FileAttachment;
90    /// use agent_diva_files::{FileManager, FileConfig};
91    /// use agent_diva_files::handle::FileMetadata;
92    /// use std::path::PathBuf;
93    ///
94    /// let config = FileConfig::with_path(PathBuf::from("./data"));
95    /// let manager = FileManager::new(config).await?;
96    /// let metadata = FileMetadata {
97    ///     name: "document.pdf".to_string(),
98    ///     size: 1024,
99    ///     mime_type: Some("application/pdf".to_string()),
100    ///     source: Some("telegram".to_string()),
101    ///     created_at: chrono::Utc::now(),
102    ///     last_accessed_at: None,
103    ///     preview: None,
104    /// };
105    /// let handle = manager.store(b"dummy content", metadata).await?;
106    /// let attachment = FileAttachment::from_handle(handle, "telegram", Some("123456"));
107    /// ```
108    pub fn from_handle(handle: FileHandle, channel: &str, message_id: Option<&str>) -> Self {
109        // Get values before consuming the handle
110        let ref_count = handle.ref_count();
111        let file_id = handle.id;
112        let metadata = &handle.metadata;
113
114        Self {
115            file_id,
116            filename: metadata.name.clone(),
117            size: metadata.size,
118            mime_type: metadata.mime_type.clone(),
119            channel: channel.to_string(),
120            message_id: message_id.map(String::from),
121            uploaded_by: metadata.source.clone(),
122            stored_at: metadata.created_at,
123            ref_count,
124        }
125    }
126
127    /// Create a `FileAttachment` from stored metadata
128    ///
129    /// This is useful when reconstructing an attachment from the database
130    /// without needing the full FileHandle.
131    pub fn from_metadata(
132        file_id: &str,
133        metadata: &FileMetadata,
134        channel: &str,
135        message_id: Option<&str>,
136        ref_count: usize,
137    ) -> Self {
138        Self {
139            file_id: file_id.to_string(),
140            filename: metadata.name.clone(),
141            size: metadata.size,
142            mime_type: metadata.mime_type.clone(),
143            channel: channel.to_string(),
144            message_id: message_id.map(String::from),
145            uploaded_by: metadata.source.clone(),
146            stored_at: metadata.created_at,
147            ref_count,
148        }
149    }
150
151    /// Get a display string for the attachment
152    ///
153    /// # Example
154    /// ```ignore
155    /// use agent_diva_core::attachment::FileAttachment;
156    /// use agent_diva_files::{FileManager, FileConfig};
157    /// use agent_diva_files::handle::FileMetadata;
158    /// use std::path::PathBuf;
159    ///
160    /// let config = FileConfig::with_path(PathBuf::from("./data"));
161    /// let manager = FileManager::new(config).await?;
162    /// let metadata = FileMetadata {
163    ///     name: "document.pdf".to_string(),
164    ///     size: 1024 * 1024, // 1 MB
165    ///     mime_type: Some("application/pdf".to_string()),
166    ///     source: Some("telegram".to_string()),
167    ///     created_at: chrono::Utc::now(),
168    ///     last_accessed_at: None,
169    ///     preview: None,
170    /// };
171    /// let handle = manager.store(b"dummy content", metadata).await?;
172    /// let attachment = FileAttachment::from_handle(handle, "telegram", Some("msg_123"));
173    /// let display = attachment.display();
174    /// // "document.pdf (1.0 MB) from telegram"
175    /// ```
176    pub fn display(&self) -> String {
177        let size_str = Self::format_size(self.size);
178        format!("{} ({}) from {}", self.filename, size_str, self.channel)
179    }
180
181    /// Format file size in human-readable format
182    fn format_size(bytes: u64) -> String {
183        const KB: u64 = 1024;
184        const MB: u64 = KB * 1024;
185        const GB: u64 = MB * 1024;
186
187        if bytes >= GB {
188            format!("{:.1} GB", bytes as f64 / GB as f64)
189        } else if bytes >= MB {
190            format!("{:.1} MB", bytes as f64 / MB as f64)
191        } else if bytes >= KB {
192            format!("{:.1} KB", bytes as f64 / KB as f64)
193        } else {
194            format!("{} B", bytes)
195        }
196    }
197
198    /// Check if this attachment is an image
199    pub fn is_image(&self) -> bool {
200        self.mime_type
201            .as_ref()
202            .map(|m| m.starts_with("image/"))
203            .unwrap_or(false)
204    }
205
206    /// Check if this attachment is a video
207    pub fn is_video(&self) -> bool {
208        self.mime_type
209            .as_ref()
210            .map(|m| m.starts_with("video/"))
211            .unwrap_or(false)
212    }
213
214    /// Check if this attachment is audio
215    pub fn is_audio(&self) -> bool {
216        self.mime_type
217            .as_ref()
218            .map(|m| m.starts_with("audio/"))
219            .unwrap_or(false)
220    }
221
222    /// Check if this attachment is a document
223    pub fn is_document(&self) -> bool {
224        self.mime_type
225            .as_ref()
226            .map(|m| {
227                m.starts_with("application/pdf")
228                    || m.starts_with("application/")
229                    || m.starts_with("text/")
230            })
231            .unwrap_or(false)
232    }
233}
234
235impl std::fmt::Display for FileAttachment {
236    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237        write!(f, "{}", self.display())
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use std::path::PathBuf;
245
246    fn create_test_metadata() -> FileMetadata {
247        FileMetadata {
248            name: "test_document.pdf".to_string(),
249            size: 1024 * 1024, // 1 MB
250            mime_type: Some("application/pdf".to_string()),
251            source: Some("telegram".to_string()),
252            created_at: Utc::now(),
253            last_accessed_at: None,
254            preview: None,
255        }
256    }
257
258    fn create_test_handle() -> FileHandle {
259        let metadata = create_test_metadata();
260        FileHandle::new(
261            "sha256:abc123def456".to_string(),
262            PathBuf::from("ab/c123def456"),
263            metadata,
264        )
265    }
266
267    #[test]
268    fn test_from_handle() {
269        let handle = create_test_handle();
270        let attachment = FileAttachment::from_handle(handle, "telegram", Some("msg_789"));
271
272        assert_eq!(attachment.file_id, "sha256:abc123def456");
273        assert_eq!(attachment.filename, "test_document.pdf");
274        assert_eq!(attachment.size, 1024 * 1024);
275        assert_eq!(attachment.mime_type, Some("application/pdf".to_string()));
276        assert_eq!(attachment.channel, "telegram");
277        assert_eq!(attachment.message_id, Some("msg_789".to_string()));
278        assert_eq!(attachment.uploaded_by, Some("telegram".to_string()));
279    }
280
281    #[test]
282    fn test_display() {
283        let handle = create_test_handle();
284        let attachment = FileAttachment::from_handle(handle, "discord", None);
285
286        let display = attachment.display();
287        assert!(display.contains("test_document.pdf"));
288        assert!(display.contains("discord"));
289        assert!(display.contains("MB")); // 1 MB formatted
290    }
291
292    #[test]
293    fn test_is_image() {
294        let mut handle = create_test_handle();
295        handle.metadata.mime_type = Some("image/png".to_string());
296        let attachment = FileAttachment::from_handle(handle, "telegram", None);
297
298        assert!(attachment.is_image());
299        assert!(!attachment.is_video());
300        assert!(!attachment.is_audio());
301        assert!(!attachment.is_document());
302    }
303
304    #[test]
305    fn test_is_video() {
306        let mut handle = create_test_handle();
307        handle.metadata.mime_type = Some("video/mp4".to_string());
308        let attachment = FileAttachment::from_handle(handle, "discord", None);
309
310        assert!(!attachment.is_image());
311        assert!(attachment.is_video());
312    }
313
314    #[test]
315    fn test_format_size() {
316        assert_eq!(FileAttachment::format_size(500), "500 B");
317        assert_eq!(FileAttachment::format_size(1024), "1.0 KB");
318        assert_eq!(FileAttachment::format_size(1024 * 512), "512.0 KB");
319        assert_eq!(FileAttachment::format_size(1024 * 1024), "1.0 MB");
320        assert_eq!(FileAttachment::format_size(1024 * 1024 * 50), "50.0 MB");
321        assert_eq!(FileAttachment::format_size(1024 * 1024 * 1024), "1.0 GB");
322    }
323
324    #[test]
325    fn test_serialize() {
326        let handle = create_test_handle();
327        let attachment = FileAttachment::from_handle(handle, "slack", Some("ts_123"));
328
329        let json = serde_json::to_string(&attachment).unwrap();
330        assert!(json.contains("test_document.pdf"));
331        assert!(json.contains("slack"));
332        assert!(json.contains("sha256:abc123def456"));
333    }
334
335    #[test]
336    fn test_deserialize() {
337        let json = r#"{
338            "file_id": "sha256:test123",
339            "filename": "doc.pdf",
340            "size": 2048,
341            "mime_type": "application/pdf",
342            "channel": "telegram",
343            "message_id": "msg_456",
344            "uploaded_by": "user123",
345            "stored_at": "2024-01-15T10:30:00Z",
346            "ref_count": 3
347        }"#;
348
349        let attachment: FileAttachment = serde_json::from_str(json).unwrap();
350        assert_eq!(attachment.file_id, "sha256:test123");
351        assert_eq!(attachment.filename, "doc.pdf");
352        assert_eq!(attachment.size, 2048);
353        assert_eq!(attachment.channel, "telegram");
354    }
355}