File size: 4,889 Bytes
f67a19b ec0130c f67a19b ec0130c f67a19b ec0130c f67a19b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """Open typed-decision server with a Jev-compatible API.
POST /v1/systemone {"state": ..., "questions": {id: {"type": "choice|score|noul", "instructions": ..., "criteria": ...}},
"model": "kivok" | "migom" | "auto"}
-> {"model": ..., "answers": {id: {...TypeSafe answer shape...}}, "usage": {...}}
GET /v1/models
Models:
kivok - small (0.3B) and fast, runs on CPU;
migom - larger (2B), better at reasoning; GPU recommended;
auto - cascade: Kivok answers, questions where its top probability is below AUTO_THRESHOLD go to Migom.
Env: KIVOK_PATH, MIGOM_PATH (either may be empty), DEVICE (cpu|cuda|xpu), DTYPE (fp16|bf16|fp32; use fp16 on Intel XPU),
GRAPHS (0 disables graphs), AUTO_THRESHOLD (default 0.8),
API_KEY (if set, require "Authorization: Bearer <key>"), HOST (default 127.0.0.1), PORT (default 8010).
"""
import hmac, os, threading, time
from typing import Any, Dict, Optional
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
KIVOK_PATH = os.environ.get('KIVOK_PATH', '')
MIGOM_PATH = os.environ.get('MIGOM_PATH', '')
DEVICE = os.environ.get('DEVICE') or None
THRESHOLD = float(os.environ.get('AUTO_THRESHOLD', '0.8'))
API_KEY = os.environ.get('API_KEY', '')
DTYPE = os.environ.get('DTYPE', '') # fp16 on Intel XPU: bf16 gives wrong answers there
GRAPHS = os.environ.get('GRAPHS', '') # '0' disables CUDA/XPU graphs
app = FastAPI(title='typed-decision server', version='0.1')
_lock = threading.Lock()
_models: Dict[str, Any] = {}
def kivok():
if 'kivok' not in _models:
if not KIVOK_PATH:
raise HTTPException(503, 'kivok is not configured (KIVOK_PATH)')
import laya
_models['kivok'] = laya.load(KIVOK_PATH, device=DEVICE) if DEVICE else laya.load(KIVOK_PATH)
return _models['kivok']
def migom():
if 'migom' not in _models:
if not MIGOM_PATH:
raise HTTPException(503, 'migom is not configured (MIGOM_PATH)')
from decider.infer import Decider
import torch
dt = {'fp16': torch.float16, 'bf16': torch.bfloat16, 'fp32': torch.float32}.get(DTYPE)
_models['migom'] = Decider(MIGOM_PATH, device=DEVICE, dtype=dt, use_graphs=False if GRAPHS == '0' else None)
return _models['migom']
def top_prob(ans: dict) -> float:
if 'probabilities' in ans and ans['probabilities']:
return max(ans['probabilities'].values())
if 'noul' in ans:
p = float(ans['noul']); return max(p, 1 - p)
return 1.0
def ask_kivok(state, questions):
return kivok().predict(state, questions)
def ask_migom(state, questions):
return migom().system_one(state, questions)
class Request_(BaseModel):
state: Any
questions: Dict[str, Dict[str, Any]]
model: Optional[str] = 'auto'
@app.middleware('http')
async def auth(request: Request, call_next):
if API_KEY and request.url.path.startswith('/v1') and not hmac.compare_digest(
request.headers.get('authorization', ''), f'Bearer {API_KEY}'):
return JSONResponse({'detail': 'missing or invalid API key'}, 401)
return await call_next(request)
@app.get('/v1/models')
def models():
return {'data': [{'id': m, 'configured': bool(p)} for m, p in
(('kivok', KIVOK_PATH), ('migom', MIGOM_PATH), ('auto', KIVOK_PATH or MIGOM_PATH))]}
@app.post('/v1/systemone')
def systemone(req: Request_):
if not req.questions:
raise HTTPException(422, 'questions must not be empty')
name = (req.model or 'auto').lower()
t0 = time.time()
with _lock:
try:
if name == 'kivok' or (name == 'auto' and not MIGOM_PATH):
out = ask_kivok(req.state, req.questions); used = {k: 'kivok' for k in req.questions}
elif name == 'migom' or (name == 'auto' and not KIVOK_PATH):
out = ask_migom(req.state, req.questions); used = {k: 'migom' for k in req.questions}
elif name == 'auto':
out = ask_kivok(req.state, req.questions)
used = {k: 'kivok' for k in req.questions}
hard = {k: q for k, q in req.questions.items() if top_prob(out['answers'][k]) < THRESHOLD}
if hard:
second = ask_migom(req.state, hard)
for k in hard:
out['answers'][k] = second['answers'][k]; used[k] = 'migom'
else:
raise HTTPException(404, f'unknown model {name!r}')
except ValueError as e:
raise HTTPException(422, str(e))
return {'model': name, 'answers': out['answers'], 'routed': used,
'usage': {'latency_ms': round(1000 * (time.time() - t0), 1)}}
if __name__ == '__main__':
import uvicorn
uvicorn.run(app, host=os.environ.get('HOST', '127.0.0.1'), port=int(os.environ.get('PORT', '8010')))
|