kiapi 0.1.0__py3-none-any.whl

kiapi/core/net/_helpers/verify_public_url.py

@@ -0,0 +1,68 @@
+"""SSRF guard for user-supplied URLs.
+
+Several capabilities materialize inputs from a client-provided ``http(s)`` URL
+(chat media, embedding image, generic file refs). Without a check, a client could
+point the server at ``http://localhost:...``, ``http://169.254.169.254/`` (cloud
+metadata), or any internal service — a server-side request forgery. Call
+:func:`verify_public_url` before fetching.
+
+The check resolves the host and rejects the URL if *any* resolved address is not
+a public unicast address. It is a best-effort mitigation, not airtight: a DNS
+name that re-resolves to a private address between this check and the actual
+fetch (DNS rebinding) is not covered. The ``KIAPI_NET_ALLOW_PRIVATE_URLS``
+setting disables the guard for trusted local development.
+"""
+
+import ipaddress
+import socket
+import urllib.parse
+
+from .._exceptions.unsafe_url_error import UnsafeURLError
+from .._settings import settings_manager
+
+
+def _is_public(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
+    mapped = getattr(addr, "ipv4_mapped", None)
+    if mapped is not None:
+        addr = mapped
+    return not (
+        addr.is_private
+        or addr.is_loopback
+        or addr.is_link_local
+        or addr.is_multicast
+        or addr.is_reserved
+        or addr.is_unspecified
+    )
+
+
+def verify_public_url(url: str, *, kind: str = "url") -> None:
+    """Raise :class:`UnsafeURLError` unless ``url`` is a fetchable public URL.
+
+    Accepts only ``http``/``https`` whose host resolves entirely to public
+    addresses. ``kind`` is used in the error message (e.g. ``"image"``).
+    """
+    parsed = urllib.parse.urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        raise UnsafeURLError(
+            f"{kind} url must be an http:// or https:// URL, got {parsed.scheme!r}"
+        )
+
+    host = parsed.hostname
+    if not host:
+        raise UnsafeURLError(f"{kind} url has no host")
+
+    if settings_manager.get_settings().allow_private_urls:
+        return
+
+    # A bare IP literal still parses as hostname; getaddrinfo handles both.
+    try:
+        infos = socket.getaddrinfo(host, parsed.port, proto=socket.IPPROTO_TCP)
+    except socket.gaierror as exc:
+        raise UnsafeURLError(f"cannot resolve {kind} url host {host!r}: {exc}") from exc
+
+    for info in infos:
+        ip = info[4][0]
+        if not _is_public(ipaddress.ip_address(ip)):
+            raise UnsafeURLError(
+                f"{kind} url host {host!r} resolves to non-public address {ip}"
+            )

kiapi/core/net/_settings.py

@@ -0,0 +1,28 @@
+"""Networking guard defaults, read from the environment (``KIAPI_NET_``)."""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class NetSettings(BaseSettings):
+    """Network safety settings for accessing user-provided URLs."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_NET_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    allow_private_urls: bool = Field(
+        default=False,
+        title="Allow private URLs",
+        description=(
+            "When true, disables the SSRF guard and allows URLs that resolve to "
+            "private, loopback, or link-local addresses.\n"
+            "Enable this only in a trusted local development environment."
+        ),
+    )
+
+
+settings_manager = SettingsManager(NetSettings)

kiapi/core/setup/__init__.py

@@ -0,0 +1,25 @@
+from ._exceptions.setup_required_error import SetupRequiredError
+from ._schemas.docker_image_resource import DockerImageResource
+from ._schemas.hf_snapshot_resource import HfSnapshotResource
+from ._schemas.local_path_resource import LocalPathResource
+from ._schemas.python_package_resource import PythonPackageResource
+from ._schemas.python_venv_resource import PythonVenvResource
+from ._schemas.setup_status import SetupStatus
+from ._schemas.url_file_resource import UrlFileResource
+from ._services.setup_manager import SetupManager
+from ._types.setup_resource import SetupResource
+from ._types.setup_target import SetupTarget
+
+__all__ = [
+    "DockerImageResource",
+    "HfSnapshotResource",
+    "LocalPathResource",
+    "PythonPackageResource",
+    "PythonVenvResource",
+    "SetupManager",
+    "SetupRequiredError",
+    "SetupResource",
+    "SetupStatus",
+    "SetupTarget",
+    "UrlFileResource",
+]

kiapi/core/setup/_exceptions/setup_required_error.py

@@ -0,0 +1,2 @@
+class SetupRequiredError(RuntimeError):
+    """Raised when a model is requested before its setup resources are ready."""

kiapi/core/setup/_schemas/docker_image_resource.py

@@ -0,0 +1,37 @@
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+
+
+class DockerImageResource(BaseModel):
+    """A Docker image that must be pulled before the model (backend) runs."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["docker_image"] = Field(
+        default="docker_image",
+        description="Resource kind discriminator.",
+        examples=["docker_image"],
+    )
+    image: str = Field(
+        ...,
+        description="Docker image reference to pull.",
+        examples=["searxng/searxng:latest"],
+    )
+    disk_gb: float | None = Field(
+        default=None,
+        ge=0.0,
+        description="Approximate on-disk size of the image in GB.",
+        examples=[0.4],
+    )
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Label accepted by kiapi activate/deactivate --repo.",
+    )
+    @property
+    def label(self) -> str:
+        return self.image
+
+    @property
+    def key(self) -> str:
+        return f"{self.kind}:{self.image}"

