dragon-cli 0.2.0__py3-none-any.whl

dragon_cli/model_hub.py

@@ -0,0 +1,252 @@
+"""Model Hub integration for Dragon CLI.
+
+Integrates with Hugging Face, Ollama, and other model hubs to:
+- Query model catalogs
+- Filter by hardware compatibility
+- Download models
+"""
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from .config import MODELS_DIR
+
+
+class ModelHub:
+    """Integration layer for model hubs (Hugging Face, Ollama, etc.)."""
+    
+    def __init__(self):
+        self.catalog_cache_file = MODELS_DIR / "model_catalog_cache.json"
+    
+    def fetch_huggingface_models(
+        self, 
+        search_query: str = "GGUF",
+        limit: int = 50
+    ) -> List[Dict[str, Any]]:
+        """Fetch models from Hugging Face.
+        
+        Args:
+            search_query: Search term (default: "GGUF" for quantized models)
+            limit: Maximum number of results
+            
+        Returns:
+            List of model metadata dictionaries
+        """
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            
+            # Search for GGUF models. ``direction`` is not supported in some
+            # huggingface_hub versions, so fall back to a signature-compatible call.
+            try:
+                models = api.list_models(
+                    search=search_query,
+                    limit=limit,
+                    sort="downloads",
+                    direction=-1,
+                )
+            except TypeError:
+                models = api.list_models(
+                    search=search_query,
+                    limit=limit,
+                    sort="downloads",
+                )
+            
+            results = []
+            for model in models:
+                # Only include if it has GGUF files
+                model_info = {
+                    "name": model.id,
+                    "repo_id": model.id,
+                    "downloads": getattr(model, "downloads", 0),
+                    "likes": getattr(model, "likes", 0),
+                    "source": "huggingface",
+                }
+                results.append(model_info)
+            
+            return results
+            
+        except ImportError:
+            print("huggingface_hub not installed. Install with: pip install huggingface_hub")
+            return []
+        except Exception as e:
+            print(f"Error fetching from Hugging Face: {e}")
+            return []
+    
+    def get_model_files(self, repo_id: str) -> List[str]:
+        """Get list of GGUF files in a repository.
+        
+        Args:
+            repo_id: Hugging Face repository ID (e.g., "Qwen/Qwen2.5-7B-Instruct-GGUF")
+            
+        Returns:
+            List of filenames ending with .gguf
+        """
+        try:
+            from huggingface_hub import HfApi
+            api = HfApi()
+            
+            files = api.list_repo_files(repo_id=repo_id)
+            gguf_files = [f for f in files if f.endswith(".gguf")]
+            
+            return gguf_files
+            
+        except Exception as e:
+            print(f"Error listing files for {repo_id}: {e}")
+            return []
+    
+    def search_models(
+        self,
+        query: str = "",
+        min_ram_gb: float = 0,
+        max_ram_gb: Optional[float] = None,
+        requires_gpu: bool = False,
+        source: str = "huggingface",
+    ) -> List[Dict[str, Any]]:
+        """Search models with filters.
+        
+        Args:
+            query: Search query string
+            min_ram_gb: Minimum RAM requirement in GB
+            max_ram_gb: Maximum RAM requirement in GB
+            requires_gpu: Whether GPU is required
+            source: Model source ("huggingface", "ollama", "all")
+            
+        Returns:
+            Filtered list of models
+        """
+        from .models.catalog import MODEL_CATALOG
+        
+        results = []
+        
+        # Search local catalog first
+        for model in MODEL_CATALOG:
+            ram_req = model.get("ram_min_gb", 0)
+            
+            # Apply filters
+            if min_ram_gb > 0 and ram_req < min_ram_gb:
+                continue
+            if max_ram_gb and ram_req > max_ram_gb:
+                continue
+            
+            # Check if matches query
+            if query and query.lower() not in model["name"].lower():
+                continue
+            
+            model_copy = model.copy()
+            model_copy["source"] = "local_catalog"
+            results.append(model_copy)
+        
+        # Optionally fetch from Hugging Face
+        if source in ("huggingface", "all") and query:
+            hf_models = self.fetch_huggingface_models(search_query=query, limit=20)
+            for model in hf_models:
+                # Check if already in results
+                if not any(m.get("repo_id") == model["repo_id"] for m in results):
+                    results.append(model)
+        
+        return results
+    
+    def get_recommended_quantization(self, vram_gb: float, ram_gb: float) -> str:
+        """Recommend quantization level based on available memory.
+        
+        Args:
+            vram_gb: Available GPU VRAM in GB
+            ram_gb: Available system RAM in GB
+            
+        Returns:
+            Recommended quantization string (e.g., "Q4_K_M", "Q8_0")
+        """
+        # If good VRAM, can use higher quantization on GPU
+        if vram_gb >= 24:
+            return "Q8_0"  # Near FP16 quality
+        elif vram_gb >= 16:
+            return "Q6_K"  # High quality
+        elif vram_gb >= 8:
+            return "Q5_K_M"  # Good balance
+        elif vram_gb >= 6:
+            return "Q4_K_M"  # Standard recommendation
+        elif ram_gb >= 16:
+            return "Q5_K_M"  # Can use more RAM
+        else:
+            return "Q4_K_M"  # Safe default
+    
+    def cache_catalog(self, models: List[Dict[str, Any]]) -> None:
+        """Cache model catalog to disk."""
+        MODELS_DIR.mkdir(exist_ok=True)
+        try:
+            with open(self.catalog_cache_file, "w") as f:
+                json.dump(models, f, indent=2)
+        except Exception as e:
+            print(f"Error caching catalog: {e}")
+    
+    def load_cached_catalog(self) -> List[Dict[str, Any]]:
+        """Load cached model catalog from disk."""
+        if not self.catalog_cache_file.exists():
+            return []
+        try:
+            with open(self.catalog_cache_file, "r") as f:
+                return json.load(f)
+        except Exception:
+            return []
+
+
+def get_model_info_from_hf(repo_id: str, filename: str) -> Optional[Dict[str, Any]]:
+    """Get detailed information about a specific model file.
+    
+    Args:
+        repo_id: Hugging Face repository ID
+        filename: GGUF filename
+        
+    Returns:
+        Model information dictionary or None
+    """
+    try:
+        from huggingface_hub import HfApi, hf_hub_download
+        api = HfApi()
+        
+        # Get file info
+        files = api.list_repo_files(repo_id=repo_id)
+        if filename not in files:
+            return None
+        
+        # Try to get file size
+        try:
+            file_info = api.hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                local_dir=str(MODELS_DIR),
+                force_download=False,
+            )
+            # This would need additional API calls for size
+        except Exception:
+            pass
+        
+        return {
+            "repo_id": repo_id,
+            "filename": filename,
+            "exists": (MODELS_DIR / filename).exists(),
+        }
+        
+    except ImportError:
+        return None
+    except Exception as e:
+        print(f"Error getting model info: {e}")
+        return None
+
+
+if __name__ == "__main__":
+    # Test the hub
+    hub = ModelHub()
+    
+    print("Searching for Qwen GGUF models...")
+    models = hub.search_models(query="Qwen GGUF", max_ram_gb=16)
+    
+    print(f"\nFound {len(models)} models:\n")
+    for i, model in enumerate(models[:10], 1):
+        name = model.get("name", model.get("repo_id", "Unknown"))
+        source = model.get("source", "unknown")
+        ram = model.get("ram_min_gb", "?")
+        print(f"  {i}. {name}")
+        print(f"     Source: {source}, Min RAM: {ram} GB\n")

