news-mcp/news_mcp/dedup/embedding_support.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone, timedelta
import json
import urllib.request
from math import sqrt
from typing import Any

from news_mcp.config import NEWS_EMBEDDINGS_ENABLED, OLLAMA_BASE_URL, OLLAMA_EMBEDDING_MODEL


@dataclass(frozen=True)
class CandidateRules:
    """Cheap, non-embedding filters before we compare vectors."""

    require_topic_match: bool = True
    require_entity_overlap: int = 1
    max_age_hours: int = 72


def cosine_similarity(a: list[float], b: list[float]) -> float:
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x * y for x, y in zip(a, b))
    na = sqrt(sum(x * x for x in a))
    nb = sqrt(sum(y * y for y in b))
    if na == 0.0 or nb == 0.0:
        return 0.0
    return dot / (na * nb)


def _to_dt(value: Any) -> datetime | None:
    if not value:
        return None
    if isinstance(value, datetime):
        return value
    try:
        s = str(value).replace("Z", "+00:00")
        dt = datetime.fromisoformat(s)
        if dt.tzinfo is None:
            return dt.replace(tzinfo=timezone.utc)
        return dt
    except Exception:
        return None


def cluster_is_candidate(
    article: dict[str, Any],
    cluster: dict[str, Any],
    *,
    rules: CandidateRules | None = None,
    article_topic: str | None = None,
) -> bool:
    rules = rules or CandidateRules()

    if rules.require_topic_match and article_topic is not None:
        if str(article_topic).strip().lower() != str(cluster.get("topic", "")).strip().lower():
            return False

    # Require some overlap in extracted entities if both sides have them.
    article_entities = {
        str(e).strip().lower()
        for e in (article.get("entities", []) or [])
        if str(e).strip()
    }
    cluster_entities = {
        str(e).strip().lower()
        for e in (cluster.get("entities", []) or [])
        if str(e).strip()
    }
    if article_entities and cluster_entities:
        overlap = len(article_entities & cluster_entities)
        if overlap < rules.require_entity_overlap:
            return False

    # Age gate: keep comparisons within a recent window.
    article_dt = _to_dt(article.get("timestamp"))
    cluster_dt = _to_dt(cluster.get("last_updated") or cluster.get("timestamp"))
    if article_dt and cluster_dt:
        age = abs(article_dt - cluster_dt)
        if age > timedelta(hours=rules.max_age_hours):
            return False

    return True


def ollama_embed(text: str, timeout: float = 20.0) -> list[float] | None:
    """Best-effort Ollama embedding call; returns None on any failure.

    Embeddings are intentionally optional. The caller should fall back to the
    heuristic path when this returns None.
    """

    if not NEWS_EMBEDDINGS_ENABLED:
        return None
    payload = json.dumps({"model": OLLAMA_EMBEDDING_MODEL, "prompt": text}).encode("utf-8")
    req = urllib.request.Request(
        f"{OLLAMA_BASE_URL.rstrip('/')}/api/embeddings",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            emb = data.get("embedding")
            if isinstance(emb, list) and emb:
                return [float(x) for x in emb]
    except Exception:
        return None
    return None