1use std::time::Duration;
9
10use fetch_kit::Client;
11use serde::{Deserialize, Serialize};
12
13use crate::clock::unix_seconds;
14use crate::error::Result;
15use crate::model::{Observation, SourceKind, SubjectRecord};
16use crate::source::{CollectCtx, CollectReport, TrendSource};
17
18const HF_API: &str = "https://huggingface.co/api/models";
19const USER_AGENT: &str = "trendkit/0.3";
20const LIMIT: u32 = 100;
21pub const HF_PACING: Duration = Duration::from_millis(1_100);
23
24#[derive(Debug, Clone)]
26pub struct HfSource {
27 client: Client,
28}
29
30impl HfSource {
31 #[must_use]
33 pub fn new() -> HfSource {
34 HfSource {
35 client: Client::builder().user_agent(USER_AGENT).build(),
36 }
37 }
38}
39
40impl Default for HfSource {
41 fn default() -> Self {
42 HfSource::new()
43 }
44}
45
46impl TrendSource for HfSource {
47 fn name(&self) -> &'static str {
48 "huggingface"
49 }
50
51 fn collect(
52 &self,
53 ctx: &CollectCtx,
54 ) -> impl std::future::Future<Output = Result<CollectReport>> {
55 let ctx = ctx.clone();
56 async move {
57 let now = unix_seconds(ctx.clock.as_ref());
58 let params = HfQuery {
59 sort: "trendingScore",
60 direction: -1,
61 limit: LIMIT,
62 };
63 let response: Vec<HfModelJson> = self
64 .client
65 .get(HF_API)
66 .query(¶ms)
67 .json_response()
68 .await?;
69 let count = response.len();
70 for model in response {
71 let observation = Observation {
72 ts: now,
73 score: model.downloads,
74 extra1: model.likes,
75 extra2: model.trending_score.map(|f| f as i64).unwrap_or(0),
76 };
77 let subject = SubjectRecord::from(model);
78 let subject_id = ctx.store.upsert_subject(&subject)?;
79 ctx.store.insert_observation(subject_id, &observation)?;
80 }
81 ctx.store.record_run(now, self.name(), 1, count as i64)?;
82 Ok(CollectReport {
83 source: self.name(),
84 items: count,
85 pages: 1,
86 })
87 }
88 }
89}
90
91#[derive(Debug, Serialize)]
92struct HfQuery<'a> {
93 sort: &'a str,
94 direction: i32,
95 limit: u32,
96}
97
98#[derive(Debug, Clone, Deserialize)]
100pub(crate) struct HfModelJson {
101 #[serde(rename = "modelId", default)]
102 pub(crate) model_id: Option<String>,
103 #[serde(rename = "id", default)]
104 pub(crate) id: String,
105 #[serde(default)]
106 pub(crate) downloads: i64,
107 #[serde(default)]
108 pub(crate) likes: i64,
109 #[serde(default)]
110 pub(crate) trending_score: Option<f64>,
111 #[serde(default)]
112 pub(crate) pipeline_tag: Option<String>,
113 #[serde(default)]
114 pub(crate) tags: Option<Vec<String>>,
115}
116
117impl From<HfModelJson> for SubjectRecord {
118 fn from(json: HfModelJson) -> Self {
119 let name = json.model_id.unwrap_or(json.id);
120 let mut topics = json.tags.unwrap_or_default();
121 if let Some(pipeline) = json.pipeline_tag {
122 if !topics.contains(&pipeline) {
123 topics.push(pipeline);
124 }
125 }
126 SubjectRecord {
127 kind: SourceKind::HfModel,
128 key: name.clone(),
129 title: None,
130 lang: None,
131 url: Some(format!("https://huggingface.co/{name}")),
132 topics,
133 created_at: None,
134 }
135 }
136}
137
138pub fn parse_models_json(payload: &str) -> Result<Vec<SubjectRecord>> {
143 let models: Vec<HfModelJson> = serde_json::from_str(payload)
144 .map_err(|e| crate::error::TrendError::Config(format!("invalid hf payload: {e}")))?;
145 Ok(models.into_iter().map(SubjectRecord::from).collect())
146}
147
148#[cfg(test)]
149mod tests {
150 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
151
152 use super::*;
153
154 const SAMPLE: &str = r#"[
155 {
156 "_id": "66f0a1",
157 "id": "org/model-a",
158 "modelId": "org/model-a",
159 "downloads": 1000000,
160 "likes": 4200,
161 "trendingScore": 97.5,
162 "pipeline_tag": "text-generation",
163 "tags": ["text-generation", "conversational"]
164 },
165 {
166 "id": "bare-model",
167 "downloads": 42
168 }
169 ]"#;
170
171 #[test]
172 fn parse_extracts_models_with_pipeline_as_topic() {
173 let subjects = parse_models_json(SAMPLE).unwrap();
174 assert_eq!(subjects.len(), 2);
175 assert_eq!(subjects[0].kind, SourceKind::HfModel);
176 assert_eq!(subjects[0].key, "org/model-a");
177 assert_eq!(
178 subjects[0].url.as_deref(),
179 Some("https://huggingface.co/org/model-a")
180 );
181 assert!(subjects[0].topics.contains(&"text-generation".to_owned()));
182 assert!(subjects[0].topics.contains(&"conversational".to_owned()));
183 assert_eq!(subjects[1].key, "bare-model");
185 }
186
187 #[test]
188 fn parse_rejects_garbage() {
189 assert!(parse_models_json("{").is_err());
190 }
191}