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