""" 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)