the-daily-epub/src/related.rs

//! Related articles under each web issue article.
//!
//! For every pick, the nearest embedded articles across the whole archive
//! (`article_embeddings`, kept forever by default). Finding them is a scan of
//! every cached vector, so it runs once when an issue is published and once
//! per issue from `features related`, and the result lives in `issue_related`;
//! reader pages only read that table.

use std::collections::{HashMap, HashSet};

use anyhow::Result;
use jiff::civil::Date;
use sqlx::Row as _;

use crate::config::VoyageConfig;
use crate::curate::embedding::{decode_blob, dot, load_article_embeddings};
use crate::db::Db;
use crate::types::ArticleId;

/// Links shown under each article.
const PER_ARTICLE: usize = 5;
/// Below this cosine a neighbour is noise, not a related article (the same
/// cut the curation prompts use for "closest rated").
const MIN_COSINE: f64 = 0.55;
/// Embedding rows decoded per statement, so memory stays flat as the archive grows.
const SCAN_PAGE: i64 = 1000;
/// Candidates kept per pick between pages: enough headroom for the title
/// de-duplication at the end to still fill [`PER_ARTICLE`].
const KEEP: usize = PER_ARTICLE * 4;

/// What [`refresh`] wrote.
#[derive(Debug, Default, PartialEq)]
pub struct Refreshed {
    pub issues: usize,
    pub links: usize,
}

/// Recompute `issue_related` for the given issues (every issue when `dates` is
/// empty) in one pass over the embedding cache.
///
/// Excluded: the pick itself, the rest of its issue, articles that only have a
/// feed excerpt (their page would be near-empty), and articles whose title
/// matches the pick's or a closer link's (syndicated copies of one story).
pub async fn refresh(db: &Db, voyage: &VoyageConfig, dates: &[Date]) -> Result<Refreshed> {
    let dates = if dates.is_empty() {
        sqlx::query_scalar::<_, String>("SELECT date FROM issues ORDER BY date")
            .fetch_all(db.pool())
            .await?
            .into_iter()
            .filter_map(|raw| raw.parse().ok())
            .collect()
    } else {
        dates.to_vec()
    };

    // Every pick of every requested issue, with its title for de-duplication.
    let mut groups: Vec<Group> = Vec::with_capacity(dates.len());
    for date in &dates {
        let rows = sqlx::query(
            "SELECT ia.article_id, a.title FROM issue_articles ia
             JOIN articles a ON a.id = ia.article_id
             WHERE ia.issue_date = ? ORDER BY ia.section, ia.position",
        )
        .bind(date.to_string())
        .fetch_all(db.pool())
        .await?;
        groups.push(Group {
            date: *date,
            members: rows.iter().map(|row| row.get("article_id")).collect(),
            picks: rows
                .iter()
                .map(|row| Target {
                    article_id: row.get("article_id"),
                    title_key: title_key(row.get::<Option<String>, _>("title").as_deref()),
                    vector: Vec::new(),
                    best: Vec::new(),
                })
                .collect(),
        });
    }
    let ids: Vec<ArticleId> = groups
        .iter()
        .flat_map(|group| group.picks.iter().map(|pick| pick.article_id))
        .collect();
    let mut vectors = load_article_embeddings(db, voyage, &ids).await?;
    for group in &mut groups {
        // A pick without a cached vector simply gets no related list.
        group
            .picks
            .retain_mut(|pick| match vectors.remove(&pick.article_id) {
                Some(vector) => {
                    pick.vector = vector;
                    true
                }
                None => false,
            });
    }

    let mut after = i64::MIN;
    loop {
        let rows = sqlx::query(
            "SELECT ae.article_id, ae.embedding, a.title
             FROM article_embeddings ae JOIN articles a ON a.id = ae.article_id
             WHERE ae.model = ? AND ae.dimension = ? AND a.excerpt_only = 0
               AND ae.article_id > ?
             ORDER BY ae.article_id LIMIT ?",
        )
        .bind(&voyage.model)
        .bind(voyage.output_dimension as i64)
        .bind(after)
        .bind(SCAN_PAGE)
        .fetch_all(db.pool())
        .await?;
        let Some(last) = rows.last() else { break };
        after = last.get("article_id");
        for row in &rows {
            let candidate: ArticleId = row.get("article_id");
            let vector = match decode_blob(
                &row.get::<Vec<u8>, _>("embedding"),
                voyage.output_dimension,
            ) {
                Ok(vector) => vector,
                Err(error) => {
                    tracing::warn!(article_id = candidate, %error, "ignoring a malformed embedding");
                    continue;
                }
            };
            let title = title_key(row.get::<Option<String>, _>("title").as_deref());
            for group in &mut groups {
                if group.members.contains(&candidate) {
                    continue;
                }
                for pick in &mut group.picks {
                    let cosine = dot(&pick.vector, &vector)?;
                    if cosine.is_finite() && cosine >= MIN_COSINE {
                        pick.best.push((candidate, cosine, title.clone()));
                    }
                }
            }
        }
        for group in &mut groups {
            for pick in &mut group.picks {
                // Highest cosine first; ties by id so reruns write the same rows.
                pick.best.sort_by(|left, right| {
                    right
                        .1
                        .total_cmp(&left.1)
                        .then_with(|| left.0.cmp(&right.0))
                });
                pick.best.truncate(KEEP);
            }
        }
    }

    let mut refreshed = Refreshed::default();
    let mut tx = db.pool().begin().await?;
    for group in &groups {
        let date = group.date.to_string();
        sqlx::query("DELETE FROM issue_related WHERE issue_date = ?")
            .bind(&date)
            .execute(&mut *tx)
            .await?;
        for pick in &group.picks {
            let mut seen: HashSet<&str> = HashSet::new();
            if !pick.title_key.is_empty() {
                seen.insert(pick.title_key.as_str());
            }
            let chosen = pick
                .best
                .iter()
                .filter(|(_, _, title)| title.is_empty() || seen.insert(title))
                .take(PER_ARTICLE);
            for (position, (related_id, cosine, _)) in chosen.enumerate() {
                sqlx::query(
                    "INSERT INTO issue_related (issue_date, article_id, position, related_id, cosine)
                     VALUES (?, ?, ?, ?, ?)",
                )
                .bind(&date)
                .bind(pick.article_id)
                .bind(position as i64)
                .bind(related_id)
                .bind(cosine)
                .execute(&mut *tx)
                .await?;
                refreshed.links += 1;
            }
        }
        refreshed.issues += 1;
    }
    tx.commit().await?;
    Ok(refreshed)
}

