kiapi 0.1.0__py3-none-any.whl

kiapi/core/file/_helpers/resolve_file_ref.py

@@ -0,0 +1,90 @@
+import base64
+import binascii
+import mimetypes
+import urllib.error
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+from kiapi.core.net import UnsafeURLError, verify_public_url
+
+from .._schemas.file_data_url_ref import FileDataURLRef
+from .._schemas.file_id_ref import FileIDRef
+from .._schemas.file_record import FileRecord
+from .._schemas.file_url_ref import FileURLRef
+from .._services.file_store import FileStore
+from .._types.file_ref import FileRef
+from .._views.resolved_file_ref import ResolvedFileRef
+
+
+def resolve_file_ref(
+    file_store: FileStore,
+    ref: FileRef,
+    *,
+    kind: str,
+) -> ResolvedFileRef:
+    if isinstance(ref, FileIDRef):
+        rec = file_store.get(ref.file_id)
+        if rec is None:
+            raise FileNotFoundError(f"unknown {kind} file_id {ref.file_id!r}")
+        return ResolvedFileRef(rec)
+
+    if isinstance(ref, FileURLRef):
+        return ResolvedFileRef(_store_url(file_store, ref, kind=kind))
+
+    if isinstance(ref, FileDataURLRef):
+        return ResolvedFileRef(_store_data_url(file_store, ref, kind=kind))
+
+
+def _store_url(file_store: FileStore, ref: FileURLRef, *, kind: str) -> FileRecord:
+    parsed = urllib.parse.urlparse(ref.url)
+    if parsed.scheme not in ("http", "https"):
+        raise ValueError(f"{kind} url must be an absolute http:// or https:// URL")
+
+    try:
+        verify_public_url(ref.url, kind=kind)
+    except UnsafeURLError as exc:
+        raise ValueError(str(exc)) from exc
+
+    try:
+        with urllib.request.urlopen(ref.url, timeout=30) as response:
+            data = response.read()
+            content_type = response.headers.get_content_type()
+    except urllib.error.URLError as exc:
+        raise ValueError(f"failed to fetch {kind} url {ref.url!r}: {exc}") from exc
+
+    filename = (
+        Path(parsed.path).name or f"{kind}{_suffix_for_content_type(content_type)}"
+    )
+    return file_store.put_bytes(
+        data,
+        filename=filename,
+        content_type=content_type,
+        meta={"source": "url", "url": ref.url, "kind": kind},
+    )
+
+
+def _store_data_url(
+    file_store: FileStore, ref: FileDataURLRef, *, kind: str
+) -> FileRecord:
+    header, sep, raw = ref.data_url.partition(",")
+    if sep != "," or not header.startswith("data:"):
+        raise ValueError(f"{kind} data_url must be a data URL")
+
+    media_type = header[5:].split(";", 1)[0] or "application/octet-stream"
+    is_base64 = ";base64" in header
+    try:
+        data = base64.b64decode(raw, validate=True) if is_base64 else raw.encode()
+    except (binascii.Error, ValueError) as exc:
+        raise ValueError(f"invalid {kind} data_url payload: {exc}") from exc
+
+    return file_store.put_bytes(
+        data,
+        filename=f"{kind}{_suffix_for_content_type(media_type)}",
+        content_type=media_type,
+        meta={"source": "data_url", "kind": kind},
+    )
+
+
+def _suffix_for_content_type(content_type: str | None) -> str:
+    return mimetypes.guess_extension(content_type or "") or ".bin"

kiapi/core/file/_schemas/file_data_url_ref.py

@@ -0,0 +1,19 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class FileDataURLRef(BaseModel):
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "type": "data_url",
+                    "data_url": "data:image/png;base64,iVBORw0KGgo...",
+                },
+            ],
+        },
+    }
+
+    type: Literal["data_url"] = "data_url"
+    data_url: str = Field(..., min_length=1)

kiapi/core/file/_schemas/file_id_ref.py

@@ -0,0 +1,16 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class FileIDRef(BaseModel):
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {"type": "file_id", "file_id": "file_0123456789abcdef"},
+            ],
+        },
+    }
+
+    type: Literal["file_id"] = "file_id"
+    file_id: str = Field(..., min_length=1)

kiapi/core/file/_schemas/file_record.py

