kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/embedding/_constants/description.py

@@ -0,0 +1,47 @@
+DESCRIPTION = """Text and multimodal embeddings for retrieval and similarity search.
+
+POST one item to `/v1/embedding` with one field per modality (`text` and/or
+`image`). Returns a single L2-normalized, last-token-pooled vector. This is **not**
+OpenAI's array `input` shape — one item per request. Sync only (no async mode).
+
+## Upstream docs
+- [mlx-embeddings](https://github.com/Blaizzy/mlx-embeddings) — the MLX embedding engine kiapi runs
+- [mlx-community/Qwen3-Embedding-8B-mxfp8](https://huggingface.co/mlx-community/Qwen3-Embedding-8B-mxfp8) — `qwen3-embedding-8b` weights
+- [mlx-community/Qwen3-VL-Embedding-2B-mxfp8](https://huggingface.co/mlx-community/Qwen3-VL-Embedding-2B-mxfp8) — `qwen3-vl-embedding-2b` weights
+
+## Choosing A Model
+`model` selects a registered embedding model (full catalog and aliases:
+`GET /v1/embedding/models`). The served models differ by input modality, so
+choose by what you send:
+- **qwen3-embedding-8b** (default) — text only. Larger, stronger text retrieval.
+- **qwen3-vl-embedding-2b** — text + image. Use it for image embeddings or a
+  shared text/image space; sending `image` to the text-only model returns 400.
+
+Vectors from different models live in different spaces and may differ in
+dimensionality — embed both your queries and your corpus with the **same** model.
+
+## Notes
+- One input item per request. To embed N items, make N requests.
+- Provide at least one modality input; an empty request returns 400.
+- Vectors are already L2-normalized, so cosine similarity reduces to a dot
+  product.
+
+## Examples
+
+### Text (default model)
+```sh
+curl -sS http://HOST:PORT/v1/embedding \\
+  -H 'Content-Type: application/json' \\
+  -d '{"text": "今日の天気は晴れです"}'
+```
+
+### Image (multimodal model)
+```sh
+curl -sS http://HOST:PORT/v1/embedding \\
+  -H 'Content-Type: application/json' \\
+  -d '{
+    "model": "qwen3-vl-embedding-2b",
+    "image": "data:image/png;base64,iVBORw0..."
+  }'
+```
+"""

kiapi/capabilities/embedding/_helpers/register.py

@@ -0,0 +1,66 @@
+"""Register embedding's models + OpenAPI description in the global registries."""
+
+from kiapi.core.capability import CapabilitySpec, capability_spec_registry
+from kiapi.core.model import ModelSpec, model_registry
+from kiapi.core.setup import HfSnapshotResource
+
+from .._constants.description import DESCRIPTION
+from .._models import qwen3_text, qwen3_vl
+
+
+def register() -> None:
+    capability_spec_registry.register(
+        CapabilitySpec(
+            name="embedding",
+            domain="embedding",
+            title="kiapi Embedding API",
+            summary="Create text and multimodal embeddings for retrieval and similarity search.",
+            description=DESCRIPTION,
+            openapi_path="/v1/embedding/openapi.json",
+            docs_path="/v1/embedding/docs",
+            redoc_path="/v1/embedding/redoc",
+            path_prefixes=("/v1/embedding",),
+        )
+    )
+
+    model_registry.register(
+        ModelSpec(
+            name="qwen3-embedding-8b",
+            family="embedding",
+            domain="embedding",
+            repo="mlx-community/Qwen3-Embedding-8B-mxfp8",
+            module=qwen3_text,
+            weight_gb=7.3,  # measured on device (estimate was 8.0)
+            peak_headroom_gb=2.0,
+            framework="mlx",
+            priority=10,
+            aliases=("text", "qwen3-embedding", "qwen3_embedding"),
+            default=True,
+            setup_resources=(
+                HfSnapshotResource(
+                    repo="mlx-community/Qwen3-Embedding-8B-mxfp8",
+                    disk_gb=7.82,
+                ),
+            ),
+        )
+    )
+    model_registry.register(
+        ModelSpec(
+            name="qwen3-vl-embedding-2b",
+            family="embedding",
+            domain="embedding",
+            repo="mlx-community/Qwen3-VL-Embedding-2B-mxfp8",
+            module=qwen3_vl,
+            weight_gb=2.4,  # measured on device (estimate was 2.0)
+            peak_headroom_gb=2.0,  # image inputs add some headroom
+            framework="mlx",
+            priority=20,  # tiny + frequently used → keep resident
+            aliases=("vl", "qwen3-vl-embedding", "qwen3_vl_embedding"),
+            setup_resources=(
+                HfSnapshotResource(
+                    repo="mlx-community/Qwen3-VL-Embedding-2B-mxfp8",
+                    disk_gb=2.59,
+                ),
+            ),
+        )
+    )

