gpc1_server/api.py

from __future__ import annotations

import asyncio
import base64
import binascii
from contextlib import asynccontextmanager
from io import BytesIO
import json
import math
import os
import secrets
import time
import uuid
from typing import Any

from fastapi import FastAPI, Header, Request
from fastapi.responses import JSONResponse
from PIL import Image
from pydantic import ValidationError
from jsonschema import Draft202012Validator

from .compiler import NumericField
from .errors import APIError
from .runtime import Backend, TransformersBackend
from .schemas import ChatRequest

MAX_BODY_BYTES = int(os.getenv("GPC1_MAX_BODY_BYTES", str(16 * 1024 * 1024)))
MAX_IMAGE_BYTES = int(os.getenv("GPC1_MAX_IMAGE_BYTES", str(3 * 1024 * 1024)))
MAX_IMAGE_PIXELS = int(os.getenv("GPC1_MAX_IMAGE_PIXELS", "16777216"))
Image.MAX_IMAGE_PIXELS = MAX_IMAGE_PIXELS
MAX_QUEUE_WAIT_SECONDS = float(os.getenv("GPC1_MAX_QUEUE_WAIT_SECONDS", "2"))
SUPPORTED_SCHEMA_KEYS = {"type", "properties", "required", "additionalProperties", "description",
                         "enum", "items", "minimum", "maximum"}


def error_response(error: APIError, request_id: str | None = None) -> JSONResponse:
    headers = {"x-request-id": request_id} if request_id else {}
    if error.status_code == 429:
        headers["Retry-After"] = "1"
    return JSONResponse(error.body(), status_code=error.status_code, headers=headers)


def validate_schema(schema: Any, path: str = "response_format.json_schema.schema", depth: int = 0) -> None:
    if depth > 16: raise APIError("JSON Schema nesting exceeds 16",param=path)
    if not isinstance(schema, dict):
        raise APIError("JSON Schema must be an object", param=path)
    unknown = set(schema) - SUPPORTED_SCHEMA_KEYS
    if unknown:
        raise APIError(f"unsupported JSON Schema keywords: {sorted(unknown)}", param=path)
    def finite(value: Any) -> bool:
        if isinstance(value,float): return math.isfinite(value)
        if isinstance(value,dict): return all(finite(k) and finite(v) for k,v in value.items())
        if isinstance(value,list): return all(finite(v) for v in value)
        return True
    if not finite(schema):
        raise APIError("JSON Schema contains a non-finite number", param=path)
    kind = schema.get("type")
    if not isinstance(kind, str) or kind not in {"object", "array", "string", "number", "integer", "boolean", "null"}:
        raise APIError("schema type must be one supported explicit type", param=f"{path}.type")
    if kind == "object":
        properties = schema.get("properties")
        required = schema.get("required")
        if not isinstance(properties, dict) or not isinstance(required, list):
            raise APIError("objects require properties and required", param=path)
        if any(not isinstance(v, str) for v in required) or len(required) != len(set(required)):
            raise APIError("required must contain unique property-name strings", param=f"{path}.required")
        if set(required) != set(properties) or schema.get("additionalProperties") is not False:
            raise APIError("strict objects require every property and additionalProperties=false", param=path)
        for name, child in properties.items(): validate_schema(child, f"{path}.properties.{name}",depth+1)
    if kind == "array":
        validate_schema(schema.get("items"), f"{path}.items",depth+1)


def validate_instance(value: Any, schema: dict[str, Any], path: str = "output") -> None:
    def finite(v: Any) -> bool:
        if isinstance(v, float): return math.isfinite(v)
        if isinstance(v, dict): return all(finite(k) and finite(x) for k, x in v.items())
        if isinstance(v, list): return all(finite(x) for x in v)
        return True
    if not finite(value): raise APIError(f"{path} contains a non-finite number", param="gpc1")
    errors=sorted(Draft202012Validator(schema).iter_errors(value),key=lambda e:list(e.path))
    if errors: raise APIError(f"{path} violates schema: {errors[0].message}",param="gpc1")


