app.py

"""app.py -- Stage 3 FastAPI backend + dashboard host.

Runs the live chat in surprise mode (our product) and shadow-runs the llm_extract WRITE
decision each turn, so the dashboard shows a real A/B write-token counter on the same
conversation. Auto-detects the LLM: a live LLMClient against whatever OpenAI-compatible
endpoint .env names, else an offline MockLLM so the demo runs with zero setup. No provider
import here -- the only provider touchpoint is llm_client (imported lazily below).
"""
from __future__ import annotations

import os
import tempfile
import threading
import uuid
from collections import OrderedDict
from pathlib import Path

from fastapi import FastAPI, Request, Response
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from agent import EXTRACT_SYSTEM, Agent
from memory.sidecar import MemorySidecar

try:  # load .env if present (creds arrive here)
    from dotenv import load_dotenv
    load_dotenv()
except Exception:
    pass

TAU = 0.40  # best fixed operating point from the benchmark sweep (see benchmarks/)
DEVICE = os.environ.get("DEVICE", "cpu")  # "cuda" on the ROCm MI300X puts the memory on-GPU
STATIC = Path(__file__).parent / "static"


class _MockResult:
    def __init__(self, text: str, total_tokens: int):
        self.text = text
        self.total_tokens = total_tokens


class MockLLM:
    """Offline stub so the dashboard runs without creds. Token counts are plausible but
    synthetic -- the headline number is only 'real' with live Fireworks."""
    def chat(self, messages, **kw):
        if "extraction" in messages[0]["content"].lower():
            user = messages[1]["content"]
            return _MockResult(user, 20 + len(user.split()))
        return _MockResult("(mock reply) Noted -- ask me anything.", 40 + len(messages[-1]["content"].split()))


def _make_llm():
    """Live LLMClient when all creds are set; otherwise MockLLM."""
    if all(os.environ.get(k) for k in ("LLM_API_KEY", "LLM_BASE_URL", "LLM_MODEL")):
        from llm_client import LLMClient  # lazy: --mock/offline needs no provider
        return LLMClient(), False
    return MockLLM(), True


class Session:
    """Single-user demo session (scope fence: no multi-user, no auth)."""
    def __init__(self):
        self.reset()

    def reset(self):
        self.llm, self.is_mock = _make_llm()
        self.sidecar = MemorySidecar("dashboard", db_path=":memory:", tau=TAU, device=DEVICE, use_nli=True)
        # Agentic read: the model drives recall via native function-calling. Write stays the token-free gate.
        self.agent = Agent(self.sidecar, self.llm, "surprise", use_tools=not self.is_mock)
        self.turns = 0
        self.llm_extract_write_tokens = 0  # shadow baseline, same conversation
        self.read_tokens = 0

    def restore(self, path: str):
        """Rebuild memory from a snapshot file (cross-session boundary): the neural memory + the
        verbatim record store come back. The live LLM client is kept; only the memory is replaced."""
        self.sidecar = MemorySidecar.load(path, db_path=":memory:", session_id=self.sidecar.session_id,
                                          tau=TAU, device=DEVICE, use_nli=True)
        self.agent = Agent(self.sidecar, self.llm, "surprise", use_tools=not self.is_mock)

    def chat(self, message: str):
        r = self.agent.turn(message)
        self.turns += 1
        self.read_tokens += r.read_tokens
        # Shadow-run the llm_extract WRITE decision to tally the baseline cost this turn.
        try:
            shadow = self.llm.chat([
                {"role": "system", "content": EXTRACT_SYSTEM},
                {"role": "user", "content": message},
            ])
            self.llm_extract_write_tokens += shadow.total_tokens
        except Exception:
            # Live LLM unreachable: fall back to the measured llm_extract average so the
            # A/B baseline counter still reflects real per-turn cost (benchmarks/results.json).
            self.llm_extract_write_tokens += 69
        return r

    def state(self) -> dict:
        snap = self.sidecar.snapshot()
        snap["tokens"] = {
            "surprise_write": self.agent.total_write_tokens,  # always 0
            "llm_extract_write": self.llm_extract_write_tokens,
            "read": self.read_tokens,
            "turns": self.turns,
        }
        snap["mock"] = self.is_mock
        snap["device"] = DEVICE
        base = os.environ.get("LLM_BASE_URL", "")
        model = (os.environ.get("LLM_MODEL", "") or "").rsplit("/", 1)[-1]
        if self.is_mock:
            snap["llm_label"] = "Mock LLM (set .env for live)"
        elif "127.0.0.1" in base or "localhost" in base:
            snap["llm_label"] = f"Gemma 4 on MI300X · {model}"
        elif "fireworks" in base:
            snap["llm_label"] = f"Fireworks · {model}"
        elif "groq" in base:
            snap["llm_label"] = f"Groq · {model}"
        elif "dashscope" in base or "aliyuncs" in base:
            snap["llm_label"] = f"Qwen · {model}"
        elif "nvidia" in base:
            snap["llm_label"] = f"NVIDIA · {model}"
        else:
            snap["llm_label"] = model or "LLM"
        return snap


app = FastAPI(title="flashbulb")

# Per-visitor isolation: each browser gets its own in-memory Session keyed by a
# cookie, so one visitor's chat is never stored into another's demo. LRU-capped
# (a client that never returns the cookie would otherwise leak a session per
# request) and lock-guarded for the threadpool the sync endpoints run in.
COOKIE = "fb_sid"
MAX_SESSIONS = 64
_sessions: OrderedDict[str, Session] = OrderedDict()
_sessions_lock = threading.Lock()
_LOCAL_SNAPS: dict[str, str] = {}  # sid -> last local snapshot path; survives reset (local-dev fallback)


