kiapi 0.1.0__py3-none-any.whl

kiapi/__init__.py

@@ -0,0 +1,3 @@
+from importlib.metadata import version
+
+__version__ = version("kiapi")

kiapi/__main__.py

@@ -0,0 +1,4 @@
+from .cli.cli import main
+
+if __name__ == "__main__":
+    main()

kiapi/api/__init__.py

@@ -0,0 +1,26 @@
+from ._constants.require_auth import REQUIRE_AUTH
+from ._helpers.build_job_responses import build_job_responses
+from ._helpers.build_openapi import build_openapi
+from ._helpers.build_train_responses import build_train_responses
+from ._helpers.get_accept import get_accept
+from ._helpers.get_ctx import get_ctx
+from ._helpers.get_worker import get_worker
+from ._helpers.register_capability_endpoints import register_capability_endpoints
+from ._helpers.submit_and_maybe_wait import submit_and_maybe_wait
+from ._settings import settings_manager
+
+__all__ = [
+    # ._constants
+    "REQUIRE_AUTH",
+    # ._helpers
+    "build_job_responses",
+    "build_openapi",
+    "build_train_responses",
+    "get_accept",
+    "get_ctx",
+    "get_worker",
+    "register_capability_endpoints",
+    # ._settings
+    "settings_manager",
+    "submit_and_maybe_wait",
+]

kiapi/api/_constants/require_auth.py

@@ -0,0 +1,14 @@
+from fastapi import Depends, HTTPException, Request
+
+from .._settings import settings_manager
+
+
+async def _require_auth(request: Request) -> None:
+    settings = settings_manager.get_settings()
+    token = settings.auth_token
+
+    if token and request.headers.get("authorization") != f"Bearer {token}":
+        raise HTTPException(status_code=401, detail="unauthorized")
+
+
+REQUIRE_AUTH = [Depends(_require_auth)]

kiapi/api/_helpers/build_job_responses.py

@@ -0,0 +1,73 @@
+from typing import Any
+
+from pydantic import BaseModel, create_model
+
+from kiapi.core.job import Job
+
+from .._views.async_job_response import AsyncJobResponse
+
+
+def build_job_responses(
+    *media_types: str,
+    result_model: type[BaseModel] | None = None,
+) -> dict[int | str, dict[str, Any]]:
+    """OpenAPI response docs for sync/async generation endpoints.
+
+    ``mode=async`` returns 202 with a lightweight job handle. ``mode=sync``
+    returns a full Job JSON when the client asks for JSON, and may return a raw
+    artifact when the endpoint produces exactly one artifact and the client did
+    not request JSON.
+
+    ``result_model`` lets a family document the capability-specific shape of
+    ``Job.result`` (otherwise a free-form object): pass the family's result view
+    and the 200 JSON schema becomes a Job whose ``result`` is that typed model.
+    """
+
+    job_model: type[BaseModel] = Job
+    if result_model is not None:
+        job_model = create_model(
+            f"Job{result_model.__name__}",
+            __base__=Job,
+            result=(result_model | None, None),
+        )
+
+    # Pass the models via FastAPI's ``model`` key (not inline ``schema``) so
+    # FastAPI hoists their definitions into ``#/components/schemas`` and rewrites
+    # refs accordingly. Inlining ``model_json_schema()`` here would leave local
+    # ``#/$defs`` refs that don't resolve at the OpenAPI document root (ReDoc
+    # raises "Invalid reference token: $defs").
+    binary_content: dict[str, Any] = {
+        media_type: {"schema": {"type": "string", "format": "binary"}}
+        for media_type in media_types
+    }
+
+    return {
+        200: {
+            "model": job_model,
+            "description": (
+                "Sync result. Returns Job JSON with Accept: application/json; "
+                "single-artifact jobs may return raw bytes otherwise."
+            ),
+            "content": binary_content,
+            "headers": {
+                "X-Kiapi-File-Id": {
+                    "description": "Produced artifact file_id when raw bytes are returned.",
+                    "schema": {"type": "string"},
+                },
+                "X-Kiapi-Job-Id": {
+                    "description": "Job id when raw bytes are returned.",
+                    "schema": {"type": "string"},
+                },
+            },
+        },
+        202: {
+            "model": AsyncJobResponse,
+            "description": "Async job accepted. Poll GET /v1/jobs/{job_id}.",
+        },
+        400: {
+            "description": "Invalid request for the selected model or file reference."
+        },
+        422: {"description": "Request schema or validation error."},
+        503: {"description": "Model setup or memory budget error."},
+        504: {"description": "Sync request exceeded the configured timeout."},
+    }

