research_agent/adapters/
arxiv_source.rs

1use async_trait::async_trait;
2use quick_xml::Reader;
3use quick_xml::events::Event;
4
5use crate::adapters::semantic_scholar_source::percent_encode;
6use crate::domain::paper::Paper;
7use crate::error::{ResearchError, Result};
8use crate::ports::paper_source::PaperSource;
9
10pub struct ArxivSource {
11    client: reqwest::Client,
12}
13
14impl ArxivSource {
15    pub fn new() -> Self {
16        Self {
17            client: reqwest::Client::new(),
18        }
19    }
20
21    /// Download a paper's PDF bytes from `https://arxiv.org/pdf/<id>`.
22    /// arXiv serves the versionless id at the latest version.
23    pub async fn download_pdf(&self, arxiv_id: &str) -> Result<Vec<u8>> {
24        let resp = self
25            .client
26            .get(format!("https://arxiv.org/pdf/{arxiv_id}"))
27            .send()
28            .await
29            .map_err(|e| ResearchError::Source(format!("arXiv PDF request failed: {e}")))?;
30        if !resp.status().is_success() {
31            let status = resp.status();
32            return Err(ResearchError::Source(format!(
33                "arXiv PDF returned HTTP {status}"
34            )));
35        }
36        let bytes = resp
37            .bytes()
38            .await
39            .map_err(|e| ResearchError::Source(format!("arXiv PDF read failed: {e}")))?;
40        Ok(bytes.to_vec())
41    }
42}
43
44impl Default for ArxivSource {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50#[async_trait]
51impl PaperSource for ArxivSource {
52    async fn fetch_papers(&self, query: &str, limit: usize) -> Result<Vec<Paper>> {
53        let encoded = percent_encode(query);
54        let url = format!(
55            "https://export.arxiv.org/api/query?search_query=all:{}&max_results={}",
56            encoded, limit
57        );
58        let resp = self
59            .client
60            .get(&url)
61            .send()
62            .await
63            .map_err(|e| ResearchError::Source(format!("arXiv request failed: {e}")))?;
64
65        if !resp.status().is_success() {
66            let status = resp.status();
67            return Err(ResearchError::Source(format!(
68                "arXiv API returned HTTP {status}"
69            )));
70        }
71
72        let body = resp
73            .text()
74            .await
75            .map_err(|e| ResearchError::Source(format!("arXiv response read failed: {e}")))?;
76
77        parse_arxiv_atom(&body)
78    }
79
80    fn name(&self) -> &str {
81        "arxiv"
82    }
83}
84
85/// Resolve an XML entity reference (`"amp"`, `"#39"`, `"#x27"`, …) to its
86/// text. Only the five predefined entities and numeric character references
87/// can occur in an arXiv Atom document — it carries no DTD. Unknown references
88/// are dropped rather than failing the whole ingest.
89fn resolve_ref(name: &str) -> Option<String> {
90    match name {
91        "amp" => Some("&".into()),
92        "lt" => Some("<".into()),
93        "gt" => Some(">".into()),
94        "quot" => Some("\"".into()),
95        "apos" => Some("'".into()),
96        other => other
97            .strip_prefix('#')
98            .and_then(|num| {
99                num.strip_prefix('x')
100                    .or_else(|| num.strip_prefix('X'))
101                    .and_then(|h| u32::from_str_radix(h, 16).ok())
102                    .or_else(|| num.parse::<u32>().ok())
103            })
104            .and_then(char::from_u32)
105            .map(|c| c.to_string()),
106    }
107}
108
109fn parse_arxiv_atom(xml: &str) -> Result<Vec<Paper>> {
110    // NB: `trim_text` stays off — it trims each text *event*, and an entity
111    // reference splits an element's text into several events, so interior
112    // whitespace would be lost ("R&D 'x'" -> "R&D'x'"). Accumulators are
113    // trimmed where they are consumed instead.
114    let mut reader = Reader::from_str(xml);
115
116    let mut papers: Vec<Paper> = Vec::new();
117
118    // Per-entry accumulators
119    let mut in_entry = false;
120    let mut current_tag = String::new();
121    let mut title = String::new();
122    let mut summary = String::new();
123    let mut arxiv_id = String::new();
124    let mut published = String::new();
125    let mut authors: Vec<String> = Vec::new();
126    let mut author_name = String::new();
127
128    loop {
129        match reader.read_event() {
130            Ok(Event::Start(ref e)) => {
131                // Strip namespace prefix (e.g. "atom:entry" → "entry")
132                let local = e.name().local_name().into_inner();
133                current_tag = local.to_string();
134                if local == "entry" {
135                    in_entry = true;
136                    title.clear();
137                    summary.clear();
138                    arxiv_id.clear();
139                    published.clear();
140                    authors.clear();
141                }
142                // Text and entity references both append, so start each
143                // element's accumulator fresh.
144                match local {
145                    "title" => title.clear(),
146                    "summary" => summary.clear(),
147                    "id" => arxiv_id.clear(),
148                    "published" => published.clear(),
149                    "name" => author_name.clear(),
150                    _ => {}
151                }
152            }
153            Ok(Event::End(ref e)) => {
154                let local = e.name().local_name().into_inner();
155                current_tag.clear();
156                if local == "name" && in_entry && !author_name.is_empty() {
157                    authors.push(author_name.trim().to_string());
158                    author_name.clear();
159                }
160                if local == "entry" && in_entry {
161                    let mut paper = Paper::new(title.trim().to_string());
162                    paper.abstract_text = summary.trim().to_string();
163                    paper.authors = authors.clone();
164                    paper.year = published
165                        .trim()
166                        .get(..4)
167                        .and_then(|y| y.parse::<u32>().ok());
168
169                    // arxiv_id from <id>http://arxiv.org/abs/2301.00234v1</id>
170                    let last = arxiv_id.split('/').next_back().unwrap_or("");
171                    // strip version suffix "v<digits>" if present
172                    let bare_id = if let Some(pos) = last.rfind('v') {
173                        let after_v = &last[pos + 1..];
174                        if after_v.chars().all(|c| c.is_ascii_digit()) && !after_v.is_empty() {
175                            last[..pos].to_string()
176                        } else {
177                            last.to_string()
178                        }
179                    } else {
180                        last.to_string()
181                    };
182                    if !bare_id.is_empty() {
183                        paper.url = Some(format!("https://arxiv.org/abs/{bare_id}"));
184                        paper.arxiv_id = Some(bare_id);
185                    }
186
187                    papers.push(paper);
188                    in_entry = false;
189                    current_tag.clear();
190                }
191            }
192            Ok(Event::Text(e)) => {
193                let text = e.xml10_content().to_string();
194                if in_entry {
195                    match current_tag.as_str() {
196                        "title" => title.push_str(&text),
197                        "summary" => summary.push_str(&text),
198                        "id" => arxiv_id.push_str(&text),
199                        "published" => published.push_str(&text),
200                        "name" => author_name.push_str(&text),
201                        _ => {}
202                    }
203                }
204            }
205            // An `&entity;` inside element text arrives as its own event; the
206            // surrounding Text pieces are delivered separately.
207            Ok(Event::GeneralRef(e)) if in_entry => {
208                let text = resolve_ref(e.xml10_content().as_ref()).unwrap_or_default();
209                match current_tag.as_str() {
210                    "title" => title.push_str(&text),
211                    "summary" => summary.push_str(&text),
212                    "id" => arxiv_id.push_str(&text),
213                    "published" => published.push_str(&text),
214                    "name" => author_name.push_str(&text),
215                    _ => {}
216                }
217            }
218            Ok(Event::Eof) => break,
219            Err(e) => return Err(ResearchError::Source(format!("arXiv XML parse error: {e}"))),
220            _ => {}
221        }
222    }
223
224    Ok(papers)
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    const ATOM_FIXTURE: &str = r#"<?xml version="1.0" encoding="UTF-8"?>
232<feed xmlns="http://www.w3.org/2005/Atom">
233  <entry>
234    <id>http://arxiv.org/abs/2301.00234v1</id>
235    <title>Async Runtimes in Rust</title>
236    <summary>A survey of async runtime design.</summary>
237    <published>2023-01-15T00:00:00Z</published>
238    <author><name>Alice Smith</name></author>
239    <author><name>Bob Jones</name></author>
240  </entry>
241  <entry>
242    <id>http://arxiv.org/abs/2302.00100v2</id>
243    <title>Tokio Internals</title>
244    <summary>Deep dive into Tokio scheduling.</summary>
245    <published>2023-02-01T00:00:00Z</published>
246    <author><name>Carol White</name></author>
247  </entry>
248</feed>"#;
249
250    #[test]
251    fn parse_atom_returns_papers() {
252        let papers = parse_arxiv_atom(ATOM_FIXTURE).unwrap();
253        assert_eq!(papers.len(), 2);
254        assert_eq!(papers[0].title, "Async Runtimes in Rust");
255        assert_eq!(papers[0].authors, vec!["Alice Smith", "Bob Jones"]);
256        assert_eq!(papers[0].year, Some(2023));
257        assert_eq!(papers[0].arxiv_id.as_deref(), Some("2301.00234"));
258        assert_eq!(papers[1].title, "Tokio Internals");
259    }
260
261    #[test]
262    fn parse_atom_empty_feed() {
263        let xml = r#"<?xml version="1.0"?><feed xmlns="http://www.w3.org/2005/Atom"></feed>"#;
264        let papers = parse_arxiv_atom(xml).unwrap();
265        assert!(papers.is_empty());
266    }
267
268    /// Entity references in text must survive the quick-xml 0.42 event model,
269    /// where `&…;` is delivered as its own `GeneralRef` event.
270    #[test]
271    fn parse_atom_unescapes_entities() {
272        let xml = r#"<?xml version="1.0"?><feed xmlns="http://www.w3.org/2005/Atom">
273  <entry>
274    <id>http://arxiv.org/abs/2303.00001v1</id>
275    <title>Q&amp;A for R&#38;D &#x27;scaling&#x27;</title>
276    <summary>AT&amp;T results &lt;published&gt; here.</summary>
277    <author><name>A &amp; B</name></author>
278  </entry>
279</feed>"#;
280        let papers = parse_arxiv_atom(xml).unwrap();
281        assert_eq!(papers.len(), 1);
282        assert_eq!(papers[0].title, "Q&A for R&D 'scaling'");
283        assert_eq!(papers[0].abstract_text, "AT&T results <published> here.");
284        assert_eq!(papers[0].authors, vec!["A & B"]);
285    }
286
287    #[test]
288    fn source_name() {
289        let source = ArxivSource::new();
290        assert_eq!(source.name(), "arxiv");
291    }
292}
research_agent/adapters/arxiv_source.rs

research_agent/adapters/
arxiv_source.rs
