| """TypeSafe SDK wire gateway to distinct shallow/full vLLM scoring engines.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| from contextlib import asynccontextmanager |
| from dataclasses import dataclass |
| import hashlib |
| import hmac |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import re |
| import uuid |
|
|
| import httpx |
| from fastapi import FastAPI, Request |
| from fastapi.responses import JSONResponse |
|
|
| from contracts import ( |
| ServingError, |
| answer, |
| compile_question, |
| prepare_request, |
| probabilities, |
| ) |
|
|
|
|
| def canonical_root(value): |
| url = httpx.URL(value) |
| path = url.path.rstrip("/").removesuffix("/v1").rstrip("/") |
| return str(url.copy_with(path=path)) |
|
|
|
|
| @dataclass(frozen=True) |
| class Settings: |
| high_url: str |
| low_url: str |
| high_model: str |
| low_model: str |
| api_key: str |
| model_name: str = "APUS-OpenJev-v1-9B" |
| max_body_bytes: int = 1_048_576 |
| max_questions: int = 16 |
| max_tokens: int = 8191 |
| max_inflight: int = 32 |
| timeout: float = 120.0 |
| high_model_path: str | None = None |
| low_model_path: str | None = None |
|
|
| def __post_init__(self): |
| for value in (self.high_url, self.low_url): |
| url = httpx.URL(value) |
| if ( |
| url.scheme not in ("http", "https") |
| or not url.host |
| or url.username |
| or url.password |
| or url.query |
| or url.fragment |
| ): |
| raise ValueError( |
| "Backend URLs must be HTTP(S) without credentials/query" |
| ) |
| if canonical_root(self.high_url) == canonical_root(self.low_url): |
| raise ValueError( |
| "Low and high must use distinct engine URLs; no silent full-depth fallback" |
| ) |
| if not all((self.high_model, self.low_model, self.api_key, self.model_name)): |
| raise ValueError( |
| "Both model aliases and a nonempty local API key are required" |
| ) |
| if self.high_model == self.low_model: |
| raise ValueError("Low and high engine model aliases must be distinct") |
|
|
| @classmethod |
| def from_env(cls): |
| return cls( |
| high_url=os.environ["JEV_HIGH_URL"], |
| low_url=os.environ["JEV_LOW_URL"], |
| high_model=os.environ["JEV_HIGH_MODEL"], |
| low_model=os.environ["JEV_LOW_MODEL"], |
| api_key=os.environ["JEV_API_KEY"], |
| high_model_path=os.environ["JEV_HIGH_MODEL_PATH"], |
| low_model_path=os.environ["JEV_LOW_MODEL_PATH"], |
| model_name=os.environ.get("JEV_MODEL_NAME", "APUS-OpenJev-v1-9B"), |
| ) |
|
|
|
|
| def unique_object(pairs): |
| result = {} |
| for key, value in pairs: |
| if key in result: |
| raise ValueError("Duplicate JSON keys") |
| result[key] = value |
| return result |
|
|
|
|
| def reject_constant(value): |
| raise ValueError("Nonfinite JSON values are forbidden") |
|
|
|
|
| def parse_scores(data, tokens): |
| try: |
| choices = data["choices"] |
| if len(choices) != 1 or len(choices[0]["logprobs"]["top_logprobs"]) != 1: |
| raise ValueError |
| values = choices[0]["logprobs"]["top_logprobs"][0] |
| if not isinstance(values, dict): |
| raise ValueError |
| by_id = {} |
| for key, value in values.items(): |
| if ( |
| not re.fullmatch(r"token_id:(0|[1-9]\d*)", key) |
| or type(value) not in (float, int) |
| or not math.isfinite(value) |
| or value > 0 |
| ): |
| raise ValueError |
| by_id[int(key.split(":")[1])] = value |
| return [by_id[token] for token in tokens] |
| except (KeyError, TypeError, ValueError, IndexError): |
| raise ServingError( |
| 502, |
| "invalid_backend_scores", |
| "Backend omitted valid candidate token logprobs", |
| ) from None |
|
|
|
|
| def endpoint(url, path): |
| root = url.rstrip("/") |
| return root + (path.removeprefix("/v1") if root.endswith("/v1") else path) |
|
|
|
|
| def file_sha(path): |
| return hashlib.sha256(Path(path).read_bytes()).hexdigest() |
|
|
|
|
| def verify_artifacts(settings): |
| """Bind configured backend roots to real model configs and export identity.""" |
| if not settings.high_model_path or not settings.low_model_path: |
| raise ValueError("Both JEV_HIGH_MODEL_PATH and JEV_LOW_MODEL_PATH are required") |
| roots = { |
| "high": Path(settings.high_model_path).resolve(), |
| "low": Path(settings.low_model_path).resolve(), |
| } |
| for effort, depth in (("high", 32), ("low", 16)): |
| root = roots[effort] |
| cfg = json.loads((root / "config.json").read_text()) |
| text = cfg.get("text_config", cfg) |
| if ( |
| cfg.get("model_type") not in ("qwen3_5", "qwen3_5_text") |
| or text.get("num_hidden_layers") != depth |
| or len(text.get("layer_types", [])) != depth |
| or text.get("hidden_size") != 4096 |
| ): |
| raise ValueError( |
| "Backend artifact does not declare the required execution depth" |
| ) |
| if (root / "adapter_config.json").exists() or ( |
| root / "EXPORT_INCOMPLETE" |
| ).exists(): |
| raise ValueError("Backend artifact is not a completed merged model") |
| if not (root / "model.safetensors.index.json").is_file(): |
| raise ValueError("Backend artifact has no merged weight index") |
| low, high = roots["low"], roots["high"] |
| proof = json.loads((low / "export-low-provenance.json").read_text()) |
| if ( |
| proof.get("status") != "verified" |
| or proof.get("actual_depth") != 16 |
| or proof.get("all_retained_tensors_bit_identical") is not True |
| ): |
| raise ValueError("Low artifact has not passed prefix export verification") |
| for key, path in { |
| "export_config_sha256": low / "config.json", |
| "export_index_sha256": low / "model.safetensors.index.json", |
| "source_config_sha256": high / "config.json", |
| "source_index_sha256": high / "model.safetensors.index.json", |
| "source_release_manifest_sha256": high / "release-manifest.json", |
| }.items(): |
| if proof.get(key) != file_sha(path): |
| raise ValueError( |
| "Low export provenance does not match the configured source or export metadata" |
| ) |
| return roots |
|
|
|
|
| def create_app(settings, tokenizer, transport=None, check_backends=True): |
| @asynccontextmanager |
| async def lifespan(app): |
| async with httpx.AsyncClient( |
| timeout=settings.timeout, |
| trust_env=False, |
| follow_redirects=False, |
| transport=transport, |
| ) as client: |
| app.state.client = client |
| if check_backends: |
| roots = verify_artifacts(settings) |
| for effort, url, alias in ( |
| ("low", settings.low_url, settings.low_model), |
| ("high", settings.high_url, settings.high_model), |
| ): |
| response = await client.get(endpoint(url, "/v1/models")) |
| response.raise_for_status() |
| matched = [m for m in response.json()["data"] if m["id"] == alias] |
| if ( |
| len(matched) != 1 |
| or not matched[0].get("root") |
| or Path(matched[0]["root"]).resolve() != roots[effort] |
| ): |
| raise RuntimeError( |
| "Engine alias/root does not match the verified depth artifact" |
| ) |
| app.state.ready = True |
| yield |
| app.state.ready = False |
|
|
| app = FastAPI(lifespan=lifespan) |
| app.state.ready = False |
| app.state.inflight = 0 |
|
|
| def authorize(request): |
| authorization = request.headers.get("authorization", "") |
| if not hmac.compare_digest( |
| authorization.encode(), ("Bearer " + settings.api_key).encode() |
| ): |
| raise ServingError(401, "unauthorized", "Invalid local API key") |
|
|
| @app.middleware("http") |
| async def request_id_header(request, call_next): |
| response = await call_next(request) |
| response.headers.setdefault("x-typesafe-request-id", str(uuid.uuid4())) |
| return response |
|
|
| @app.exception_handler(ServingError) |
| async def serving_error(_request, exc): |
| return JSONResponse( |
| status_code=exc.status, |
| content={"error": {"code": exc.code, "message": exc.message}}, |
| ) |
|
|
| @app.get("/health") |
| async def health(): |
| return JSONResponse( |
| status_code=200 if app.state.ready else 503, |
| content={ |
| "ready": app.state.ready, |
| "model": settings.model_name, |
| "efforts": {"low": 16, "high": 32}, |
| "routing": "separate-vllm-engines", |
| "automatic_routing": False, |
| }, |
| ) |
|
|
| @app.get("/v1/models") |
| async def models(request: Request): |
| authorize(request) |
| return { |
| "models": [ |
| {"name": name, "description": description, "release_date": "2026-09-21"} |
| for name, description in [ |
| ( |
| settings.model_name, |
| "APUS local decision model; explicit low16/high32; uncalibrated candidate scores.", |
| ), |
| ( |
| "jev-latest", |
| "Compatibility alias for " |
| + settings.model_name |
| + "; NOT the official Jev model.", |
| ), |
| ] |
| ] |
| } |
|
|
| @app.post("/v1/systemone") |
| async def systemone(request: Request): |
| authorize(request) |
| if not app.state.ready: |
| raise ServingError(503, "not_ready", "Backends are not ready") |
| if app.state.inflight >= settings.max_inflight: |
| raise ServingError(429, "busy", "Gateway in-flight budget exceeded") |
| app.state.inflight += 1 |
| try: |
| async with asyncio.timeout(settings.timeout): |
| raw = bytearray() |
| async for chunk in request.stream(): |
| raw.extend(chunk) |
| if len(raw) > settings.max_body_bytes: |
| raise ServingError( |
| 413, "request_too_large", "Request byte budget exceeded" |
| ) |
| try: |
| body = json.loads( |
| raw, |
| object_pairs_hook=unique_object, |
| parse_constant=reject_constant, |
| ) |
| if not isinstance(body, dict): |
| raise ValueError |
| except (ValueError, TypeError, RecursionError, UnicodeError): |
| raise ServingError( |
| 422, "invalid_json", "Malformed JSON object" |
| ) from None |
| effort = body.get("effort", "high") |
| if effort not in ("low", "high"): |
| raise ServingError( |
| 422, "invalid_effort", "effort must be low or high" |
| ) |
| questions = prepare_request( |
| body, |
| model_name=settings.model_name, |
| max_questions=settings.max_questions, |
| allow_score=True, |
| ) |
| compiled = [ |
| compile_question(q, tokenizer, max_tokens=settings.max_tokens) |
| for q in questions |
| ] |
| url, model, depth = ( |
| (settings.low_url, settings.low_model, 16) |
| if effort == "low" |
| else (settings.high_url, settings.high_model, 32) |
| ) |
| answers = {} |
| for item in compiled: |
| if not item.prompt_ids: |
| answers[item.question.key] = answer(item.question, [1.0]) |
| continue |
| payload = { |
| "model": model, |
| "prompt": item.prompt_ids, |
| "max_tokens": 1, |
| "n": 1, |
| "stream": False, |
| "echo": False, |
| "add_special_tokens": False, |
| "logprobs": len(item.candidate_token_ids), |
| "logprob_token_ids": item.candidate_token_ids, |
| "return_tokens_as_token_ids": True, |
| "temperature": 1.0, |
| "top_p": 1.0, |
| "top_k": -1, |
| "min_p": 0.0, |
| "frequency_penalty": 0.0, |
| "presence_penalty": 0.0, |
| "repetition_penalty": 1.0, |
| } |
| response = await app.state.client.post( |
| endpoint(url, "/v1/completions"), json=payload |
| ) |
| if response.status_code != 200: |
| raise ServingError( |
| 502, |
| "backend_error", |
| f"Inference backend returned HTTP {response.status_code}", |
| ) |
| try: |
| data = json.loads( |
| response.content, |
| object_pairs_hook=unique_object, |
| parse_constant=reject_constant, |
| ) |
| except (ValueError, TypeError, UnicodeError, RecursionError): |
| raise ServingError( |
| 502, "backend_json", "Backend returned invalid JSON" |
| ) from None |
| answers[item.question.key] = answer( |
| item.question, |
| probabilities(parse_scores(data, item.candidate_token_ids)), |
| ) |
| result = { |
| "model": settings.model_name, |
| "answers": answers, |
| "usage": { |
| "input_tokens": sum(len(c.prompt_ids) for c in compiled), |
| "output_tokens": sum(bool(c.prompt_ids) for c in compiled), |
| "billing_units": len(compiled), |
| }, |
| } |
| return JSONResponse( |
| result, |
| headers={ |
| "x-openjev-effort": effort, |
| "x-openjev-executed-layers": str( |
| depth if any(c.prompt_ids for c in compiled) else 0 |
| ), |
| "x-openjev-engine-model": model, |
| "x-typesafe-request-id": str(uuid.uuid4()), |
| }, |
| ) |
| except (TimeoutError, httpx.TimeoutException): |
| raise ServingError( |
| 504, "deadline_exceeded", "Inference request exceeded timeout" |
| ) from None |
| except httpx.HTTPError: |
| raise ServingError( |
| 502, "backend_transport", "Inference backend connection failed" |
| ) from None |
| finally: |
| app.state.inflight -= 1 |
|
|
| return app |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--host", default="127.0.0.1") |
| parser.add_argument("--port", type=int, default=30013) |
| args = parser.parse_args() |
| from transformers import AutoTokenizer |
| import uvicorn |
|
|
| settings = Settings.from_env() |
| if ( |
| Path(os.environ["JEV_MODEL_PATH"]).resolve() |
| != Path(settings.high_model_path).resolve() |
| ): |
| raise ValueError("Tokenizer must come from the configured full-depth model") |
| tokenizer = AutoTokenizer.from_pretrained( |
| os.environ["JEV_MODEL_PATH"], local_files_only=True |
| ) |
| uvicorn.run(create_app(settings, tokenizer), host=args.host, port=args.port) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