kiapi/api/_helpers/build_openapi.py

@@ -0,0 +1,55 @@
+from fastapi import FastAPI
+from fastapi.openapi.utils import get_openapi
+from fastapi.routing import APIRoute
+
+from kiapi import __version__
+from kiapi.core.capability import CapabilitySpec
+
+
+def build_openapi(
+    app: FastAPI,
+    *,
+    title: str,
+    path_prefixes: tuple[str, ...],
+    include_paths: tuple[str, ...] = (),
+    description: str | None = None,
+    capability: CapabilitySpec | None = None,
+) -> dict:
+    routes = [
+        route
+        for route in app.routes
+        if isinstance(route, APIRoute)
+        and _path_matches(
+            route.path,
+            path_prefixes=path_prefixes,
+            include_paths=include_paths,
+        )
+    ]
+    schema = get_openapi(
+        title=title,
+        version=__version__,
+        description=description,
+        routes=routes,
+    )
+
+    if capability is not None:
+        schema["x-kiapi-capability"] = capability.name
+        schema["x-kiapi-domain"] = capability.domain
+        schema["x-kiapi-root-openapi"] = "/openapi.json"
+
+    return schema
+
+
+def _path_matches(
+    path: str,
+    *,
+    path_prefixes: tuple[str, ...],
+    include_paths: tuple[str, ...],
+) -> bool:
+    return any(_matches_prefix(path, prefix) for prefix in path_prefixes) or path in (
+        include_paths
+    )
+
+
+def _matches_prefix(path: str, prefix: str) -> bool:
+    return path == prefix or path.startswith(f"{prefix}/")

kiapi/api/_helpers/build_train_responses.py

@@ -0,0 +1,21 @@
+from typing import Any
+
+from .._views.async_job_response import AsyncJobResponse
+
+
+def build_train_responses() -> dict[int | str, dict[str, Any]]:
+    return {
+        # Pass the model via FastAPI's ``model`` key (not inline ``schema``) so
+        # FastAPI hoists nested ``$defs`` (e.g. JobStatus) into
+        # ``#/components/schemas`` and rewrites refs. Inlining
+        # ``model_json_schema()`` leaves dangling ``#/$defs/...`` refs that don't
+        # resolve at the document root (ReDoc raises "Invalid reference token:
+        # $defs").
+        202: {
+            "model": AsyncJobResponse,
+            "description": "Training job accepted. Poll GET /v1/jobs/{job_id}.",
+        },
+        400: {"description": "Invalid model or dataset FileRef."},
+        422: {"description": "Request schema or validation error."},
+        503: {"description": "Model setup or memory budget error."},
+    }

kiapi/api/_helpers/get_accept.py

@@ -0,0 +1,20 @@
+from typing import Annotated
+
+from fastapi import Header
+
+
+def get_accept(
+    accept: Annotated[
+        str | None,
+        Header(
+            alias="Accept",
+            description=(
+                "Response media type preference. application/json returns the Job "
+                "JSON; otherwise sync requests with one artifact return raw bytes "
+                "when possible."
+            ),
+            examples=["application/json", "image/png", "audio/wav", "video/mp4"],
+        ),
+    ] = None,
+) -> str | None:
+    return accept