@@ -0,0 +1,46 @@
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from .._types.file_id import FileID
+
+
+class FileRecord(BaseModel):
+    """A stored file's metadata. Doubles as the Files API response schema.
+
+    ``path`` is the on-disk location and is internal to the file store; it is
+    excluded from serialization so ``model_dump()`` yields the API/sidecar shape.
+    """
+
+    file_id: FileID = Field(
+        ...,
+        description="Persistent file id. Use this in FileRef inputs or download it with GET /v1/files/{file_id}/download.",
+        examples=["file_0123456789abcdef"],
+    )
+    filename: str = Field(
+        ...,
+        description="Original or generated filename associated with the stored bytes.",
+        examples=["image_1766200000.png"],
+    )
+    content_type: str = Field(
+        ...,
+        description="Media type of the stored file.",
+        examples=["image/png"],
+    )
+    size: int = Field(
+        ...,
+        ge=0,
+        description="File size in bytes.",
+        examples=[123456],
+    )
+    created_at: float = Field(
+        ...,
+        description="Unix timestamp when the file was stored.",
+        examples=[1766200000.0],
+    )
+    meta: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Capability-specific metadata such as seed, request params, timings, dimensions, or artifact kind.",
+        examples=[{"model": "turbo", "width": 1024, "height": 1024}],
+    )
+    path: str = Field(..., exclude=True)

kiapi/core/file/_schemas/file_url_ref.py

@@ -0,0 +1,16 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class FileURLRef(BaseModel):
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {"type": "url", "url": "https://example.com/input.png"},
+            ],
+        },
+    }
+
+    type: Literal["url"] = "url"
+    url: str = Field(..., min_length=1)

kiapi/core/file/_services/file_store.py

@@ -0,0 +1,138 @@
+"""Unified file store. Artifacts (and uploads) are addressed by ``file_id``.
+
+Files outlive jobs while their files remain on disk: jobs are in-memory and reset
+on process restart, while stored files stay self-describing via a sidecar
+``<file_id>.json`` holding the original filename, content type, size, and an
+optional ``meta`` blob. This
+mirrors the per-server Files APIs (ace-step / audiocraft / video) unified into
+one store under ``files_root``.
+"""
+
+import json
+import mimetypes
+import shutil
+import time
+import uuid
+from pathlib import Path
+from typing import Any
+
+from .._schemas.file_record import FileRecord
+from .._settings import FileSettings
+from .._types.file_id import FileID
+
+
+class FileStore:
+    def __init__(self, settings: FileSettings) -> None:
+        self.root = Path(settings.files_root).expanduser()
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def _blob_path(self, file_id: FileID, suffix: str) -> Path:
+        return self.root / f"{file_id}{suffix}"
+
+    def _sidecar(self, file_id: FileID) -> Path:
+        return self.root / f"{file_id}.json"
+
+    # -- write ----------------------------------------------------------------
+
+    def put_bytes(
+        self,
+        data: bytes,
+        *,
+        filename: str,
+        content_type: str | None = None,
+        meta: dict[str, Any] | None = None,
+    ) -> FileRecord:
+        file_id = _generate_file_id()
+        suffix = Path(filename).suffix
+        path = self._blob_path(file_id, suffix)
+        path.write_bytes(data)
+        rec = FileRecord(
+            file_id=file_id,
+            filename=filename,
+            content_type=_resolve_content_type(filename, content_type),
+            size=len(data),
+            created_at=time.time(),
+            meta=meta or {},
+            path=str(path),
+        )
+        self._write_sidecar(rec, suffix)
+        return rec
+
+    def put_path(
+        self,
+        src: str | Path,
+        *,
+        filename: str | None = None,
+        content_type: str | None = None,
+        meta: dict[str, Any] | None = None,
+        move: bool = False,
+    ) -> FileRecord:
+        src = Path(src)
+        file_id = _generate_file_id()
+        suffix = src.suffix
+        filename = filename or src.name
+        path = self._blob_path(file_id, suffix)
+        if move:
+            shutil.move(str(src), str(path))
+        else:
+            shutil.copyfile(src, path)
+        rec = FileRecord(
+            file_id=file_id,
+            filename=filename,
+            content_type=_resolve_content_type(filename, content_type),
+            size=path.stat().st_size,
+            created_at=time.time(),
+            meta=meta or {},
+            path=str(path),
+        )
+        self._write_sidecar(rec, suffix)
+        return rec
+
+    def _write_sidecar(self, rec: FileRecord, suffix: str) -> None:
+        payload = rec.model_dump()
+        payload["_suffix"] = suffix
+        self._sidecar(rec.file_id).write_text(json.dumps(payload, ensure_ascii=False))
+
+    # -- read -----------------------------------------------------------------
+
+    def get(self, file_id: FileID) -> FileRecord | None:
+        sidecar = self._sidecar(file_id)
+        if not sidecar.exists():
+            return None
+        d = json.loads(sidecar.read_text())
+        suffix = d.get("_suffix", "")
+        return FileRecord(
+            file_id=d["file_id"],
+            filename=d["filename"],
+            content_type=d["content_type"],
+            size=d["size"],
+            created_at=d["created_at"],
+            meta=d.get("meta", {}),
+            path=str(self._blob_path(file_id, suffix)),
+        )
+
+    def list(self) -> list[FileRecord]:
+        recs = []
+        for sidecar in self.root.glob("*.json"):
+            rec = self.get(sidecar.stem)
+            if rec is not None:
+                recs.append(rec)
+        return sorted(recs, key=lambda r: r.created_at, reverse=True)
+
+    def delete(self, file_id: FileID) -> bool:
+        rec = self.get(file_id)
+        if rec is None:
+            return False
+        Path(rec.path).unlink(missing_ok=True)
+        self._sidecar(file_id).unlink(missing_ok=True)
+        return True
+
+
+def _generate_file_id() -> FileID:
+    return f"file_{uuid.uuid4().hex}"
+
+
+def _resolve_content_type(filename: str, content_type: str | None) -> str:
+    return (
+        content_type or mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    )

