kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/chat/_operations/parse_messages.py

@@ -0,0 +1,326 @@
+"""Shared helper: parse OpenAI-style multimodal ``messages`` into the Qwen
+chat-template shape. Model handlers call this with the set of modalities they
+support (``allow=``); anything outside that set is rejected with 400.
+
+Recognised content parts (lenient — aliases accepted):
+
+  text   : {"type": "text", "text": "..."}
+  image  : {"type": "image_url",  "image_url": {"url": "data:image/...;base64,..." | "http..." }}
+           aliases: {"type": "image", "image": <url|data|path>}, {"type": "input_image", ...}
+  audio  : {"type": "input_audio", "input_audio": {"data": "<base64>", "format": "wav"}}
+           aliases: {"type": "audio", "audio": <url|data|path>}, {"type": "audio_url", ...}
+  video  : {"type": "video_url",  "video_url": {"url": "data:video/...;base64,..." | "http..." }}
+           aliases: {"type": "video", "video": <url|data|path>}, {"type": "input_video", ...}
+
+Each media source (data URL, http(s) URL, or bare base64) is materialized to a
+file under ``tmp_dir``. Server-side local paths are intentionally not accepted,
+so a chat message cannot read arbitrary files off the host. The function returns:
+
+  - ``template_messages``: same conversation, but every media part rewritten to
+    the Qwen-native form ({"type": "image"/"audio"/"video", <key>: <path>}) so the
+    processor's chat_template inserts placeholders in document order.
+  - ``images`` / ``audios`` / ``videos``: ordered lists of file paths, collected in
+    the SAME walk order, to hand to ``generate(image=, audio=, video=)``.
+"""
+
+import base64
+import binascii
+import json
+import re
+import subprocess
+import urllib.request
+import uuid
+from pathlib import Path
+from typing import Any
+
+from kiapi.core.net import UnsafeURLError, verify_public_url
+
+_DATA_URL = re.compile(r"^data:(?P<mime>[\w./+-]+)?;base64,(?P<data>.*)$", re.DOTALL)
+
+_EXT_BY_MIME = {
+    "image/png": ".png",
+    "image/jpeg": ".jpg",
+    "image/jpg": ".jpg",
+    "image/webp": ".webp",
+    "image/gif": ".gif",
+    "audio/wav": ".wav",
+    "audio/x-wav": ".wav",
+    "audio/mpeg": ".mp3",
+    "audio/mp3": ".mp3",
+    "audio/mp4": ".m4a",
+    "audio/ogg": ".ogg",
+    "audio/flac": ".flac",
+    "video/mp4": ".mp4",
+    "video/quicktime": ".mov",
+    "video/webm": ".webm",
+}
+_DEFAULT_EXT = {"image": ".png", "audio": ".wav", "video": ".mp4"}
+
+
+class MediaError(ValueError):
+    """Bad/undecodable media in a message; mapped to HTTP 400 by the API layer."""
+
+
+class CapabilityError(MediaError):
+    """A modality the selected model does not support; mapped to HTTP 400."""
+
+
+def parse_messages(
+    messages: list[dict[str, Any]],
+    tmp_dir: Path,
+    *,
+    allow: set[str] | None = None,
+    use_audio_in_video: bool = True,
+) -> tuple[list[dict[str, Any]], list[str], list[str], list[str]]:
+    """Parse messages into (template_messages, images, audios, videos).
+
+    ``allow`` is the set of modalities the selected model supports
+    (e.g. ``{"text", "image"}``). A part outside it raises :class:`CapabilityError`.
+    ``None`` allows everything (used by tests).
+    """
+    template_messages: list[dict[str, Any]] = []
+    images: list[str] = []
+    audios: list[str] = []
+    videos: list[str] = []
+
+    def _check(kind: str) -> None:
+        if allow is not None and kind not in allow:
+            raise CapabilityError(
+                f"{kind} input is not supported by the selected model"
+            )
+
+    for raw_msg in messages:
+        msg = _normalize_tool_calls_for_template(raw_msg)
+        role = msg.get("role", "user")
+        content = msg.get("content")
+
+        if content is None or isinstance(content, str):
+            # Pass plain text / tool fields (tool_calls, tool_call_id, name) through.
+            template_messages.append(dict(msg))
+            continue
+
+        if not isinstance(content, list):
+            raise MediaError(f"message.content must be a string or list (role={role})")
+
+        new_content: list[dict[str, Any]] = []
+        for part in content:
+            if not isinstance(part, dict):
+                raise MediaError("each content part must be an object")
+            kind, source = _extract_part(part)
+
+            if kind == "text":
+                new_content.append({"type": "text", "text": source or ""})
+                continue
+
+            _check(kind)
+
+            fmt = None
+            if isinstance(source, tuple):  # (data, format) from input_audio/input_video
+                source, fmt = source
+            path = _materialize(source, kind, tmp_dir, fmt=fmt)
+
+            if kind == "image":
+                images.append(path)
+                new_content.append({"type": "image", "image": path})
+            elif kind == "audio":
+                audios.append(path)
+                new_content.append({"type": "audio", "audio": path})
+            else:  # video
+                videos.append(path)
+                new_content.append({"type": "video", "video": path})
+                # A video that carries sound: demux its audio track and feed it as
+                # an audio input too (placeholder right after the video), so the
+                # model hears the clip, not just sees it.
+                if use_audio_in_video and (allow is None or "audio" in allow):
+                    apath = _extract_audio(path, tmp_dir)
+                    if apath is not None:
+                        audios.append(apath)
+                        new_content.append({"type": "audio", "audio": apath})
+
+        out = dict(msg)
+        out["content"] = new_content
+        template_messages.append(out)
+
+    return template_messages, images, audios, videos
+
+
+def _template_tool_arguments(arguments: Any) -> dict[str, Any]:
+    """Convert OpenAI's JSON-string tool args to the mapping Qwen templates expect."""
+    if isinstance(arguments, dict):
+        return arguments
+    if isinstance(arguments, str):
+        try:
+            parsed = json.loads(arguments)
+        except json.JSONDecodeError:
+            return {"_raw": arguments}
+        if isinstance(parsed, dict):
+            return parsed
+        return {"value": parsed}
+    if arguments is None:
+        return {}
+    return {"value": arguments}
+
+
+def _normalize_tool_calls_for_template(msg: dict[str, Any]) -> dict[str, Any]:
+    """Return a copy of msg with tool-call history shaped for HF chat templates."""
+    out = dict(msg)
+    calls = out.get("tool_calls")
+    if not isinstance(calls, list):
+        return out
+
+    if out.get("content") is None:
+        out["content"] = ""
+
+    normalized = []
+    for call in calls:
+        if not isinstance(call, dict):
+            normalized.append(call)
+            continue
+        ncall = dict(call)
+        fn = ncall.get("function")
+        if isinstance(fn, dict):
+            nfn = dict(fn)
+            if "arguments" in nfn:
+                nfn["arguments"] = _template_tool_arguments(nfn["arguments"])
+            ncall["function"] = nfn
+        elif "arguments" in ncall:
+            ncall["arguments"] = _template_tool_arguments(ncall["arguments"])
+        normalized.append(ncall)
+
+    out["tool_calls"] = normalized
+    return out
+
+
+def _write_bytes(data: bytes, tmp_dir: Path, kind: str, ext: str) -> str:
+    path = tmp_dir / f"{kind}_{uuid.uuid4().hex}{ext}"
+    path.write_bytes(data)
+    return str(path)
+
+
+def _materialize(
+    source: str, kind: str, tmp_dir: Path, *, fmt: str | None = None
+) -> str:
+    """Turn a data URL / http(s) URL / bare base64 into a local file path."""
+    if not isinstance(source, str) or not source:
+        raise MediaError(f"{kind}: empty or non-string source")
+
+    m = _DATA_URL.match(source)
+    if m:
+        mime = m.group("mime") or ""
+        ext = _EXT_BY_MIME.get(mime, _DEFAULT_EXT[kind])
+        try:
+            return _write_bytes(base64.b64decode(m.group("data")), tmp_dir, kind, ext)
+        except (binascii.Error, ValueError) as exc:
+            raise MediaError(f"{kind}: invalid base64 data URL ({exc})")  # noqa: B904
+
+    if source.startswith(("http://", "https://")):
+        try:
+            verify_public_url(source, kind=kind)
+        except UnsafeURLError as exc:
+            raise MediaError(f"{kind}: {exc}")  # noqa: B904
+        try:
+            with urllib.request.urlopen(source, timeout=30) as resp:
+                data = resp.read()
+        except Exception as exc:
+            raise MediaError(f"{kind}: failed to fetch {source!r} ({exc})")  # noqa: B904
+        ext = Path(source.split("?")[0]).suffix or _DEFAULT_EXT[kind]
+        return _write_bytes(data, tmp_dir, kind, ext)
+
+    # Bare base64 (e.g. input_audio.data without a data: prefix). Server-side
+    # local paths are intentionally NOT accepted: a chat message must not be able
+    # to read arbitrary files off the host.
+    try:
+        data = base64.b64decode(source, validate=True)
+    except (binascii.Error, ValueError):
+        raise MediaError(f"{kind}: source is not a data URL, URL, or base64")  # noqa: B904
+    ext = (
+        _EXT_BY_MIME.get(f"{kind}/{fmt}", _DEFAULT_EXT[kind])
+        if fmt
+        else _DEFAULT_EXT[kind]
+    )
+    return _write_bytes(data, tmp_dir, kind, ext)
+
+
+def _extract_part(part: dict[str, Any]) -> tuple[str, Any]:
+    """Return (kind, source) for one content part. kind ∈ text|image|audio|video."""
+    t = part.get("type")
+
+    if t == "text":
+        return "text", part.get("text", "")
+
+    if (
+        t in ("image_url", "image", "input_image")
+        or "image_url" in part
+        or "image" in part
+    ):
+        if "image_url" in part:
+            iu = part["image_url"]
+            return "image", iu.get("url") if isinstance(iu, dict) else iu
+        return "image", part.get("image") or part.get("input_image")
+
+    if (
+        t in ("input_audio", "audio", "audio_url")
+        or "input_audio" in part
+        or "audio_url" in part
+    ):
+        if "input_audio" in part:
+            ia = part["input_audio"]
+            if isinstance(ia, dict):
+                return "audio", (ia.get("data"), ia.get("format"))
+            return "audio", ia
+        if "audio_url" in part:
+            au = part["audio_url"]
+            return "audio", au.get("url") if isinstance(au, dict) else au
+        return "audio", part.get("audio")
+
+    if (
+        t in ("video_url", "video", "input_video")
+        or "video_url" in part
+        or "video" in part
+    ):
+        if "video_url" in part:
+            vu = part["video_url"]
+            return "video", vu.get("url") if isinstance(vu, dict) else vu
+        if "input_video" in part:
+            iv = part["input_video"]
+            if isinstance(iv, dict):
+                return "video", (iv.get("data"), iv.get("format"))
+            return "video", iv
+        return "video", part.get("video")
+
+    raise MediaError(f"unsupported content part type: {t!r}")
+
+
+def _extract_audio(video_path: str, tmp_dir: Path) -> str | None:
+    """Demux a video's audio track to mono 16 kHz wav. None if it has no audio."""
+    wav = tmp_dir / f"audio_{uuid.uuid4().hex}.wav"
+    try:
+        # If the video has no audio stream, ffmpeg will fail/exit with an error,
+        # and either not create the file or leave it empty.
+        subprocess.run(
+            [
+                "ffmpeg",
+                "-v",
+                "error",
+                "-y",
+                "-i",
+                video_path,
+                "-vn",
+                "-ac",
+                "1",
+                "-ar",
+                "16000",
+                str(wav),
+            ],
+            capture_output=True,
+            timeout=120,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+
+    if wav.exists() and wav.stat().st_size > 0:
+        return str(wav)
+
+    if wav.exists():
+        wav.unlink(missing_ok=True)
+    return None

kiapi/capabilities/chat/_operations/resolve_chat_params.py

@@ -0,0 +1,50 @@
+"""Merge a chat request with settings defaults into the complete ChatParams.
+
+Resolves and caps the sampling knobs and fills multimodal defaults once, here,
+so the per-model ``run`` works purely from :class:`ChatParams`. Model-specific
+template switches (e.g. Qwen3.6's ``enable_thinking`` default) stay in the model;
+this only passes ``chat_template_kwargs`` through verbatim.
+"""
+
+from .._settings import ChatSettings
+from .._views.chat_params import ChatParams
+from .._views.chat_request import ChatRequest
+
+
+def resolve_chat_params(
+    settings: ChatSettings,
+    req: ChatRequest,
+    *,
+    variant: str,
+) -> ChatParams:
+    requested_max_tokens = req.max_completion_tokens
+    if requested_max_tokens is None:
+        requested_max_tokens = settings.default_max_tokens
+    max_tokens = min(requested_max_tokens, settings.max_tokens_cap)
+
+    temperature = (
+        req.temperature if req.temperature is not None else settings.default_temperature
+    )
+    top_p = req.top_p if req.top_p is not None else settings.default_top_p
+    fps = req.fps if req.fps is not None else settings.default_fps
+    use_audio_in_video = (
+        req.use_audio_in_video
+        if req.use_audio_in_video is not None
+        else settings.use_audio_in_video
+    )
+
+    return ChatParams(
+        model=variant,
+        messages=req.messages,
+        tools=req.tools,
+        tool_choice=req.tool_choice,
+        parallel_tool_calls=req.parallel_tool_calls,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        seed=req.seed,
+        fps=fps,
+        use_audio_in_video=use_audio_in_video,
+        chat_template_kwargs=req.chat_template_kwargs,
+        stream=req.stream,
+    )

kiapi/capabilities/chat/_operations/stream_text_from_tokens.py

@@ -0,0 +1,32 @@
+from collections.abc import Iterable, Iterator
+from copy import copy
+
+
+def stream_text_from_tokens(processor, chunks: Iterable) -> Iterator:  # type: ignore
+    """Yield chunks whose text is decoded from token ids when mlx-vlm supports it."""
+    try:
+        from mlx_vlm.tokenizer_utils import (  # type: ignore
+            _ServerTokenStreamer,
+            make_streaming_detokenizer,
+        )
+    except Exception:
+        yield from chunks
+        return
+
+    tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor
+    streamer = _ServerTokenStreamer(tokenizer, make_streaming_detokenizer(processor))
+
+    for chunk in chunks:
+        token = getattr(chunk, "token", None)
+        finish_reason = getattr(chunk, "finish_reason", None)
+
+        if token is None:
+            text = str(getattr(chunk, "text", "") or "")
+        elif finish_reason is None:
+            text = streamer.advance(token, None)
+        else:
+            text = streamer.finalize()
+
+        out = copy(chunk)
+        out.text = text
+        yield out

kiapi/capabilities/chat/_operations/strip_tool_calls.py

@@ -0,0 +1,15 @@
+"""Shared helper: the assistant preamble before any tool call.
+
+Tool calls are emitted at the end of the turn, so everything before the first
+``<tool_call>`` tag is the natural-language preamble (often empty). Close-tag
+independent and format-independent — both the JSON (Qwen3-Omni) and Hermes/XML
+(Qwen3.6) formats use the same ``<tool_call>`` opening tag — so it is shared.
+"""
+
+_TAG = "<tool_call>"
+
+
+def strip_tool_calls(full_text: str) -> str:
+    """Return the text before the first ``<tool_call>`` tag (stripped)."""
+    tag = full_text.find(_TAG)
+    return (full_text if tag == -1 else full_text[:tag]).strip()

kiapi/capabilities/chat/_settings.py

@@ -0,0 +1,74 @@
+"""Chat-capability generation defaults, read from the environment (``KIAPI_CHAT_``).
+
+Separate from the global infra settings: each capability owns its own knobs. A
+request may override these per call.
+"""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class ChatSettings(BaseSettings):
+    """Settings for Chat Completions generation defaults and input-processing limits."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_CHAT_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    default_max_tokens: int = Field(
+        default=512,
+        title="Default maximum generated tokens",
+        description="Generation token limit used when a request omits max_tokens.",
+    )
+
+    default_temperature: float = Field(
+        default=0.7,
+        title="Default temperature",
+        description=(
+            "Sampling temperature used when a request omits temperature.\n"
+            "Lower values are more stable; higher values are more diverse."
+        ),
+    )
+
+    default_top_p: float = Field(
+        default=1.0,
+        title="Default top_p",
+        description=(
+            "Nucleus sampling threshold used when a request omits top_p.\n"
+            "Use a value greater than 0 and less than or equal to 1.0."
+        ),
+    )
+
+    default_fps: float = Field(
+        default=1.0,
+        title="Default video input sampling FPS",
+        description=(
+            "Default FPS for extracting frames from video input before passing them to the model.\n"
+            "Higher values provide finer video understanding, but increase "
+            "processing time and input size."
+        ),
+    )
+
+    use_audio_in_video: bool = Field(
+        default=True,
+        title="Use audio in video",
+        description=(
+            "When true, audio tracks in video files are demuxed and passed to "
+            "the model as audio input."
+        ),
+    )
+
+    max_tokens_cap: int = Field(
+        default=4096,
+        title="Maximum generated token cap",
+        description=(
+            "Server-side upper limit for max_tokens accepted in a request.\n"
+            "This protects the single worker queue from being blocked for a long time."
+        ),
+    )
+
+
+settings_manager = SettingsManager(ChatSettings)

kiapi/capabilities/chat/_utils/apply_parallel_tool_call_policy.py

@@ -0,0 +1,17 @@
+"""Apply OpenAI's ``parallel_tool_calls`` response contract.
+
+Qwen templates do not expose a native switch for this, so each model enforces it
+after parsing the generated tool calls: when ``parallel_tool_calls`` is false, at
+most one tool call is returned.
+"""
+
+from typing import Any
+
+
+def apply_parallel_tool_call_policy(
+    tool_calls: list[dict[str, Any]],
+    parallel_tool_calls: bool,
+) -> list[dict[str, Any]]:
+    if parallel_tool_calls is False:
+        return tool_calls[:1]
+    return tool_calls

kiapi/capabilities/chat/_utils/apply_seed.py

@@ -0,0 +1,8 @@
+"""Seed MLX's RNG for a reproducible generation (no-op when ``seed`` is None)."""
+
+
+def apply_seed(seed: int | None) -> None:
+    if seed is not None:
+        import mlx.core as mx
+
+        mx.random.seed(seed)

kiapi/capabilities/chat/_utils/load_audio_mono.py

@@ -0,0 +1,31 @@
+"""Load an audio file as a mono float32 array at the target sampling rate.
+
+Replaces ``mlx_vlm.utils.load_audio`` for the Qwen3-Omni audio path. That helper
+has a bug: ``read_audio`` returns audio channel-last ``(samples, channels)``, but
+``load_audio`` then resamples with ``resample_audio(..., axis=-1)`` (the *channel*
+axis) while downmixing with ``mean(axis=1)`` (also the channel axis). The two are
+inconsistent, so a stereo clip whose rate differs from the target (e.g. a 48 kHz
+stereo wav) is never actually resampled on the time axis — its 48 kHz samples are
+handed to the model as if they were 16 kHz, and Omni can't recognize the speech.
+Mono clips happen to work because their only axis is the time axis.
+
+Here we downmix to mono *first* (so the array is 1-D) and only then resample, so
+``resample_audio``'s default ``axis=-1`` is the time axis and the rate conversion
+is correct. Drop this once mlx-vlm fixes ``load_audio``.
+"""
+
+import numpy as np
+
+
+def load_audio_mono(path: str, sr: int) -> np.ndarray:
+    """Read ``path`` and return a mono float32 array resampled to ``sr`` Hz."""
+    from mlx_audio.audio_io import read as read_audio  # type: ignore
+    from mlx_audio.utils import resample_audio  # type: ignore
+
+    audio, orig_sr = read_audio(path, dtype="float32")
+    audio = np.asarray(audio, dtype=np.float32)
+    if audio.ndim > 1:  # (samples, channels) → mono, before any resample
+        audio = audio.mean(axis=1)
+    if orig_sr != sr:
+        audio = resample_audio(audio, orig_sr, sr)  # 1-D, so axis=-1 is time
+    return np.asarray(audio, dtype=np.float32)

kiapi/capabilities/chat/_utils/load_mlx_vlm.py

@@ -0,0 +1,12 @@
+"""Load an mlx-vlm model + processor. Shared by every chat model's ``load()``."""
+
+from types import SimpleNamespace
+
+from kiapi.core.model import ModelSpec
+
+
+def load_mlx_vlm(spec: ModelSpec) -> SimpleNamespace:
+    from mlx_vlm import load  # type: ignore
+
+    model, processor = load(spec.repo)
+    return SimpleNamespace(model=model, processor=processor)

kiapi/capabilities/chat/_utils/warmup_params.py

@@ -0,0 +1,26 @@
+"""Build the tiny :class:`ChatParams` used to prime MLX kernels at startup.
+
+A 1-token, deterministic request for the given model variant — each model's
+``warmup(payload)`` feeds it to ``run`` so the first real request isn't slowed by
+kernel compilation.
+"""
+
+from .._views.chat_params import ChatParams
+
+
+def warmup_params(model: str) -> ChatParams:
+    return ChatParams(
+        model=model,
+        messages=[{"role": "user", "content": "hi"}],
+        tools=None,
+        tool_choice=None,
+        parallel_tool_calls=True,
+        max_tokens=1,
+        temperature=0.0,
+        top_p=1.0,
+        seed=None,
+        fps=1.0,
+        use_audio_in_video=False,
+        chat_template_kwargs=None,
+        stream=False,
+    )

kiapi/capabilities/chat/_views/chat_params.py

@@ -0,0 +1,34 @@
+"""The complete contract for one chat completion run.
+
+Built from settings + request by ``resolve_chat_params``; a model's ``run`` needs
+nothing else (no ``req`` / ``settings``) to generate the completion. ``model`` is
+the canonical model name to echo back; sampling knobs are already resolved and
+capped. ``messages`` / ``tools`` / ``tool_choice`` are the request's raw OpenAI
+shapes (parsed per-model via ``parse_messages`` + each model's tool format).
+"""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class ChatParams(BaseModel):
+    model: str
+
+    messages: list[dict[str, Any]]
+    tools: list[dict[str, Any]] | None
+    tool_choice: Any | None
+    parallel_tool_calls: bool
+
+    # Resolved + capped sampling knobs.
+    max_tokens: int
+    temperature: float
+    top_p: float
+    seed: int | None
+
+    # Multimodal knobs (resolved defaults).
+    fps: float
+    use_audio_in_video: bool
+
+    chat_template_kwargs: dict[str, Any] | None
+    stream: bool