dragon_cli/models/__init__.py

@@ -0,0 +1,11 @@
+from .catalog import MODEL_CATALOG, add_model_to_catalog, available_ram_gb, list_catalog
+from .loader import download_model, load_llm
+
+__all__ = [
+    "MODEL_CATALOG",
+    "add_model_to_catalog",
+    "available_ram_gb",
+    "list_catalog",
+    "download_model",
+    "load_llm",
+]

dragon_cli/models/catalog.py

@@ -0,0 +1,205 @@
+"""Catalog of verified GGUF models.
+
+Each entry carries the metadata the runtime needs to choose a sensible GPU
+offload without crashing on small VRAM:
+
+    ram_min_gb   — system RAM (GB) required to run the model
+    vram_min_gb  — GPU VRAM (GB) required to offload ALL layers (full GPU mode)
+    n_layers     — transformer block count, used to compute partial offload and
+                   to back off gracefully when full offload runs out of VRAM
+    ctx          — default context window
+"""
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+from ..config import C_DIM, C_GREEN, C_RED, C_RESET, MODELS_DIR
+
+MODEL_CATALOG = [
+    {
+        "name": "Qwen2.5-1.5B-Instruct (Q4_K_M) — tiny, GPU-friendly",
+        "repo_id": "Qwen/Qwen2.5-1.5B-Instruct-GGUF",
+        "filename": "qwen2.5-1.5b-instruct-q4_k_m.gguf",
+        "ram_min_gb": 3, "vram_min_gb": 1.5, "n_layers": 28, "ctx": 8192,
+    },
+    {
+        "name": "DeepSeek-R1-Distill-Qwen-1.5B (Q4_K_M) — light, for testing",
+        "repo_id": "bartowski/DeepSeek-R1-Distill-Qwen-1.5B-GGUF",
+        "filename": "DeepSeek-R1-Distill-Qwen-1.5B-Q4_K_M.gguf",
+        "ram_min_gb": 3, "vram_min_gb": 1.5, "n_layers": 28, "ctx": 4096,
+    },
+    {
+        "name": "Llama-3.2-3B-Instruct (Q4_K_M) — small, GPU-friendly",
+        "repo_id": "bartowski/Llama-3.2-3B-Instruct-GGUF",
+        "filename": "Llama-3.2-3B-Instruct-Q4_K_M.gguf",
+        "ram_min_gb": 4, "vram_min_gb": 2.5, "n_layers": 28, "ctx": 8192,
+    },
+    {
+        "name": "Qwen2.5-3B-Instruct (Q4_K_M) — fast, good multilingual quality",
+        "repo_id": "Qwen/Qwen2.5-3B-Instruct-GGUF",
+        "filename": "qwen2.5-3b-instruct-q4_k_m.gguf",
+        "ram_min_gb": 4, "vram_min_gb": 2.5, "n_layers": 36, "ctx": 8192,
+    },
+    {
+        "name": "Qwen2.5-7B-Instruct (Q4_K_M) — general use, great for tools",
+        "repo_id": "Qwen/Qwen2.5-7B-Instruct-GGUF",
+        "filename": "qwen2.5-7b-instruct-q4_k_m.gguf",
+        "ram_min_gb": 8, "vram_min_gb": 5.0, "n_layers": 28, "ctx": 8192,
+    },
+    {
+        "name": "DeepSeek-R1-Distill-Qwen-7B (Q4_K_M) — reasoning/code",
+        "repo_id": "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF",
+        "filename": "DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf",
+        "ram_min_gb": 8, "vram_min_gb": 5.0, "n_layers": 28, "ctx": 8192,
+    },
+    {
+        "name": "Meta-Llama-3.1-8B-Instruct (Q4_K_M) — general use",
+        "repo_id": "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF",
+        "filename": "Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
+        "ram_min_gb": 9, "vram_min_gb": 5.5, "n_layers": 32, "ctx": 8192,
+    },
+    {
+        "name": "Qwen2.5-14B-Instruct (Q4_K_M) — heavy, best quality",
+        "repo_id": "Qwen/Qwen2.5-14B-Instruct-GGUF",
+        "filename": "qwen2.5-14b-instruct-q4_k_m.gguf",
+        "ram_min_gb": 14, "vram_min_gb": 9.0, "n_layers": 48, "ctx": 8192,
+    },
+]
+
+CUSTOM_CATALOG_FILE = MODELS_DIR / "custom_model_catalog.json"
+
+
+def _model_key(model: Dict[str, Any]) -> Tuple[str, str]:
+    repo_id = str(model.get("repo_id", "")).strip().lower()
+    filename = str(model.get("filename", "")).strip().lower()
+    return repo_id, filename
+
+
+def _merge_model(target: Dict[str, Any], source: Dict[str, Any]) -> Dict[str, Any]:
+    merged = dict(target)
+    for key in ("name", "repo_id", "filename", "ram_min_gb", "vram_min_gb", "n_layers", "ctx"):
+        if key in source:
+            merged[key] = source[key]
+    return merged
+
+
+def load_custom_catalog() -> List[Dict[str, Any]]:
+    """Load custom model entries from disk, returning validated dicts."""
+    if not CUSTOM_CATALOG_FILE.exists():
+        return []
+    try:
+        data = json.loads(CUSTOM_CATALOG_FILE.read_text(encoding="utf-8"))
+    except Exception:
+        return []
+    if not isinstance(data, list):
+        return []
+
+    items: List[Dict[str, Any]] = []
+    for raw in data:
+        if not isinstance(raw, dict):
+            continue
+        repo_id = str(raw.get("repo_id", "")).strip()
+        filename = str(raw.get("filename", "")).strip()
+        if not repo_id or not filename:
+            continue
+        items.append(
+            {
+                "name": str(raw.get("name", filename)).strip() or filename,
+                "repo_id": repo_id,
+                "filename": filename,
+                "ram_min_gb": int(raw.get("ram_min_gb", 4) or 4),
+                "vram_min_gb": float(raw.get("vram_min_gb", 0) or 0),
+                "n_layers": int(raw.get("n_layers", 32) or 32),
+                "ctx": int(raw.get("ctx", 8192) or 8192),
+            }
+        )
+    return items
+
+
+def save_custom_catalog(models: List[Dict[str, Any]]) -> None:
+    """Persist custom model entries under models/custom_model_catalog.json."""
+    MODELS_DIR.mkdir(exist_ok=True)
+    payload: List[Dict[str, Any]] = []
+    seen = set()
+    for model in models:
+        key = _model_key(model)
+        if key in seen:
+            continue
+        seen.add(key)
+        payload.append(
+            {
+                "name": str(model.get("name", "")).strip(),
+                "repo_id": str(model.get("repo_id", "")).strip(),
+                "filename": str(model.get("filename", "")).strip(),
+                "ram_min_gb": int(model.get("ram_min_gb", 4) or 4),
+                "vram_min_gb": float(model.get("vram_min_gb", 0) or 0),
+                "n_layers": int(model.get("n_layers", 32) or 32),
+                "ctx": int(model.get("ctx", 8192) or 8192),
+            }
+        )
+    CUSTOM_CATALOG_FILE.write_text(
+        json.dumps(payload, ensure_ascii=True, indent=2) + "\n", encoding="utf-8"
+    )
+
+
+def add_model_to_catalog(model: Dict[str, Any], persist: bool = False) -> Tuple[int, bool]:
+    """Add a model to runtime catalog and optionally persist it.
+
+    Returns ``(index, created)`` where index is 1-based in MODEL_CATALOG.
+    """
+    incoming = _merge_model({}, model)
+    incoming_key = _model_key(incoming)
+
+    for i, existing in enumerate(MODEL_CATALOG, 1):
+        if _model_key(existing) == incoming_key:
+            MODEL_CATALOG[i - 1] = _merge_model(existing, incoming)
+            if persist:
+                _persist_model_entry(MODEL_CATALOG[i - 1])
+            return i, False
+
+    MODEL_CATALOG.append(incoming)
+    if persist:
+        _persist_model_entry(incoming)
+    return len(MODEL_CATALOG), True
+
+
+def _persist_model_entry(model: Dict[str, Any]) -> None:
+    items = load_custom_catalog()
+    model_key = _model_key(model)
+    for i, existing in enumerate(items):
+        if _model_key(existing) == model_key:
+            items[i] = _merge_model(existing, model)
+            save_custom_catalog(items)
+            return
+    items.append(model)
+    save_custom_catalog(items)
+
+
+def _load_custom_into_runtime_catalog() -> None:
+    for model in load_custom_catalog():
+        add_model_to_catalog(model, persist=False)
+
+
+_load_custom_into_runtime_catalog()
+
+
+def available_ram_gb() -> float:
+    try:
+        import psutil
+        return psutil.virtual_memory().available / (1024 ** 3)
+    except ImportError:
+        return -1
+
+
+def list_catalog() -> None:
+    ram = available_ram_gb()
+    if ram > 0:
+        print(f"{C_DIM}Available RAM: {ram:.1f} GB{C_RESET}")
+    for i, m in enumerate(MODEL_CATALOG, 1):
+        downloaded = (MODELS_DIR / m["filename"]).exists()
+        tag = f"{C_GREEN}[downloaded]{C_RESET}" if downloaded else f"{C_DIM}[cloud]{C_RESET}"
+        warning = f" {C_RED}[!] RAM may be too low right now{C_RESET}" if 0 < ram < m["ram_min_gb"] else ""
+        gpu = m.get("vram_min_gb")
+        gpu_tag = f" {C_DIM}(GPU full ~{gpu:.0f}GB){C_RESET}" if gpu else ""
+        print(f"  {i}. {m['name']}  {tag}{gpu_tag}{warning}")