kiapi/core/file/_settings.py

@@ -0,0 +1,22 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class FileSettings(BaseSettings):
+    """Settings for the file store that keeps uploads and generated artifacts."""
+
+    model_config = SettingsConfigDict(env_prefix="KIAPI_")
+
+    files_root: str = Field(
+        default="/tmp/kiapi/files",
+        title="File storage directory",
+        description=(
+            "Directory where uploaded files and generated artifacts are stored.\n"
+            "The default is under /tmp so files may disappear after reboot or tmp cleanup.\n"
+            "~ is expanded as the home directory of the running user."
+        ),
+    )
+
+
+settings_manager = SettingsManager(FileSettings)

kiapi/core/file/_types/file_id.py

@@ -0,0 +1,2 @@
+type FileID = str
+"""The address of a stored file while it remains available in the Files API."""

kiapi/core/file/_types/file_ref.py

@@ -0,0 +1,12 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from .._schemas.file_data_url_ref import FileDataURLRef
+from .._schemas.file_id_ref import FileIDRef
+from .._schemas.file_url_ref import FileURLRef
+
+type FileRef = Annotated[
+    FileIDRef | FileURLRef | FileDataURLRef,
+    Field(discriminator="type"),
+]

kiapi/core/file/_views/resolved_file_ref.py

@@ -0,0 +1,16 @@
+from dataclasses import dataclass
+
+from .._schemas.file_record import FileRecord
+
+
+@dataclass(frozen=True)
+class ResolvedFileRef:
+    record: FileRecord
+
+    @property
+    def file_id(self) -> str:
+        return self.record.file_id
+
+    @property
+    def path(self) -> str:
+        return self.record.path

kiapi/core/job/__init__.py

@@ -0,0 +1,24 @@
+"""Unified job model + in-memory job store. Every generation (sync or async) is
+a Job; sync waits for it, async returns its id. Jobs reset on restart.
+
+Public surface (was ``core/jobs.py``): :class:`Job`, :class:`JobStore`,
+:class:`JobStatus`, :class:`ProgressReporter`, :func:`creep_progress`.
+"""
+
+from ._enums.job_status import JobStatus
+from ._helpers.creep_progress import creep_progress
+from ._models.job import Job
+from ._services.job_store import JobStore
+from ._services.progress_reporter import ProgressReporter
+from ._types.job_id import JobID
+from ._types.job_result import JobResult
+
+__all__ = [
+    "Job",
+    "JobID",
+    "JobResult",
+    "JobStatus",
+    "JobStore",
+    "ProgressReporter",
+    "creep_progress",
+]

kiapi/core/job/_enums/job_status.py

@@ -0,0 +1,9 @@
+from enum import StrEnum
+
+
+class JobStatus(StrEnum):
+    QUEUED = "queued"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELED = "canceled"

kiapi/core/job/_helpers/creep_progress.py