def _decode_image(url: str) -> Image.Image:
    if not url.startswith("data:image/"):
        raise APIError("only inline data:image URIs are accepted; remote URLs are disabled",
                       param="messages.content.image_url.url", code="remote_image_disabled")
    try:
        header, encoded = url.split(",", 1)
        media = header[5:].split(";", 1)[0].lower()
        if header != f"data:{media};base64" or media not in {"image/png", "image/jpeg", "image/webp"}:
            raise ValueError
        raw = base64.b64decode(encoded, validate=True)
    except (ValueError, binascii.Error) as error:
        raise APIError("invalid base64 image data URI", param="messages.content.image_url.url") from error
    if len(raw) > MAX_IMAGE_BYTES:
        raise APIError("decoded image exceeds byte limit", status_code=413, code="image_too_large")
    try:
        image = Image.open(BytesIO(raw))
        if image.width * image.height > MAX_IMAGE_PIXELS:
            raise APIError("decoded image exceeds pixel limit", status_code=413, code="image_too_large")
        image.verify()
        image = Image.open(BytesIO(raw)).convert("RGB")
    except APIError:
        raise
    except Exception as error:
        raise APIError("image payload failed safe decode", param="messages.content.image_url.url") from error
    return image


def flatten_messages(messages: list[Any]) -> tuple[str, Image.Image | None]:
    texts=[]; image=None
    for message in messages:
        if isinstance(message.content, str):
            texts.append(f"{message.role.upper()}: {message.content}")
            continue
        for part in message.content:
            if set(part) == {"type", "text"} and part["type"] == "text" and isinstance(part["text"], str):
                texts.append(f"{message.role.upper()}: {part['text']}")
            elif set(part) == {"type", "image_url"} and part["type"] == "image_url" \
                    and isinstance(part["image_url"], dict) and set(part["image_url"]) <= {"url", "detail"}:
                if image is not None:
                    raise APIError("at most one image is supported", param="messages")
                image = _decode_image(str(part["image_url"].get("url", "")))
            else:
                raise APIError("unsupported message content part", param="messages.content")
    context="\n".join(texts)
    if not context.strip(): raise APIError("messages must include text context", param="messages")
    return context, image


