Skip to main content

trendkit_core/
hf.rs

1//! HuggingFace trending-models source (public Hub API, no auth).
2//!
3//! The AI wave deserves a dedicated signal: trending *models* lead the
4//! repos and stories that follow them. Score = downloads (cumulative),
5//! `extra1` = likes; HF tags flow into subjects.topics, so the personal
6//! profile scores model trends the same way as repo topics.
7
8use 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;
21/// Politeness pacing; single request per run today, kept for pagination.
22pub const HF_PACING: Duration = Duration::from_millis(1_100);
23
24/// HuggingFace trending-models source.
25#[derive(Debug, Clone)]
26pub struct HfSource {
27    client: Client,
28}
29
30impl HfSource {
31    /// Build a source.
32    #[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(&params)
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/// The subset of a Hub model entry trendkit consumes.
99#[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
138/// Parse a Hub models response body into subjects.
139///
140/// # Errors
141/// Fails when the payload is not a Hub API response.
142pub 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        // Degenerate entry falls back to `id`.
184        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}