"""HTTP API for THX-01: calibrated typed decisions over REST. THX01_MODEL=doofz/THX-01 THX01_API_KEY=... python -m thx01.server --port 8095 Endpoints GET /health liveness GET /v1/models the loaded model GET /v1/presets ready-made question schemas POST /v1/decide {state, questions} -> answers (+ detected language, latency) POST /v1/decide/batch {items: [{state, questions}]} -> list of results, one batched pass POST /v1/systemone TypeSafe-compatible alias of /v1/decide (same request and response shape) POST /v1/systemone/batch alias of /v1/decide/batch Question types: choice, noul, score, number (a value stated in the document; number.py), excerpt (a verbatim part of the document; excerpt.py). Any question may add "cite": true to get the parts of the document that support it. """ import argparse import asyncio import json import os import time from typing import Any, Dict, List, Optional, Union from fastapi import Depends, FastAPI, Header, HTTPException from pydantic import BaseModel from . import __version__, presets from . import number as num from . import excerpt as exc from . import compat from .agent import Agent from .lang import analyse MODEL = os.environ.get("THX01_MODEL", os.environ.get("SYSTEM01_MODEL", "doofz/THX-01")) API_KEY = os.environ.get("THX01_API_KEY", os.environ.get("SYSTEM01_API_KEY")) app = FastAPI(title="THX-01 API", version=__version__, description="THX-01: non-autoregressive, calibrated decisions in one forward pass.") agent: Optional[Agent] = None BATCH_WINDOW_MS = float(os.environ.get("THX01_BATCH_WINDOW_MS", os.environ.get("SYSTEM01_BATCH_WINDOW_MS", "3"))) BATCH_MAX = int(os.environ.get("THX01_BATCH_MAX", os.environ.get("SYSTEM01_BATCH_MAX", "256"))) _queue: Optional[asyncio.Queue] = None async def _batcher(): """Micro-batching: requests arriving within a few ms share one forward pass on the GPU.""" loop = asyncio.get_running_loop() while True: first = await _queue.get() items = [first] deadline = loop.time() + BATCH_WINDOW_MS / 1000 while len(items) < BATCH_MAX: timeout = deadline - loop.time() if timeout <= 0: break try: items.append(await asyncio.wait_for(_queue.get(), timeout)) except asyncio.TimeoutError: break try: outs = await loop.run_in_executor(None, agent.predict_batch, [i[0] for i in items], [i[1] for i in items]) for (_, _, fut), o in zip(items, outs): fut.set_result(o) except Exception as e: # a bad request must not poison its batch-mates for st, qs, fut in items: try: fut.set_result(await loop.run_in_executor(None, agent.predict, st, qs)) except Exception as e2: fut.set_exception(e2) def auth(authorization: Optional[str] = Header(default=None)): if API_KEY and authorization != "Bearer " + API_KEY: raise HTTPException(401, "missing or invalid bearer token") class DecideRequest(BaseModel): state: Union[str, Dict[str, Any], List[Any]] questions: Dict[str, Dict[str, Any]] model: Optional[str] = None # accepted for compatibility; one model is served class BatchItem(BaseModel): state: Union[str, Dict[str, Any], List[Any]] questions: Dict[str, Dict[str, Any]] class BatchRequest(BaseModel): items: List[BatchItem] model: Optional[str] = None @app.on_event("startup") async def _startup(): global agent, _queue agent = Agent(MODEL, device=os.environ.get("THX01_DEVICE", os.environ.get("SYSTEM01_DEVICE", "cuda"))) _queue = asyncio.Queue() asyncio.get_running_loop().create_task(_batcher()) @app.get("/health") def health(): return {"status": "ok", "model": "THX-01", "version": __version__} @app.get("/v1/models", dependencies=[Depends(auth)]) def models(): return {"object": "list", "data": [{"id": "THX-01", "source": MODEL, "device": str(agent.device)}]} @app.get("/v1/presets", dependencies=[Depends(auth)]) def get_presets(): return {"triage": presets.triage_questions(), "email": presets.email_questions(), "guard": presets.guard_questions(), "moderation": presets.moderation_questions(), "router": presets.router_questions()} def _lang(state): d = analyse(state) return {"script": d["script"], "language": d["language"]} def _predict(states, qs): return [o["answers"] for o in agent.predict_batch(states, qs)] def _extras(state, questions, answers, adapters=None): """excerpt questions, compatibility adapters and cite flags: extra THX-01 passes driven by the service layer.""" for qid, (kind, info) in (adapters or {}).items(): answers[qid] = compat.answer(_predict, state, questions[qid], kind, info) for qid, q in questions.items(): if q.get("type") == "excerpt": answers[qid] = exc.excerpt(_predict, state, q) if q.get("cite") and qid in answers: text = q.get("instructions") or q.get("question") or "" text = text if isinstance(text, str) else json.dumps(text, ensure_ascii=False) answers[qid]["citations"] = exc.cite(_predict, state, text) return answers async def _ask(state, questions): fut = asyncio.get_running_loop().create_future() await _queue.put((state, questions, fut)) return await fut @app.post("/v1/decide", dependencies=[Depends(auth)]) @app.post("/v1/systemone", dependencies=[Depends(auth)]) async def decide(req: DecideRequest): t0 = time.perf_counter() try: adapters = compat.route(req.state, req.questions) regular = {k: v for k, v in req.questions.items() if v.get("type") != "excerpt" and k not in adapters} plain, numq, meta = num.split(req.state, regular) parts = [] if plain: parts.append(_ask(req.state, plain)) if numq: # number questions read the document text itself parts.append(_ask(num.document_text(req.state), numq)) res = await asyncio.gather(*parts) p_out = res[0] if plain else None n_out = res[-1] if numq else None out = dict(p_out or n_out or {"model": "THX-01", "usage": {"input_tokens": 0, "output_tokens": 0}}) out["answers"] = num.merge(regular, p_out, n_out, meta) if plain and numq: out["usage"] = {k: p_out["usage"].get(k, 0) + n_out["usage"].get(k, 0) for k in p_out["usage"]} if len(regular) != len(req.questions) or any(q.get("cite") for q in req.questions.values()): ans = await asyncio.get_running_loop().run_in_executor(None, _extras, req.state, req.questions, out["answers"], adapters) out["answers"] = {k: ans[k] for k in req.questions} except (ValueError, KeyError, TypeError, AttributeError) as e: raise HTTPException(400, str(e)) out["detected"] = _lang(req.state) out["latency_ms"] = round((time.perf_counter() - t0) * 1000, 2) return out @app.post("/v1/decide/batch", dependencies=[Depends(auth)]) @app.post("/v1/systemone/batch", dependencies=[Depends(auth)]) def decide_batch(req: BatchRequest): t0 = time.perf_counter() try: adps = [compat.route(it.state, it.questions) for it in req.items] regs = [{k: v for k, v in it.questions.items() if v.get("type") != "excerpt" and k not in ad} for it, ad in zip(req.items, adps)] splits = [num.split(it.state, rq) for it, rq in zip(req.items, regs)] states, qs, where = [], [], [] for i, (it, (plain, numq, _)) in enumerate(zip(req.items, splits)): if plain: states.append(it.state); qs.append(plain); where.append((i, "p")) if numq: states.append(num.document_text(it.state)); qs.append(numq); where.append((i, "n")) flat = agent.predict_batch(states, qs) if states else [] got = {} for (i, kind), o in zip(where, flat): got[(i, kind)] = o outs = [] for i, (it, (plain, numq, meta)) in enumerate(zip(req.items, splits)): p_out, n_out = got.get((i, "p")), got.get((i, "n")) o = dict(p_out or n_out or {"model": "THX-01", "usage": {"input_tokens": 0, "output_tokens": 0}}) o["answers"] = num.merge(regs[i], p_out, n_out, meta) if len(regs[i]) != len(it.questions) or any(q.get("cite") for q in it.questions.values()): ans = _extras(it.state, it.questions, o["answers"], adps[i]) o["answers"] = {k: ans[k] for k in it.questions} outs.append(o) except (ValueError, KeyError, TypeError, AttributeError) as e: raise HTTPException(400, str(e)) return {"results": outs, "latency_ms": round((time.perf_counter() - t0) * 1000, 2)} def main(): import uvicorn ap = argparse.ArgumentParser() ap.add_argument("--host", default="0.0.0.0") ap.add_argument("--port", type=int, default=8095) a = ap.parse_args() uvicorn.run(app, host=a.host, port=a.port) if __name__ == "__main__": main()