File size: 8,505 Bytes
dbb7b0f | 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """Serve the trained checkpoint through the TypeSafe decision API."""
from __future__ import annotations
import base64
import binascii
import hmac
import os
import threading
import time
import uuid
from _thread import LockType
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Literal, cast
from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import HTMLResponse, JSONResponse, Response
from PIL import Image, UnidentifiedImageError
from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
from starlette.concurrency import run_in_threadpool
from starlette.middleware.base import RequestResponseEndpoint
from autojev.types import Answer, DecisionInput, DecisionResponse, JSONValue, Question as DecisionQuestion
if TYPE_CHECKING:
from autojev.model import DecisionModel
type Content = str | dict[str, JsonValue] | list[JsonValue]
DEFAULT_MODEL = "autojev-qwen3.8-27b"
ALIASES = {"autojev", "jev-latest", "jev-preview", "jev-1.13.0", DEFAULT_MODEL}
@dataclass
class Service:
model: DecisionModel | None = None
name: str = DEFAULT_MODEL
checkpoint: str = "checkpoints/selected"
release_date: str = ""
lock: LockType = field(default_factory=threading.Lock)
service = Service()
class Question(BaseModel):
model_config = ConfigDict(extra="forbid")
instructions: Content | None = None
class Choice(Question):
type: Literal["choice"]
criteria: dict[str, Content | None] = Field(min_length=1, max_length=255)
class Score(Question):
type: Literal["score"]
criteria: list[Content] = Field(min_length=2, max_length=10)
class Noul(Question):
type: Literal["noul"]
criteria: dict[Literal["true", "false"], Content | None] | None = None
class EvaluationRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
model: str
state: Content
questions: dict[str, Annotated[Choice | Score | Noul, Field(discriminator="type")]] = Field(min_length=1)
images: list[str] = Field(default_factory=list, max_length=4)
@field_validator("model")
@classmethod
def known_model(cls, value: str) -> str:
if value not in ALIASES | {service.name}:
raise ValueError(f"Unknown model. Use {service.name} or jev-latest.")
return value
@field_validator("images")
@classmethod
def valid_images(cls, values: list[str]) -> list[str]:
for value in values:
if len(value) > 12_000_000:
raise ValueError("Each image must be at most 8 MB before base64 encoding.")
header, separator, encoded = value.partition(",")
if not separator or header not in {
"data:image/png;base64", "data:image/jpeg;base64", "data:image/webp;base64",
}:
raise ValueError("Images must be base64 PNG, JPEG, or WebP data URLs.")
try:
content = base64.b64decode(encoded, validate=True)
if len(content) > 8_000_000:
raise ValueError("Each image must be at most 8 MB.")
with Image.open(BytesIO(content)) as image:
if image.width * image.height > 16_000_000:
raise ValueError("Each image must have at most 16 million pixels.")
if image.format not in {"PNG", "JPEG", "WEBP"}:
raise ValueError("Unsupported image format.")
image.verify()
except (binascii.Error, OSError, SyntaxError, UnidentifiedImageError, Image.DecompressionBombError) as error:
raise ValueError("Invalid image data.") from error
return values
def authenticate(authorization: str | None = Header(default=None)) -> None:
key = os.getenv("AUTOJEV_API_KEY")
if key and not hmac.compare_digest((authorization or "").encode(), f"Bearer {key}".encode()):
raise HTTPException(401, "Missing or invalid API key.", headers={"WWW-Authenticate": "Bearer"})
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
from autojev.model import DecisionModel
service.checkpoint = os.getenv("AUTOJEV_CHECKPOINT", "checkpoints/selected")
service.model = await run_in_threadpool(DecisionModel, checkpoint=service.checkpoint)
service.name = f"autojev-{service.model.base_model.rsplit('/', 1)[-1].lower()}"
modified = (Path(service.checkpoint) / "decision_config.json").stat().st_mtime
service.release_date = datetime.fromtimestamp(modified, timezone.utc).date().isoformat()
try:
yield
finally:
service.model = None
app = FastAPI(title="AutoJev", version="0.2.0", lifespan=lifespan)
@app.middleware("http")
async def request_metadata(request: Request, call_next: RequestResponseEndpoint) -> Response:
started, identifier = time.perf_counter(), uuid.uuid4().hex
response = await call_next(request)
response.headers["x-typesafe-request-id"] = identifier
response.headers["x-request-id"] = identifier
response.headers["server-timing"] = f"total;dur={(time.perf_counter() - started) * 1000:.1f}"
return response
@app.exception_handler(RequestValidationError)
async def validation_error(request: Request, error: RequestValidationError) -> JSONResponse:
return JSONResponse(status_code=422, content={"detail": [
{"loc": item["loc"], "msg": item["msg"], "type": item["type"]}
for item in error.errors()
]})
@app.get("/", response_class=HTMLResponse, include_in_schema=False)
def playground() -> str:
return Path(__file__).with_name("playground.html").read_text()
@app.get("/health", response_model=None)
def health() -> dict[str, JSONValue]:
return {"status": "ready" if service.model is not None else "loading", "model": service.name,
"checkpoint": service.checkpoint, "authentication": bool(os.getenv("AUTOJEV_API_KEY")),
"modalities": ["text", "image"]}
@app.get("/v1/models", dependencies=[Depends(authenticate)], response_model=None)
def models() -> dict[str, JSONValue]:
return {"models": [
{"name": name, "description": "Local AutoJev text and image decisions.", "release_date": service.release_date}
for name in sorted(ALIASES | {service.name})
]}
def predict(model: DecisionModel, body: EvaluationRequest) -> DecisionResponse:
import torch
from autojev.model import answer
questions = {key: cast(DecisionQuestion, question.model_dump(exclude_none=True))
for key, question in body.questions.items()}
identifiers = list(questions)
rows: list[DecisionInput] = [{"state": body.state, "question": question, "images": list(body.images)}
for question in questions.values()]
answers: dict[str, Answer] = {}
input_tokens = 0
with torch.inference_mode():
for start in range(0, len(rows), 8):
batch = model.prepare(rows[start:start + 8])
distributions: list[list[float]] = (model(batch) / model.temperature).softmax(-1).cpu().tolist()
for identifier, values, count in zip(identifiers[start:start + 8], distributions, batch.counts, strict=True):
answers[identifier] = answer(questions[identifier], values[:count])
input_tokens += batch.input_tokens
return {"model": service.name, "answers": answers, "usage": {"input_tokens": input_tokens, "output_tokens": 0}}
@app.post("/v1/systemone", dependencies=[Depends(authenticate)], response_model=None)
async def system_one(body: EvaluationRequest) -> DecisionResponse:
model = service.model
if model is None:
raise HTTPException(503, "The model is not ready.")
if not service.lock.acquire(blocking=False):
raise HTTPException(529, "The model is busy. Retry shortly.", headers={"Retry-After": "1"})
try:
return await run_in_threadpool(predict, model, body)
except ValueError as error:
raise HTTPException(422, str(error)) from error
finally:
service.lock.release()
def main() -> None:
import uvicorn
uvicorn.run("autojev.server:app", host=os.getenv("AUTOJEV_HOST", "127.0.0.1"), port=int(os.getenv("PORT", "8000")))
if __name__ == "__main__":
main()
|