dragon_cli/models/loader.py

@@ -0,0 +1,187 @@
+"""Model download and loading (huggingface_hub + llama-cpp-python)."""
+
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from .._cuda import preload_cuda_runtime as _preload_cuda_runtime
+from ..config import C_GREEN, C_RED, C_RESET, C_YELLOW, MODELS_DIR
+
+
+
+def download_model(model: dict) -> Path:
+    from huggingface_hub import HfApi, hf_hub_download
+
+    repo_id = model["repo_id"]
+    filename = model["filename"]
+
+    # HF file paths are case-sensitive. If catalog data was saved in lowercase,
+    # resolve the exact-cased filename from the repository before downloading.
+    try:
+        api = HfApi()
+        files = api.list_repo_files(repo_id=repo_id)
+        if filename not in files:
+            by_lower = {f.lower(): f for f in files}
+            resolved = by_lower.get(filename.lower())
+            if resolved:
+                print(f"{C_YELLOW}[download]{C_RESET} Using canonical HF filename: {resolved}")
+                filename = resolved
+                model["filename"] = resolved
+    except Exception:
+        # If listing fails (network/rate limit), keep original filename and let
+        # hf_hub_download surface a clear error.
+        pass
+
+    MODELS_DIR.mkdir(exist_ok=True)
+    destination = MODELS_DIR / filename
+    if destination.exists():
+        print(f"{C_GREEN}[ok]{C_RESET} Model already saved at {destination}")
+        return destination
+    print(f"{C_YELLOW}[download]{C_RESET} Downloading {filename} "
+          f"(a few GB, only happens once)...")
+    path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=str(MODELS_DIR))
+    print(f"{C_GREEN}[ok]{C_RESET} Saved to {path}")
+    return Path(path)
+
+
+def load_llm(
+    path: Path,
+    ctx: int,
+    model_info: Optional[Dict[str, Any]] = None,
+    hardware_profile: Optional[Dict[str, Any]] = None,
+    gpu_layers: Optional[int] = None,
+    force_cpu: bool = False,
+    force_gpu: bool = False,
+):
+    """Load LLM with optimized runtime parameters.
+
+    GPU offload precedence (first match wins):
+      1. ``gpu_layers`` — explicit layer count (0 = CPU, -1/<0 = all layers);
+      2. ``force_cpu``  — disables the GPU entirely;
+      3. ``force_gpu``  — offloads all layers when a GPU is present;
+      4. otherwise      — auto, from the hardware-aware ``RuntimeConfig``.
+
+    Args:
+        path: Path to the GGUF model file
+        ctx: Context window size (fallback if not auto-detected)
+        model_info: Optional model metadata for optimization
+        hardware_profile: Optional hardware profile for optimization
+        gpu_layers: Explicit number of layers to offload to GPU (None = auto)
+        force_cpu: Force CPU-only execution (ignore any GPU)
+        force_gpu: Force GPU offload of all layers
+
+    Returns:
+        Loaded Llama instance
+    """
+    _preload_cuda_runtime()
+    from llama_cpp import Llama
+
+    # Try to use runtime configuration for optimal params
+    try:
+        from ..runtime import RuntimeConfig
+        config = RuntimeConfig(hardware_profile=hardware_profile, model_info=model_info)
+        params = config.get_runtime_params()
+
+        # Override context if explicitly provided
+        if ctx:
+            params["n_ctx"] = ctx
+
+        threads = params["n_threads"]
+        batch_size = params["n_batch"]
+        gpu_available = config.hardware.get("gpu_available", False)
+
+        # Resolve the number of GPU layers according to the precedence above.
+        if force_cpu:
+            n_gpu_layers = 0
+            exec_mode = "CPU"
+        elif gpu_layers is not None:
+            n_gpu_layers = gpu_layers
+            exec_mode = "CPU" if gpu_layers == 0 else "GPU"
+        elif force_gpu:
+            if gpu_available:
+                n_gpu_layers = -1  # llama.cpp convention: -1 == all layers
+                exec_mode = "GPU"
+            else:
+                print(f"{C_YELLOW}[gpu]{C_RESET} --gpu requested but no GPU was "
+                      f"detected; falling back to CPU.")
+                n_gpu_layers = 0
+                exec_mode = "CPU"
+        else:
+            n_gpu_layers = params["n_gpu_layers"]
+            exec_mode = config.decide_execution_mode().upper()
+
+        mode_msg = f"mode={exec_mode} n_gpu_layers={n_gpu_layers}"
+
+    except ImportError:
+        # Fallback to simple configuration
+        threads = max(1, (os.cpu_count() or 4) - 2)
+        n_gpu_layers = 0
+        batch_size = 512
+        exec_mode = "CPU"
+        mode_msg = "mode=CPU (fallback)"
+
+    print(f"{C_YELLOW}[load]{C_RESET} Loading model (threads={threads}, "
+          f"ctx={params.get('n_ctx', ctx)}, {mode_msg})...")
+
+    # Full GPU offload can run out of VRAM on small cards (the whole model +
+    # the display's memory). Rather than crash, try the requested layer count,
+    # then progressively fewer GPU layers, finally falling back to CPU.
+    total_layers = int((model_info or {}).get("n_layers", 32))
+    if n_gpu_layers == 0:
+        candidates = [0]
+    elif n_gpu_layers > 0:
+        candidates = []
+        for n in (n_gpu_layers, n_gpu_layers // 2, 0):
+            if n not in candidates:
+                candidates.append(n)
+    else:  # -1 / "all layers"
+        candidates = []
+        for frac in (1.0, 0.66, 0.5, 0.33, 0.0):
+            n = int(round(total_layers * frac))
+            if n not in candidates:
+                candidates.append(n)
+
+    n_ctx = params.get("n_ctx", ctx)
+    llm = None
+    last_err = None
+    for ng in candidates:
+        try:
+            llm = Llama(
+                model_path=str(path),
+                n_ctx=n_ctx,
+                n_threads=threads,
+                n_gpu_layers=ng,
+                n_batch=batch_size,
+                verbose=False,
+            )
+            used_layers = ng
+            break
+        except Exception as exc:  # noqa: BLE001 - ggml surfaces OOM as ValueError
+            last_err = exc
+            llm = None
+            if ng == 0:
+                break
+            print(f"{C_RED}[gpu]{C_RESET} Offloading {ng} layers failed (likely "
+                  f"out of VRAM). Retrying with fewer layers...")
+            # A failed constructor can strand VRAM until GC runs; reclaim it
+            # before the next attempt so retries don't cascade into failure.
+            import gc
+            gc.collect()
+
+    if llm is None:
+        # Even CPU (0 layers) failed — this is a real error (corrupt file, etc.).
+        raise last_err
+
+    if used_layers > 0 and exec_mode != "CPU":
+        exec_mode = "GPU"
+        if used_layers < total_layers:
+            print(f"{C_YELLOW}[gpu]{C_RESET} Running with {used_layers}/{total_layers} "
+                  f"layers on GPU (partial offload).")
+    else:
+        exec_mode = "CPU"
+        print(f"{C_YELLOW}[gpu]{C_RESET} GPU offload unavailable; running on CPU.")
+
+    # Expose the chosen backend so the UI (chat banner, server) can display it.
+    llm.exec_mode = exec_mode
+    llm.n_gpu_layers = used_layers
+    return llm