ragrig_gpui/rag/
attach.rs

1//! Context formatting for the Sciscraper co-writer prompt.
2//!
3//! Ragrig's built-in strategies emit plain `[Attached: …]` blocks. Sciscraper's
4//! prompt instead documents two marked containers — `source="creator"` for
5//! papers the author attached and `source="corpus"` for retrieval results.
6//! [`SciscraperAttach`] 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. [`SciscraperAttach::with_context_tokens`]
12//! therefore splits that budget in favour of explicit attachments and marks
13//! every excerpt with the prompt's `<sciscraper: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 Sciscraper's context containers.
25///
26/// ```text
27/// <sciscraper:v1:context source="creator">
28/// <sciscraper:v1:document name="@Kano2025">…full text…</sciscraper:v1:document>
29/// </sciscraper:v1:context>
30///
31/// <sciscraper:v1:context source="corpus">
32/// …retrieved passages…
33/// </sciscraper: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 SciscraperAttach {
42    context_tokens: Option<usize>,
43}
44
45impl SciscraperAttach {
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 SciscraperAttach {
85    fn merge(&self, retrieved: &str, attachments: &[AttachedDocument]) -> String {
86        let mut context = String::new();
87        if !attachments.is_empty() {
88            context.push_str("<sciscraper: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                let limit = if index + 1 == attachments.len() {
93                    remaining
94                } else {
95                    share.min(remaining)
96                };
97                let total = attachment.content.chars().count();
98                let (content, truncated) = truncate(&attachment.content, limit);
99                remaining = remaining.saturating_sub(content.chars().count());
100                context.push_str(&format!(
101                    "<sciscraper:v1:document name=\"{}\">{content}",
102                    attribute(&attachment.name)
103                ));
104                if truncated {
105                    context.push_str(&format!(
106                        "<sciscraper:v1:truncated shown=\"{}\" total=\"{total}\"/>",
107                        content.chars().count()
108                    ));
109                }
110                context.push_str("</sciscraper:v1:document>\n");
111            }
112            context.push_str("</sciscraper:v1:context>");
113        }
114
115        let retrieved = retrieved.trim_end_matches('\n');
116        let (retrieved, _) = truncate(retrieved, self.retrieval_budget(attachments));
117        if !retrieved.trim().is_empty() {
118            if !context.is_empty() {
119                context.push_str("\n\n");
120            }
121            context.push_str("<sciscraper:v1:context source=\"corpus\">\n");
122            context.push_str(&retrieved);
123            context.push_str("\n</sciscraper:v1:context>");
124        }
125        context
126    }
127
128    fn name(&self) -> &'static str {
129        "sciscraper"
130    }
131}
132
133/// Return `text` shortened to at most `max` characters, and whether it was cut.
134fn truncate(text: &str, max: usize) -> (String, bool) {
135    let count = text.chars().count();
136    if count <= max {
137        return (text.to_owned(), false);
138    }
139    (text.chars().take(max).collect(), true)
140}
141
142/// Escape the few characters that would end or corrupt an XML-like attribute.
143fn attribute(value: &str) -> String {
144    value
145        .replace('&', "&amp;")
146        .replace('"', "&quot;")
147        .replace('<', "&lt;")
148}
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153
154    fn document(name: &str, content: &str) -> AttachedDocument {
155        AttachedDocument {
156            name: name.to_owned(),
157            content: content.to_owned(),
158        }
159    }
160
161    #[test]
162    fn attachments_and_retrieval_use_the_documented_containers() {
163        let merged = SciscraperAttach::new().merge(
164            "[doc: Kano2025]\nretrieved passage",
165            &[
166                document("@Kano2025", "PAPER BODY"),
167                document("@Kobayashi2011", "OTHER BODY"),
168            ],
169        );
170
171        assert_eq!(
172            merged,
173            "<sciscraper:v1:context source=\"creator\">\n\
174             <sciscraper:v1:document name=\"@Kano2025\">PAPER BODY</sciscraper:v1:document>\n\
175             <sciscraper:v1:document name=\"@Kobayashi2011\">OTHER BODY</sciscraper:v1:document>\n\
176             </sciscraper:v1:context>\n\n\
177             <sciscraper:v1:context source=\"corpus\">\n\
178             [doc: Kano2025]\nretrieved passage\n\
179             </sciscraper:v1:context>"
180        );
181    }
182
183    #[test]
184    fn empty_parts_do_not_leave_empty_containers() {
185        assert_eq!(SciscraperAttach::new().merge("", &[]), "");
186        assert_eq!(
187            SciscraperAttach::new().merge("chunks", &[]),
188            "<sciscraper:v1:context source=\"corpus\">\nchunks\n</sciscraper:v1:context>"
189        );
190        assert_eq!(
191            SciscraperAttach::new().merge("", &[document("@Key", "BODY")]),
192            "<sciscraper:v1:context source=\"creator\">\n\
193             <sciscraper:v1:document name=\"@Key\">BODY</sciscraper:v1:document>\n\
194             </sciscraper:v1:context>"
195        );
196    }
197
198    #[test]
199    fn attribute_values_are_escaped() {
200        let merged = SciscraperAttach::new().merge("", &[document("a\"<&", "BODY")]);
201        assert!(merged.contains("name=\"a&quot;&lt;&amp;\""));
202    }
203
204    #[test]
205    fn a_token_budget_keeps_the_whole_context_within_the_window() {
206        // 1,024 reserved + 100 context tokens at three characters each.
207        let merged = SciscraperAttach::with_context_tokens(1124)
208            .merge(&"r".repeat(1_000), &[document("@Key", &"a".repeat(1_000))]);
209
210        // 70% of the 300-character budget goes to the paper, the rest to retrieval.
211        assert!(merged.contains(
212            "<sciscraper:v1:truncated shown=\"210\" total=\"1000\"/></sciscraper:v1:document>"
213        ));
214        assert!(merged.contains(&format!("source=\"corpus\">\n{}\n", "r".repeat(90))));
215        let body = merged
216            .split_once("name=\"@Key\">")
217            .expect("document body")
218            .1
219            .split_once("<sciscraper:v1:truncated")
220            .expect("truncation marker")
221            .0;
222        assert_eq!(body.chars().count(), 210);
223    }
224
225    #[test]
226    fn short_papers_do_not_get_a_marker_and_leave_room_for_retrieval() {
227        let merged = SciscraperAttach::with_context_tokens(1124)
228            .merge(&"r".repeat(1_000), &[document("@Key", "SHORT")]);
229        assert!(!merged.contains("<sciscraper:v1:truncated"));
230        assert!(merged.contains(">SHORT</sciscraper:v1:document>"));
231        // The unused attachment share does not starve retrieval.
232        assert!(merged.contains(&"r".repeat(90)));
233    }
234}
ragrig_gpui/rag/attach.rs

ragrig_gpui/rag/
attach.rs
