ragrig_gpui/rag/
attach.rs

1//! Context formatting for the chat prompt.
2//!
3//! Ragrig's built-in strategies emit plain `[Attached: …]` blocks. The chat
4//! prompt instead documents two marked containers — `source="creator"` for
5//! papers the author attached and `source="corpus"` for retrieval results.
6//! [`ScicoAttach`] produces exactly that markup, so the model can tell an
7//! explicitly attached `@Key` paper from a retrieved chunk.
8//!
9//! A paper easily exceeds a local model's window: the default 4096-token
10//! context leaves about 9,000 characters for the whole context block, while a
11//! cited PDF can carry 100,000. [`ScicoAttach::with_context_tokens`]
12//! therefore splits that budget in favour of explicit attachments and marks
13//! every excerpt with the prompt's `<scico:v1:truncated …/>` element.
14
15use ragrig::{AttachStrategy, AttachedDocument};
16
17/// Tokens Ragrig reserves for the system prompt, transcript, and query.
18const RESERVED_PROMPT_TOKENS: usize = 1024;
19/// Ragrig's character-per-token estimate for prompt assembly.
20const CHARS_PER_TOKEN: usize = 3;
21/// Percentage of the context block reserved for explicit attachments.
22const ATTACHMENT_PERCENT: usize = 70;
23
24/// Merge attachments and retrieved chunks into Scico's context containers.
25///
26/// ```text
27/// <scico:v1:context source="creator">
28/// <scico:v1:document name="@Kano2025">…full text…</scico:v1:document>
29/// </scico:v1:context>
30///
31/// <scico:v1:context source="corpus">
32/// …retrieved passages…
33/// </scico:v1:context>
34/// ```
35///
36/// Without a budget the content passes through unchanged. With
37/// [`with_context_tokens`](Self::with_context_tokens), explicit attachments
38/// receive most of the context window and every shortened paper carries a
39/// truncation marker.
40#[derive(Clone, Debug, Default)]
41pub struct ScicoAttach {
42    context_tokens: Option<usize>,
43}
44
45impl ScicoAttach {
46    /// Format without a character budget.
47    pub fn new() -> Self {
48        Self::default()
49    }
50
51    /// Keep the context block within a model window of `tokens`.
52    ///
53    /// Mirrors Ragrig's own retrieval budget: its fixed prompt reserve and
54    /// three characters per token.
55    pub fn with_context_tokens(tokens: usize) -> Self {
56        Self {
57            context_tokens: Some(tokens),
58        }
59    }
60
61    /// The content budget in characters, when one is configured.
62    pub fn budget_chars(&self) -> Option<usize> {
63        self.context_tokens.map(|tokens| {
64            tokens
65                .saturating_sub(RESERVED_PROMPT_TOKENS)
66                .saturating_mul(CHARS_PER_TOKEN)
67        })
68    }
69
70    fn attachment_budget(&self) -> usize {
71        self.budget_chars()
72            .map_or(usize::MAX, |budget| budget * ATTACHMENT_PERCENT / 100)
73    }
74
75    fn retrieval_budget(&self, attachments: &[AttachedDocument]) -> usize {
76        match self.budget_chars() {
77            None => usize::MAX,
78            Some(budget) if attachments.is_empty() => budget,
79            Some(budget) => budget - budget * ATTACHMENT_PERCENT / 100,
80        }
81    }
82}
83
84impl AttachStrategy for ScicoAttach {
85    fn merge(&self, retrieved: &str, attachments: &[AttachedDocument]) -> String {
86        let mut context = String::new();
87        if !attachments.is_empty() {
88            context.push_str("<scico:v1:context source=\"creator\">\n");
89            let mut remaining = self.attachment_budget();
90            let share = remaining / attachments.len();
91            for (index, attachment) in attachments.iter().enumerate() {
92                // A pre-built protocol container (the document buffer, the
93                // bibliography, a marked selection) is inserted verbatim: the
94                // host already attached the attributes the prompt documents
95                // (cursor, range, digest), and wrapping it again would hide
96                // them inside a second `document` element.
97                if is_container(&attachment.content) {
98                    let content = attachment.content.trim_end();
99                    remaining = remaining.saturating_sub(content.chars().count());
100                    context.push_str(content);
101                    context.push('\n');
102                    continue;
103                }
104                let limit = if index + 1 == attachments.len() {
105                    remaining
106                } else {
107                    share.min(remaining)
108                };
109                let total = attachment.content.chars().count();
110                let (content, truncated) = truncate(&attachment.content, limit);
111                remaining = remaining.saturating_sub(content.chars().count());
112                context.push_str(&format!(
113                    "<scico:v1:document name=\"{}\">{content}",
114                    attribute(&attachment.name)
115                ));
116                if truncated {
117                    context.push_str(&format!(
118                        "<scico:v1:truncated shown=\"{}\" total=\"{total}\"/>",
119                        content.chars().count()
120                    ));
121                }
122                context.push_str("</scico:v1:document>\n");
123            }
124            context.push_str("</scico:v1:context>");
125        }
126
127        let retrieved = retrieved.trim_end_matches('\n');
128        let (retrieved, _) = truncate(retrieved, self.retrieval_budget(attachments));
129        if !retrieved.trim().is_empty() {
130            if !context.is_empty() {
131                context.push_str("\n\n");
132            }
133            context.push_str("<scico:v1:context source=\"corpus\">\n");
134            context.push_str(&retrieved);
135            context.push_str("\n</scico:v1:context>");
136        }
137        context
138    }
139
140    fn name(&self) -> &'static str {
141        "scico"
142    }
143}
144
145/// Whether `text` already is one of the protocol's containers.
146///
147/// Hosts build richer containers than [`ScicoAttach`] itself does — the
148/// document buffer adds `cursor`, a selection adds `range` and `digest` — and
149/// hand them in as attachment content. Those must pass through untouched.
150fn is_container(text: &str) -> bool {
151    text.trim_start().starts_with("<scico:v1")
152}
153
154/// Return `text` shortened to at most `max` characters, and whether it was cut.
155fn truncate(text: &str, max: usize) -> (String, bool) {
156    let count = text.chars().count();
157    if count <= max {
158        return (text.to_owned(), false);
159    }
160    (text.chars().take(max).collect(), true)
161}
162
163/// Escape the few characters that would end or corrupt an XML-like attribute.
164fn attribute(value: &str) -> String {
165    value
166        .replace('&', "&amp;")
167        .replace('"', "&quot;")
168        .replace('<', "&lt;")
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    fn document(name: &str, content: &str) -> AttachedDocument {
176        AttachedDocument {
177            name: name.to_owned(),
178            content: content.to_owned(),
179        }
180    }
181
182    #[test]
183    fn attachments_and_retrieval_use_the_documented_containers() {
184        let merged = ScicoAttach::new().merge(
185            "[doc: Kano2025]\nretrieved passage",
186            &[
187                document("@Kano2025", "PAPER BODY"),
188                document("@Kobayashi2011", "OTHER BODY"),
189            ],
190        );
191
192        assert_eq!(
193            merged,
194            "<scico:v1:context source=\"creator\">\n\
195             <scico:v1:document name=\"@Kano2025\">PAPER BODY</scico:v1:document>\n\
196             <scico:v1:document name=\"@Kobayashi2011\">OTHER BODY</scico:v1:document>\n\
197             </scico:v1:context>\n\n\
198             <scico:v1:context source=\"corpus\">\n\
199             [doc: Kano2025]\nretrieved passage\n\
200             </scico:v1:context>"
201        );
202    }
203
204    #[test]
205    fn empty_parts_do_not_leave_empty_containers() {
206        assert_eq!(ScicoAttach::new().merge("", &[]), "");
207        assert_eq!(
208            ScicoAttach::new().merge("chunks", &[]),
209            "<scico:v1:context source=\"corpus\">\nchunks\n</scico:v1:context>"
210        );
211        assert_eq!(
212            ScicoAttach::new().merge("", &[document("@Key", "BODY")]),
213            "<scico:v1:context source=\"creator\">\n\
214             <scico:v1:document name=\"@Key\">BODY</scico:v1:document>\n\
215             </scico:v1:context>"
216        );
217    }
218
219    #[test]
220    fn prebuilt_containers_pass_through_verbatim() {
221        let selection = "<scico:v1:selection name=\"draft.md\" range=\"2..8\" chars=\"6\" digest=\"b3:abc\">\nmarked\n</scico:v1:selection>";
222        let merged = ScicoAttach::new().merge(
223            "",
224            &[document("ignored", selection), document("@Key", "PAPER")],
225        );
226
227        assert!(
228            merged.contains(selection),
229            "the richer container is kept: {merged}"
230        );
231        assert!(
232            !merged.contains("<scico:v1:document name=\"ignored\">"),
233            "a container is not wrapped twice: {merged}"
234        );
235        assert!(merged.contains("<scico:v1:document name=\"@Key\">PAPER</scico:v1:document>"));
236    }
237
238    #[test]
239    fn attribute_values_are_escaped() {
240        let merged = ScicoAttach::new().merge("", &[document("a\"<&", "BODY")]);
241        assert!(merged.contains("name=\"a&quot;&lt;&amp;\""));
242    }
243
244    #[test]
245    fn a_token_budget_keeps_the_whole_context_within_the_window() {
246        // 1,024 reserved + 100 context tokens at three characters each.
247        let merged = ScicoAttach::with_context_tokens(1124)
248            .merge(&"r".repeat(1_000), &[document("@Key", &"a".repeat(1_000))]);
249
250        // 70% of the 300-character budget goes to the paper, the rest to retrieval.
251        assert!(
252            merged
253                .contains("<scico:v1:truncated shown=\"210\" total=\"1000\"/></scico:v1:document>")
254        );
255        assert!(merged.contains(&format!("source=\"corpus\">\n{}\n", "r".repeat(90))));
256        let body = merged
257            .split_once("name=\"@Key\">")
258            .expect("document body")
259            .1
260            .split_once("<scico:v1:truncated")
261            .expect("truncation marker")
262            .0;
263        assert_eq!(body.chars().count(), 210);
264    }
265
266    #[test]
267    fn short_papers_do_not_get_a_marker_and_leave_room_for_retrieval() {
268        let merged = ScicoAttach::with_context_tokens(1124)
269            .merge(&"r".repeat(1_000), &[document("@Key", "SHORT")]);
270        assert!(!merged.contains("<scico:v1:truncated"));
271        assert!(merged.contains(">SHORT</scico:v1:document>"));
272        // The unused attachment share does not starve retrieval.
273        assert!(merged.contains(&"r".repeat(90)));
274    }
275}
ragrig_gpui/rag/attach.rs

ragrig_gpui/rag/
attach.rs