kiapi/api/_helpers/get_ctx.py

@@ -0,0 +1,7 @@
+from fastapi import Request
+
+from kiapi.core.app import AppContext
+
+
+def get_ctx(request: Request) -> AppContext:
+    return request.app.state.ctx  # type: ignore

kiapi/api/_helpers/get_worker.py

@@ -0,0 +1,7 @@
+from fastapi import Request
+
+from kiapi.core.worker import Worker
+
+
+def get_worker(request: Request) -> Worker:
+    return request.app.state.worker  # type: ignore

kiapi/api/_helpers/register_capability_endpoints.py

@@ -0,0 +1,70 @@
+from fastapi import APIRouter, Request
+from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
+from fastapi.responses import HTMLResponse
+
+from kiapi.api._helpers.build_openapi import build_openapi
+from kiapi.api._views.capability_model_spec import CapabilityModelSpec
+from kiapi.core.capability import CapabilitySpec, capability_spec_registry
+from kiapi.core.model import model_registry
+
+
+def register_capability_endpoints(
+    router: APIRouter,
+    *,
+    name: str,
+    base_path: str,
+) -> None:
+    """Register the shared per-capability tail endpoints on ``router``.
+
+    Every capability exposes the same four boilerplate endpoints that differ
+    only by capability ``name`` (== family == model-registry token) and the
+    ``base_path`` under which they live:
+
+    - ``GET {base_path}/models``        — public model catalog for ``name``
+    - ``GET {base_path}/openapi.json``  — capability OpenAPI (operation layer)
+    - ``GET {base_path}/docs``          — Swagger UI for that OpenAPI
+    - ``GET {base_path}/redoc``         — ReDoc for that OpenAPI
+
+    The capability spec is resolved per-request (it may not be registered when
+    the router module is imported), mirroring the previous hand-written handlers.
+    """
+
+    def _spec() -> CapabilitySpec:
+        spec = capability_spec_registry.get(name)
+        if spec is None:
+            raise RuntimeError(f"kiapi capability {name!r} is not registered")
+        return spec
+
+    @router.get(f"{base_path}/models", response_model=list[CapabilityModelSpec])
+    async def list_models() -> list[CapabilityModelSpec]:
+        """List the servable models for this capability.
+
+        Returns the public catalog of every variant selectable via the ``model``
+        field on this capability's endpoints.
+        """
+        return [
+            CapabilityModelSpec.from_model_spec(spec)
+            for spec in model_registry.list_specs(name)
+        ]
+
+    @router.get(f"{base_path}/openapi.json", include_in_schema=False)
+    async def capability_openapi(request: Request) -> dict:
+        spec = _spec()
+        return build_openapi(
+            request.app,
+            title=spec.title,
+            description=spec.description,
+            path_prefixes=spec.path_prefixes,
+            include_paths=spec.include_paths,
+            capability=spec,
+        )
+
+    @router.get(f"{base_path}/docs", include_in_schema=False)
+    async def capability_docs() -> HTMLResponse:
+        spec = _spec()
+        return get_swagger_ui_html(openapi_url=spec.openapi_path, title=spec.title)
+
+    @router.get(f"{base_path}/redoc", include_in_schema=False)
+    async def capability_redoc() -> HTMLResponse:
+        spec = _spec()
+        return get_redoc_html(openapi_url=spec.openapi_path, title=spec.title)

kiapi/api/_helpers/submit_and_maybe_wait.py