kiapi/capabilities/embedding/_models/qwen3_text.py

@@ -0,0 +1,53 @@
+"""Handler for Qwen3-Embedding-8B (text embeddings; ``model_type: qwen3``).
+
+A text-only embedding model. Owns its own load + embed flow; shares the
+input-parsing and response-shaping operations with the VL handler.
+"""
+
+import tempfile
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+from kiapi.core.model import ModelSpec
+
+from .._operations.format_response import format_response
+from .._operations.parse_inputs import parse_inputs
+from .._utils.to_vector import to_vector
+from .._views.embed_params import EmbedParams
+
+FEATURES = {"text"}
+
+
+def load(spec: ModelSpec) -> SimpleNamespace:
+    from mlx_embeddings import load as _load  # type: ignore
+
+    model, processor = _load(spec.repo)
+    return SimpleNamespace(model=model, processor=processor)
+
+
+def warmup(payload: SimpleNamespace) -> None:
+    run(
+        payload,
+        EmbedParams(model="qwen3-embedding-8b", inputs={"text": "hi"}, max_length=512),
+    )
+
+
+def run(payload: SimpleNamespace, params: EmbedParams) -> dict[str, Any]:
+    from mlx_embeddings import generate
+
+    model, processor = payload.model, payload.processor
+    # text-only: parse_inputs validates capabilities; no media files are written,
+    # so the tmp_dir is never actually touched.
+    parsed = parse_inputs(params.inputs, Path(tempfile.gettempdir()), allow=FEATURES)
+
+    t0 = time.time()
+    output = generate(model, processor, parsed["text"], max_length=params.max_length)
+    elapsed = time.time() - t0
+
+    return format_response(
+        model_name=params.model,
+        embedding=to_vector(output),
+        elapsed=elapsed,
+    )

kiapi/capabilities/embedding/_models/qwen3_vl.py

@@ -0,0 +1,97 @@
+"""Handler for Qwen3-VL-Embedding-2B (text + image embeddings; ``model_type: qwen3_vl``).
+
+Owns the full embed flow for the multimodal embedding model. The MLX-converted
+processor may be missing the ``chat_template`` and token-id attributes that
+Transformers normally sets up, so :func:`load` patches them in after loading (a
+known Qwen3-VL quirk, carried over from the original single-purpose server).
+
+Embedding is done via the model's high-level ``process([item], processor=...)``
+API, where ``item`` is ``{"text": ..., "image": <path>}`` (either key optional).
+"""
+
+import shutil
+import time
+from types import SimpleNamespace
+from typing import Any
+
+from kiapi.core.model import ModelSpec
+from kiapi.core.workdir import create_work_dir
+
+from .._operations.format_response import format_response
+from .._operations.parse_inputs import parse_inputs
+from .._utils.to_vector import to_vector
+from .._views.embed_params import EmbedParams
+
+FEATURES = {"text", "image"}
+
+# Source repo for the chat template, when the converted processor lacks one.
+TEMPLATE_SOURCE_MODEL_ID = "Qwen/Qwen3-VL-Embedding-2B"
+
+
+def load(spec: ModelSpec) -> SimpleNamespace:
+    from mlx_embeddings import load as _load  # type: ignore
+
+    model, processor = _load(spec.repo)
+    payload = SimpleNamespace(model=model, processor=processor)
+    _patch_processor(payload.processor)
+    return payload
+
+
+def warmup(payload: SimpleNamespace) -> None:
+    run(
+        payload,
+        EmbedParams(
+            model="qwen3-vl-embedding-2b", inputs={"text": "hi"}, max_length=512
+        ),
+    )
+
+
+def _patch_processor(processor: Any) -> None:
+    """Patch in the chat template / token-id attributes the converted processor
+    may be missing (Qwen3-VL quirk). Best-effort and idempotent."""
+    inner = getattr(processor, "processor", None)
+    tokenizer = getattr(processor, "tokenizer", None)
+
+    if getattr(processor, "chat_template", None) is None:
+        try:
+            chat_template = processor._load_chat_template(TEMPLATE_SOURCE_MODEL_ID)
+            if inner is not None:
+                inner.chat_template = chat_template
+            if tokenizer is not None:
+                tokenizer.chat_template = chat_template
+        except Exception:
+            pass
+
+    if inner is not None:
+        image_token_id = getattr(processor, "image_token_id", None)
+        video_token_id = getattr(processor, "video_token_id", None)
+        if image_token_id is not None:
+            inner.image_ids = [image_token_id]
+        if video_token_id is not None:
+            inner.video_ids = [video_token_id]
+        inner.audio_ids = []
+
+
+def run(payload: SimpleNamespace, params: EmbedParams) -> dict[str, Any]:
+    model, processor = payload.model, payload.processor
+    tmp_dir = create_work_dir("embedding/qwen3_vl")
+    try:
+        parsed = parse_inputs(params.inputs, tmp_dir, allow=FEATURES)
+
+        item: dict[str, Any] = {}
+        if "text" in parsed:
+            item["text"] = parsed["text"]
+        if "image" in parsed:
+            item["image"] = parsed["image"]  # local file path
+
+        t0 = time.time()
+        embeddings = model.process([item], processor=processor)
+        elapsed = time.time() - t0
+
+        return format_response(
+            model_name=params.model,
+            embedding=to_vector(embeddings),
+            elapsed=elapsed,
+        )
+    finally:
+        shutil.rmtree(tmp_dir, ignore_errors=True)

