news/dedup/deduper.py

"""三层去重主流程。

调用顺序:check / ingest 内部按 L1 -> L2 -> L3 顺序判定,任意层命中即返回。

Deduper 不要求线程安全;批处理串行调用即可。
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path

from loguru import logger

from extractor import Article

from .hasher import (
    DEFAULT_HAMMING_THRESHOLD,
    content_hash,
    hamming,
    simhash64,
)
from .models import DedupLayer, DedupResult, DedupStats, Fingerprint
from .store import DEFAULT_DB_PATH, FingerprintStore

# 默认时间窗口(±N 天)
DEFAULT_TIME_WINDOW_DAYS = 30


def _publish_date(article: Article) -> str | None:
    """从 Article.publish_time 取 YYYY-MM-DD 字符串。"""
    if article.publish_time is None:
        return None
    return article.publish_time.strftime("%Y-%m-%d")


def article_to_fingerprint(article: Article) -> Fingerprint:
    """构造 Fingerprint(用于 ingest 写入或对外只读)。

    source_ids 初始为 [article.source_id],后续重复文章命中时由 ingest 合并。
    """
    return Fingerprint(
        url_hash=article.url_hash,
        content_hash=content_hash(article.content),
        simhash=simhash64(article.content),
        source_id=article.source_id,
        url=article.url,
        title=article.title,
        publish_date=_publish_date(article),
        ingested_at=datetime.now(),
        source_ids=[article.source_id],
    )


class Deduper:
    """三层去重器。

    构造完毕后:
        - check(article)  仅判断,不写入;
        - ingest(article) 判断,不重复则写入指纹库,返回结果。
    """

    def __init__(
        self,
        db_path: str | Path = DEFAULT_DB_PATH,
        simhash_threshold: int = DEFAULT_HAMMING_THRESHOLD,
        time_window_days: int = DEFAULT_TIME_WINDOW_DAYS,
    ) -> None:
        self.store = FingerprintStore(db_path)
        self.simhash_threshold = simhash_threshold
        self.time_window_days = time_window_days

    def close(self) -> None:
        self.store.close()

    def __enter__(self) -> Deduper:
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    # ------------------------------------------------------------------ #
    # 公共 API
    # ------------------------------------------------------------------ #

    def check(self, article: Article) -> DedupResult:
        """三层判重(只读)。命中时附带匹配指纹的多源信息(all_source_ids)。"""
        fp = article_to_fingerprint(article)

        # L1: URL hash
        existing = self.store.get_by_url_hash(fp.url_hash)
        if existing is not None:
            return DedupResult(
                url_hash=fp.url_hash,
                is_duplicate=True,
                matched_layer=DedupLayer.URL,
                matched_url_hash=existing.url_hash,
                matched_url=existing.url,
                matched_title=existing.title,
                matched_source_id=existing.source_id,
                all_source_ids=existing.source_ids,
            )

        # L2: 内容 hash
        existing = self.store.find_by_content_hash(fp.content_hash)
        if existing is not None:
            return DedupResult(
                url_hash=fp.url_hash,
                is_duplicate=True,
                matched_layer=DedupLayer.CONTENT,
                matched_url_hash=existing.url_hash,
                matched_url=existing.url,
                matched_title=existing.title,
                matched_source_id=existing.source_id,
                all_source_ids=existing.source_ids,
            )

        # L3: SimHash 模糊
        candidates = self.store.candidates_for_simhash(
            fp.publish_date, self.time_window_days
        )
        best_dist: int | None = None
        best_match: Fingerprint | None = None
        for c in candidates:
            d = hamming(fp.simhash, c.simhash)
            if d <= self.simhash_threshold and (best_dist is None or d < best_dist):
                best_dist = d
                best_match = c
                if d == 0:  # 不可能更近,提前结束
                    break

        if best_match is not None:
            return DedupResult(
                url_hash=fp.url_hash,
                is_duplicate=True,
                matched_layer=DedupLayer.SIMHASH,
                matched_url_hash=best_match.url_hash,
                matched_url=best_match.url,
                matched_title=best_match.title,
                matched_source_id=best_match.source_id,
                all_source_ids=best_match.source_ids,
                hamming_distance=best_dist,
            )

        return DedupResult(url_hash=fp.url_hash, is_duplicate=False)

    def ingest(self, article: Article) -> DedupResult:
        """判重 + 不重复则入库。

        命中重复时,把当前文章的 source_id 合并进匹配指纹的 source_ids
        (记录同一内容组的全部来源),并更新 all_source_ids 后返回。
        """
        result = self.check(article)
        if not result.is_duplicate:
            fp = article_to_fingerprint(article)
            self.store.upsert(fp)
            logger.debug("入库: {} {}", fp.url_hash, fp.title[:30])
            return result

        # 重复:合并来源到匹配指纹(主源保持首位,Fingerprint validator 负责去重)
        if result.matched_url_hash and article.source_id not in result.all_source_ids:
            matched = self.store.get_by_url_hash(result.matched_url_hash)
            if matched is not None:
                merged = [*matched.source_ids, article.source_id]
                self.store.upsert(matched.model_copy(update={"source_ids": merged}))
                result = result.model_copy(
                    update={"all_source_ids": merged}
                )
                logger.debug(
                    "合并来源 {} -> {} ({} 个源)",
                    article.source_id, result.matched_url_hash, len(merged),
                )
        logger.debug("命中重复: {}", result.short_summary())
        return result

    def stats(self) -> DedupStats:
        lo, hi = self.store.date_range()
        return DedupStats(
            total=self.store.count(),
            by_source=self.store.count_by_source(),
            earliest=lo,
            latest=hi,
        )