/// The stored related links for one pick, closest first.
pub async fn for_article(db: &Db, date: Date, article_id: ArticleId) -> Result<Vec<ArticleId>> {
    Ok(sqlx::query_scalar(
        "SELECT related_id FROM issue_related
         WHERE issue_date = ? AND article_id = ? ORDER BY position",
    )
    .bind(date.to_string())
    .bind(article_id)
    .fetch_all(db.pool())
    .await?)
}

/// The most recent issue each of `ids` was published in, for linking a related
/// article to its issue page rather than the standalone one.
pub async fn latest_issues(db: &Db, ids: &[ArticleId]) -> Result<HashMap<ArticleId, Date>> {
    let mut issues = HashMap::new();
    for id in ids {
        let date: Option<String> =
            sqlx::query_scalar("SELECT MAX(issue_date) FROM issue_articles WHERE article_id = ?")
                .bind(id)
                .fetch_one(db.pool())
                .await?;
        if let Some(date) = date.and_then(|raw| raw.parse().ok()) {
            issues.insert(*id, date);
        }
    }
    Ok(issues)
}

struct Group {
    date: Date,
    /// Every pick of the issue, embedded or not: none of them is "related".
    members: HashSet<ArticleId>,
    picks: Vec<Target>,
}

struct Target {
    article_id: ArticleId,
    title_key: String,
    vector: Vec<f32>,
    /// (candidate, cosine, title key), closest first after each scan page.
    best: Vec<(ArticleId, f64, String)>,
}

