Skip to main content

systemprompt_cli/commands/core/files/
upload.rs

1use std::path::{Path, PathBuf};
2
3use anyhow::{Result, anyhow};
4use base64::Engine;
5use base64::engine::general_purpose::STANDARD;
6use clap::Args;
7use sha2::{Digest, Sha256};
8use systemprompt_files::{FileUploadRequest, FileUploadService, FilesConfig};
9use systemprompt_identifiers::{ContextId, SessionId, UserId};
10use systemprompt_runtime::AppContext;
11use tokio::fs;
12
13use super::types::FileUploadOutput;
14use crate::CliConfig;
15use crate::shared::CommandResult;
16
17#[derive(Debug, Clone, Args)]
18pub struct UploadArgs {
19    #[arg(help = "Path to file to upload")]
20    pub file_path: PathBuf,
21
22    #[arg(long, help = "Context ID (required)")]
23    pub context: String,
24
25    #[arg(long, help = "User ID")]
26    pub user: Option<String>,
27
28    #[arg(long, help = "Session ID")]
29    pub session: Option<String>,
30
31    #[arg(long, help = "Mark as AI-generated content")]
32    pub ai: bool,
33}
34
35pub async fn execute(
36    args: UploadArgs,
37    _config: &CliConfig,
38) -> Result<CommandResult<FileUploadOutput>> {
39    let ctx = AppContext::new().await?;
40    let files_config = FilesConfig::get()?;
41    let service = FileUploadService::new(ctx.db_pool(), files_config.clone())?;
42
43    if !service.is_enabled() {
44        return Err(anyhow!("File uploads are disabled in configuration"));
45    }
46
47    let file_path = args
48        .file_path
49        .canonicalize()
50        .map_err(|e| anyhow!("File not found: {} - {}", args.file_path.display(), e))?;
51
52    let bytes = fs::read(&file_path).await?;
53    let bytes_base64 = STANDARD.encode(&bytes);
54    let digest = Sha256::digest(&bytes);
55    let checksum_sha256 = digest.iter().fold(String::with_capacity(64), |mut acc, b| {
56        acc.push_str(&format!("{b:02x}"));
57        acc
58    });
59    let size_bytes = bytes.len() as i64;
60
61    let mime_type = detect_mime_type(&file_path);
62    let filename = file_path
63        .file_name()
64        .and_then(|n| n.to_str())
65        .map(String::from);
66
67    let context_id = ContextId::new(args.context);
68
69    let request = FileUploadRequest {
70        name: filename,
71        mime_type: mime_type.clone(),
72        bytes_base64,
73        context_id,
74        user_id: args.user.map(UserId::new),
75        session_id: args.session.map(SessionId::new),
76        trace_id: None,
77    };
78
79    let result = service.upload_file(request).await?;
80
81    let output = FileUploadOutput {
82        file_id: result.file_id,
83        path: result.path,
84        public_url: result.public_url,
85        size_bytes,
86        mime_type,
87        checksum_sha256,
88    };
89
90    Ok(CommandResult::card(output).with_title("File Uploaded"))
91}
92
93const EXTENSION_MIME_TABLE: &[(&[&str], &str)] = &[
94    (&["jpg", "jpeg"], "image/jpeg"),
95    (&["png"], "image/png"),
96    (&["gif"], "image/gif"),
97    (&["webp"], "image/webp"),
98    (&["svg"], "image/svg+xml"),
99    (&["bmp"], "image/bmp"),
100    (&["tiff", "tif"], "image/tiff"),
101    (&["ico"], "image/x-icon"),
102    (&["pdf"], "application/pdf"),
103    (&["doc"], "application/msword"),
104    (
105        &["docx"],
106        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
107    ),
108    (&["xls"], "application/vnd.ms-excel"),
109    (
110        &["xlsx"],
111        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
112    ),
113    (&["ppt"], "application/vnd.ms-powerpoint"),
114    (
115        &["pptx"],
116        "application/vnd.openxmlformats-officedocument.presentationml.presentation",
117    ),
118    (&["txt"], "text/plain"),
119    (&["csv"], "text/csv"),
120    (&["md"], "text/markdown"),
121    (&["html", "htm"], "text/html"),
122    (&["json"], "application/json"),
123    (&["xml"], "application/xml"),
124    (&["rtf"], "application/rtf"),
125    (&["mp3"], "audio/mpeg"),
126    (&["wav"], "audio/wav"),
127    (&["ogg"], "audio/ogg"),
128    (&["aac"], "audio/aac"),
129    (&["flac"], "audio/flac"),
130    (&["m4a"], "audio/mp4"),
131    (&["mp4"], "video/mp4"),
132    (&["webm"], "video/webm"),
133    (&["mov"], "video/quicktime"),
134    (&["avi"], "video/x-msvideo"),
135    (&["mkv"], "video/x-matroska"),
136];
137
138pub fn detect_mime_type(path: &Path) -> String {
139    let extension = path
140        .extension()
141        .and_then(|e| e.to_str())
142        .map(str::to_lowercase);
143    let Some(ext) = extension.as_deref() else {
144        return "application/octet-stream".to_string();
145    };
146    EXTENSION_MIME_TABLE
147        .iter()
148        .find(|(exts, _)| exts.contains(&ext))
149        .map_or("application/octet-stream", |(_, mime)| *mime)
150        .to_string()
151}