Skip to main content

paper_gap/
repos.rs

1use crate::arxiv::Paper;
2use crate::{Result, provider::ProviderOutcome};
3use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue, USER_AGENT};
4use serde::{Deserialize, Serialize};
5use std::time::Duration;
6
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
8pub struct Repository {
9    pub forge: String,
10    pub url: String,
11    pub owner: String,
12    pub name: String,
13    pub description: Option<String>,
14    pub primary_language: Option<String>,
15    pub stars: Option<u64>,
16    pub forks: Option<u64>,
17    pub license: Option<String>,
18    pub archived: Option<bool>,
19    pub default_branch: Option<String>,
20    pub last_commit_date: Option<String>,
21    pub source_provider: String,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
25pub struct RepositoryMatch {
26    pub repository: Repository,
27    pub confidence_score: u32,
28    pub evidence: Vec<MatchEvidence>,
29}
30
31#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
32pub struct MatchEvidence {
33    pub kind: String,
34    pub value: String,
35    pub points: u32,
36}
37
38const MIN_MATCH_CONFIDENCE: u32 = 30;
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
41pub struct RepositoryRef {
42    pub forge: String,
43    pub owner: String,
44    pub name: String,
45    pub url: String,
46}
47
48pub fn repository_from_url(url: &str) -> Option<Repository> {
49    let normalized = url.trim().trim_end_matches('/').trim_end_matches(".git");
50    let rest = normalized
51        .strip_prefix("https://")
52        .or_else(|| normalized.strip_prefix("http://"))?;
53    let (host, path) = rest.split_once('/')?;
54    let forge = match host.to_ascii_lowercase().as_str() {
55        "github.com" => "github",
56        "gitlab.com" => "gitlab",
57        "codeberg.org" => "codeberg",
58        _ => return None,
59    };
60    let parts: Vec<_> = path.split('/').filter(|part| !part.is_empty()).collect();
61    if parts.len() < 2 || (forge != "gitlab" && parts.len() != 2) {
62        return None;
63    }
64    let name = parts.last()?.to_string();
65    let owner = parts[..parts.len() - 1].join("/");
66    Some(Repository {
67        forge: forge.into(),
68        url: format!("https://{host}/{path}")
69            .trim_end_matches('/')
70            .trim_end_matches(".git")
71            .to_string(),
72        owner,
73        name,
74        description: None,
75        primary_language: None,
76        stars: None,
77        forks: None,
78        license: None,
79        archived: None,
80        default_branch: None,
81        last_commit_date: None,
82        source_provider: "papers-with-code".into(),
83    })
84}
85
86#[derive(Debug, Clone)]
87pub struct RepoSearchQuery {
88    pub terms: Vec<String>,
89    pub limit: usize,
90}
91
92pub fn search_terms(paper: &Paper) -> Vec<String> {
93    let mut terms = vec![format!("\"{}\"", paper.title)];
94    if let Some(arxiv_id) = paper.arxiv_id.as_deref() {
95        terms.push(trim_arxiv_version(arxiv_id).to_string());
96    }
97    if let Some(doi) = paper.doi.as_deref() {
98        terms.push(doi.to_string());
99    }
100    terms.extend(
101        paper
102            .extracted_methods
103            .iter()
104            .filter(|term| is_specific_method_term(term))
105            .cloned(),
106    );
107    dedup_strings(terms)
108}
109
110pub async fn search_for_paper(paper: &Paper, limit: usize) -> Result<Vec<Repository>> {
111    let query = RepoSearchQuery {
112        // ponytail: cap live forge searches; search_terms still exposes the full list for reports/tests.
113        terms: search_terms(paper).into_iter().take(3).collect(),
114        limit,
115    };
116    let mut repos = Vec::new();
117    repos.extend(github_search(&query).await.unwrap_or_default());
118    repos.extend(gitlab_search(&query).await.unwrap_or_default());
119    repos.extend(
120        forgejo_search("codeberg", "https://codeberg.org/api/v1", &query)
121            .await
122            .unwrap_or_default(),
123    );
124    Ok(filter_repositories(paper, dedup_repositories(repos)))
125}
126
127pub async fn search_for_paper_with_outcome(
128    paper: &Paper,
129    limit: usize,
130) -> (Vec<Repository>, ProviderOutcome) {
131    search_for_paper_with_providers_and_outcome(paper, limit, &["github", "gitlab", "codeberg"])
132        .await
133}
134
135pub async fn search_for_paper_with_providers_and_outcome(
136    paper: &Paper,
137    limit: usize,
138    providers: &[&str],
139) -> (Vec<Repository>, ProviderOutcome) {
140    let query = RepoSearchQuery {
141        terms: search_terms(paper).into_iter().take(3).collect(),
142        limit,
143    };
144    let mut repositories = Vec::new();
145    let mut outcomes = Vec::new();
146
147    for provider in providers {
148        let result = match *provider {
149            "github" => github_search(&query).await,
150            "gitlab" => gitlab_search(&query).await,
151            "codeberg" => forgejo_search("codeberg", "https://codeberg.org/api/v1", &query).await,
152            _ => continue,
153        };
154        match result {
155            Ok(found) => {
156                repositories.extend(found);
157                outcomes.push(ProviderOutcome::success(*provider));
158            }
159            Err(error) => outcomes.push(ProviderOutcome::from_error(*provider, &error)),
160        }
161    }
162
163    (
164        dedup_repositories(repositories),
165        ProviderOutcome::aggregate("repository-search", &outcomes),
166    )
167}
168
169pub fn filter_repositories(paper: &Paper, repositories: Vec<Repository>) -> Vec<Repository> {
170    rank_repositories(paper, repositories, &[])
171        .into_iter()
172        .map(|matched| matched.repository)
173        .collect()
174}
175
176pub fn rank_repositories(
177    paper: &Paper,
178    repositories: Vec<Repository>,
179    official_urls: &[String],
180) -> Vec<RepositoryMatch> {
181    let mut matches: Vec<_> = repositories
182        .into_iter()
183        .filter_map(|repository| {
184            let matched = match_repository(paper, repository, official_urls);
185            (matched.confidence_score >= MIN_MATCH_CONFIDENCE).then_some(matched)
186        })
187        .collect();
188    matches.sort_by_key(|matched| std::cmp::Reverse(matched.confidence_score));
189    matches
190}
191
192fn match_repository(
193    paper: &Paper,
194    repository: Repository,
195    official_urls: &[String],
196) -> RepositoryMatch {
197    let mut evidence = Vec::new();
198    let identity = normalize(&format!(
199        "{} {} {}",
200        repository.owner, repository.name, repository.url
201    ));
202    let description = normalize(repository.description.as_deref().unwrap_or_default());
203    let whole = format!("{identity}{description}");
204
205    if official_urls
206        .iter()
207        .any(|url| normalized_repo_url(url) == normalized_repo_url(&repository.url))
208    {
209        push_evidence(
210            &mut evidence,
211            "official_pwc_link",
212            repository.url.clone(),
213            100,
214        );
215    }
216    if let Some(id) = paper.arxiv_id.as_deref().map(trim_arxiv_version) {
217        if whole.contains(&normalize(id)) {
218            push_evidence(&mut evidence, "exact_arxiv_id", id.into(), 100);
219        }
220    }
221    if let Some(doi) = paper.doi.as_deref() {
222        if whole.contains(&normalize(doi)) {
223            push_evidence(&mut evidence, "exact_doi", doi.into(), 100);
224        }
225    }
226
227    let title = normalize(&paper.title);
228    if !title.is_empty() && whole.contains(&title) {
229        push_evidence(&mut evidence, "exact_title", paper.title.clone(), 60);
230    } else {
231        let title_terms = meaningful_title_terms(&paper.title);
232        let identity_terms: Vec<_> = title_terms
233            .iter()
234            .filter(|term| identity.contains(term.as_str()))
235            .cloned()
236            .collect();
237        let description_terms: Vec<_> = title_terms
238            .iter()
239            .filter(|term| !identity_terms.contains(term) && description.contains(term.as_str()))
240            .cloned()
241            .collect();
242        if !identity_terms.is_empty() {
243            push_evidence(
244                &mut evidence,
245                "title_identity_overlap",
246                identity_terms.join(", "),
247                (identity_terms.len() as u32 * 15).min(45),
248            );
249        }
250        if !description_terms.is_empty() {
251            push_evidence(
252                &mut evidence,
253                "description_overlap",
254                description_terms.join(", "),
255                (description_terms.len() as u32 * 10).min(30),
256            );
257        }
258    }
259
260    let methods: Vec<_> = paper
261        .extracted_methods
262        .iter()
263        .filter(|term| is_specific_method_term(term) && whole.contains(&normalize(term)))
264        .cloned()
265        .collect();
266    if !methods.is_empty() {
267        push_evidence(
268            &mut evidence,
269            "method_overlap",
270            methods.join(", "),
271            (methods.len() as u32 * 10).min(20),
272        );
273    }
274
275    let confidence_score = evidence
276        .iter()
277        .map(|item| item.points)
278        .sum::<u32>()
279        .min(100);
280    RepositoryMatch {
281        repository,
282        confidence_score,
283        evidence,
284    }
285}
286
287fn push_evidence(evidence: &mut Vec<MatchEvidence>, kind: &str, value: String, points: u32) {
288    evidence.push(MatchEvidence {
289        kind: kind.into(),
290        value,
291        points,
292    });
293}
294
295fn normalized_repo_url(url: &str) -> String {
296    url.trim_end_matches('/')
297        .trim_end_matches(".git")
298        .to_ascii_lowercase()
299}
300
301fn meaningful_title_terms(title: &str) -> Vec<String> {
302    let stop = [
303        "with", "from", "into", "using", "based", "towards", "toward", "paper", "study", "method",
304        "methods", "system", "systems", "program", "programs",
305    ];
306    title
307        .split(|c: char| !c.is_alphanumeric())
308        .map(str::to_ascii_lowercase)
309        .filter(|term| term.len() >= 4 && !stop.contains(&term.as_str()))
310        .collect()
311}
312
313fn is_specific_method_term(term: &str) -> bool {
314    let broad = [
315        "ai", "ml", "rl", "api", "cpu", "gpu", "fft", "pde", "ode", "chc", "nlp", "llm", "cnn",
316        "rnn", "gcn", "sql", "http", "json", "xml",
317    ];
318    let lower = term.to_ascii_lowercase();
319    term.len() > 3 && !broad.contains(&lower.as_str())
320}
321
322fn normalize(text: &str) -> String {
323    text.chars()
324        .filter(|c| c.is_ascii_alphanumeric())
325        .flat_map(char::to_lowercase)
326        .collect()
327}
328
329pub fn github_search_url(term: &str, limit: usize) -> String {
330    format!(
331        "https://api.github.com/search/repositories?q={}&per_page={}",
332        urlencoding::encode(term),
333        limit
334    )
335}
336
337pub fn gitlab_search_url(term: &str, limit: usize) -> String {
338    format!(
339        "https://gitlab.com/api/v4/projects?search={}&per_page={}",
340        urlencoding::encode(term),
341        limit
342    )
343}
344
345pub fn forgejo_search_url(base_url: &str, term: &str, limit: usize) -> String {
346    format!(
347        "{}/repos/search?q={}&limit={}",
348        base_url.trim_end_matches('/'),
349        urlencoding::encode(term),
350        limit
351    )
352}
353
354#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
355pub struct RepositoryFile {
356    pub path: String,
357}
358
359pub fn github_tree_url(owner: &str, name: &str, branch: &str) -> String {
360    format!(
361        "https://api.github.com/repos/{}/{}/git/trees/{}?recursive=1",
362        owner,
363        name,
364        urlencoding::encode(branch)
365    )
366}
367
368pub fn gitlab_tree_url(owner: &str, name: &str, branch: &str) -> String {
369    format!(
370        "https://gitlab.com/api/v4/projects/{}/repository/tree?recursive=true&per_page=100&ref={}",
371        urlencoding::encode(&format!("{owner}/{name}")),
372        urlencoding::encode(branch)
373    )
374}
375
376pub fn forgejo_tree_url(base_url: &str, owner: &str, name: &str, branch: &str) -> String {
377    format!(
378        "{}/repos/{}/{}/git/trees/{}?recursive=true",
379        base_url.trim_end_matches('/'),
380        owner,
381        name,
382        urlencoding::encode(branch)
383    )
384}
385
386pub async fn file_paths(repo: &Repository) -> Result<Vec<String>> {
387    let client = client(match repo.forge.as_str() {
388        "github" => Some("GITHUB_TOKEN"),
389        "gitlab" => Some("GITLAB_TOKEN"),
390        _ => None,
391    });
392    let branch = repo.default_branch.as_deref().unwrap_or("main");
393    let url = match repo.forge.as_str() {
394        "github" => github_tree_url(&repo.owner, &repo.name, branch),
395        "gitlab" => gitlab_tree_url(&repo.owner, &repo.name, branch),
396        "codeberg" => forgejo_tree_url(
397            "https://codeberg.org/api/v1",
398            &repo.owner,
399            &repo.name,
400            branch,
401        ),
402        _ => return Ok(Vec::new()),
403    };
404    let response = client.get(url).send().await?;
405    if response.status() == reqwest::StatusCode::NOT_FOUND {
406        return Ok(Vec::new());
407    }
408    let text = response.error_for_status()?.text().await?;
409    match repo.forge.as_str() {
410        "github" => parse_github_tree(&text),
411        "gitlab" => parse_gitlab_tree(&text),
412        _ => parse_forgejo_tree(&text),
413    }
414}
415
416pub fn parse_github_tree(json: &str) -> Result<Vec<String>> {
417    #[derive(Deserialize)]
418    struct Tree {
419        tree: Vec<TreeItem>,
420    }
421    #[derive(Deserialize)]
422    struct TreeItem {
423        path: String,
424    }
425    Ok(serde_json::from_str::<Tree>(json)?
426        .tree
427        .into_iter()
428        .map(|item| item.path)
429        .collect())
430}
431
432pub fn parse_gitlab_tree(json: &str) -> Result<Vec<String>> {
433    #[derive(Deserialize)]
434    struct TreeItem {
435        path: String,
436    }
437    Ok(serde_json::from_str::<Vec<TreeItem>>(json)?
438        .into_iter()
439        .map(|item| item.path)
440        .collect())
441}
442
443pub fn parse_forgejo_tree(json: &str) -> Result<Vec<String>> {
444    parse_github_tree(json)
445}
446
447pub fn github_inspect_url(owner: &str, name: &str) -> String {
448    format!("https://api.github.com/repos/{}/{}", owner, name)
449}
450
451pub fn gitlab_inspect_url(owner: &str, name: &str) -> String {
452    format!(
453        "https://gitlab.com/api/v4/projects/{}",
454        urlencoding::encode(&format!("{owner}/{name}"))
455    )
456}
457
458pub fn forgejo_inspect_url(base_url: &str, owner: &str, name: &str) -> String {
459    format!(
460        "{}/repos/{}/{}",
461        base_url.trim_end_matches('/'),
462        owner,
463        name
464    )
465}
466
467pub async fn github_search(query: &RepoSearchQuery) -> Result<Vec<Repository>> {
468    let client = client(Some("GITHUB_TOKEN"));
469    let mut repos = Vec::new();
470    for term in &query.terms {
471        let response = client
472            .get(github_search_url(term, query.limit))
473            .send()
474            .await?;
475        if response.status() == reqwest::StatusCode::FORBIDDEN
476            || response.status() == reqwest::StatusCode::UNAUTHORIZED
477        {
478            continue;
479        }
480        repos.extend(parse_github_search(
481            &response.error_for_status()?.text().await?,
482        )?);
483    }
484    Ok(dedup_repositories(repos))
485}
486
487pub async fn gitlab_search(query: &RepoSearchQuery) -> Result<Vec<Repository>> {
488    let client = client(Some("GITLAB_TOKEN"));
489    let mut repos = Vec::new();
490    for term in &query.terms {
491        let response = client
492            .get(gitlab_search_url(term, query.limit))
493            .send()
494            .await?;
495        if response.status() == reqwest::StatusCode::FORBIDDEN
496            || response.status() == reqwest::StatusCode::UNAUTHORIZED
497        {
498            continue;
499        }
500        repos.extend(parse_gitlab_search(
501            &response.error_for_status()?.text().await?,
502        )?);
503    }
504    Ok(dedup_repositories(repos))
505}
506
507pub async fn forgejo_search(
508    provider: &str,
509    base_url: &str,
510    query: &RepoSearchQuery,
511) -> Result<Vec<Repository>> {
512    let client = client(None);
513    let mut repos = Vec::new();
514    for term in &query.terms {
515        let response = client
516            .get(forgejo_search_url(base_url, term, query.limit))
517            .send()
518            .await?;
519        if response.status() == reqwest::StatusCode::FORBIDDEN
520            || response.status() == reqwest::StatusCode::UNAUTHORIZED
521        {
522            continue;
523        }
524        repos.extend(parse_forgejo_search(
525            provider,
526            &response.error_for_status()?.text().await?,
527        )?);
528    }
529    Ok(dedup_repositories(repos))
530}
531
532pub async fn inspect(repo: &RepositoryRef) -> Result<Option<Repository>> {
533    let client = client(match repo.forge.as_str() {
534        "github" => Some("GITHUB_TOKEN"),
535        "gitlab" => Some("GITLAB_TOKEN"),
536        _ => None,
537    });
538    let url = match repo.forge.as_str() {
539        "github" => github_inspect_url(&repo.owner, &repo.name),
540        "gitlab" => gitlab_inspect_url(&repo.owner, &repo.name),
541        "codeberg" => forgejo_inspect_url("https://codeberg.org/api/v1", &repo.owner, &repo.name),
542        _ => return Ok(None),
543    };
544    let response = client.get(url).send().await?;
545    if response.status() == reqwest::StatusCode::NOT_FOUND {
546        return Ok(None);
547    }
548    let text = response.error_for_status()?.text().await?;
549    Ok(Some(match repo.forge.as_str() {
550        "github" => parse_github_repo(&text)?,
551        "gitlab" => parse_gitlab_repo(&text)?,
552        _ => parse_forgejo_repo(&repo.forge, &text)?,
553    }))
554}
555
556pub fn parse_github_search(json: &str) -> Result<Vec<Repository>> {
557    #[derive(Deserialize)]
558    struct Search {
559        items: Vec<GithubRepo>,
560    }
561    Ok(serde_json::from_str::<Search>(json)?
562        .items
563        .into_iter()
564        .map(Into::into)
565        .collect())
566}
567
568pub fn parse_github_repo(json: &str) -> Result<Repository> {
569    Ok(serde_json::from_str::<GithubRepo>(json)?.into())
570}
571
572pub fn parse_gitlab_search(json: &str) -> Result<Vec<Repository>> {
573    Ok(serde_json::from_str::<Vec<GitlabRepo>>(json)?
574        .into_iter()
575        .map(Into::into)
576        .collect())
577}
578
579pub fn parse_gitlab_repo(json: &str) -> Result<Repository> {
580    Ok(serde_json::from_str::<GitlabRepo>(json)?.into())
581}
582
583pub fn parse_forgejo_search(provider: &str, json: &str) -> Result<Vec<Repository>> {
584    #[derive(Deserialize)]
585    struct Search {
586        data: Vec<ForgejoRepo>,
587    }
588    Ok(serde_json::from_str::<Search>(json)?
589        .data
590        .into_iter()
591        .map(|r| r.into_repo(provider))
592        .collect())
593}
594
595pub fn parse_forgejo_repo(provider: &str, json: &str) -> Result<Repository> {
596    Ok(serde_json::from_str::<ForgejoRepo>(json)?.into_repo(provider))
597}
598
599fn client(token_env: Option<&str>) -> reqwest::Client {
600    let mut headers = HeaderMap::new();
601    headers.insert(USER_AGENT, HeaderValue::from_static("paper-gap"));
602    if let Some(token_env) = token_env.and_then(|name| std::env::var(name).ok()) {
603        if let Ok(value) = HeaderValue::from_str(&format!("Bearer {token_env}")) {
604            headers.insert(AUTHORIZATION, value);
605        }
606    }
607    reqwest::Client::builder()
608        .default_headers(headers)
609        .timeout(Duration::from_secs(8))
610        .build()
611        .expect("valid reqwest client")
612}
613
614#[derive(Deserialize)]
615struct GithubRepo {
616    html_url: String,
617    name: String,
618    owner: GithubOwner,
619    description: Option<String>,
620    language: Option<String>,
621    stargazers_count: Option<u64>,
622    forks_count: Option<u64>,
623    license: Option<GithubLicense>,
624    archived: Option<bool>,
625    default_branch: Option<String>,
626    pushed_at: Option<String>,
627}
628#[derive(Deserialize)]
629struct GithubOwner {
630    login: String,
631}
632#[derive(Deserialize)]
633struct GithubLicense {
634    spdx_id: Option<String>,
635    name: Option<String>,
636}
637
638impl From<GithubRepo> for Repository {
639    fn from(repo: GithubRepo) -> Self {
640        Repository {
641            forge: "github".into(),
642            url: repo.html_url,
643            owner: repo.owner.login,
644            name: repo.name,
645            description: repo.description,
646            primary_language: repo.language,
647            stars: repo.stargazers_count,
648            forks: repo.forks_count,
649            license: repo.license.and_then(|l| l.spdx_id.or(l.name)),
650            archived: repo.archived,
651            default_branch: repo.default_branch,
652            last_commit_date: repo.pushed_at,
653            source_provider: "github".into(),
654        }
655    }
656}
657
658#[derive(Deserialize)]
659struct GitlabRepo {
660    web_url: String,
661    path: String,
662    path_with_namespace: String,
663    description: Option<String>,
664    star_count: Option<u64>,
665    forks_count: Option<u64>,
666    archived: Option<bool>,
667    default_branch: Option<String>,
668    last_activity_at: Option<String>,
669}
670
671impl From<GitlabRepo> for Repository {
672    fn from(repo: GitlabRepo) -> Self {
673        let owner = repo
674            .path_with_namespace
675            .rsplit_once('/')
676            .map_or("", |(owner, _)| owner)
677            .to_string();
678        Repository {
679            forge: "gitlab".into(),
680            url: repo.web_url,
681            owner,
682            name: repo.path,
683            description: repo.description,
684            primary_language: None,
685            stars: repo.star_count,
686            forks: repo.forks_count,
687            license: None,
688            archived: repo.archived,
689            default_branch: repo.default_branch,
690            last_commit_date: repo.last_activity_at,
691            source_provider: "gitlab".into(),
692        }
693    }
694}
695
696#[derive(Deserialize)]
697struct ForgejoRepo {
698    html_url: String,
699    name: String,
700    owner: ForgejoOwner,
701    description: Option<String>,
702    language: Option<String>,
703    stars_count: Option<u64>,
704    forks_count: Option<u64>,
705    archived: Option<bool>,
706    default_branch: Option<String>,
707    updated_at: Option<String>,
708}
709#[derive(Deserialize)]
710struct ForgejoOwner {
711    login: String,
712}
713
714impl ForgejoRepo {
715    fn into_repo(self, provider: &str) -> Repository {
716        Repository {
717            forge: provider.into(),
718            url: self.html_url,
719            owner: self.owner.login,
720            name: self.name,
721            description: self.description,
722            primary_language: self.language,
723            stars: self.stars_count,
724            forks: self.forks_count,
725            license: None,
726            archived: self.archived,
727            default_branch: self.default_branch,
728            last_commit_date: self.updated_at,
729            source_provider: provider.into(),
730        }
731    }
732}
733
734fn dedup_repositories(repos: Vec<Repository>) -> Vec<Repository> {
735    let mut out = Vec::new();
736    for repo in repos {
737        if !out
738            .iter()
739            .any(|existing: &Repository| existing.url == repo.url)
740        {
741            out.push(repo);
742        }
743    }
744    out
745}
746
747fn dedup_strings(values: Vec<String>) -> Vec<String> {
748    let mut out = Vec::new();
749    for value in values.into_iter().filter(|s| !s.trim().is_empty()) {
750        if !out.iter().any(|existing| existing == &value) {
751            out.push(value);
752        }
753    }
754    out
755}
756
757fn trim_arxiv_version(arxiv_id: &str) -> &str {
758    arxiv_id.rsplit_once('v').map_or(arxiv_id, |(id, version)| {
759        if version.chars().all(|c| c.is_ascii_digit()) {
760            id
761        } else {
762            arxiv_id
763        }
764    })
765}