@@ -0,0 +1,123 @@
+"""Synthetic 'liveness' progress for jobs whose backend exposes no per-step hook.
+
+Some generation backends (notably ltx2's ``mlx_video.generate_video``, which only
+drives its own internal rich progress bar) give us no way to observe denoising
+steps — unlike mflux, which offers a callback registry (see capabilities'
+``attach_mflux_progress``). A polling client, typically an LLM agent, reads a job
+whose ``progress`` never moves as possibly hung and may stop waiting.
+
+For those backends we creep the reported fraction forward on a piecewise-linear
+schedule (see :func:`_scheduled_fraction`) — constant, easy-to-read rates rather
+than a smooth curve. This is a deliberate fiction: it signals 'alive and working',
+not real completion; the genuine terminal 1.0 still comes only from
+:meth:`Job.mark_succeeded`.
+
+The job's worker thread is blocked inside the backend call while this runs, so a
+daemon thread does the ticking. ContextVars don't cross threads, so it captures
+the worker thread's bound :class:`ProgressReporter` up front and pushes updates
+through it — the same throttling and job wiring apply.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+from .._services.progress_reporter import ProgressReporter
+
+
+def _scheduled_fraction(
+    elapsed: float,
+    *,
+    eta_s: float,
+    cap: float,
+    ramp_frac: float,
+    ramp_rate: float,
+    mid_frac: float,
+    mid_fallback_rate: float,
+    tail_rate: float,
+) -> float:
+    """The reported fraction at ``elapsed`` seconds — three constant-rate phases:
+
+    1. **ramp** — linear from 0 to ``ramp_frac`` at a *fixed* ``ramp_rate`` per
+       second (e.g. 0 → 0.20 at 0.5%/s, reaching 0.20 after 40s). The rate is
+       always honoured, even for a short ``eta_s``, so the start is always a gentle
+       crawl — easier to wait on than an immediate jump.
+    2. **mid** — linear from ``ramp_frac`` to ``mid_frac``. Normally paced to hit
+       ``mid_frac`` exactly at ``eta_s``; but if the ramp has already consumed
+       ``eta_s`` (``eta_s <= ramp duration``), there is no room left to pace to, so
+       it falls back to the fixed ``mid_fallback_rate`` per second (e.g. 1%/s).
+    3. **tail** — linear from ``mid_frac`` upward at the slow ``tail_rate`` per
+       second, clamped at ``cap`` — the overrun crawl.
+    """
+    knee_t = ramp_frac / ramp_rate  # fixed: guarantees the slow start
+    if elapsed <= knee_t:
+        return min(cap, ramp_rate * elapsed)
+
+    if eta_s > knee_t:
+        mid_end_t = eta_s  # room to pace: land on mid_frac exactly at eta
+    else:
+        mid_end_t = knee_t + (mid_frac - ramp_frac) / mid_fallback_rate
+
+    if elapsed <= mid_end_t:
+        fraction = ramp_frac + (mid_frac - ramp_frac) * (elapsed - knee_t) / (
+            mid_end_t - knee_t
+        )
+    else:
+        fraction = mid_frac + tail_rate * (elapsed - mid_end_t)
+    return min(cap, max(0.0, fraction))
+
+
+@contextmanager
+def creep_progress(
+    *,
+    eta_s: float,
+    cap: float = 0.999,
+    label: str = "generating",
+    interval_s: float = 1.0,
+    ramp_frac: float = 0.20,
+    ramp_rate: float = 0.005,
+    mid_frac: float = 0.80,
+    mid_fallback_rate: float = 0.01,
+    tail_rate: float = 0.0005,
+) -> Iterator[None]:
+    """Tick synthetic progress on a piecewise-linear schedule while the block runs.
+
+    The fraction follows :func:`_scheduled_fraction`: a fixed-rate ramp to
+    ``ramp_frac`` (always a gentle start), an even climb to ``mid_frac`` timed to
+    land at ``eta_s`` (or a fixed ``mid_fallback_rate`` if ``eta_s`` is already
+    spent), then a slow ``tail_rate`` crawl up to ``cap``. ``eta_s`` only sets the
+    pace, never the ceiling, so a job that overruns keeps inching forward in the
+    tail phase rather than stalling. ``eta_s <= 0`` disables creeping entirely.
+    """
+    reporter = ProgressReporter.current()
+    if eta_s <= 0:
+        yield
+        return
+
+    stop = threading.Event()
+    t0 = time.monotonic()
+
+    def _tick() -> None:
+        while not stop.wait(interval_s):
+            fraction = _scheduled_fraction(
+                time.monotonic() - t0,
+                eta_s=eta_s,
+                cap=cap,
+                ramp_frac=ramp_frac,
+                ramp_rate=ramp_rate,
+                mid_frac=mid_frac,
+                mid_fallback_rate=mid_fallback_rate,
+                tail_rate=tail_rate,
+            )
+            reporter.update(fraction, label)
+
+    thread = threading.Thread(target=_tick, name="kiapi-progress-creep", daemon=True)
+    thread.start()
+    try:
+        yield
+    finally:
+        stop.set()
+        thread.join(timeout=2.0)

kiapi/core/job/_models/job.py

@@ -0,0 +1,121 @@
+"""Unified job model.
+
+Every generation in kiapi — sync or async, chat or video — is a Job. Sync
+requests create a job, wait for it, and return its result; async requests create
+a job and return its id immediately. This gives one consistent place to inspect
+work in flight and one consistent result shape.
+
+``result`` is a free-form dict whose shape depends on ``type`` (the caller
+dispatches on ``type``); ``artifacts`` lists file_ids produced (see core/file).
+Jobs are in-memory only and reset on restart; artifacts are Files API entries
+that can outlive jobs while their files remain on disk.
+"""
+
+import time
+import uuid
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from kiapi.core.file import FileID
+
+from .._enums.job_status import JobStatus
+from .._types.job_id import JobID
+from .._types.job_result import JobResult
+from .._types.job_type import JobType
+
+
+class Job(BaseModel):
+    type: JobType = Field(
+        ...,
+        description="Job type. Use this to interpret the capability-specific result payload.",
+        examples=["zimage"],
+    )
+    params: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Request parameters captured for inspection and reproducibility. Secret or large media payloads may be omitted or redacted by endpoints.",
+    )
+    id: JobID = Field(
+        default_factory=lambda: f"job_{uuid.uuid4().hex}",
+        description="In-memory job id. Jobs are cleared when the kiapi process restarts.",
+        examples=["job_0123456789abcdef"],
+    )
+    status: JobStatus = Field(
+        default=JobStatus.QUEUED,
+        description="Job lifecycle state: queued, running, succeeded, failed, or canceled.",
+        examples=["succeeded"],
+    )
+    result: JobResult | None = Field(
+        default=None,
+        description="Capability-specific result metadata. Completed generation jobs usually also expose produced file ids in artifacts.",
+        examples=[{"file_id": "file_0123456789abcdef", "image_bytes": 123456}],
+    )
+    artifacts: list[FileID] = Field(
+        default_factory=list,
+        description="File ids produced by the job. Use GET /v1/files/{file_id} for metadata or /download for bytes.",
+        examples=[["file_0123456789abcdef"]],
+    )
+    error: str | None = Field(
+        default=None,
+        description="Error message when status is failed; otherwise null.",
+        examples=["model 'turbo' is not activated"],
+    )
+    created_at: float = Field(
+        default_factory=lambda: time.time(),
+        description="Unix timestamp when the job was created.",
+        examples=[1766200000.0],
+    )
+    started_at: float | None = Field(
+        default=None,
+        description="Unix timestamp when the worker started the job, or null while queued.",
+        examples=[1766200001.0],
+    )
+    finished_at: float | None = Field(
+        default=None,
+        description="Unix timestamp when the job reached a terminal state, or null while queued/running.",
+        examples=[1766200030.0],
+    )
+    progress: float | None = Field(
+        default=None,
+        ge=0.0,
+        le=1.0,
+        description="Best-effort completion fraction in [0.0, 1.0]. Null means the job has not reported progress.",
+        examples=[0.42],
+    )
+    progress_label: str = Field(
+        default="queued",
+        description="Short human-readable phase label such as queued, running, denoising, saving, or done.",
+        examples=["denoising"],
+    )
+
+    def to_dict(self) -> dict[str, Any]:
+        return self.model_dump()
+
+    def mark_running(self) -> None:
+        self.status = JobStatus.RUNNING
+        self.started_at = time.time()
+        self.progress = 0.0
+        self.progress_label = "running"
+
+    def update_progress(self, fraction: float, label: str | None = None) -> None:
+        self.progress = min(1.0, max(0.0, fraction))
+
+        if label is not None:
+            self.progress_label = label
+
+    def mark_succeeded(self, result: JobResult, artifacts: list[FileID]) -> None:
+        self.status = JobStatus.SUCCEEDED
+        self.result = result
+        self.artifacts = artifacts
+        self.progress = 1.0
+        self.progress_label = "done"
+        self.finished_at = time.time()
+
+    def mark_failed(self, error: str) -> None:
+        self.status = JobStatus.FAILED
+        self.error = error
+        self.finished_at = time.time()
+
+    def mark_canceled(self) -> None:
+        self.status = JobStatus.CANCELED
+        self.finished_at = time.time()