news-mcp/news_mcp/enrichment/llm_enrich.py

from __future__ import annotations

from typing import Any, Dict

from news_mcp.config import NEWS_ENTITY_BLACKLIST
from news_mcp.llm import call_extraction, call_summary


def _filter_entities(entities, blacklist=None):
    banned = set(x.strip().lower() for x in (blacklist if blacklist is not None else NEWS_ENTITY_BLACKLIST))
    out = []
    for ent in entities or []:
        key = str(ent).strip().lower()
        if not key or key in banned:
            continue
        out.append(ent)
    return out


async def classify_cluster_groq(cluster: Dict[str, Any]) -> Dict[str, Any]:
    parsed = await call_extraction(cluster)
    out = dict(cluster)
    out.update({
        "topic": parsed.get("topic", cluster.get("topic")),
        "entities": _filter_entities(parsed.get("entities", [])),
        "sentiment": parsed.get("sentiment", "neutral"),
        "sentimentScore": parsed.get("sentimentScore"),
        "keywords": parsed.get("keywords", []),
    })
    return out


async def summarize_cluster_groq(cluster: Dict[str, Any]) -> Dict[str, Any]:
    parsed = await call_summary(cluster)
    return parsed