news-mcp/news_mcp/dedup/embedding_support.py

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass
from datetime import datetime, timezone, timedelta
from math import sqrt
from typing import Any

import httpx

from news_mcp.config import (
    NEWS_EMBEDDINGS_ENABLED,
    OLLAMA_BASE_URL,
    OLLAMA_EMBEDDING_MODEL,
    _NEEDLE_OLLAMA_MAX_CONCURRENCY,
)


_ollama_semaphore = asyncio.Semaphore(_NEEDLE_OLLAMA_MAX_CONCURRENCY)


@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


async def ollama_embed(text: str, timeout: float = 20.0) -> list[float] | None:
    """Async Ollama embedding call with concurrency limiting.

    Returns None on any failure so the caller falls back to heuristic clustering.
    """
    if not NEWS_EMBEDDINGS_ENABLED:
        return None

    payload = json.dumps({"model": OLLAMA_EMBEDDING_MODEL, "prompt": text}).encode("utf-8")
    url = f"{OLLAMA_BASE_URL.rstrip('/')}/api/embeddings"

    async with _ollama_semaphore:
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                resp = await client.post(
                    url,
                    content=payload,
                    headers={"Content-Type": "application/json"},
                )
                resp.raise_for_status()
                data = resp.json()
                emb = data.get("embedding")
                if isinstance(emb, list) and emb:
                    return [float(x) for x in emb]
        except Exception:
            return None
    return None