attachments.rs

crates/cortex-core/src/attachments.rs

1. //! Attachments. Images are downscaled before they are stored, because a vision
2. //! model bills by image size; text files are read into the prompt as a fenced
3. //! block. Files are stored by content hash, so the same image twice is one file.
5. use std::path::{Path, PathBuf};
7. use base64::Engine;
8. use image::GenericImageView;
9. use sha2::{Digest, Sha256};
11. /// Long edge, in pixels. Larger images cost more and rarely read better.
12. pub const MAX_IMAGE_EDGE: u32 = 1568;
13. /// Text files beyond this are refused rather than silently cut.
14. pub const MAX_TEXT_BYTES: usize = 256 * 1024;
15. /// Images beyond this are refused before decoding.
16. pub const MAX_IMAGE_BYTES: usize = 40 * 1024 * 1024;
18. #[derive(Debug, thiserror::Error, PartialEq)]
19. pub enum AttachError {
20. #[error("{0} is larger than {1} KB")]
21. TooLarge(String, usize),
22. #[error("{0} is not a readable image")]
23. BadImage(String),
24. #[error("{0} is not UTF-8 text")]
25. NotText(String),
26. #[error("could not save the attachment: {0}")]
27. Io(String),
28. }
30. #[derive(Debug)]
31. pub struct ProcessedImage {
32. pub bytes: Vec<u8>,
33. pub mime: &'static str,
34. pub ext: &'static str,
35. pub width: u32,
36. pub height: u32,
37. pub sha256: String,
38. }
40. pub fn sha256_hex(bytes: &[u8]) -> String {
41. Sha256::digest(bytes).iter().map(|b| format!("{b:02x}")).collect()
42. }
44. /// The image types Cortex accepts, by extension.
45. pub const IMAGE_EXTENSIONS: [&str; 5] = ["png", "jpg", "jpeg", "gif", "webp"];
47. pub fn is_image_name(name: &str) -> bool {
48. let n = name.to_lowercase();
49. n.rsplit_once('.').is_some_and(|(_, ext)| IMAGE_EXTENSIONS.contains(&ext))
50. }
52. /// Decode, shrink to MAX_IMAGE_EDGE, re-encode: PNG when it has transparency, else JPEG.
53. pub fn process_image(filename: &str, bytes: &[u8]) -> Result<ProcessedImage, AttachError> {
54. if bytes.len() > MAX_IMAGE_BYTES {
55. return Err(AttachError::TooLarge(filename.into(), MAX_IMAGE_BYTES / 1024));
56. }
57. let img = image::load_from_memory(bytes).map_err(|_| AttachError::BadImage(filename.into()))?;
58. let (w, h) = img.dimensions();
59. let img = if w.max(h) > MAX_IMAGE_EDGE {
60. img.resize(MAX_IMAGE_EDGE, MAX_IMAGE_EDGE, image::imageops::FilterType::Lanczos3)
61. } else {
62. img
63. };
64. let (width, height) = img.dimensions();
65. let mut out = std::io::Cursor::new(Vec::new());
66. let (mime, ext) = if img.color().has_alpha() {
67. img.write_to(&mut out, image::ImageFormat::Png).map_err(|e| AttachError::Io(e.to_string()))?;
68. ("image/png", "png")
69. } else {
70. let rgb = image::DynamicImage::ImageRgb8(img.to_rgb8());
71. let enc = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut out, 85);
72. rgb.write_with_encoder(enc).map_err(|e| AttachError::Io(e.to_string()))?;
73. ("image/jpeg", "jpg")
74. };
75. let bytes = out.into_inner();
76. Ok(ProcessedImage { sha256: sha256_hex(&bytes), bytes, mime, ext, width, height })
77. }
79. /// Save under `dir/<sha256>.<ext>`; returns the path.
80. pub fn store(dir: &Path, img: &ProcessedImage) -> Result<PathBuf, AttachError> {
81. std::fs::create_dir_all(dir).map_err(|e| AttachError::Io(e.to_string()))?;
82. let path = dir.join(format!("{}.{}", img.sha256, img.ext));
83. if !path.exists() {
84. std::fs::write(&path, &img.bytes).map_err(|e| AttachError::Io(e.to_string()))?;
85. }
86. Ok(path)
87. }
89. pub fn read_text(filename: &str, bytes: &[u8]) -> Result<String, AttachError> {
90. if bytes.len() > MAX_TEXT_BYTES {
91. return Err(AttachError::TooLarge(filename.into(), MAX_TEXT_BYTES / 1024));
92. }
93. String::from_utf8(bytes.to_vec()).map_err(|_| AttachError::NotText(filename.into()))
94. }
96. /// How a text file travels in the prompt.
97. pub fn text_block(filename: &str, text: &str) -> String {
98. let fence = if text.contains("```") { "~~~~" } else { "```" };
99. format!("File: {filename}\n{fence}\n{}\n{fence}", text.trim_end())
100. }
102. pub fn data_url(mime: &str, bytes: &[u8]) -> String {
103. format!("data:{mime};base64,{}", base64::engine::general_purpose::STANDARD.encode(bytes))
104. }
106. #[cfg(test)]
107. mod tests {
108. use super::*;
110. fn png(w: u32, h: u32, alpha: bool) -> Vec<u8> {
111. let img = if alpha {
112. image::DynamicImage::ImageRgba8(image::RgbaImage::from_pixel(w, h, image::Rgba([10, 20, 30, 128])))
113. } else {
114. image::DynamicImage::ImageRgb8(image::RgbImage::from_pixel(w, h, image::Rgb([200, 100, 50])))
115. };
116. let mut out = std::io::Cursor::new(Vec::new());
117. img.write_to(&mut out, image::ImageFormat::Png).unwrap();
118. out.into_inner()
119. }
121. #[test]
122. fn large_images_shrink_and_keep_aspect() {
123. let p = process_image("big.png", &png(3136, 1000, false)).unwrap();
124. assert_eq!((p.width, p.height), (1568, 500));
125. assert_eq!(p.mime, "image/jpeg");
126. let small = process_image("s.png", &png(300, 200, true)).unwrap();
127. assert_eq!((small.width, small.height, small.mime), (300, 200, "image/png"));
128. }
130. #[test]
131. fn stored_by_hash_once() {
132. let dir = tempfile::tempdir().unwrap();
133. let p = process_image("a.png", &png(10, 10, false)).unwrap();
134. let a = store(dir.path(), &p).unwrap();
135. let b = store(dir.path(), &p).unwrap();
136. assert_eq!(a, b);
137. assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 1);
138. }
140. #[test]
141. fn rejects_bad_input() {
142. assert_eq!(process_image("x.png", b"nope").unwrap_err(), AttachError::BadImage("x.png".into()));
143. assert!(matches!(read_text("big.txt", &vec![b'a'; MAX_TEXT_BYTES + 1]), Err(AttachError::TooLarge(..))));
144. assert_eq!(read_text("bin", &[0xff, 0xfe]).unwrap_err(), AttachError::NotText("bin".into()));
145. }
147. #[test]
148. fn text_blocks_and_data_urls() {
149. assert_eq!(text_block("a.md", "hi\n"), "File: a.md\n```\nhi\n```");
150. assert!(text_block("b.md", "```rust\n```").starts_with("File: b.md\n~~~~\n"));
151. assert_eq!(data_url("image/png", b"hi"), "data:image/png;base64,aGk=");
152. assert!(is_image_name("Photo.JPEG") && !is_image_name("notes.txt"));
153. }
154. }