kiapi 0.1.0__py3-none-any.whl

kiapi/api/chat/router.py

@@ -0,0 +1,253 @@
+"""Chat API: OpenAI-compatible ``POST /v1/chat/completions``.
+
+Internally still a job (so it shows up in /v1/jobs and shares the single-flight
+worker). Non-streaming callers wait for the completion dict. Streaming callers
+receive OpenAI-style SSE chunks while the worker thread generates; the final job
+result remains inspectable through /v1/jobs.
+"""
+
+import asyncio
+import base64
+import binascii
+import json
+import logging
+import re
+from collections.abc import AsyncIterator
+from copy import deepcopy
+from typing import Any
+
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import StreamingResponse
+
+from kiapi.api import (
+    REQUIRE_AUTH,
+    get_ctx,
+    get_worker,
+    register_capability_endpoints,
+)
+from kiapi.api._settings import settings_manager
+from kiapi.capabilities.chat import ChatRequest, handle_chat
+from kiapi.core.app import AppContext
+from kiapi.core.memory import MemoryBudgetError
+from kiapi.core.model import UnknownModelError, model_registry
+from kiapi.core.setup import SetupRequiredError
+from kiapi.core.worker import Worker
+
+from ._views.chat_completion_response import ChatCompletionResponse
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(dependencies=REQUIRE_AUTH)
+
+_BASE64_PLACEHOLDER = "<base64_encoded_data>"
+
+_DATA_URL_BASE64 = re.compile(
+    r"^(?P<prefix>data:[^,]*;base64,)(?P<data>.*)$",
+    re.DOTALL,
+)
+
+
+@router.post(
+    "/v1/chat/completions",
+    response_model=None,
+    responses={
+        200: {
+            "model": ChatCompletionResponse,
+            "description": (
+                "Non-streaming: the full `chat.completion` object. "
+                "Streaming (`stream: true`): an OpenAI `text/event-stream` of "
+                "`chat.completion.chunk` events, each `data: {...}`, terminated "
+                "by `data: [DONE]`. Tool calls arrive as `delta.tool_calls`; the "
+                "final chunk carries `finish_reason`."
+            ),
+            "content": {
+                "text/event-stream": {"schema": {"type": "string", "format": "binary"}},
+            },
+        },
+        400: {"description": "Invalid request, unknown model, or bad input/modality."},
+        503: {"description": "Model not set up, or memory budget exceeded."},
+        504: {"description": "Sync timeout exceeded; the job keeps running."},
+    },
+)
+async def chat_completions(
+    req: ChatRequest,
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> dict | StreamingResponse:
+    """Generate a chat completion (OpenAI-compatible).
+
+    Accepts the OpenAI `chat.completions` request shape with multimodal
+    `messages` (text / image / audio / video), function `tools`, and
+    `tool_choice`. The resolved `model` (see GET /v1/models) determines which
+    input modalities are accepted.
+
+    Every request runs as a single-flight job, so it appears in /v1/jobs and is
+    serialized with all other generation. Non-streaming callers wait up to the
+    sync timeout and receive the full `chat.completion`; on timeout the job
+    keeps running and can be polled at /v1/jobs/{id}. Set `stream: true` to
+    receive incremental `chat.completion.chunk` SSE events instead.
+    """
+    settings = settings_manager.get_settings()
+
+    logger.debug(
+        "chat/completions request validated: %s",
+        _redacted_chat_request_dump(req),
+    )
+
+    if req.stream:
+        try:
+            model_registry.resolve("chat", req.model)
+        except UnknownModelError as exc:
+            raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+
+        loop = asyncio.get_running_loop()
+        queue: asyncio.Queue[dict | None] = asyncio.Queue()
+
+        def emit(chunk: dict) -> None:
+            loop.call_soon_threadsafe(queue.put_nowait, chunk)
+
+        def finish_stream() -> None:
+            loop.call_soon_threadsafe(queue.put_nowait, None)
+
+        def thunk():  # type: ignore
+            try:
+                return handle_chat(ctx, req, emit=emit)
+            finally:
+                finish_stream()
+
+        job = ctx.job_store.create(
+            type="chat", params={"model": req.model, "stream": True}
+        )
+        fut = await worker.submit(job, thunk)
+
+        async def events() -> AsyncIterator[str]:
+            try:
+                while True:
+                    item = await queue.get()
+                    if item is None:
+                        break
+                    yield _sse(item)
+
+                try:
+                    await asyncio.shield(fut)
+                except Exception as exc:
+                    yield _sse(_stream_error(exc))
+                yield _sse("[DONE]")
+            except asyncio.CancelledError:
+                # The running worker job is not preemptible; let it finish and
+                # keep its final state in the job store.
+                return
+
+        return StreamingResponse(
+            events(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
+    job = ctx.job_store.create(type="chat", params={"model": req.model})
+    fut = await worker.submit(job, lambda: handle_chat(ctx, req))
+
+    try:
+        return await asyncio.wait_for(fut, timeout=settings.sync_timeout_s)
+    except TimeoutError:
+        raise HTTPException(  # noqa: B904
+            status_code=504,
+            detail=f"chat 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:
+        # MediaError/CapabilityError (bad input / unsupported modality) → 400.
+        if exc.__class__.__name__ in ("MediaError", "CapabilityError"):
+            raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+        raise HTTPException(status_code=500, detail=f"generation failed: {exc}")  # noqa: B904
+
+
+register_capability_endpoints(router, name="chat", base_path="/v1/chat")
+
+
+def _redact_data_url(value: str) -> str:
+    m = _DATA_URL_BASE64.match(value)
+    if not m:
+        return value
+    return f"{m.group('prefix')}{_BASE64_PLACEHOLDER}"
+
+
+def _looks_like_base64(value: str) -> bool:
+    """Heuristic for bare media payloads; avoids redacting ordinary short text."""
+    s = value.strip()
+    if len(s) < 128:
+        return False
+    try:
+        base64.b64decode(s, validate=True)
+    except (binascii.Error, ValueError):
+        return False
+    return True
+
+
+def _redact_media_source(value: Any, *, bare_base64: bool = False) -> Any:
+    if not isinstance(value, str):
+        return value
+    redacted = _redact_data_url(value)
+    if redacted != value:
+        return redacted
+    if bare_base64 and _looks_like_base64(value):
+        return _BASE64_PLACEHOLDER
+    return value
+
+
+def _redact_content_part(part: Any) -> Any:
+    if not isinstance(part, dict):
+        return part
+
+    out = deepcopy(part)
+    for key in ("image_url", "audio_url", "video_url"):
+        container = out.get(key)
+        if isinstance(container, dict) and "url" in container:
+            container["url"] = _redact_media_source(container["url"])
+        elif isinstance(container, str):
+            out[key] = _redact_media_source(container)
+
+    for key in ("input_audio", "input_video"):
+        container = out.get(key)
+        if isinstance(container, dict) and "data" in container:
+            container["data"] = _BASE64_PLACEHOLDER
+        elif isinstance(container, str):
+            out[key] = _redact_media_source(container, bare_base64=True)
+
+    for key in ("image", "audio", "video"):
+        if key in out:
+            out[key] = _redact_media_source(out[key], bare_base64=True)
+
+    return out
+
+
+def _redacted_chat_request_dump(req: ChatRequest) -> str:
+    data = req.model_dump(mode="json")
+    messages = data.get("messages")
+    if isinstance(messages, list):
+        for msg in messages:
+            if not isinstance(msg, dict):
+                continue
+            content = msg.get("content")
+            if isinstance(content, list):
+                msg["content"] = [_redact_content_part(part) for part in content]
+    return json.dumps(data, ensure_ascii=False, indent=2)
+
+
+def _sse(payload) -> str:  # type: ignore
+    if payload == "[DONE]":
+        return "data: [DONE]\n\n"
+    return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
+
+
+def _stream_error(exc: BaseException) -> dict:
+    return {"error": {"message": str(exc), "type": exc.__class__.__name__}}

kiapi/api/embedding/_views/embedding_response.py

@@ -0,0 +1,35 @@
+"""``/v1/embedding`` response model.
+
+Mirrors the dict shaped by ``capabilities/embedding/_operations/format_response``
+so the response is self-describing in OpenAPI. It is an output-side projection of
+the model result (not a model-coupled view), so it lives under
+``api/embedding/_views``.
+
+``EmbeddingResponse`` is the sole public element; the nested pieces are private
+since nothing references them on their own.
+"""
+
+from pydantic import BaseModel, Field
+
+
+class _Usage(BaseModel):
+    prompt_tokens: int = Field(
+        description="Tokens in the embedded input (0 when the model does not report it)."
+    )
+    total_tokens: int = Field(
+        description="Total tokens accounted for; equals `prompt_tokens` (no generation)."
+    )
+
+
+class _Timings(BaseModel):
+    total_s: float = Field(description="Wall-clock embedding time in seconds.")
+
+
+class EmbeddingResponse(BaseModel):
+    model: str = Field(description="Resolved model name that produced the vector.")
+    embedding: list[float] = Field(
+        description="L2-normalized, last-token-pooled embedding vector for the item."
+    )
+    dimension: int = Field(description="Length of `embedding` (vector dimensionality).")
+    usage: _Usage = Field(description="Token accounting for the request.")
+    timings: _Timings = Field(description="kiapi extension: server-side timing.")

kiapi/api/embedding/router.py

@@ -0,0 +1,82 @@
+"""Embedding API: ``POST /v1/embedding`` (sync only).
+
+Like chat, still a job internally (shows up in /v1/jobs, shares the single-flight
+worker), but no async mode: the caller always gets the vector back. A sync wait
+that exceeds ``sync_timeout_s`` is an error (504).
+"""
+
+import asyncio
+
+from fastapi import APIRouter, Depends, HTTPException
+
+from kiapi.api import (
+    REQUIRE_AUTH,
+    get_ctx,
+    get_worker,
+    register_capability_endpoints,
+)
+from kiapi.api._settings import settings_manager
+from kiapi.capabilities.embedding import EmbedRequest, handle_embed
+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 Worker
+
+from ._views.embedding_response import EmbeddingResponse
+
+router = APIRouter(dependencies=REQUIRE_AUTH)
+
+
+@router.post(
+    "/v1/embedding",
+    response_model=EmbeddingResponse,
+    responses={
+        400: {"description": "Unknown model, empty input, or unsupported modality."},
+        503: {"description": "Model not set up, or memory budget exceeded."},
+        504: {"description": "Sync timeout exceeded; the job keeps running."},
+    },
+)
+async def embedding(
+    req: EmbedRequest,
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> dict:
+    """Embed a single item (one field per modality).
+
+    Send `text` and/or `image` in one request (one item, not OpenAI's array
+    `input`). The resolved `model` (see GET /v1/embedding/models) determines which
+    modalities are accepted; provide at least one input it supports. Returns one
+    L2-normalized, last-token-pooled vector with its `dimension`.
+
+    Embedding runs as a single-flight job, so it appears in /v1/jobs and is
+    serialized with all other generation, but there is no async mode: the caller
+    always waits for the vector. A wait that exceeds the sync timeout returns 504
+    while the job keeps running and can be polled at /v1/jobs/{id}.
+    """
+    settings = settings_manager.get_settings()
+
+    job = ctx.job_store.create(type="embedding", params={"model": req.model})
+    fut = await worker.submit(job, lambda: handle_embed(ctx, req))
+
+    try:
+        return await asyncio.wait_for(fut, timeout=settings.sync_timeout_s)
+    except TimeoutError:
+        raise HTTPException(  # noqa: B904
+            status_code=504,
+            detail=f"embedding job {job.id} exceeded sync timeout "
+            f"({settings.sync_timeout_s}s); 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"):
+            raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+        raise HTTPException(status_code=500, detail=f"embedding failed: {exc}")  # noqa: B904
+
+
+register_capability_endpoints(router, name="embedding", base_path="/v1/embedding")

kiapi/api/files/_operations/get_file.py

@@ -0,0 +1,14 @@
+from typing import Annotated
+
+from fastapi import File, UploadFile
+
+
+def _get_file(
+    file: Annotated[
+        UploadFile,
+        File(
+            description="File bytes to store and reference from inference requests.",
+        ),
+    ],
+) -> UploadFile:
+    return file

kiapi/api/files/_operations/get_file_id.py

@@ -0,0 +1,18 @@
+from typing import Annotated
+
+from fastapi import Path
+
+from kiapi.core.file import FileID
+
+
+def _get_file_id(
+    file_id: Annotated[
+        str,
+        Path(
+            pattern=r"^file_[0-9a-f]+$",
+            description="Persistent file id returned by the Files API.",
+            examples=["file_0123456789abcdef"],
+        ),
+    ],
+) -> FileID:
+    return file_id

kiapi/api/files/_views/file_delete_response.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+
+
+class FileDeleteResponse(BaseModel):
+    deleted: bool = Field(
+        ...,
+        description="True when the file record and stored bytes were deleted.",
+        examples=[True],
+    )
+    file_id: str = Field(
+        ...,
+        description="Deleted file id.",
+        examples=["file_0123456789abcdef"],
+    )

kiapi/api/files/_views/file_list_response.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, Field
+
+from kiapi.core.file import FileRecord
+
+
+class FileListResponse(BaseModel):
+    object: str = Field(
+        default="list",
+        description="OpenAI-style list envelope marker.",
+        examples=["list"],
+    )
+    data: list[FileRecord] = Field(
+        default_factory=list,
+        description="Stored files, including uploaded inputs and generated artifacts.",
+    )

kiapi/api/files/router.py

@@ -0,0 +1,105 @@
+"""Files API: upload / list / get / download / delete. Unified across capabilities."""
+
+from fastapi import APIRouter, Depends, HTTPException, UploadFile
+from fastapi.responses import FileResponse
+
+from kiapi.api import REQUIRE_AUTH, get_ctx
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileID, FileRecord
+
+from ._operations.get_file import _get_file
+from ._operations.get_file_id import _get_file_id
+from ._views.file_delete_response import FileDeleteResponse
+from ._views.file_list_response import FileListResponse
+
+router = APIRouter(prefix="/v1/files", dependencies=REQUIRE_AUTH)
+
+
+@router.post("", response_model=FileRecord)
+async def upload(
+    file: UploadFile = Depends(_get_file),
+    ctx: AppContext = Depends(get_ctx),
+) -> FileRecord:
+    """Store a file and return a reusable file_id.
+
+    Uploaded files become persistent Files API records. Use the returned
+    file_id as a FileRef input for generation/editing endpoints, or download
+    the bytes later with GET /v1/files/{file_id}/download.
+    """
+    data = await file.read()
+    return ctx.file_store.put_bytes(
+        data,
+        filename=file.filename or "upload",
+        content_type=file.content_type,
+    )
+
+
+@router.get("", response_model=FileListResponse)
+async def list_files(ctx: AppContext = Depends(get_ctx)) -> FileListResponse:
+    """List stored files.
+
+    Files include explicit uploads and generated artifacts. File records are
+    disk-backed and can outlive in-memory jobs across process restarts.
+    """
+    return FileListResponse(data=ctx.file_store.list())
+
+
+@router.get("/{file_id}", response_model=FileRecord)
+async def get_file(
+    file_id: FileID = Depends(_get_file_id),
+    ctx: AppContext = Depends(get_ctx),
+) -> FileRecord:
+    """Get metadata for a stored file.
+
+    This returns metadata only, not file bytes. Use the download endpoint when
+    the client needs the artifact body.
+    """
+    rec = ctx.file_store.get(file_id)
+    if rec is None:
+        raise HTTPException(status_code=404, detail=f"unknown file {file_id!r}")
+    return rec
+
+
+@router.get(
+    "/{file_id}/download",
+    responses={
+        200: {
+            "description": "Stored file bytes.",
+            "content": {
+                "application/octet-stream": {
+                    "schema": {"type": "string", "format": "binary"}
+                }
+            },
+        },
+        404: {"description": "Unknown file_id."},
+    },
+)
+async def download_file(
+    file_id: FileID = Depends(_get_file_id),
+    ctx: AppContext = Depends(get_ctx),
+) -> FileResponse:
+    """Download stored file bytes.
+
+    The response media type is the stored file's content_type, such as
+    image/png, audio/wav, video/mp4, text/markdown, or application/pdf.
+    """
+    rec = ctx.file_store.get(file_id)
+    if rec is None:
+        raise HTTPException(status_code=404, detail=f"unknown file {file_id!r}")
+    return FileResponse(rec.path, media_type=rec.content_type, filename=rec.filename)
+
+
+@router.delete("/{file_id}", response_model=FileDeleteResponse)
+async def delete_file(
+    file_id: FileID = Depends(_get_file_id),
+    ctx: AppContext = Depends(get_ctx),
+) -> dict:
+    """Delete a stored file.
+
+    Deleting a file removes the stored bytes and metadata. Existing job records
+    may still mention the deleted file_id in artifacts, but the file can no
+    longer be downloaded.
+    """
+    if not ctx.file_store.delete(file_id):
+        raise HTTPException(status_code=404, detail=f"unknown file {file_id!r}")
+    return {"deleted": True, "file_id": file_id}

kiapi/api/health/_views/health_response.py

@@ -0,0 +1,26 @@
+from pydantic import BaseModel, Field
+
+from kiapi.core.memory import MemoryStats
+
+
+class HealthResponse(BaseModel):
+    status: str = Field(
+        ...,
+        description="Service health marker. The value is ok when the API process is accepting requests.",
+        examples=["ok"],
+    )
+    warm: bool = Field(
+        ...,
+        description="Whether startup warmup has finished. Warmup may skip models that are not activated.",
+        examples=[True],
+    )
+    queue_len: int = Field(
+        ...,
+        ge=0,
+        description="Number of jobs waiting in the single-flight worker queue.",
+        examples=[0],
+    )
+    memory: MemoryStats = Field(
+        ...,
+        description="Resident model memory budget and loaded model snapshot.",
+    )

kiapi/api/health/router.py

@@ -0,0 +1,28 @@
+from fastapi import APIRouter, Depends
+
+from kiapi.api import get_ctx, get_worker
+from kiapi.core.app import AppContext
+from kiapi.core.worker import Worker
+
+from ._views.health_response import HealthResponse
+
+router = APIRouter()
+
+
+@router.get("/health", response_model=HealthResponse)
+async def health(
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> dict:
+    """Inspect server readiness, queue depth, and resident model memory.
+
+    Use this endpoint to see whether warmup has finished, how many jobs are
+    queued behind the single-flight worker, and which resident models currently
+    count against the global memory budget.
+    """
+    return {
+        "status": "ok",
+        "warm": worker.warm,
+        "queue_len": worker.queue.qsize(),
+        "memory": ctx.memory_manager.stats(),
+    }

kiapi/api/image/depthpro/_views/estimate_response.py

@@ -0,0 +1,70 @@
+"""``/v1/image/depthpro/estimate`` result payload (the Job ``result`` shape).
+
+Mirrors the dict returned by ``capabilities/depthpro/_models/depthpro._estimate``
+so the depth-map result is self-describing in OpenAPI. It is an output-side
+projection of the model result (not a model-coupled view), so it lives under
+``api/image/depthpro/_views``.
+
+``EstimateResponse`` is the sole public element; nested pieces stay private.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class _Timings(BaseModel):
+    total_s: float = Field(
+        description="Wall-clock depth-estimation time in seconds (model run only)."
+    )
+
+
+class EstimateResponse(BaseModel):
+    """Capability-specific ``result`` for a succeeded depthpro estimate job."""
+
+    model: str = Field(
+        description="Resolved model variant that produced the depth map."
+    )
+
+    depth_image_file_id: str = Field(
+        description=(
+            "Files-API id of the grayscale depth PNG (near = bright, far = dark). "
+            "Fetch metadata at GET /v1/files/{id} or bytes at /download. This is "
+            "also the artifact returned as raw bytes by a single-artifact sync call."
+        )
+    )
+    depth_image_bytes: int = Field(description="Size of the depth PNG in bytes.")
+    depth_data_file_id: str | None = Field(
+        description=(
+            "Files-API id of the compressed NPZ holding the raw float depth array "
+            "plus min/max depth, or null when `include_depth_data` was false."
+        )
+    )
+    depth_data_bytes: int | None = Field(
+        description="Size of the NPZ in bytes, or null when no depth data was stored."
+    )
+
+    input_width: int = Field(description="Width in pixels of the submitted image.")
+    input_height: int = Field(description="Height in pixels of the submitted image.")
+    input_mode: str = Field(
+        description="PIL mode of the submitted image (e.g. RGB, L, RGBA)."
+    )
+
+    width: int = Field(description="Width in pixels of the produced depth map.")
+    height: int = Field(description="Height in pixels of the produced depth map.")
+    mode: str = Field(description="PIL mode of the produced depth PNG (typically L).")
+
+    array_shape: list[int] = Field(
+        description="Shape of the raw depth array stored in the NPZ (rows, cols)."
+    )
+    min_depth: float = Field(
+        description="Minimum depth value in the raw array, in meters."
+    )
+    max_depth: float = Field(
+        description="Maximum depth value in the raw array, in meters."
+    )
+
+    params: dict[str, Any] = Field(
+        description="Resolved parameters used for the run (model, quantize, etc.)."
+    )
+    timings: _Timings = Field(description="kiapi extension: server-side timing.")