/// Lowercased letters and digits only, so "Foo: Bar" and "foo — bar" collide.
fn title_key(title: Option<&str>) -> String {
    title
        .unwrap_or_default()
        .chars()
        .filter(|c| c.is_alphanumeric())
        .flat_map(char::to_lowercase)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::curate::embedding::encode_blob;

    fn voyage() -> VoyageConfig {
        VoyageConfig {
            model: "voyage-4-lite".into(),
            output_dimension: 2,
            ..VoyageConfig::default()
        }
    }

    async fn harness() -> (tempfile::TempDir, Db) {
        let dir = tempfile::tempdir().unwrap();
        let db = Db::open_and_migrate(&dir.path().join("related.db"))
            .await
            .unwrap();
        (dir, db)
    }

    async fn article(db: &Db, id: ArticleId, title: &str, vector: &[f32], excerpt_only: bool) {
        sqlx::query(
            "INSERT INTO articles (id, canonical_url, title, excerpt_only, first_seen)
             VALUES (?, ?, ?, ?, '2026-08-15T00:00:00Z')",
        )
        .bind(id)
        .bind(format!("https://example.com/{id}"))
        .bind(title)
        .bind(excerpt_only)
        .execute(db.pool())
        .await
        .unwrap();
        let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
        let unit: Vec<f32> = vector.iter().map(|v| v / norm).collect();
        sqlx::query(
            "INSERT INTO article_embeddings
                 (article_id, model, dimension, input_hash, embedding, created_at)
             VALUES (?, 'voyage-4-lite', 2, 'h', ?, '2026-08-15T00:00:00Z')",
        )
        .bind(id)
        .bind(encode_blob(&unit).unwrap())
        .execute(db.pool())
        .await
        .unwrap();
    }

    async fn issue(db: &Db, date: &str, picks: &[ArticleId]) {
        sqlx::query("INSERT INTO issues (date, issue_number, generated_at) VALUES (?, 1, ?)")
            .bind(date)
            .bind("2026-09-01T00:00:00Z")
            .execute(db.pool())
            .await
            .unwrap();
        for (position, id) in picks.iter().enumerate() {
            sqlx::query(
                "INSERT INTO issue_articles (issue_date, article_id, section, position)
                 VALUES (?, ?, 'Top Stories', ?)",
            )
            .bind(date)
            .bind(id)
            .bind(position as i64)
            .execute(db.pool())
            .await
            .unwrap();
        }
    }

    fn day(raw: &str) -> Date {
        raw.parse().unwrap()
    }

    #[tokio::test]
    async fn links_skip_the_issue_excerpts_duplicates_and_far_articles() {
        let (_dir, db) = harness().await;
        article(&db, 1, "Rust borrow checker", &[1.0, 0.0], false).await;
        article(&db, 2, "Same issue neighbour", &[1.0, 0.01], false).await;
        article(&db, 3, "Lifetimes explained", &[1.0, 0.1], false).await;
        article(&db, 4, "Paywalled", &[1.0, 0.05], true).await;
        article(&db, 5, "Rust Borrow-Checker!", &[1.0, 0.02], false).await;
        article(&db, 6, "Ownership deep dive", &[1.0, 0.3], false).await;
        article(&db, 7, "Ownership: deep dive", &[1.0, 0.31], false).await;
        article(&db, 8, "Gardening", &[0.0, 1.0], false).await;
        issue(&db, "2026-09-01", &[1, 2]).await;

        let refreshed = refresh(&db, &voyage(), &[day("2026-09-01")]).await.unwrap();
        let one = for_article(&db, day("2026-09-01"), 1).await.unwrap();
        assert_eq!(one, vec![3, 6], "closest first, one copy of each story");
        // 5 only repeats article 1's title, so it is fair game for article 2.
        let two = for_article(&db, day("2026-09-01"), 2).await.unwrap();
        assert_eq!(two, vec![5, 3, 6]);
        assert_eq!(
            refreshed,
            Refreshed {
                issues: 1,
                links: 5
            }
        );
    }

    #[tokio::test]
    async fn refreshing_every_issue_replaces_rows_and_links_earlier_picks() {
        let (_dir, db) = harness().await;
        article(&db, 1, "One", &[1.0, 0.0], false).await;
        article(&db, 2, "Two", &[1.0, 0.1], false).await;
        article(&db, 3, "Three", &[1.0, 0.3], false).await;
        issue(&db, "2026-09-01", &[1]).await;
        issue(&db, "2026-09-02", &[2]).await;

        assert_eq!(refresh(&db, &voyage(), &[]).await.unwrap().issues, 2);
        // A rerun rewrites rather than duplicates.
        refresh(&db, &voyage(), &[]).await.unwrap();
        let ids = for_article(&db, day("2026-09-02"), 2).await.unwrap();
        assert_eq!(ids, vec![1, 3]);
        let issues = latest_issues(&db, &[1, 3]).await.unwrap();
        assert_eq!(issues.get(&1), Some(&day("2026-09-01")));
        assert_eq!(issues.get(&3), None);
    }

    #[test]
    fn title_keys_ignore_case_and_punctuation() {
        assert_eq!(title_key(Some("Foo: Bar — Baz!")), "foobarbaz");
        assert_eq!(title_key(None), "");
    }
}