rig_tools/
attach.rs

1//! Files a message names as `@path`, read to go with it: an image as an
2//! image, any other file as its text, capped the way `read` shows it, so
3//! the model needs no `read` call for it.
4
5use std::path::{Path, PathBuf};
6
7use base64::Engine;
8use rig_core::completion::Message;
9use rig_core::conversation::Conversation;
10use rig_core::message::{ImageMediaType, UserContent};
11
12use crate::workspace::{Stat, Workspace};
13use crate::{MAX_LINES, show_lines};
14
15/// The largest image attached: providers refuse bigger ones (Anthropic
16/// takes 5 MB per image).
17const MOST_BYTES: u64 = 5 * 1024 * 1024;
18
19/// The most bytes of a text file attached; the model reads on with `read`.
20const MOST_TEXT_BYTES: usize = 32 * 1024;
21
22/// How an attached text file starts, followed by its label and `">`.
23const FILE_OPEN: &str = "<file path=\"";
24
25/// How an attached text file ends.
26const FILE_CLOSE: &str = "</file>";
27
28/// The file extensions looked at as images.
29const EXTENSIONS: [&str; 5] = ["png", "jpg", "jpeg", "gif", "webp"];
30
31/// Whether `path` names an image by its extension.
32pub fn is_image_path(path: &Path) -> bool {
33    path.extension()
34        .and_then(|extension| extension.to_str())
35        .is_some_and(|extension| {
36            EXTENSIONS
37                .iter()
38                .any(|known| extension.eq_ignore_ascii_case(known))
39        })
40}
41
42/// What a message's `@path` references came to.
43#[derive(Clone, Debug, Default, PartialEq)]
44pub struct Attached {
45    /// The files read, each with its label: its path as the message named
46    /// it, `~/` made the home directory.
47    pub files: Vec<(String, UserContent)>,
48    /// A line for the user about each file that is there but could not be
49    /// attached, and why.
50    pub notes: Vec<String>,
51    /// The references that name a path with no file there, as the message
52    /// spelled them.
53    pub missing: Vec<String>,
54}
55
56/// The files `text` names with `@path`, each once, resolved from `base` in
57/// `workspace` and read from it. Every `@` token that names a file there is
58/// attached, whatever its shape (`@notes.md`, `@src/lib.rs`, `@Makefile`,
59/// `@.env`); one with no file there is [missing](Attached::missing) when it
60/// [looks like a path](is_path_reference), and otherwise stays text, such
61/// as `@alice` or `@alice.smith`. A token's trailing punctuation is not
62/// part of its path. A leading `~/` is the user's home on this machine; in
63/// another workspace it is not attached, with a note.
64///
65/// Reads `workspace`: on a thread that must not wait, call it only for a
66/// local one ([`Workspace::host`] is `None`).
67pub fn attachments(text: &str, base: &Path, workspace: &dyn Workspace) -> Attached {
68    let mut attached = Attached::default();
69    let mut seen: Vec<PathBuf> = Vec::new();
70    for token in text.split_whitespace() {
71        let Some(reference) = token.strip_prefix('@') else {
72            continue;
73        };
74        let trimmed = reference.trim_end_matches([',', '.', ';', ':', '!', '?', ')', '"', '\'']);
75        if trimmed.starts_with("~/") && workspace.host().is_some() {
76            attached.notes.push(format!(
77                "{trimmed} is not attached: home paths are this machine's."
78            ));
79            continue;
80        }
81        let found = [reference, trimmed].into_iter().find_map(|named| {
82            let label = expand(named);
83            let path = base.join(&label);
84            matches!(workspace.stat(&path), Ok(Stat::File { .. }))
85                .then(|| (label.display().to_string(), path))
86        });
87        let Some((label, path)) = found else {
88            if is_path_reference(trimmed) && !attached.missing.iter().any(|seen| seen == trimmed) {
89                attached.missing.push(trimmed.to_owned());
90            }
91            continue;
92        };
93        if seen.contains(&path) {
94            continue;
95        }
96        let read = if is_image_path(&path) {
97            read_image(workspace, &path)
98        } else {
99            read_file(workspace, &path, &label)
100        };
101        match read {
102            Ok(content) => attached.files.push((label, content)),
103            Err(why) => attached
104                .notes
105                .push(format!("{label} is not attached: {why}.")),
106        }
107        seen.push(path);
108    }
109    attached
110}
111
112/// Whether `reference`, an `@` token without its `@`, looks like a path by
113/// its shape alone: it holds a `/`, or ends in a dot and one to four
114/// letters or digits. [`attachments`] reports such a token as missing when
115/// no file is there; any other token (a mention, an address, or a file
116/// with no extension) is attached only when its file is there, and is
117/// otherwise left as text.
118///
119/// ```
120/// use rig_tools::attach::is_path_reference;
121///
122/// assert!(is_path_reference("src/main.rs"));
123/// assert!(is_path_reference("notes.md"));
124/// assert!(!is_path_reference("reviewer"));
125/// ```
126pub fn is_path_reference(reference: &str) -> bool {
127    reference.contains('/')
128        || reference.rsplit_once('.').is_some_and(|(name, extension)| {
129            !name.is_empty()
130                && (1..=4).contains(&extension.len())
131                && extension.chars().all(|c| c.is_ascii_alphanumeric())
132        })
133}
134
135/// The label and the contents of an attached text file, when `text` is
136/// one, so a view can show it without the markup around it.
137pub fn attached_file(text: &str) -> Option<(&str, &str)> {
138    let rest = text.strip_prefix(FILE_OPEN)?.strip_suffix(FILE_CLOSE)?;
139    rest.split_once("\">\n")
140}
141
142/// `path` with a leading `~/` made the home directory.
143fn expand(path: &str) -> PathBuf {
144    match (path.strip_prefix("~/"), std::env::home_dir()) {
145        (Some(rest), Some(home)) => home.join(rest),
146        _ => PathBuf::from(path),
147    }
148}
149
150/// The image at `path` in `workspace` as message content, its type told
151/// by its first bytes; why not, for a file that cannot be read, is larger
152/// than a provider takes, or is not a PNG, JPEG, GIF or WebP image.
153pub fn read_image(workspace: &dyn Workspace, path: &Path) -> Result<UserContent, String> {
154    let contents = workspace
155        .read(path, MOST_BYTES)
156        .map_err(|failure| failure.to_string())?;
157    if contents.size > MOST_BYTES {
158        return Err(format!(
159            "it has {} KB, more than the {} KB a provider takes",
160            contents.size / 1024,
161            MOST_BYTES / 1024
162        ));
163    }
164    let bytes = contents.bytes;
165    let media_type = ImageMediaType::sniff(&bytes)
166        .ok_or("it is not a PNG, JPEG, GIF or WebP image".to_owned())?;
167    let data = base64::engine::general_purpose::STANDARD.encode(&bytes);
168    Ok(UserContent::image_base64(data, Some(media_type), None))
169}
170
171/// The text file at `path` as message content: its lines the way `read`
172/// shows them, at most [`MOST_TEXT_BYTES`], with a line saying how
173/// many were left out, between a `<file path="label">` line and `</file>`.
174fn read_file(workspace: &dyn Workspace, path: &Path, label: &str) -> Result<UserContent, String> {
175    let text = workspace
176        .read_text(path)
177        .map_err(|failure| failure.to_string())?;
178    if text.contains('\0') {
179        return Err("it is not a text file".to_owned());
180    }
181    let lines = show_lines(&text, 1, MAX_LINES, MOST_TEXT_BYTES);
182    Ok(UserContent::text(format!(
183        "{FILE_OPEN}{label}\">\n{lines}{FILE_CLOSE}"
184    )))
185}
186
187/// What the user typed in message `at` of `conversation`: the text items
188/// of a message of the user's own
189/// ([`EntryRef::is_users_own`](rig_core::conversation::EntryRef::is_users_own)), never of
190/// one an agent or a plugin delivered, such as a note, nor a file
191/// attached, each trimmed, joined by a blank line; `None` when it has
192/// none.
193pub fn typed(conversation: &Conversation, at: usize) -> Option<String> {
194    let entry = conversation.entries().nth(at)?;
195    let Message::User { content } = entry.message else {
196        return None;
197    };
198    if !entry.is_users_own() {
199        return None;
200    }
201    let typed: Vec<&str> = content
202        .iter()
203        .filter_map(|content| match content {
204            UserContent::Text(text) if attached_file(&text.text).is_none() => {
205                Some(text.text.trim())
206            }
207            _ => None,
208        })
209        .filter(|text| !text.is_empty())
210        .collect();
211    (!typed.is_empty()).then(|| typed.join("\n\n"))
212}
213
214/// What the user typed first in `conversation` ([`typed`]).
215pub fn first_typed(conversation: &Conversation) -> Option<String> {
216    (0..conversation.messages().len()).find_map(|at| typed(conversation, at))
217}
218
219#[cfg(test)]
220mod tests;
rig_tools/attach.rs

rig_tools/
attach.rs