@@ -0,0 +1,113 @@
+"""Shared sync/async job flow for generation capabilities (se, video, music).
+
+Chat and embedding are sync-only and return their provider-native payload
+directly. The generation capabilities instead:
+  - accept a ``mode`` ("sync" | "async") on a single endpoint,
+  - always create a Job (so it shows up in /v1/jobs and shares the single-flight
+    worker), and
+  - produce file artifacts referenced by ``file_id`` (see files.py).
+
+``submit_and_maybe_wait`` centralizes that: enqueue the job's thunk, then either
+return the job immediately (async, HTTP 202) or wait for it (sync), mapping the
+usual failures to HTTP status codes. A sync wait that exceeds ``sync_timeout_s``
+is an error (504) — no fallback to async.
+
+**Sync content negotiation**: a sync request that produces exactly
+one artifact returns that artifact's *raw bytes* by default — convenient for
+``curl -o out.png .../generate`` and direct clients (``Accept: */*``). The file_id
+and job_id ride along in response headers (``X-Kiapi-File-Id`` / ``X-Kiapi-Job-Id``)
+so the full metadata (seed, params, timings) stays fetchable via
+``GET /v1/files/{id}``. A client that sets ``Accept: application/json`` instead
+gets the Job dict — as do jobs with zero or multiple artifacts (e.g. acestep
+extract), and all async (202) responses.
+"""
+
+import asyncio
+
+from fastapi import HTTPException
+from fastapi.responses import FileResponse, JSONResponse, Response
+
+from kiapi.core.app import AppContext
+from kiapi.core.memory import MemoryBudgetError
+from kiapi.core.model import UnknownModelError
+from kiapi.core.setup import SetupRequiredError
+from kiapi.core.worker import JobThunk, Worker
+
+from .._settings import settings_manager
+
+
+async def submit_and_maybe_wait(
+    ctx: AppContext,
+    worker: Worker,
+    *,
+    type: str,
+    params: dict,
+    thunk: JobThunk,
+    mode: str,
+    prefer_raw: bool = True,
+    accept: str | None = None,
+) -> Response:
+    """Run a generation job. Returns the artifact bytes (sync, single artifact,
+    non-JSON Accept), the job dict (sync otherwise), or {job_id,...} (async)."""
+    settings = settings_manager.get_settings()
+
+    job = ctx.job_store.create(type=type, params=params)
+    fut = await worker.submit(job, thunk)
+
+    if mode == "async":
+        return JSONResponse(
+            status_code=202,
+            content={"job_id": job.id, "type": job.type, "status": job.status},
+        )
+
+    try:
+        await asyncio.wait_for(fut, timeout=settings.sync_timeout_s)
+    except TimeoutError:
+        raise HTTPException(  # noqa: B904
+            status_code=504,
+            detail=f"{type} job {job.id} exceeded sync timeout "
+            f"({settings.sync_timeout_s}s); it keeps running — poll /v1/jobs/{job.id}",
+        )
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    except SetupRequiredError as exc:
+        raise HTTPException(status_code=503, detail=str(exc))  # noqa: B904
+    except MemoryBudgetError as exc:
+        raise HTTPException(status_code=503, detail=str(exc))  # noqa: B904
+    except Exception as exc:
+        if exc.__class__.__name__ in (
+            "MediaError",
+            "CapabilityError",
+            "ValidationError",
+        ):
+            raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+        raise HTTPException(status_code=500, detail=f"{type} failed: {exc}")  # noqa: B904
+
+    # The worker mutates this same Job instance in place, so it is now updated.
+
+    # Sync content negotiation: return the single artifact's raw bytes unless the
+    # client asked for JSON (or there isn't exactly one artifact).
+    if prefer_raw and not _wants_json(accept) and len(job.artifacts) == 1:
+        raw = _raw_response(ctx, job, job.artifacts[0])
+        if raw is not None:
+            return raw
+
+    return JSONResponse(status_code=200, content=job.to_dict())
+
+
+def _wants_json(accept: str | None) -> bool:
+    return "application/json" in (accept or "").lower()
+
+
+def _raw_response(ctx: AppContext, job, file_id: str) -> Response | None:  # type: ignore
+    """Build a raw-bytes FileResponse for a single artifact, or None if the file
+    can't be resolved (caller falls back to JSON)."""
+    rec = ctx.file_store.get(file_id)
+    if rec is None:
+        return None
+    return FileResponse(
+        rec.path,
+        media_type=rec.content_type,
+        filename=rec.filename,
+        headers={"X-Kiapi-File-Id": rec.file_id, "X-Kiapi-Job-Id": job.id},
+    )

