FabledGrimoire/src/fabledassistant/services/rss_classifier.py

"""
RSS item topic classifier.

Classifies RSS items into topic tags using a fast non-streaming LLM call.
Called from rss.py after new items are stored — fire-and-forget.
"""

import json
import logging
import re
from datetime import datetime, timezone

import httpx

from fabledassistant.config import Config

logger = logging.getLogger(__name__)

STANDARD_TOPICS = [
    "technology", "science", "politics", "business",
    "health", "environment", "local", "entertainment", "sports", "other",
]

_CLASSIFY_PROMPT = """\
Classify each news item into 1-3 topics. Use only topics from this list: {vocab}.
Return ONLY a JSON object mapping item_id (as string) to a list of topics.
Example: {{"1": ["technology", "ai"], "2": ["politics"]}}

Items:
{items_block}"""


async def _llm_classify(prompt: str, model: str) -> str:
    """Make a fast non-streaming LLM call and return the raw text response."""
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "stream": False,
        "options": {"num_ctx": 2048, "temperature": 0.0},
    }
    async with httpx.AsyncClient(timeout=120.0) as client:
        resp = await client.post(f"{Config.OLLAMA_URL}/api/chat", json=payload)
        resp.raise_for_status()
        return resp.json().get("message", {}).get("content", "")


async def classify_items_batch(
    items: list[dict],
    user_include_topics: list[str],
    model: str | None = None,
) -> dict[int, list[str]]:
    """
    Classify a batch of RSS items into topic tags.

    Args:
        items: list of dicts with 'id', 'title', 'content'
        user_include_topics: extra topics from user preferences to add to vocabulary
        model: Ollama model name; defaults to Config.OLLAMA_MODEL

    Returns:
        dict mapping item_id (int) -> list of topic strings.
        Items not returned had classification fail; callers should leave classified_at=NULL.
    """
    if not items:
        return {}

    if model is None:
        model = Config.OLLAMA_BACKGROUND_MODEL

    vocab = STANDARD_TOPICS + [t for t in user_include_topics if t not in STANDARD_TOPICS]
    items_block = "\n".join(
        f"[{item['id']}] {item['title']} — {item.get('content', '')[:300]}"
        for item in items
    )
    prompt = _CLASSIFY_PROMPT.format(vocab=", ".join(vocab), items_block=items_block)

    try:
        raw = await _llm_classify(prompt, model)
        # Strip <think>...</think> blocks emitted by reasoning models (e.g. qwen3)
        raw = re.sub(r"<think>.*?</think>", "", raw, flags=re.DOTALL)
        # Extract JSON from response (LLM may wrap it in markdown)
        raw = raw.strip()
        if raw.startswith("```"):
            raw = raw.split("```")[1]
            if raw.startswith("json"):
                raw = raw[4:]
        raw = raw.strip()
        # Allow control characters that may appear in LLM-generated JSON strings
        parsed = json.loads(raw, strict=False)
        return {int(k): v for k, v in parsed.items() if isinstance(v, list)}
    except Exception:
        logger.warning("RSS classification failed", exc_info=True)
        return {}


async def classify_and_store(
    item_ids: list[int],
    user_id: int,
) -> None:
    """
    Classify unclassified RSS items and write results to DB.
    Called as a fire-and-forget task from rss.py.
    """
    from sqlalchemy import select

    from fabledassistant.models import async_session
    from fabledassistant.models.rss_feed import RssItem
    from fabledassistant.services.settings import get_setting

    if not item_ids:
        return

    # Load the items
    async with async_session() as session:
        result = await session.execute(
            select(RssItem).where(RssItem.id.in_(item_ids))
        )
        items = list(result.scalars().all())

    if not items:
        return

    # Get user's include topics to extend vocabulary
    raw_include = await get_setting(user_id, "briefing_include_topics", "[]")
    try:
        include_topics = json.loads(raw_include) if isinstance(raw_include, str) else []
    except Exception:
        include_topics = []

    model = await get_setting(user_id, "background_model", Config.OLLAMA_BACKGROUND_MODEL)

    # Classify in batches of 10
    batch_size = 10
    all_results: dict[int, list[str]] = {}
    for i in range(0, len(items), batch_size):
        batch = items[i: i + batch_size]
        batch_dicts = [{"id": it.id, "title": it.title, "content": it.content} for it in batch]
        results = await classify_items_batch(batch_dicts, include_topics, model=model)
        all_results.update(results)

    # Write back to DB
    now = datetime.now(timezone.utc)
    async with async_session() as session:
        for item in items:
            item_db = await session.get(RssItem, item.id)
            if item_db is None:
                continue
            topics = all_results.get(item.id)
            if topics is not None:
                item_db.topics = topics
                item_db.classified_at = now
        await session.commit()