core/llm.py

"""Aurelius core — optional Gemini LLM layer.

This is the entire LLM surface. It is deliberately isolated so the rest of
the engine never imports an LLM and never depends on one: every function
here returns a plain result dict and **never raises**, so callers can treat
"the LLM said X" and "the LLM is unavailable" as ordinary data.

Contract (the load-bearing requirement):
  * No key            → available False, reason "no_key"; no network call.
  * Over soft budget  → reason "rate_limited"; no network call (we never
                        cause an upstream 429).
  * Cooling down      → reason "cooling_down" after an upstream 429/5xx.
  * Timeout / error   → ok False; swallowed, never propagated.

The model only ever NARRATES already-computed evidence — the prompt
builders below feed it the graph's own output (titles, edge phrases,
correlations, bridges), so it explains real structure rather than inventing
facts.

Backend (Google AI Studio Generative Language API):
  POST .../v1beta/models/{model}:generateContent?key={KEY}
  { "contents":[{"parts":[{"text": prompt}]}], "generationConfig": {...} }
  → candidates[0].content.parts[0].text
"""

from __future__ import annotations

import asyncio
import time
from collections import deque

import httpx

from config import (
    GEMINI_API_KEY, GEMINI_MODEL, GEMINI_RPM, GEMINI_COOLDOWN_S,
    GEMINI_TIMEOUT_S, GEMINI_MAX_CONCURRENCY, CONTACT_EMAIL,
)

_ENDPOINT = ("https://generativelanguage.googleapis.com/v1beta/"
             "models/{model}:generateContent")
_HEADERS = {"User-Agent": f"Aurelius/1.0 ({CONTACT_EMAIL})",
            "Content-Type": "application/json"}

# ── rate-limit / cooldown state (process-local, single event loop) ───────
_call_times: deque[float] = deque()     # timestamps of recent Gemini calls
_cooldown_until = 0.0
_sema = asyncio.Semaphore(max(1, GEMINI_MAX_CONCURRENCY))


def _budget_reason() -> str | None:
    """Why we must NOT call Gemini right now, or None if we may."""
    if not GEMINI_API_KEY:
        return "no_key"
    now = time.time()
    if now < _cooldown_until:
        return "cooling_down"
    while _call_times and now - _call_times[0] > 60.0:
        _call_times.popleft()
    if len(_call_times) >= GEMINI_RPM:
        return "rate_limited"
    return None


def llm_status() -> dict:
    """Whether AI features can run right now — booleans only, never the key.
    Safe to expose to the browser via /api/llm/status."""
    reason = _budget_reason()
    return {"available": reason is None, "reason": reason or "ok",
            "model": GEMINI_MODEL if GEMINI_API_KEY else None}


def available() -> bool:
    return _budget_reason() is None


def _trip_cooldown():
    global _cooldown_until
    _cooldown_until = time.time() + GEMINI_COOLDOWN_S


async def generate(prompt: str, *, system: str | None = None,
                   max_output_tokens: int = 512,
                   temperature: float = 0.4) -> dict:
    """Run one Gemini completion. Returns {ok, text, reason} and never raises.

    ok False + reason in {no_key, rate_limited, cooling_down, timeout, error,
    empty} means the caller should fall back to its non-LLM output.
    """
    reason = _budget_reason()
    if reason is not None:
        return {"ok": False, "text": None, "reason": reason}

    parts = (f"{system}\n\n{prompt}" if system else prompt)
    body = {
        "contents": [{"parts": [{"text": parts}]}],
        "generationConfig": {
            "temperature": temperature,
            "maxOutputTokens": max_output_tokens,
        },
    }
    url = _ENDPOINT.format(model=GEMINI_MODEL)

    _call_times.append(time.time())   # count the attempt against the budget
    try:
        async with _sema:
            async with httpx.AsyncClient(timeout=GEMINI_TIMEOUT_S) as client:
                r = await client.post(url, params={"key": GEMINI_API_KEY},
                                       headers=_HEADERS, json=body)
    except (httpx.TimeoutException, httpx.HTTPError):
        return {"ok": False, "text": None, "reason": "timeout"}
    except Exception:
        return {"ok": False, "text": None, "reason": "error"}

    if r.status_code == 429 or r.status_code >= 500:
        _trip_cooldown()   # back off — don't hammer the quota/service
        return {"ok": False, "text": None,
                "reason": "rate_limited" if r.status_code == 429 else "error"}
    if r.status_code != 200:
        return {"ok": False, "text": None, "reason": "error"}

    try:
        data = r.json()
        text = (data["candidates"][0]["content"]["parts"][0]["text"]).strip()
    except (KeyError, IndexError, ValueError, TypeError):
        return {"ok": False, "text": None, "reason": "empty"}
    if not text:
        return {"ok": False, "text": None, "reason": "empty"}
    return {"ok": True, "text": text, "reason": "ok"}