kiapi/core/setup/_schemas/hf_snapshot_resource.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+
+
+class HfSnapshotResource(BaseModel):
+    """A Hugging Face snapshot that must be downloaded before the model runs."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["hf_snapshot"] = Field(
+        default="hf_snapshot",
+        description="Resource kind discriminator.",
+        examples=["hf_snapshot"],
+    )
+    repo: str = Field(
+        ...,
+        description="Hugging Face repository id to snapshot-download.",
+        examples=["mlx-community/Qwen3-Omni-30B-A3B-Instruct-4bit"],
+    )
+    disk_gb: float | None = Field(
+        default=None,
+        ge=0.0,
+        description="Approximate on-disk size of the snapshot in GB.",
+        examples=[21.8],
+    )
+    revision: str | None = Field(
+        default=None,
+        description="Git revision (branch, tag, or commit) to pin; None means main.",
+        examples=["main"],
+    )
+    local_dir: str | None = Field(
+        default=None,
+        description="If set, materialize into this directory instead of the HF cache.",
+        examples=[None],
+    )
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Label accepted by kiapi activate/deactivate --repo.",
+    )
+    @property
+    def label(self) -> str:
+        return self.repo
+
+    @property
+    def key(self) -> str:
+        revision = self.revision or "main"
+        local_dir = str(Path(self.local_dir).expanduser()) if self.local_dir else ""
+        return f"{self.kind}:{self.repo}:{revision}:{local_dir}"

kiapi/core/setup/_schemas/local_path_resource.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+
+
+class LocalPathResource(BaseModel):
+    """A local filesystem path that must already exist before the model runs."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["local_path"] = Field(
+        default="local_path",
+        description="Resource kind discriminator.",
+        examples=["local_path"],
+    )
+    path: str = Field(
+        ...,
+        description="Filesystem path that must exist for the model to run.",
+        examples=["./checkpoints/ace-step-v1.5-xl-base"],
+    )
+    label_name: str | None = Field(
+        default=None,
+        description="Optional friendly label; defaults to the path when unset.",
+        examples=["ace-step-xl-base"],
+    )
+    disk_gb: float | None = Field(
+        default=None,
+        ge=0.0,
+        description="Approximate on-disk size of the path in GB.",
+        examples=[12.0],
+    )
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Label accepted by kiapi activate/deactivate --repo.",
+    )
+    @property
+    def label(self) -> str:
+        return self.label_name or self.path
+
+    @property
+    def key(self) -> str:
+        return f"{self.kind}:{Path(self.path).expanduser()}"

kiapi/core/setup/_schemas/python_package_resource.py

