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