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