@@ -0,0 +1,58 @@
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+
+
+class PythonPackageResource(BaseModel):
+    """A Python package installed into the current kiapi environment."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["python_package"] = Field(
+        default="python_package",
+        description="Resource kind discriminator.",
+        examples=["python_package"],
+    )
+    package: str = Field(
+        ...,
+        description="Distribution name used when uninstalling the package.",
+        examples=["mlx-video"],
+    )
+    spec: str = Field(
+        ...,
+        description="Package spec used when installing the package.",
+        examples=["mlx-video @ git+https://github.com/Blaizzy/mlx-video.git@<commit>"],
+    )
+    import_name: str = Field(
+        ...,
+        description="Module imported to verify that the package is ready.",
+        examples=["mlx_video.models.ltx_2.generate"],
+    )
+    verify_attrs: tuple[str, ...] = Field(
+        default_factory=tuple,
+        description="Attributes that must exist on the imported module.",
+        examples=[["PipelineType", "generate_video"]],
+    )
+    label_name: str | None = Field(
+        default=None,
+        description="Optional friendly label; defaults to the package name.",
+        examples=["mlx-video-ltx2"],
+    )
+    disk_gb: float | None = Field(
+        default=None,
+        ge=0.0,
+        description="Approximate on-disk size of the installed package in GB.",
+        examples=[0.5],
+    )
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Label accepted by kiapi activate/deactivate --repo.",
+    )
+    @property
+    def label(self) -> str:
+        return self.label_name or self.package
+
+    @property
+    def key(self) -> str:
+        attrs = ",".join(self.verify_attrs)
+        return f"{self.kind}:{self.package}:{self.spec}:{self.import_name}:{attrs}"

kiapi/core/setup/_schemas/python_venv_resource.py

@@ -0,0 +1,63 @@
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+
+
+class PythonVenvResource(BaseModel):
+    """A Python virtual environment that must be built before the model runs."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["python_venv"] = Field(
+        default="python_venv",
+        description="Resource kind discriminator.",
+        examples=["python_venv"],
+    )
+    path: str = Field(
+        ...,
+        description="Virtual environment directory.",
+        examples=[".venv-acestep"],
+    )
+    python: str = Field(
+        default="3.12",
+        description="Python version or executable used to create the virtual environment.",
+        examples=["3.12"],
+    )
+    packages: tuple[str, ...] = Field(
+        ...,
+        min_length=1,
+        description="Pip package specs installed into the virtual environment.",
+        examples=[["ace-step @ git+https://github.com/ace-step/ACE-Step-1.5.git"]],
+    )
+    import_name: str = Field(
+        ...,
+        description="Module imported to verify that the environment is ready.",
+        examples=["acestep"],
+    )
+    label_name: str | None = Field(
+        default=None,
+        description="Optional friendly label; defaults to the venv path when unset.",
+        examples=["acestep-venv"],
+    )
+    disk_gb: float | None = Field(
+        default=None,
+        ge=0.0,
+        description="Approximate on-disk size of the virtual environment in GB.",
+        examples=[8.0],
+    )
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Label accepted by kiapi activate/deactivate --repo.",
+    )
+    @property
+    def label(self) -> str:
+        return self.label_name or self.path
+
+    @property
+    def key(self) -> str:
+        packages = ",".join(self.packages)
+        return (
+            f"{self.kind}:{Path(self.path).expanduser()}:{self.python}:"
+            f"{packages}:{self.import_name}"
+        )

kiapi/core/setup/_schemas/setup_status.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class SetupStatus:
+    ready: bool
+    detail: str

kiapi/core/setup/_schemas/url_file_resource.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field, computed_field
+
+
+class UrlFileResource(BaseModel):
+    """A single file fetched from a URL and cached at a local path."""
+
+    model_config = ConfigDict(frozen=True)
+
+    kind: Literal["url_file"] = Field(
+        default="url_file",
+        description="Resource kind discriminator.",
+        examples=["url_file"],
+    )
+    url: str = Field(
+        ...,
+        description="Source URL the file is downloaded from.",
+        examples=["https://ml-site.cdn-apple.com/models/depth-pro/depth_pro.pt"],
+    )
+    path: str = Field(
+        ...,
+        description="Local path the downloaded file is stored at.",
+        examples=["./checkpoints/depth-pro/depth_pro.pt"],
+    )
+    disk_gb: float | None = Field(
+        default=None,
+        ge=0.0,
+        description="Approximate on-disk size of the file in GB.",
+        examples=[1.9],
+    )
+
+    @computed_field(  # type: ignore[prop-decorator]
+        description="Label accepted by kiapi activate/deactivate --repo.",
+    )
+    @property
+    def label(self) -> str:
+        return self.url
+
+    @property
+    def key(self) -> str:
+        return f"{self.kind}:{self.url}:{Path(self.path).expanduser()}"