def _session_for(request: Request, response: Response) -> Session:
    sid = request.cookies.get(COOKIE)
    with _sessions_lock:
        if sid and sid in _sessions:
            _sessions.move_to_end(sid)
        else:
            sid = uuid.uuid4().hex
            _sessions[sid] = Session()
            while len(_sessions) > MAX_SESSIONS:
                _sessions.popitem(last=False)  # evict least-recently-used
        sess = _sessions[sid]
    sess.sid = sid  # visitor key for per-session cloud snapshots
    response.set_cookie(COOKIE, sid, max_age=86_400, httponly=True, samesite="lax")
    return sess


class ChatIn(BaseModel):
    message: str


@app.post("/chat")
def chat(inp: ChatIn, request: Request, response: Response):
    r = _session_for(request, response).chat(inp.message)
    return {
        "answer": r.answer,
        "write_tokens": r.write_tokens,
        "read_tokens": r.read_tokens,
        "stored": [rec.text for rec in r.stored],
        "skipped": r.skipped,
        "tool_calls": r.tool_calls,
    }


@app.get("/memory/state")
def memory_state(request: Request, response: Response):
    return _session_for(request, response).state()


@app.post("/reset")
def reset(request: Request, response: Response):
    _session_for(request, response).reset()
    return {"ok": True}


def _snap_key(sess) -> str:
    return f"snapshots/{sess.sid}.pt"


@app.post("/snapshot")
def snapshot(request: Request, response: Response):
    """Persist this session's learned memory -- the first half of the cross-session boundary.
    Uses Alibaba Cloud OSS when the oss2 SDK and OSS env vars are both present, and falls back to
    a local file otherwise. The response always names the backend it actually used, so the demo
    never claims a cloud it is not on. See cloud/oss_store.py."""
    sess = _session_for(request, response)
    n = len(sess.sidecar._active_records())
    path = str(Path(tempfile.gettempdir()) / f"flashbulb_{sess.sid}_{uuid.uuid4().hex}.pt")
    sess.sidecar.save(path)
    from cloud import oss_store  # lazy: local/offline dev needs no oss2
    if oss_store.oss_configured():
        key = _snap_key(sess)
        oss_store.upload_snapshot(path, key)
        return {"ok": True, "backend": "alibaba-oss", "object": key,
                "bucket": os.environ.get("OSS_BUCKET"), "records": n}
    _LOCAL_SNAPS[sess.sid] = path
    return {"ok": True, "backend": "local-only", "path": path, "records": n,
            "note": "set ALIBABA_CLOUD_ACCESS_KEY_ID/SECRET + OSS_BUCKET to persist to Alibaba Cloud"}


@app.post("/restore")
def restore(request: Request, response: Response):
    """Rebuild this session's memory from its snapshot -- the SECOND half of the cross-session
    boundary. Pulls the object back down from Alibaba Cloud OSS when configured; otherwise reloads
    the last local snapshot. Demonstrates that memory survives 'close -> reopen'."""
    sess = _session_for(request, response)
    from cloud import oss_store
    if oss_store.oss_configured():
        key = _snap_key(sess)
        if not oss_store.snapshot_exists(key):
            return {"ok": False, "error": "no snapshot in Alibaba OSS for this session yet -- Snapshot first"}
        path = str(Path(tempfile.gettempdir()) / f"flashbulb_restore_{sess.sid}_{uuid.uuid4().hex}.pt")
        oss_store.download_snapshot(key, path)
        sess.restore(path)
        return {"ok": True, "backend": "alibaba-oss", "object": key,
                "records": len(sess.sidecar._active_records())}
    local = _LOCAL_SNAPS.get(sess.sid)
    if local and Path(local).exists():
        sess.restore(local)
        return {"ok": True, "backend": "local-only", "records": len(sess.sidecar._active_records())}
    return {"ok": False, "error": "no snapshot found -- click Snapshot first"}


PRESETS = {
    "duplicate": [
        "I live in Ames, Iowa.",
        "My dog is named Rex.",
        "I live in Ames, Iowa.",     # exact duplicate -> gated, 0 tokens
        "I reside in Ames, Iowa.",   # paraphrase -> gated
    ],
    "knowledge_update": [
        "I live in Ames, Iowa.",
        "I work as a data engineer.",
        "My dog is named Rex.",
        "I moved from Ames to Chicago last month.",  # supersedes the stale location
    ],
}


@app.post("/preset/{name}")
def preset(name: str, request: Request, response: Response):
    if name not in PRESETS:
        return {"error": f"unknown preset '{name}'", "available": list(PRESETS)}
    sess = _session_for(request, response)
    sess.reset()
    turns = [{"message": m, **_turn_summary(sess.chat(m))} for m in PRESETS[name]]
    return {"preset": name, "turns": turns}


def _turn_summary(r) -> dict:
    return {"answer": r.answer, "stored": [x.text for x in r.stored], "skipped": r.skipped,
            "tool_calls": r.tool_calls}


app.mount("/static", StaticFiles(directory=str(STATIC)), name="static")


@app.get("/")
def index(request: Request):
    resp = FileResponse(str(STATIC / "index.html"))
    _session_for(request, resp)  # allocate the visitor's session + set its cookie on load
    return resp