kiapi/api/_settings.py

@@ -0,0 +1,50 @@
+from pydantic import Field, SecretStr
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class APISettings(BaseSettings):
+    """Settings for API server binding, authentication, and sync request waits."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_",
+        extra="ignore",
+    )
+
+    host: str = Field(
+        default="127.0.0.1",
+        title="Bind host",
+        description=(
+            "Hostname or IP address that the kiapi API server binds to.\n"
+            "Use 127.0.0.1 for local-only access, or 0.0.0.0 in a trusted "
+            "environment when other devices on the same network should connect."
+        ),
+    )
+
+    port: int = Field(
+        default=8000,
+        title="Bind port",
+        description="TCP port that the kiapi API server binds to.\nThe usual value is 8000.",
+    )
+
+    auth_token: SecretStr | None = Field(
+        default=None,
+        title="API auth token",
+        description=(
+            "When set, API calls must use Bearer token authentication.\n"
+            "When unset, requests are accepted without authentication."
+        ),
+    )
+
+    sync_timeout_s: float = Field(
+        default=600.0,
+        title="Sync request wait seconds",
+        description=(
+            "Maximum number of seconds a sync-mode API call waits for job completion.\n"
+            "If this is exceeded, kiapi returns a timeout error instead of "
+            "switching the request to async mode."
+        ),
+    )
+
+
+settings_manager = SettingsManager(APISettings)

kiapi/api/_views/async_job_response.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel, Field
+
+from kiapi.core.job import JobStatus
+
+
+class AsyncJobResponse(BaseModel):
+    job_id: str = Field(
+        ...,
+        description="In-memory job id. Poll GET /v1/jobs/{job_id} to inspect status, progress, result, and artifacts.",
+        examples=["job_0123456789abcdef"],
+    )
+    type: str = Field(
+        ...,
+        description="Job type. Generation APIs use values such as zimage, flux2-edit, or acestep-extract.",
+        examples=["zimage"],
+    )
+    status: JobStatus = Field(
+        ...,
+        description="Initial job status. Async responses are normally queued unless the worker starts immediately.",
+        examples=["queued"],
+    )

kiapi/api/_views/capability_model_spec.py

@@ -0,0 +1,49 @@
+from pydantic import BaseModel, Field
+
+from kiapi.core.model import ModelSpec
+
+
+class CapabilityModelSpec(BaseModel):
+    """Public model discovery entry for capability-specific model lists."""
+
+    name: str = Field(
+        ...,
+        description="Model variant name to pass in the request model field.",
+        examples=["turbo"],
+    )
+    family: str = Field(
+        ...,
+        description="Capability family that resolves this model variant.",
+        examples=["zimage"],
+    )
+    domain: str = Field(
+        ...,
+        description="Capability domain used for grouping model lists.",
+        examples=["image"],
+    )
+    aliases: tuple[str, ...] = Field(
+        default_factory=tuple,
+        description="Alternative names that also resolve to this model.",
+        examples=[["omni", "qwen3-omni-30b"]],
+    )
+    default: bool = Field(
+        default=False,
+        description="Whether this is the default model when the request omits model.",
+        examples=[True],
+    )
+    features: tuple[str, ...] = Field(
+        default_factory=tuple,
+        description="Handler-declared modalities and features supported by this model.",
+        examples=[["text", "image"]],
+    )
+
+    @classmethod
+    def from_model_spec(cls, spec: ModelSpec) -> "CapabilityModelSpec":
+        return cls(
+            name=spec.name,
+            family=spec.family,
+            domain=spec.domain,
+            aliases=spec.aliases,
+            default=spec.default,
+            features=tuple(sorted(spec.features)),
+        )