llama_core/
files.rs

1//! Define APIs for file operations.
2
3use crate::{error::LlamaCoreError, ARCHIVES_DIR};
4use base64::{engine::general_purpose, Engine as _};
5use endpoints::files::{DeleteFileStatus, FileObject, ListFilesResponse};
6use serde_json::{json, Value};
7use std::{
8    fs::{self, File},
9    io::Read,
10    path::Path,
11};
12use walkdir::{DirEntry, WalkDir};
13
14/// Remove the target file by id.
15///
16/// # Arguments
17///
18/// * `id`: The id of the target file.
19///
20/// # Returns
21///
22/// A `DeleteFileStatus` instance.
23pub fn remove_file(id: impl AsRef<str>) -> Result<DeleteFileStatus, LlamaCoreError> {
24    let root = format!("{}/{}", ARCHIVES_DIR, id.as_ref());
25    let status = match fs::remove_dir_all(root) {
26        Ok(_) => {
27            #[cfg(feature = "logging")]
28            info!(target: "stdout", "Successfully deleted the target file with id {}.", id.as_ref());
29
30            DeleteFileStatus {
31                id: id.as_ref().into(),
32                object: "file".to_string(),
33                deleted: true,
34            }
35        }
36        Err(e) => {
37            let err_msg = format!(
38                "Failed to delete the target file with id {}. {}",
39                id.as_ref(),
40                e
41            );
42
43            // log
44            #[cfg(feature = "logging")]
45            error!(target: "stdout", "{}", &err_msg);
46
47            DeleteFileStatus {
48                id: id.as_ref().into(),
49                object: "file".to_string(),
50                deleted: false,
51            }
52        }
53    };
54
55    Ok(status)
56}
57
58/// List all files in the archives directory.
59///
60/// # Returns
61///
62/// A `ListFilesResponse` instance.
63pub fn list_files() -> Result<ListFilesResponse, LlamaCoreError> {
64    #[cfg(feature = "logging")]
65    info!(target: "stdout", "Listing all archive files");
66
67    let mut file_objects: Vec<FileObject> = Vec::new();
68    for entry in WalkDir::new(ARCHIVES_DIR)
69        .into_iter()
70        .filter_map(|e| e.ok())
71    {
72        if !is_hidden(&entry) && entry.path().is_file() {
73            #[cfg(feature = "logging")]
74            info!(target: "stdout", "archive file: {}", entry.path().display());
75
76            let id = entry
77                .path()
78                .parent()
79                .and_then(|p| p.file_name())
80                .unwrap()
81                .to_str()
82                .unwrap()
83                .to_string();
84
85            let filename = entry
86                .path()
87                .file_name()
88                .and_then(|n| n.to_str())
89                .unwrap()
90                .to_string();
91
92            let metadata = entry.path().metadata().unwrap();
93
94            let created_at = metadata
95                .created()
96                .unwrap()
97                .duration_since(std::time::UNIX_EPOCH)
98                .unwrap()
99                .as_secs();
100
101            let bytes = metadata.len();
102
103            let fo = FileObject {
104                id,
105                bytes,
106                created_at,
107                filename,
108                object: "file".to_string(),
109                purpose: "assistants".to_string(),
110            };
111
112            file_objects.push(fo);
113        }
114    }
115
116    #[cfg(feature = "logging")]
117    info!(target: "stdout", "Found {} archive files", file_objects.len());
118
119    let file_objects = ListFilesResponse {
120        object: "list".to_string(),
121        data: file_objects,
122    };
123
124    Ok(file_objects)
125}
126
127/// Retrieve information about a specific file by id.
128///
129/// # Arguments
130///
131/// * `id`: The id of the target file.
132///
133/// # Returns
134///
135/// A `FileObject` instance.
136pub fn retrieve_file(id: impl AsRef<str>) -> Result<FileObject, LlamaCoreError> {
137    #[cfg(feature = "logging")]
138    info!(target: "stdout", "Retrieving the target file with id {}", id.as_ref());
139
140    let root = format!("{}/{}", ARCHIVES_DIR, id.as_ref());
141    for entry in WalkDir::new(root).into_iter().filter_map(|e| e.ok()) {
142        if !is_hidden(&entry) && entry.path().is_file() {
143            #[cfg(feature = "logging")]
144            info!(target: "stdout", "archive file: {}", entry.path().display());
145
146            let filename = entry
147                .path()
148                .file_name()
149                .and_then(|n| n.to_str())
150                .unwrap()
151                .to_string();
152
153            let metadata = entry.path().metadata().unwrap();
154
155            let created_at = metadata
156                .created()
157                .unwrap()
158                .duration_since(std::time::UNIX_EPOCH)
159                .unwrap()
160                .as_secs();
161
162            let bytes = metadata.len();
163
164            return Ok(FileObject {
165                id: id.as_ref().into(),
166                bytes,
167                created_at,
168                filename,
169                object: "file".to_string(),
170                purpose: "assistants".to_string(),
171            });
172        }
173    }
174
175    Err(LlamaCoreError::FileNotFound)
176}
177
178/// Retrieve the content of a specific file by id.
179///
180/// # Arguments
181///
182/// * `id`: The id of the target file.
183///
184/// # Returns
185///
186/// A `Value` instance.
187pub fn retrieve_file_content(id: impl AsRef<str>) -> Result<Value, LlamaCoreError> {
188    #[cfg(feature = "logging")]
189    info!(target: "stdout", "Retrieving the content of the target file with id {}", id.as_ref());
190
191    let file_object = retrieve_file(id)?;
192    let file_path = Path::new(ARCHIVES_DIR)
193        .join(&file_object.id)
194        .join(&file_object.filename);
195
196    let base64_content = file_to_base64(&file_path)?;
197
198    Ok(json!({
199        "id": file_object.id,
200        "bytes": file_object.bytes,
201        "created_at": file_object.created_at,
202        "filename": file_object.filename,
203        "content": base64_content,
204    }))
205}
206
207/// Download a specific file by id.
208///
209/// # Arguments
210///
211/// * `id`: The id of the target file.
212///
213/// # Returns
214///
215/// A tuple of `(String, Vec<u8>)`. The first element is the filename, and the second element is the file content.
216pub fn download_file(id: impl AsRef<str>) -> Result<(String, Vec<u8>), LlamaCoreError> {
217    #[cfg(feature = "logging")]
218    info!(target: "stdout", "Downloading the target file with id {}", id.as_ref());
219
220    let file_object = retrieve_file(id)?;
221    let file_path = Path::new(ARCHIVES_DIR)
222        .join(&file_object.id)
223        .join(&file_object.filename);
224
225    if !file_path.exists() {
226        return Err(LlamaCoreError::FileNotFound);
227    }
228
229    // Open the file
230    let mut file = match File::open(file_path) {
231        Ok(file) => file,
232        Err(e) => {
233            let err_msg = format!("Failed to open the target file. {}", e);
234            return Err(LlamaCoreError::Operation(err_msg));
235        }
236    };
237
238    // read the file content as bytes
239    let mut buffer = Vec::new();
240    match file.read_to_end(&mut buffer) {
241        Ok(_) => Ok((file_object.filename.clone(), buffer)),
242        Err(e) => {
243            let err_msg = format!("Failed to read the content of the target file. {}", e);
244
245            // log
246            #[cfg(feature = "logging")]
247            error!(target: "stdout", "{}", &err_msg);
248
249            Err(LlamaCoreError::Operation(err_msg))
250        }
251    }
252}
253
254fn is_hidden(entry: &DirEntry) -> bool {
255    entry
256        .file_name()
257        .to_str()
258        .map(|s| s.starts_with("."))
259        .unwrap_or(false)
260}
261
262fn file_to_base64(file_path: impl AsRef<Path>) -> Result<String, LlamaCoreError> {
263    if !file_path.as_ref().exists() {
264        return Err(LlamaCoreError::FileNotFound);
265    }
266
267    // Open the file
268    let mut file = match File::open(file_path) {
269        Ok(file) => file,
270        Err(e) => {
271            let err_msg = format!("Failed to open the target file. {}", e);
272            return Err(LlamaCoreError::Operation(err_msg));
273        }
274    };
275
276    // read the file content as bytes
277    let mut buffer = Vec::new();
278    file.read_to_end(&mut buffer).unwrap();
279
280    Ok(general_purpose::STANDARD.encode(&buffer))
281}