kiapi/capabilities/embedding/_operations/format_response.py

@@ -0,0 +1,24 @@
+"""Shape an embedding vector into the server's ``/v1/embedding`` response dict.
+
+Identical across handlers, so it lives here and the per-model ``run()`` calls it
+after computing the vector.
+"""
+
+from typing import Any
+
+
+def format_response(
+    *,
+    model_name: str,
+    embedding: list[float],
+    elapsed: float,
+    prompt_tokens: int = 0,
+) -> dict[str, Any]:
+    """Build the ``/v1/embedding`` response dict for a single embedding."""
+    return {
+        "model": model_name,
+        "embedding": embedding,
+        "dimension": len(embedding),
+        "usage": {"prompt_tokens": prompt_tokens, "total_tokens": prompt_tokens},
+        "timings": {"total_s": round(elapsed, 3)},
+    }

kiapi/capabilities/embedding/_operations/handle_embed.py

@@ -0,0 +1,24 @@
+"""Embedding service entry: resolve → acquire (memory budget) → run.
+
+Body of the embedding job's worker-thread thunk. Mirrors mlx-embedding-server's
+Engine.run, but loads through the *global* memory manager so embedding models
+share the budget with every other capability.
+"""
+
+from kiapi.core.app import AppContext
+from kiapi.core.model import model_registry
+
+from .._settings import settings_manager
+from .._views.embed_request import EmbedRequest
+from .resolve_embed_params import resolve_embed_params
+
+
+def handle_embed(ctx: AppContext, req: EmbedRequest) -> tuple[dict, list[str]]:
+    """Run one embedding. Returns (embedding_result, artifact_file_ids)."""
+    settings = settings_manager.get_settings()
+    spec = model_registry.resolve("embedding", req.model)
+    ctx.ensure_model_ready(spec)
+    params = resolve_embed_params(settings, req, variant=spec.name)
+    payload = ctx.memory_manager.acquire(spec)
+    result = spec.module.run(payload, params)
+    return result, []  # embeddings produce no file artifacts

kiapi/capabilities/embedding/_operations/parse_inputs.py

