paper_gap/
pwc.rs

1use crate::{Result, arxiv::Paper, provider};
2use serde::{Deserialize, Serialize};
3use std::time::Duration;
4
5const BASE_URL: &str = "https://paperswithcode.com/api/v1";
6
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
8pub struct CodeLink {
9    pub source_provider: String,
10    pub paper_id: Option<String>,
11    pub repository_url: String,
12    pub repository_name: Option<String>,
13    pub is_official: Option<bool>,
14}
15
16#[derive(Debug, Deserialize)]
17struct List<T> {
18    results: Vec<T>,
19}
20
21#[derive(Debug, Deserialize)]
22struct PwcPaper {
23    id: String,
24}
25
26#[derive(Debug, Deserialize)]
27struct PwcRepository {
28    url: Option<String>,
29    name: Option<String>,
30    is_official: Option<bool>,
31}
32
33pub fn paper_search_url(paper: &Paper) -> String {
34    if let Some(arxiv_id) = paper.arxiv_id.as_deref() {
35        format!(
36            "{BASE_URL}/papers/?arxiv_id={}",
37            urlencoding::encode(trim_arxiv_version(arxiv_id))
38        )
39    } else {
40        format!("{BASE_URL}/papers/?q={}", urlencoding::encode(&paper.title))
41    }
42}
43
44pub fn paper_repositories_url(pwc_paper_id: &str) -> String {
45    format!(
46        "{BASE_URL}/papers/{}/repositories/",
47        urlencoding::encode(pwc_paper_id)
48    )
49}
50
51pub async fn code_links(paper: &Paper) -> Result<Vec<CodeLink>> {
52    let client = reqwest::Client::builder()
53        .timeout(Duration::from_secs(8))
54        .build()?;
55    let paper_response = provider::send_with_retry(client.get(paper_search_url(paper))).await?;
56    if paper_response.status() == reqwest::StatusCode::NOT_FOUND {
57        return Ok(Vec::new());
58    }
59    let paper_response = paper_response.error_for_status()?;
60    let papers = parse_paper_search(&paper_response.text().await?)?;
61    let Some(pwc_paper) = papers.first() else {
62        return Ok(Vec::new());
63    };
64
65    let repo_response =
66        provider::send_with_retry(client.get(paper_repositories_url(&pwc_paper.id))).await?;
67    if repo_response.status() == reqwest::StatusCode::NOT_FOUND {
68        return Ok(Vec::new());
69    }
70    let repo_response = repo_response.error_for_status()?;
71    parse_repositories(&pwc_paper.id, &repo_response.text().await?)
72}
73
74fn parse_paper_search(json: &str) -> Result<Vec<PwcPaper>> {
75    Ok(serde_json::from_str::<List<PwcPaper>>(json)?.results)
76}
77
78pub fn parse_repositories(paper_id: &str, json: &str) -> Result<Vec<CodeLink>> {
79    let repos = serde_json::from_str::<List<PwcRepository>>(json)?;
80    Ok(repos
81        .results
82        .into_iter()
83        .filter_map(|repo| {
84            Some(CodeLink {
85                source_provider: "papers-with-code".into(),
86                paper_id: Some(paper_id.into()),
87                repository_url: repo.url?,
88                repository_name: repo.name,
89                is_official: repo.is_official,
90            })
91        })
92        .collect())
93}
94
95fn trim_arxiv_version(arxiv_id: &str) -> &str {
96    arxiv_id.rsplit_once('v').map_or(arxiv_id, |(id, version)| {
97        if version.chars().all(|c| c.is_ascii_digit()) {
98            id
99        } else {
100            arxiv_id
101        }
102    })
103}
paper_gap/pwc.rs

paper_gap/
pwc.rs
