apsimo 1.3.0__py3-none-any.whl

apsimo/vector/safety_image.py

@@ -0,0 +1,159 @@
+"""Colony Vector — image content classifier.
+
+Validates images before embedding: format, dimensions, and optional
+content classification.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class ImageCheckLevel(str, Enum):
+    """Image validation strictness."""
+
+    OFF = "off"        # No checks
+    BASIC = "basic"    # Format, dimension, and size validation only
+    STRICT = "strict"  # Format + dimensions + content classification
+
+
+@dataclass
+class ImageCheckResult:
+    """Result of an image content check."""
+
+    safe: bool = True
+    level: ImageCheckLevel = ImageCheckLevel.BASIC
+    errors: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+    classification: Optional[dict[str, float]] = None  # content scores for strict mode
+
+    @property
+    def reason(self) -> str:
+        return "; ".join(self.errors) if self.errors else ""
+
+
+async def check_image(
+    data: bytes,
+    mime_type: str = "",
+    width: int = 0,
+    height: int = 0,
+    level: ImageCheckLevel = ImageCheckLevel.BASIC,
+) -> ImageCheckResult:
+    """Check image content.
+
+    Parameters
+    ----------
+    data : bytes
+        Raw image bytes.
+    mime_type : str
+        Detected or provided MIME type.
+    width, height : int
+        Image dimensions (0 if unknown).
+    level : ImageCheckLevel
+        How strict to check.
+
+    Returns
+    -------
+    ImageCheckResult
+    """
+    result = ImageCheckResult(level=level)
+
+    if level == ImageCheckLevel.OFF:
+        return result
+
+    # --- Basic checks ---
+
+    # Format
+    valid_mimes = {"image/jpeg", "image/png", "image/webp", "image/gif"}
+    if mime_type and mime_type not in valid_mimes:
+        result.errors.append(f"Unsupported image format: {mime_type}")
+        result.safe = False
+
+    # Size
+    max_size = 20 * 1024 * 1024  # 20 MB
+    if len(data) > max_size:
+        result.errors.append(f"Image size {len(data)} exceeds 20MB limit")
+        result.safe = False
+
+    if len(data) == 0:
+        result.errors.append("Empty image data")
+        result.safe = False
+
+    # Dimensions (if known)
+    max_dim = 4096
+    if width > max_dim or height > max_dim:
+        result.warnings.append(f"Image dimensions {width}x{height} exceed {max_dim}px — will be resized")
+
+    # Minimum dimensions (avoid 1px images, etc.)
+    if width > 0 and height > 0 and (width < 16 or height < 16):
+        result.warnings.append(f"Very small image: {width}x{height}px")
+
+    if level == ImageCheckLevel.BASIC:
+        return result
+
+    # --- Strict checks (content classification) ---
+
+    try:
+        classification = await _classify_content(data, mime_type)
+        result.classification = classification
+
+        # Flag content with high explicit scores
+        explicit_score = classification.get("explicit", 0.0)
+        if explicit_score > 0.8:
+            result.errors.append(f"Image classified as explicit content (score: {explicit_score:.2f})")
+            result.safe = False
+        elif explicit_score > 0.5:
+            result.warnings.append(f"Image may contain explicit content (score: {explicit_score:.2f})")
+
+    except Exception as exc:
+        logger.warning("Content classification failed: %s — allowing image (basic checks passed)", exc)
+        result.warnings.append("Content classification unavailable")
+
+    return result
+
+
+async def _classify_content(data: bytes, mime_type: str) -> dict[str, float]:
+    """Classify image content.
+
+    Tries a local content classifier (falconsai/nsfw_image_detection via
+    transformers). Falls back to safe scores if the model isn't available.
+    """
+    try:
+        import asyncio
+        from PIL import Image as PILImage
+        import io
+
+        img = PILImage.open(io.BytesIO(data))
+        if img.mode in ("RGBA", "P"):
+            img = img.convert("RGB")
+
+        def _run_classifier():
+            from transformers import pipeline as hf_pipeline
+            classifier = hf_pipeline("image-classification", model="falconsai/nsfw_image_detection", device=-1)
+            return classifier(img)
+
+        results = await asyncio.get_running_loop().run_in_executor(None, _run_classifier)
+
+        # Parse results into score dict
+        scores = {"explicit": 0.0, "safe": 1.0}
+        for r in results:
+            label = r.get("label", "").lower()
+            score = r.get("score", 0.0)
+            if "nsfw" in label or "explicit" in label:
+                scores["explicit"] = score
+            if "safe" in label or "normal" in label:
+                scores["safe"] = score
+
+        return scores
+
+    except ImportError:
+        logger.debug("transformers not available for image classification — returning safe defaults")
+        return {"explicit": 0.0, "safe": 1.0}
+    except Exception as exc:
+        logger.debug("Image classification failed: %s — returning safe defaults", exc)
+        return {"explicit": 0.0, "safe": 1.0}

apsimo/vector/scanner.py

@@ -0,0 +1,197 @@
+"""System scanner — detect GPU, VRAM, RAM to select embedding tier.
+
+Scans the host hardware and returns a ``HardwareProfile`` that the
+tier selector uses to pick the right models.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import platform
+import shutil
+import subprocess
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class HardwareProfile:
+    """Detected hardware capabilities."""
+
+    gpu_type: str           # "cuda", "mlx", "none"
+    gpu_name: str           # e.g. "NVIDIA GB10"
+    vram_gb: int            # GPU memory in GB (0 if no GPU)
+    ram_gb: int             # System RAM in GB
+    cuda_version: str | None = None
+    os: str = ""            # "linux", "darwin", "windows"
+
+
+def scan() -> HardwareProfile:
+    """Scan the system and return a HardwareProfile."""
+    gpu_type = "none"
+    gpu_name = ""
+    vram_gb = 0
+    cuda_version = None
+
+    # Detect CUDA
+    if _has_cuda():
+        gpu_type = "cuda"
+        gpu_name, vram_gb, cuda_version = _cuda_info()
+
+    # Detect MLX (Apple Silicon)
+    if gpu_type == "none" and _has_mlx():
+        gpu_type = "mlx"
+        gpu_name, vram_gb = _mlx_info()
+
+    ram_gb = _ram_gb()
+    system = platform.system().lower()
+
+    profile = HardwareProfile(
+        gpu_type=gpu_type,
+        gpu_name=gpu_name,
+        vram_gb=vram_gb,
+        ram_gb=ram_gb,
+        cuda_version=cuda_version,
+        os=system,
+    )
+
+    logger.info(
+        "Hardware scan: gpu=%s (%s, %dGB VRAM), ram=%dGB, cuda=%s, os=%s",
+        profile.gpu_type, profile.gpu_name, profile.vram_gb,
+        profile.ram_gb, profile.cuda_version, profile.os,
+    )
+
+    return profile
+
+
+def _has_cuda() -> bool:
+    """Check if CUDA is available."""
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
+            capture_output=True, text=True, timeout=5,
+        )
+        return result.returncode == 0
+    except Exception:
+        return False
+
+
+def _cuda_info() -> tuple[str, int, str | None]:
+    """Get CUDA GPU name, VRAM, and driver version.
+
+    Handles unified memory GPUs (Grace Blackwell, Grace Hopper, Tegra)
+    where nvidia-smi reports [N/A] for memory.total.  On these chips
+    the GPU uses system RAM, so we fall back to total system RAM.
+    """
+    name = "Unknown NVIDIA GPU"
+    vram = 0
+    cuda_ver = None
+
+    try:
+        result = subprocess.run(
+            ["nvidia-smi",
+             "--query-gpu=name,memory.total",
+             "--format=csv,noheader,nounits"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0:
+            parts = result.stdout.strip().split(",")
+            if len(parts) >= 2:
+                name = parts[0].strip()
+                mem_str = parts[1].strip()
+                # Handle [N/A] for unified memory GPUs
+                if mem_str.lower() in ("[n/a]", "n/a", ""):
+                    vram = _ram_gb()  # Unified memory: GPU RAM == system RAM
+                else:
+                    vram = int(float(mem_str))
+    except Exception:
+        pass
+
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0:
+            cuda_ver = result.stdout.strip()
+    except Exception:
+        pass
+
+    return name, vram, cuda_ver
+
+
+def _has_mlx() -> bool:
+    """Check if Apple MLX is available (M-series chip)."""
+    if platform.system() != "Darwin":
+        return False
+    try:
+        # Check for Apple Silicon
+        result = subprocess.run(
+            ["sysctl", "-n", "machdep.cpu.brand_string"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0 and "Apple" in result.stdout:
+            return True
+    except Exception:
+        pass
+    return False
+
+
+def _mlx_info() -> tuple[str, int]:
+    """Get Apple Silicon GPU name and unified memory."""
+    name = "Apple Silicon"
+    vram = 0
+
+    try:
+        result = subprocess.run(
+            ["sysctl", "-n", "machdep.cpu.brand_string"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0:
+            name = result.stdout.strip()
+    except Exception:
+        pass
+
+    # On Apple Silicon, VRAM == RAM (unified memory)
+    vram = _ram_gb()
+
+    return name, vram
+
+
+def _ram_gb() -> int:
+    """Get system RAM in GB."""
+    system = platform.system().lower()
+
+    try:
+        if system == "linux":
+            result = subprocess.run(
+                ["free", "-g", "--si"],
+                capture_output=True, text=True, timeout=5,
+            )
+            if result.returncode == 0:
+                for line in result.stdout.splitlines():
+                    if line.startswith("Mem:"):
+                        return int(line.split()[1])
+        elif system == "darwin":
+            result = subprocess.run(
+                ["sysctl", "-n", "hw.memsize"],
+                capture_output=True, text=True, timeout=5,
+            )
+            if result.returncode == 0:
+                return int(result.stdout.strip()) // (1024 ** 3)
+    except Exception:
+        pass
+
+    # Fallback: parse /proc/meminfo
+    try:
+        with open("/proc/meminfo") as f:
+            for line in f:
+                if line.startswith("MemTotal:"):
+                    kb = int(line.split()[1])
+                    return kb // (1024 * 1024) + 1
+    except Exception:
+        pass
+
+    return 8  # Safe default

apsimo/vector/setup.py

@@ -0,0 +1,289 @@
+"""Colony Vector Store — hardware detection, config recommendation, and CLI.
+
+Provides:
+  - detect_hardware()    — probe CUDA, MLX, CPU resources
+  - recommend_config()   — select best EmbeddingConfig for detected hardware
+  - ensure_embed_config() — resolve config from env or auto-detect
+  - ensure_model_ready()  — download model weights if missing
+  - CLI entry point for ``colony vector setup`` and ``colony vector rebuild``
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import sys
+from typing import Optional
+
+from apsimo.vector.config import EmbeddingConfig, HardwareProfile, VectorStoreConfig
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Hardware detection
+# ---------------------------------------------------------------------------
+
+def detect_hardware() -> HardwareProfile:
+    """Probe available compute resources.  Returns a structured profile."""
+    profile = HardwareProfile()
+
+    # CUDA
+    try:
+        import torch
+        if torch.cuda.is_available():
+            profile.cuda_available = True
+            profile.cuda_device_count = torch.cuda.device_count()
+            profile.cuda_vram_gb = sum(
+                torch.cuda.get_device_properties(i).total_memory
+                for i in range(profile.cuda_device_count)
+            ) / 1e9
+    except ImportError:
+        pass
+
+    # Apple Silicon (MLX)
+    try:
+        import mlx.core as mx  # noqa: F401
+        profile.mlx_available = True
+        profile.mlx_memory_gb = _probe_mlx_memory()
+    except ImportError:
+        pass
+
+    # CPU fallback — always available
+    try:
+        import psutil
+        profile.cpu_ram_gb = psutil.virtual_memory().total / 1e9
+        profile.cpu_cores = psutil.cpu_count(logical=False) or 1
+    except ImportError:
+        import multiprocessing
+        profile.cpu_cores = multiprocessing.cpu_count() or 1
+
+    return profile
+
+
+def _probe_mlx_memory() -> float:
+    """Attempt to detect unified memory on Apple Silicon."""
+    try:
+        import subprocess
+        result = subprocess.run(
+            ["sysctl", "-n", "hw.memsize"],
+            capture_output=True, text=True, timeout=5,
+        )
+        if result.returncode == 0:
+            return int(result.stdout.strip()) / 1e9
+    except Exception:
+        pass
+    return 0.0
+
+
+# ---------------------------------------------------------------------------
+# Model recommendation
+# ---------------------------------------------------------------------------
+
+# Tiered recommendation matrix from spec §4.4.
+# Each entry: (provider, model_id, dimensions, quantization)
+_TIER_MATRIX = {
+    "gpu-high": ("cuda", "Qwen/Qwen3-Embedding-8B", 4096, "fp8"),
+    "gpu-balanced": ("cuda", "BAAI/bge-m3", 1024, None),
+    "gpu-fast": ("cuda", "nomic-ai/nomic-embed-text-v1.5", 768, None),
+    "native-mlx-high": ("native_mlx", "Qwen/Qwen3-Embedding-8B", 4096, None),
+    "native-mlx-balanced": ("native_mlx", "BAAI/bge-m3", 1024, None),
+    "mlx": ("mlx", "nomic-ai/nomic-embed-text-v1.5", 768, None),
+    "cpu-quality": ("cpu", "nomic-ai/nomic-embed-text-v1.5", 768, "int8"),
+    "cpu-lightweight": ("cpu", "BAAI/bge-small-en-v1.5", 384, None),
+}
+
+
+def _detect_tier(profile: HardwareProfile) -> str:
+    """Select the best tier for the given hardware profile."""
+    if profile.cuda_available:
+        if profile.cuda_vram_gb >= 24:
+            return "gpu-high"
+        if profile.cuda_vram_gb >= 8:
+            return "gpu-balanced"
+        return "gpu-fast"
+    if profile.mlx_available:
+        # Use high tier for machines with 64GB+ unified memory
+        if profile.mlx_memory_gb >= 64:
+            return "native-mlx-high"
+        if profile.mlx_memory_gb >= 32:
+            return "native-mlx-balanced"
+        return "mlx"
+    if profile.cpu_ram_gb >= 16:
+        return "cpu-quality"
+    return "cpu-lightweight"
+
+
+def recommend_config(profile: HardwareProfile) -> EmbeddingConfig:
+    """Return the best EmbeddingConfig for the detected hardware profile."""
+    tier = _detect_tier(profile)
+    provider, model_id, dims, quant = _TIER_MATRIX[tier]
+    return EmbeddingConfig(
+        provider=provider,
+        model_id=model_id,
+        dimensions=dims,
+        quantization=quant,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Config resolution
+# ---------------------------------------------------------------------------
+
+async def ensure_embed_config() -> Optional[EmbeddingConfig]:
+    """Resolve embedding config from env vars or auto-detect.
+
+    Returns ``None`` if no config is set and the process is non-interactive,
+    allowing Colony to start without embeddings (graceful degradation).
+    """
+    config = EmbeddingConfig.from_env()
+    if config is not None:
+        return config
+
+    # Auto-detect hardware and recommend
+    profile = detect_hardware()
+    recommended = recommend_config(profile)
+
+    # Non-interactive: degrade gracefully
+    if not sys.stdin.isatty():
+        logger.warning(
+            "Vector search disabled — no COLONY_EMBED_MODEL configured. "
+            "Run 'colony vector setup' to enable. "
+            "Recommended: COLONY_EMBED_MODEL=%s COLONY_EMBED_DIMS=%d COLONY_EMBED_PROVIDER=%s",
+            recommended.model_id,
+            recommended.dimensions,
+            recommended.provider,
+        )
+        return None
+
+    # Interactive: print recommendation
+    print("\n--- Colony Vector Store Setup ---")
+    print(f"Detected hardware: {profile}")
+    print(f"Recommended model: {recommended.model_id}")
+    print(f"Provider: {recommended.provider}, Dimensions: {recommended.dimensions}")
+    print("\nSet these environment variables to enable vector search:")
+    print(f"  COLONY_EMBED_PROVIDER={recommended.provider}")
+    print(f"  COLONY_EMBED_MODEL={recommended.model_id}")
+    print(f"  COLONY_EMBED_DIMS={recommended.dimensions}")
+    if recommended.quantization:
+        print(f"  COLONY_EMBED_QUANTIZATION={recommended.quantization}")
+    print()
+    return None
+
+
+async def ensure_model_ready(config: EmbeddingConfig) -> None:
+    """Download model weights if not present in cache_dir.  Run warmup."""
+    try:
+        from huggingface_hub import snapshot_download
+
+        cache_dir = config.cache_dir or None
+        snapshot_download(config.model_id, cache_dir=cache_dir)
+        logger.info("Model %s ready (cache=%s)", config.model_id, cache_dir or "default")
+    except ImportError:
+        logger.warning("huggingface_hub not installed — skipping model download check")
+    except Exception as exc:
+        logger.warning("Model download check failed (may already be cached): %s", exc)
+
+
+# ---------------------------------------------------------------------------
+# CLI entry points
+# ---------------------------------------------------------------------------
+
+def cli_setup(args: list[str] | None = None) -> None:
+    """CLI handler for ``colony vector setup``."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Colony vector store setup")
+    parser.add_argument("--model", help="HuggingFace model ID to use")
+    parser.add_argument(
+        "--tier",
+        choices=list(_TIER_MATRIX.keys()),
+        help="Hardware tier (auto-detected if not specified)",
+    )
+    parsed = parser.parse_args(args if args is not None else [])
+
+    profile = detect_hardware()
+    print(f"\nHardware profile:")
+    print(f"  CUDA: {'Yes' if profile.cuda_available else 'No'}"
+          + (f" ({profile.cuda_device_count} device(s), {profile.cuda_vram_gb:.1f}GB VRAM)" if profile.cuda_available else ""))
+    print(f"  MLX:  {'Yes' if profile.mlx_available else 'No'}"
+          + (f" ({profile.mlx_memory_gb:.1f}GB unified)" if profile.mlx_available else ""))
+    print(f"  CPU:  {profile.cpu_cores} cores, {profile.cpu_ram_gb:.1f}GB RAM")
+
+    if parsed.tier:
+        tier = parsed.tier
+    else:
+        tier = _detect_tier(profile)
+    print(f"\nSelected tier: {tier}")
+
+    if parsed.model:
+        provider, _, dims, quant = _TIER_MATRIX[tier]
+        config = EmbeddingConfig(
+            provider=provider,
+            model_id=parsed.model,
+            dimensions=dims,
+            quantization=quant,
+        )
+    else:
+        config = recommend_config(profile)
+
+    print(f"Model: {config.model_id}")
+    print(f"Provider: {config.provider}, Dimensions: {config.dimensions}")
+    if config.quantization:
+        print(f"Quantization: {config.quantization}")
+
+    print("\nAdd to your .env:")
+    print(f"  COLONY_EMBED_PROVIDER={config.provider}")
+    print(f"  COLONY_EMBED_MODEL={config.model_id}")
+    print(f"  COLONY_EMBED_DIMS={config.dimensions}")
+    if config.quantization:
+        print(f"  COLONY_EMBED_QUANTIZATION={config.quantization}")
+
+    # Download model
+    try:
+        print(f"\nDownloading model weights for {config.model_id}...")
+        from huggingface_hub import snapshot_download
+        cache_dir = config.cache_dir or None
+        snapshot_download(config.model_id, cache_dir=cache_dir)
+        print("Model downloaded successfully.")
+    except ImportError:
+        print("huggingface_hub not installed — install with: pip install huggingface-hub")
+    except Exception as exc:
+        print(f"Download failed (model may already be cached): {exc}")
+
+
+def cli_rebuild(args: list[str] | None = None) -> None:
+    """CLI handler for ``colony vector rebuild``."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Rebuild vector collections")
+    parser.add_argument(
+        "--collection",
+        default="all",
+        help="Collection to rebuild (default: all)",
+    )
+    parsed = parser.parse_args(args if args is not None else [])
+
+    print(f"Rebuilding vector collection(s): {parsed.collection}")
+    print("NOTE: Vector search requires COLONY_EMBED_MODEL to be set.")
+    print("      Run 'colony vector setup' first, then use scripts/vector_backfill.py")
+    print("      for the full backfill process.")
+
+
+def cli_main() -> None:
+    """Entry point for ``colony vector <subcommand>``."""
+    if len(sys.argv) < 3:
+        print("Usage: colony vector {setup|rebuild}")
+        sys.exit(1)
+
+    subcommand = sys.argv[2]
+    remaining = sys.argv[3:]
+
+    if subcommand == "setup":
+        cli_setup(remaining)
+    elif subcommand == "rebuild":
+        cli_rebuild(remaining)
+    else:
+        print(f"Unknown subcommand: {subcommand}")
+        print("Usage: colony vector {setup|rebuild}")
+        sys.exit(1)