@@ -0,0 +1,108 @@
+"""Turn an embedding request's modality inputs into usable values.
+
+The request carries one field per modality (``text``, ``image``, ...). Model
+handlers call :func:`parse_inputs` with the set of modalities they support
+(``allow=``); anything outside that set — or an empty request — is rejected with
+:class:`~kiapi.capabilities.CapabilityError` (mapped to HTTP 400 by the API layer).
+
+  - ``text``  is passed through as-is.
+  - ``image`` accepts a data URL, an http(s) URL, or bare base64, and is
+    materialized to a file under ``tmp_dir`` (the loader reads paths).
+
+Returns a dict keyed by modality, e.g. ``{"text": "...", "image": "/tmp/..png"}``.
+"""
+
+import base64
+import binascii
+import re
+import urllib.request
+import uuid
+from pathlib import Path
+from typing import Any
+
+from kiapi.capabilities import CapabilityError
+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",
+}
+_DEFAULT_EXT = {"image": ".png"}
+
+
+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) -> str:
+    """Turn a data URL / http(s) URL / bare base64 into a file path."""
+    if not isinstance(source, str) or not source:
+        raise CapabilityError(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 CapabilityError(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 CapabilityError(f"{kind}: {exc}")  # noqa: B904
+        try:
+            with urllib.request.urlopen(source, timeout=30) as resp:
+                data = resp.read()
+        except Exception as exc:
+            raise CapabilityError(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 (no data: prefix) — the common case for {"image": "<b64>"}.
+    try:
+        data = base64.b64decode(source, validate=True)
+    except (binascii.Error, ValueError):
+        raise CapabilityError(f"{kind}: source is not a data URL, URL, or base64")  # noqa: B904
+    return _write_bytes(data, tmp_dir, kind, _DEFAULT_EXT[kind])
+
+
+def parse_inputs(
+    inputs: dict[str, Any],
+    tmp_dir: Path,
+    *,
+    allow: set[str] | None = None,
+) -> dict[str, Any]:
+    """Validate + materialize request inputs into a per-modality dict.
+
+    ``allow`` is the set of modalities the selected model supports (e.g.
+    ``{"text", "image"}``). An input outside it raises :class:`CapabilityError`.
+    ``None`` allows everything (used by tests).
+    """
+    if not inputs:
+        raise CapabilityError(
+            "request has no input (provide at least one of: text, image)"
+        )
+
+    out: dict[str, Any] = {}
+    for kind, source in inputs.items():
+        if allow is not None and kind not in allow:
+            raise CapabilityError(
+                f"{kind} input is not supported by the selected model"
+            )
+        if kind == "text":
+            out["text"] = source or ""
+        elif kind == "image":
+            out["image"] = _materialize(source, "image", tmp_dir)
+        else:  # pragma: no cover — schema only exposes known modalities
+            raise CapabilityError(f"unsupported input modality: {kind!r}")
+    return out

kiapi/capabilities/embedding/_operations/resolve_embed_params.py

@@ -0,0 +1,18 @@
+"""Merge an embedding request with settings defaults into the complete EmbedParams."""
+
+from .._settings import EmbeddingSettings
+from .._views.embed_params import EmbedParams
+from .._views.embed_request import EmbedRequest
+
+
+def resolve_embed_params(
+    settings: EmbeddingSettings,
+    req: EmbedRequest,
+    *,
+    variant: str,
+) -> EmbedParams:
+    return EmbedParams(
+        model=variant,
+        inputs=req.inputs(),
+        max_length=settings.max_length,
+    )

kiapi/capabilities/embedding/_settings.py

@@ -0,0 +1,27 @@
+"""Embedding-capability defaults, read from the environment (``KIAPI_EMBEDDING_``)."""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class EmbeddingSettings(BaseSettings):
+    """Settings for Embedding input length limits."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_EMBEDDING_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    max_length: int = Field(
+        default=512,
+        title="Maximum input token length",
+        description=(
+            "Maximum token length passed to the embedding model tokenizer.\n"
+            "Inputs longer than this are truncated before model processing."
+        ),
+    )
+
+
+settings_manager = SettingsManager(EmbeddingSettings)

kiapi/capabilities/embedding/_utils/to_vector.py

@@ -0,0 +1,20 @@
+"""Extract a single embedding vector from an mlx-embeddings model output."""
+
+from typing import Any
+
+
+def to_vector(output: Any) -> list[float]:
+    """Extract a single embedding vector from an mlx-embeddings model output.
+
+    mlx-embeddings returns a model-output object with ``text_embeds`` (normalized,
+    pooled, shape ``(batch, dim)``). We embed one item per request, so take row 0.
+    Falls back to treating the output itself as the array if there's no attribute.
+    """
+    import mlx.core as mx
+
+    embeds = getattr(output, "text_embeds", None)
+    if embeds is None:
+        embeds = getattr(output, "pooler_output", output)
+    mx.eval(embeds)
+    rows = embeds.tolist()
+    return rows[0] if rows and isinstance(rows[0], list) else rows

kiapi/capabilities/embedding/_views/embed_params.py

@@ -0,0 +1,17 @@
+"""The complete contract for one embedding run.
+
+Built from settings + request by ``resolve_embed_params``; a model's ``run`` needs
+nothing else (no ``req`` / ``settings``) to produce the embedding and its response.
+``model`` is the canonical model name to echo back; ``inputs`` are the request's
+raw modality sources (materialized per-model via ``parse_inputs``).
+"""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class EmbedParams(BaseModel):
+    model: str
+    inputs: dict[str, Any]
+    max_length: int

kiapi/capabilities/embedding/_views/embed_request.py

@@ -0,0 +1,39 @@
+"""Embedding request model — one item, one field per modality."""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class EmbedRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = Field(
+        default=None,
+        description=(
+            "Registry name, alias, or repo of the embedding model (see "
+            "`GET /v1/embedding/models`). Omit to use the default model."
+        ),
+    )
+
+    text: str | None = Field(
+        default=None,
+        description="Text to embed. Supported by every embedding model.",
+    )
+    image: str | None = Field(
+        default=None,
+        description=(
+            "Image to embed, as a base64 string, data URL, http(s) URL, or local "
+            "path. Only multimodal models accept it; sending it to a text-only "
+            "model returns HTTP 400."
+        ),
+    )
+
+    def inputs(self) -> dict[str, Any]:
+        """Present modality inputs, keyed by modality name (skips unset ones)."""
+        out: dict[str, Any] = {}
+        if self.text is not None:
+            out["text"] = self.text
+        if self.image is not None:
+            out["image"] = self.image
+        return out