# ══════════════════════════════════════════════════════════════
# Prompt builders — each narrates already-computed evidence.
# Kept here so the LLM's whole footprint lives in one file.
# ══════════════════════════════════════════════════════════════

_STYLE = ("You are Aurelius, a graph-intelligence assistant. Explain clearly "
          "for a curious non-expert. Be concise and concrete. Use ONLY the "
          "evidence given — never invent facts, numbers, or links. No preamble "
          "like 'Sure' or 'Here is'; start with the explanation.")


def explain_path(source: str, nodes: list[str], edges: list[str]) -> str:
    steps = []
    for i in range(len(nodes) - 1):
        rel = edges[i] if i < len(edges) and edges[i] else "connects to"
        steps.append(f'"{nodes[i]}" --({rel})--> "{nodes[i + 1]}"')
    chain = "\n".join(steps)
    return (f"{_STYLE}\n\nThis is a real path found in the '{source}' graph "
            f"between \"{nodes[0]}\" and \"{nodes[-1]}\". Each arrow is a "
            f"verified relationship:\n{chain}\n\n"
            "In 2-3 sentences, explain how these two things connect through "
            "this chain, and why the link is interesting. Do not restate the "
            "arrows mechanically — tell the story of the connection.")


def analyze_relation(source: str, a: str, b: str, data: dict,
                     exposure: list[dict] | None = None) -> str:
    lines = [f"Connection strength: {data.get('strength')}/100."]
    if data.get("direct", {}).get("a_to_b") or data.get("direct", {}).get("b_to_a"):
        lines.append("There is a direct relationship between them.")
    if data.get("n_paths"):
        lines.append(f"{data['n_paths']} one-step paths connect them.")
    inter = [x["title"] for x in (data.get("paths_a_to_b", [])
             + data.get("paths_b_to_a", []) + data.get("co_targets", []))][:8]
    if inter:
        lines.append("Shared intermediaries: " + ", ".join(inter) + ".")
    if data.get("n_co_targets"):
        lines.append(f"They connect to {data['n_co_targets']} of the same things.")
    if data.get("similarity") is not None:
        lines.append(f"Embedding similarity: {data['similarity']}.")
    if exposure:
        exp = ", ".join(f"{e['title']} ({e['chain']})" for e in exposure[:5])
        lines.append("If one moves, exposure flows to: " + exp + ".")
    evidence = "\n".join(f"- {ln}" for ln in lines)
    return (f"{_STYLE}\n\nTwo entities in the '{source}' graph: \"{a}\" and "
            f"\"{b}\". Evidence:\n{evidence}\n\n"
            "In 3-4 sentences write an analyst-style note: how are they "
            "connected, how strongly, and what it means. If price/correlation "
            "or exposure evidence is present, interpret it plainly.")


def summarize_coverage(entity: str, headlines: list[dict]) -> str:
    lines = []
    for h in headlines[:14]:
        tone = h.get("sentiment_label", "")
        medium = h.get("medium", "news")
        lines.append(f"- [{medium}/{tone}] {h.get('title', '')}")
    body = "\n".join(lines)
    return (f"{_STYLE}\n\nRecent news and discussion headlines about "
            f"\"{entity}\":\n{body}\n\n"
            "Write a 2-3 sentence summary of what the coverage is about right "
            "now and the overall mood. Then, on a new line, output exactly one "
            "of: TONE: positive | TONE: negative | TONE: neutral — your read of "
            "the overall sentiment.")


def explain_discovery(source: str, a: str, candidates: list[dict]) -> str:
    lines = []
    for c in candidates[:6]:
        bridges = ", ".join(b["title"] for b in c.get("bridges", [])[:4])
        lines.append(f"- \"{c['title']}\": no direct link, but reached through "
                     f"{c.get('n_bridges', 0)} shared connections "
                     f"({bridges}).")
    body = "\n".join(lines)
    return (f"{_STYLE}\n\nIn the '{source}' graph these are hidden connections "
            f"from \"{a}\" — entities with strong indirect support but no "
            f"direct link:\n{body}\n\n"
            "In 2-3 sentences explain what these hidden connections suggest and "
            "why they're worth a look. Speak to the pattern, not each item.")


def summarize_entity(source: str, title: str, info: dict) -> str:
    facts = []
    feats = info.get("features", {}) or {}
    if feats.get("kind"):
        facts.append(f"kind: {feats['kind']}")
    if feats.get("sector"):
        facts.append(f"sector: {feats['sector']}")
    if info.get("summary"):
        facts.append(f"note: {info['summary']}")
    fact_str = "; ".join(facts) if facts else "(no structured facts)"
    return (f"{_STYLE}\n\nEntity in the '{source}' graph: \"{title}\" "
            f"({fact_str}). In 1-2 sentences, explain what it is in plain "
            "language for someone unfamiliar with it.")