def create_app(backend: Backend | None = None) -> FastAPI:
    selected = backend
    if int(os.getenv("GPC1_MAX_CONCURRENCY", "1")) != 1:
        raise RuntimeError("canonical request-local runtime requires GPC1_MAX_CONCURRENCY=1")
    max_pending = int(os.getenv("GPC1_MAX_PENDING_REQUESTS", "4"))
    if max_pending < 0:
        raise RuntimeError("GPC1_MAX_PENDING_REQUESTS must be nonnegative")
    if not math.isfinite(MAX_QUEUE_WAIT_SECONDS) or MAX_QUEUE_WAIT_SECONDS <= 0:
        raise RuntimeError("GPC1_MAX_QUEUE_WAIT_SECONDS must be finite and positive")
    semaphore = asyncio.Semaphore(1)
    admitted = 0

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        nonlocal selected
        if selected is None:
            if not os.getenv("GPC1_API_KEY"): raise RuntimeError("GPC1_API_KEY is required")
            selected = TransformersBackend()
        yield

    app = FastAPI(title="GPC-1 OpenAI-style API", version="0.1.0", lifespan=lifespan)

    async def run_native(call, *args):
        task=asyncio.create_task(asyncio.to_thread(call,*args))
        try: return await asyncio.shield(task)
        except asyncio.CancelledError:
            # Retain the model lease until the worker thread finishes.
            while not task.done():
                try:
                    await asyncio.shield(task)
                except asyncio.CancelledError:
                    continue
                except Exception:
                    break
            if not task.cancelled():
                task.exception()
            raise

    @app.exception_handler(APIError)
    async def api_error_handler(_: Request, error: APIError): return error_response(error)

    @app.middleware("http")
    async def admission(request: Request, call_next):
        request_id = "req_" + uuid.uuid4().hex
        if request.method == "POST" and "content-length" not in request.headers:
            return error_response(APIError("Content-Length is required", 411, code="length_required"), request_id)
        try: content_length = int(request.headers.get("content-length", "0") or 0)
        except ValueError: return error_response(APIError("invalid Content-Length", 400), request_id)
        if content_length < 0 or content_length > MAX_BODY_BYTES:
            return error_response(APIError("request body exceeds byte limit", 413, code="request_too_large"), request_id)
        configured = os.getenv("GPC1_API_KEY")
        supplied = request.headers.get("authorization", "")
        if request.url.path not in {"/healthz", "/readyz"} and configured and not secrets.compare_digest(supplied, f"Bearer {configured}"):
            return error_response(APIError("invalid API key", 401, "authentication_error", code="invalid_api_key"), request_id)
        response = await call_next(request); response.headers["x-request-id"] = request_id; return response

    @app.get("/healthz", include_in_schema=False)
    async def health():
        return {"status": "ok"}

    @app.get("/readyz", include_in_schema=False)
    async def ready():
        if selected is None:
            return JSONResponse({"status": "not_ready"}, status_code=503)
        return {"status": "ready"}

    @app.get("/v1/models")
    async def models():
        assert selected is not None
        return {"object": "list", "data": [{"id": selected.model_id, "object": "model",
                "created": 0, "owned_by": "gpc-1"}]}

    @app.post("/v1/chat/completions")
    async def completions(request: Request):
        nonlocal admitted
        if admitted >= 1 + max_pending:
            raise APIError("server is at capacity; retry with backoff", 429, "rate_limit_error", code="server_busy")
        admitted += 1
        try:
            try:
                await asyncio.wait_for(semaphore.acquire(), timeout=MAX_QUEUE_WAIT_SECONDS)
            except TimeoutError as error:
                raise APIError("server is at capacity; retry with backoff", 429, "rate_limit_error", code="server_busy") from error
            try:
                return await complete_admitted(request)
            finally:
                semaphore.release()
        finally:
            admitted -= 1

    async def complete_admitted(request: Request):
        assert selected is not None
        try:
            chunks=[]; actual=0
            async for chunk in request.stream():
                actual += len(chunk)
                if actual > MAX_BODY_BYTES: raise APIError("request body exceeds byte limit",413,code="request_too_large")
                chunks.append(chunk)
            raw=await asyncio.to_thread(json.loads, b"".join(chunks),
                                        parse_constant=lambda value: (_ for _ in ()).throw(
                                            ValueError("JSON contains a non-finite number; use finite numeric values")))
            parsed = await asyncio.to_thread(ChatRequest.model_validate, raw)
        except (json.JSONDecodeError, UnicodeDecodeError, ValueError, ValidationError) as error:
            raise APIError(str(error), param="body") from error
        if parsed.model != selected.model_id:
            raise APIError(f"model {parsed.model!r} not found", 404, "invalid_request_error", "model", "model_not_found")
        wrapper = parsed.response_format.json_schema
        if wrapper.get("strict") is not True or not isinstance(wrapper.get("name"), str) or set(wrapper) != {"name", "strict", "schema"}:
            raise APIError("json_schema requires exactly name, strict=true, and schema", param="response_format.json_schema")
        await asyncio.to_thread(validate_schema, wrapper["schema"])
        try: await asyncio.to_thread(Draft202012Validator.check_schema, wrapper["schema"])
        except Exception as error: raise APIError(f"invalid JSON Schema: {error}",param="response_format.json_schema.schema") from error
        context, image = await asyncio.to_thread(flatten_messages, parsed.messages)
        if parsed.temperature is not None or parsed.max_completion_tokens is not None:
            raise APIError("temperature and max_completion_tokens are unsupported; omit them",param="temperature")
        try:
            if parsed.gpc1.mode == "atomic":
                if image is not None: raise APIError("atomic image mode is not supported",param="messages")
                choices=parsed.gpc1.choices or []
                try:
                    for choice in choices: json.dumps(choice,allow_nan=False)
                except (TypeError,ValueError) as error: raise APIError("atomic choices must be finite JSON",param="gpc1.choices") from error
                field_key=str(parsed.gpc1.field_key)
                schema=wrapper["schema"]
                properties=schema.get("properties",{})
                if set(properties)!={field_key} or properties[field_key].get("enum") != choices:
                    raise APIError("atomic field and choices must exactly match the schema enum",param="gpc1")
                result=await run_native(selected.atomic,context,field_key,choices)
                value={field_key:result["value"]}
            elif parsed.gpc1.mode == "numeric101":
                raw_fields = parsed.gpc1.numeric_fields or []
                if not 1 <= len(raw_fields) <= 512:
                    raise APIError("numeric101 requires 1..512 fields", param="gpc1.numeric_fields")
                expected_field_keys = {"key", "description", "minimum", "maximum", "unit"}
                optional_field_keys = {"aliases", "reference"}
                if any(not expected_field_keys <= set(v) or set(v) - expected_field_keys - optional_field_keys for v in raw_fields):
                    raise APIError("numeric field keys are invalid", param="gpc1.numeric_fields")
                fields = [NumericField(
                    key=str(v["key"]), description=str(v["description"]),
                    minimum=float(v["minimum"]), maximum=float(v["maximum"]), unit=str(v["unit"]),
                    aliases=tuple(map(str, v.get("aliases", []))), reference=str(v.get("reference", "declared_by_caller")))
                    for v in raw_fields]
                if len({field.key for field in fields}) != len(fields):
                    raise APIError("numeric field keys must be unique", param="gpc1.numeric_fields")
                names=[name for field in fields for name in (field.key,*field.aliases)]
                if len(names) != len(set(names)):
                    raise APIError("numeric canonical keys and aliases must be globally unique",param="gpc1.numeric_fields")
                properties = wrapper["schema"].get("properties", {})
                if set(properties) != {field.key for field in fields} or any(
                    properties[f.key].get("type") != "number"
                    or properties[f.key].get("minimum") != f.minimum
                    or properties[f.key].get("maximum") != f.maximum for f in fields):
                    raise APIError("numeric fields and bounds must exactly match number properties", param="gpc1.numeric_fields")
                result = await run_native(selected.numeric, context, fields, image)
                value = {key: field["value"] for key, field in result["fields"].items()}
            else:
                if image is not None: raise APIError("finite_joint image mode is not supported", param="messages")
                records = parsed.gpc1.allowed_records or []
                if not 2 <= len(records) <= 255: raise APIError("finite_joint requires 2..255 allowed records", param="gpc1.allowed_records")
                try:
                    for record in records: json.dumps(record,allow_nan=False)
                except (TypeError,ValueError) as error: raise APIError("allowed records must be finite JSON",param="gpc1.allowed_records") from error
                for index, record in enumerate(records):
                    await asyncio.to_thread(validate_instance, record, wrapper["schema"], f"allowed_records[{index}]")
                encoded = [json.dumps(v, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False) for v in records]
                if len(encoded) != len(set(encoded)): raise APIError("allowed records must be unique", param="gpc1.allowed_records")
                result = await run_native(selected.finite_joint, context, wrapper["schema"], records)
                value = result["value"]
            await asyncio.to_thread(validate_instance, value, wrapper["schema"])
        except APIError: raise
        except (KeyError, TypeError, ValueError) as error:
            raise APIError(str(error), param="gpc1") from error
        content = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
        now = int(time.time()); completion_id = "chatcmpl_" + uuid.uuid4().hex
        return {"id": completion_id, "object": "chat.completion", "created": now,
                "model": selected.model_id, "choices": [{"index": 0,
                "message": {"role": "assistant", "content": content}, "finish_reason": "stop", "logprobs": None}],
                "usage": None,
                "gpc1": result}
    return app


app = create_app()