backend/llm.py

"""
LLM layer for Polis.

Wraps the OpenAI API with three safety features that matter in a portfolio piece:
  1. A hard *budget guard* — the sim refuses to spend past POLIS_BUDGET_USD.
  2. A deterministic *mock backend* — if there is no API key, agents still think,
     so the Space boots and demos with zero cost.
  3. A tiny in-process *embedding cache* — identical strings are never re-embedded.

Nothing here is OpenAI-specific beyond the client construction, so swapping in a
local model later is a one-file change.
"""
from __future__ import annotations

import hashlib
import math
import os
import random
import threading
from dataclasses import dataclass, field
from typing import List, Optional

# ---- pricing (USD per 1M tokens), used only for the budget guard ------------
# Kept intentionally conservative; update if you change models.
PRICING = {
    "gpt-4o-mini":        {"in": 0.15, "out": 0.60},
    "gpt-4o":             {"in": 2.50, "out": 10.00},
    "text-embedding-3-small": {"in": 0.02, "out": 0.0},
}

CHAT_MODEL = os.getenv("POLIS_CHAT_MODEL", "gpt-4o-mini")
EMBED_MODEL = os.getenv("POLIS_EMBED_MODEL", "text-embedding-3-small")
BUDGET_USD = float(os.getenv("POLIS_BUDGET_USD", "1.00"))


@dataclass
class Ledger:
    """Tracks spend so a runaway loop can't drain the account."""
    spent_usd: float = 0.0
    calls: int = 0
    tokens_in: int = 0
    tokens_out: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def charge(self, model: str, tin: int, tout: int) -> None:
        p = PRICING.get(model, {"in": 0.0, "out": 0.0})
        cost = (tin / 1e6) * p["in"] + (tout / 1e6) * p["out"]
        with self._lock:
            self.spent_usd += cost
            self.calls += 1
            self.tokens_in += tin
            self.tokens_out += tout

    def remaining(self) -> float:
        return max(0.0, BUDGET_USD - self.spent_usd)

    def as_dict(self) -> dict:
        return {
            "spent_usd": round(self.spent_usd, 4),
            "budget_usd": BUDGET_USD,
            "remaining_usd": round(self.remaining(), 4),
            "calls": self.calls,
            "tokens_in": self.tokens_in,
            "tokens_out": self.tokens_out,
        }


class BudgetExceeded(RuntimeError):
    pass


class LLM:
    """Unified chat + embedding client with a mock fallback."""

    def __init__(self) -> None:
        self.ledger = Ledger()
        self._embed_cache: dict[str, List[float]] = {}
        self._client = None
        self.live = False
        key = os.getenv("OPENAI_API_KEY", "").strip()
        if key:
            try:
                from openai import OpenAI  # imported lazily so mock mode needs no dep
                self._client = OpenAI(api_key=key)
                self.live = True
            except Exception as exc:  # pragma: no cover - defensive
                print(f"[llm] OpenAI unavailable, using mock backend: {exc}")

    # -- chat -----------------------------------------------------------------
    def chat(self, system: str, user: str, *, temperature: float = 0.8,
             max_tokens: int = 220) -> str:
        if not self.live:
            return self._mock_chat(system, user)
        if self.ledger.remaining() <= 0:
            raise BudgetExceeded(
                f"Budget of ${BUDGET_USD:.2f} exhausted; refusing to spend more."
            )
        resp = self._client.chat.completions.create(
            model=CHAT_MODEL,
            messages=[{"role": "system", "content": system},
                      {"role": "user", "content": user}],
            temperature=temperature,
            max_tokens=max_tokens,
        )
        usage = resp.usage
        self.ledger.charge(CHAT_MODEL, usage.prompt_tokens, usage.completion_tokens)
        return (resp.choices[0].message.content or "").strip()

    # -- embeddings -----------------------------------------------------------
    def embed(self, text: str) -> List[float]:
        h = hashlib.sha1(text.encode("utf-8")).hexdigest()
        if h in self._embed_cache:
            return self._embed_cache[h]
        if not self.live or self.ledger.remaining() <= 0:
            vec = self._mock_embed(text)
        else:
            resp = self._client.embeddings.create(model=EMBED_MODEL, input=text)
            self.ledger.charge(EMBED_MODEL, resp.usage.prompt_tokens, 0)
            vec = resp.data[0].embedding
        self._embed_cache[h] = vec
        return vec

    # -- deterministic mock backend ------------------------------------------
    def _mock_embed(self, text: str, dim: int = 64) -> List[float]:
        """Hash-seeded pseudo-embedding. Not semantic, but stable and cheap —
        enough to make retrieval do *something* sensible offline."""
        seed = int(hashlib.sha1(text.lower().encode()).hexdigest()[:8], 16)
        rng = random.Random(seed)
        # bias vector by simple word hashing so related strings cluster a bit
        vec = [0.0] * dim
        for tok in text.lower().split():
            t = int(hashlib.sha1(tok.encode()).hexdigest()[:8], 16)
            vec[t % dim] += 1.0
        # add small noise for uniqueness
        vec = [v + rng.uniform(-0.05, 0.05) for v in vec]
        norm = math.sqrt(sum(v * v for v in vec)) or 1.0
        return [v / norm for v in vec]

    _MOCK_ACTIONS = [
        "heads to the plaza to see who is around",
        "strikes up a conversation about the day's news",
        "tends to work, humming quietly",
        "shares a rumor they overheard this morning",
        "invites a neighbor to the evening gathering",
        "reflects on a recent argument and softens",
        "sketches a small plan for tomorrow",
        "offers to help someone carry supplies",
    ]

    def _mock_chat(self, system: str, user: str) -> str:
        seed = int(hashlib.sha1((system + user).encode()).hexdigest()[:8], 16)
        rng = random.Random(seed)
        if "reflect" in user.lower() or "insight" in user.lower():
            return rng.choice([
                "I value the people who show up for me.",
                "Small kindnesses seem to matter more than grand plans.",
                "I am becoming more curious about the newcomers.",
            ])
        if "dialogue" in user.lower() or "say to" in user.lower():
            return rng.choice([
                "\"Have you heard? Something is stirring near the market.\"",
                "\"Come by tonight — we could use another set of hands.\"",
                "\"I've been thinking about what you said yesterday.\"",
            ])
        return rng.choice(self._MOCK_ACTIONS)


# module-level singleton
llm = LLM()


def cosine(a: List[float], b: List[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a)) or 1.0
    nb = math.sqrt(sum(x * x for x in b)) or 1.0
    return dot / (na * nb)