upbeatBytes/goodnews/llm.py

from __future__ import annotations

import json
import os
import sqlite3
import urllib.error
import urllib.request
from dataclasses import dataclass


DEFAULT_BASE_URL = "http://127.0.0.1:1234/v1"
DEFAULT_MODEL = "gpt-oss"


SYSTEM_PROMPT = """You classify article metadata for a calm constructive-news digest.

Judge emotional aftertaste, not simple positivity. Accept stories that leave a reader informed without feeling drained, especially when they include repair, progress, agency, resilience, human benefit, scientific discovery, environmental improvement, community action, or useful perspective.

Reject stories centered on fear, outrage, partisan conflict, crime, tragedy, disaster repetition, celebrity drama, market panic, or corporate PR without clear public benefit.

Return only JSON with this exact shape:
{
  "constructive_score": 0,
  "cortisol_score": 0,
  "ragebait_score": 0,
  "agency_score": 0,
  "human_benefit_score": 0,
  "novelty_score": 0,
  "pr_risk_score": 0,
  "accepted": false,
  "reason_code": "short_snake_case",
  "reason_text": "one concise sentence"
}
"""


@dataclass
class LocalModelClient:
    base_url: str
    model: str
    api_key: str | None = None
    timeout: int = 90

    @classmethod
    def from_env(cls) -> "LocalModelClient":
        return cls(
            base_url=os.environ.get("GOODNEWS_LLM_BASE_URL", DEFAULT_BASE_URL).rstrip("/"),
            model=os.environ.get("GOODNEWS_LLM_MODEL", DEFAULT_MODEL),
            api_key=os.environ.get("GOODNEWS_LLM_API_KEY"),
        )

    def classify(self, article: sqlite3.Row) -> dict:
        payload = {
            "model": self.model,
            "temperature": 0.1,
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": _article_prompt(article)},
            ],
            "response_format": {"type": "json_object"},
        }
        try:
            return self._chat(payload)
        except RuntimeError as exc:
            if "HTTP 400" not in str(exc):
                raise
            payload.pop("response_format", None)
            return self._chat(payload)

    def list_models(self) -> list[str]:
        headers = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        request = urllib.request.Request(f"{self.base_url}/models", headers=headers)
        try:
            with urllib.request.urlopen(request, timeout=10) as response:
                data = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"HTTP {exc.code} from local model: {detail}") from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(f"could not reach local model at {self.base_url}: {exc.reason}") from exc

        models = data.get("data", [])
        names = []
        for model in models:
            if isinstance(model, dict) and model.get("id"):
                names.append(str(model["id"]))
        return names

    def _chat(self, payload: dict) -> dict:
        body = json.dumps(payload).encode("utf-8")
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        request = urllib.request.Request(
            f"{self.base_url}/chat/completions",
            data=body,
            headers=headers,
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                data = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"HTTP {exc.code} from local model: {detail}") from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(f"could not reach local model at {self.base_url}: {exc.reason}") from exc

        try:
            content = data["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError) as exc:
            raise RuntimeError(f"unexpected local model response: {data}") from exc
        return parse_classifier_json(content)


def classify_articles(
    conn: sqlite3.Connection,
    client: LocalModelClient,
    limit: int,
    include_rejected: bool = False,
    dry_run: bool = False,
) -> list[tuple[int, dict]]:
    rows = _classification_candidates(conn, limit=limit, include_rejected=include_rejected)
    results = []
    for row in rows:
        scores = client.classify(row)
        scores = normalize_scores(scores, model_name=client.model)
        results.append((row["id"], scores))
        if not dry_run:
            upsert_article_score(conn, row["id"], scores)
    if not dry_run:
        conn.commit()
    return results


def parse_classifier_json(content: str) -> dict:
    content = content.strip()
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        start = content.find("{")
        end = content.rfind("}")
        if start == -1 or end == -1 or end <= start:
            raise RuntimeError(f"model did not return JSON: {content}")
        return json.loads(content[start : end + 1])


def normalize_scores(data: dict, model_name: str) -> dict:
    return {
        "constructive_score": _bounded_int(data.get("constructive_score")),
        "cortisol_score": _bounded_int(data.get("cortisol_score")),
        "ragebait_score": _bounded_int(data.get("ragebait_score")),
        "agency_score": _bounded_int(data.get("agency_score")),
        "human_benefit_score": _bounded_int(data.get("human_benefit_score")),
        "novelty_score": _bounded_int(data.get("novelty_score")),
        "pr_risk_score": _bounded_int(data.get("pr_risk_score")),
        "accepted": 1 if bool(data.get("accepted")) else 0,
        "reason_code": str(data.get("reason_code") or "model_no_reason")[:120],
        "reason_text": str(data.get("reason_text") or "")[:1000],
        "model_name": model_name,
    }


def upsert_article_score(conn: sqlite3.Connection, article_id: int, scores: dict) -> None:
    conn.execute(
        """
        INSERT INTO article_scores (
            article_id, constructive_score, cortisol_score, ragebait_score,
            agency_score, human_benefit_score, novelty_score, pr_risk_score,
            accepted, reason_code, reason_text, model_name, scored_at
        )
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
        ON CONFLICT(article_id) DO UPDATE SET
            constructive_score = excluded.constructive_score,
            cortisol_score = excluded.cortisol_score,
            ragebait_score = excluded.ragebait_score,
            agency_score = excluded.agency_score,
            human_benefit_score = excluded.human_benefit_score,
            novelty_score = excluded.novelty_score,
            pr_risk_score = excluded.pr_risk_score,
            accepted = excluded.accepted,
            reason_code = excluded.reason_code,
            reason_text = excluded.reason_text,
            model_name = excluded.model_name,
            scored_at = CURRENT_TIMESTAMP
        """,
        (
            article_id,
            scores["constructive_score"],
            scores["cortisol_score"],
            scores["ragebait_score"],
            scores["agency_score"],
            scores["human_benefit_score"],
            scores["novelty_score"],
            scores["pr_risk_score"],
            scores["accepted"],
            scores["reason_code"],
            scores["reason_text"],
            scores["model_name"],
        ),
    )


def _classification_candidates(
    conn: sqlite3.Connection,
    limit: int,
    include_rejected: bool,
) -> list[sqlite3.Row]:
    where = "" if include_rejected else "WHERE s.accepted = 1 OR s.constructive_score >= 4"
    return conn.execute(
        f"""
        SELECT
            a.id,
            a.title,
            a.description,
            a.published_at,
            a.canonical_url,
            src.name AS source_name,
            src.default_category,
            src.trust_score AS source_trust_score,
            src.pr_risk_score AS source_pr_risk_score,
            s.constructive_score,
            s.cortisol_score,
            s.ragebait_score,
            s.agency_score,
            s.human_benefit_score,
            s.pr_risk_score,
            s.accepted,
            s.reason_code
        FROM articles a
        JOIN sources src ON src.id = a.source_id
        LEFT JOIN article_scores s ON s.article_id = a.id
        {where}
        ORDER BY
            CASE WHEN s.model_name LIKE 'heuristic-%' THEN 0 ELSE 1 END,
            COALESCE(a.published_at, a.discovered_at) DESC
        LIMIT ?
        """,
        (limit,),
    ).fetchall()


def _article_prompt(article: sqlite3.Row) -> str:
    return "\n".join(
        [
            f"Source: {article['source_name']}",
            f"Source category: {article['default_category'] or 'unknown'}",
            f"Source trust score: {article['source_trust_score']}/10",
            f"Source PR risk score: {article['source_pr_risk_score']}/10",
            f"Published: {article['published_at'] or 'unknown'}",
            f"Title: {article['title']}",
            f"Snippet: {article['description'] or ''}",
            f"URL: {article['canonical_url']}",
        ]
    )


def _bounded_int(value: object) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        parsed = 0
    return max(0, min(10, parsed))