parishad 0.1.0__py3-none-any.whl

parishad/models/backends/llama_cpp.py

@@ -0,0 +1,239 @@
+"""
+Llama.cpp backend for GGUF models.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Optional
+
+from .base import BackendConfig, BackendError, BackendResult, BaseBackend
+
+logger = logging.getLogger(__name__)
+
+# Suppress verbose llama.cpp logging (ggml_metal_init messages)
+os.environ.setdefault("GGML_METAL_LOG_LEVEL", "0")
+os.environ.setdefault("LLAMA_CPP_LOG_LEVEL", "0")
+
+# Lazy imports
+_llama_cpp = None
+_suppress_output = None
+
+
+def _get_llama_cpp():
+    """Lazy import of llama-cpp-python."""
+    global _llama_cpp, _suppress_output
+    if _llama_cpp is None:
+        try:
+            import llama_cpp
+            from llama_cpp import suppress_stdout_stderr
+            _llama_cpp = llama_cpp
+            _suppress_output = suppress_stdout_stderr
+        except ImportError:
+            raise ImportError(
+                "llama-cpp-python is required for LlamaCppBackend. "
+                "Install with: pip install llama-cpp-python"
+            )
+    return _llama_cpp
+
+
+def resolve_model_path(model_id: str) -> Optional[Path]:
+    """Resolve a model ID to a file path."""
+    # Direct path
+    direct = Path(model_id)
+    if direct.exists():
+        return direct
+    
+    # Try model manager
+    try:
+        from ..downloader import ModelManager, get_default_model_dir
+        manager = ModelManager()
+        
+        # Check registry
+        path = manager.get_model_path(model_id)
+        if path and path.exists():
+            return path
+        
+        # Check ollama symlinks
+        ollama_dir = get_default_model_dir() / "ollama"
+        if ollama_dir.exists():
+            safe_name = model_id.replace(":", "_").replace("/", "_")
+            for suffix in [".gguf", ""]:
+                candidate = ollama_dir / f"{safe_name}{suffix}"
+                if candidate.exists():
+                    return candidate
+        
+        # Search for matching GGUF files
+        for model in manager.list_models():
+            if model.format.value == "gguf":
+                if model_id in model.name or model_id in str(model.path):
+                    return model.path
+    except Exception as e:
+        logger.debug(f"Model manager lookup failed: {e}")
+    
+    # Common locations
+    search_paths = [
+        Path.cwd() / "models",
+        Path.home() / ".cache" / "parishad" / "models",
+        Path.home() / ".local" / "share" / "parishad" / "models",
+    ]
+    
+    for search_dir in search_paths:
+        if search_dir.exists():
+            candidate = search_dir / model_id
+            if candidate.exists():
+                return candidate
+            for gguf_file in search_dir.rglob("*.gguf"):
+                if model_id in gguf_file.name:
+                    return gguf_file
+    
+    return None
+
+
+class LlamaCppBackend(BaseBackend):
+    """Backend for GGUF models using llama-cpp-python."""
+    
+    _name = "llama_cpp"
+    
+    def __init__(self):
+        """Initialize LlamaCppBackend."""
+        super().__init__()
+        self._llm = None
+    
+    def load(self, config: BackendConfig) -> None:
+        """Load a GGUF model."""
+        llama_cpp = _get_llama_cpp()
+        
+        model_path = resolve_model_path(config.model_id)
+        
+        if model_path is None:
+            raise BackendError(
+                f"Model not found: {config.model_id}. "
+                "Download with: parishad download <model_name>",
+                backend_name=self._name,
+                model_id=config.model_id,
+            )
+        
+        extra = config.extra or {}
+        n_gpu_layers = extra.get("n_gpu_layers", -1)
+        n_ctx = extra.get("n_ctx", config.context_length)
+        n_batch = extra.get("n_batch", 512)
+        verbose = extra.get("verbose", False)
+        chat_format = extra.get("chat_format", None)
+        
+        try:
+            suppress_ctx = _suppress_output(disable=False) if _suppress_output else None
+            if suppress_ctx:
+                with suppress_ctx:
+                    self._llm = llama_cpp.Llama(
+                        model_path=str(model_path),
+                        n_gpu_layers=n_gpu_layers,
+                        n_ctx=n_ctx,
+                        n_batch=n_batch,
+                        verbose=verbose,
+                        chat_format=chat_format,
+                    )
+            else:
+                self._llm = llama_cpp.Llama(
+                    model_path=str(model_path),
+                    n_gpu_layers=n_gpu_layers,
+                    n_ctx=n_ctx,
+                    n_batch=n_batch,
+                    verbose=verbose,
+                    chat_format=chat_format,
+                )
+            self._config = config
+            self._model_id = config.model_id
+            self._loaded = True
+            
+        except Exception as e:
+            raise BackendError(
+                f"Failed to load model: {e}",
+                backend_name=self._name,
+                model_id=config.model_id,
+                original_error=e,
+            )
+    
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        stop: list[str] | None = None,
+    ) -> BackendResult:
+        """Generate text using llama.cpp."""
+        if not self._loaded or self._llm is None:
+            raise BackendError(
+                "Model not loaded",
+                backend_name=self._name,
+                model_id=self._model_id,
+            )
+        
+        start_time = time.perf_counter()
+        
+        try:
+            logger.debug(f"Calling llama_cpp with prompt len={len(prompt)}, max_tokens={max_tokens}, temp={temperature}")
+            
+            result = self._llm(
+                prompt,
+                max_tokens=max_tokens,
+                temperature=max(temperature, 0.01),
+                top_p=top_p,
+                stop=stop or [],
+                echo=False,
+            )
+            
+            logger.debug(f"llama_cpp raw result keys: {result.keys()}")
+            if "choices" in result and result["choices"]:
+                logger.debug(f"First choice keys: {result['choices'][0].keys()}")
+                logger.debug(f"Finish reason: {result['choices'][0].get('finish_reason')}")
+            else:
+                logger.error(f"No choices in result: {result}")
+            
+            text = result["choices"][0]["text"]
+            finish_reason = result["choices"][0].get("finish_reason", "stop")
+            
+            usage = result.get("usage", {})
+            tokens_in = usage.get("prompt_tokens", self._estimate_tokens(prompt))
+            tokens_out = usage.get("completion_tokens", self._estimate_tokens(text))
+            
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            
+            return BackendResult(
+                text=text,
+                tokens_in=tokens_in,
+                tokens_out=tokens_out,
+                model_id=self._model_id,
+                finish_reason=finish_reason,
+                latency_ms=latency_ms,
+                extra={"total_tokens": usage.get("total_tokens", tokens_in + tokens_out)},
+            )
+            
+        except Exception as e:
+            raise BackendError(
+                f"Generation failed: {e}",
+                backend_name=self._name,
+                model_id=self._model_id,
+                original_error=e,
+            )
+    
+    def unload(self) -> None:
+        """Unload the model to free memory."""
+        if self._llm is not None:
+            del self._llm
+            self._llm = None
+        
+        super().unload()
+        gc.collect()
+    
+    @property
+    def context_length(self) -> int:
+        """Get the model's context length."""
+        if self._llm is not None:
+            return self._llm.n_ctx()
+        return self._config.context_length if self._config else 4096

parishad/models/backends/mlx_lm.py

@@ -0,0 +1,141 @@
+"""
+MLX backend for Apple Silicon.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from pathlib import Path
+
+from .base import BackendConfig, BackendError, BackendResult, BaseBackend
+
+logger = logging.getLogger(__name__)
+
+# Lazy imports
+_mlx_lm = None
+
+
+def _get_mlx_lm():
+    """Lazy import of mlx-lm."""
+    global _mlx_lm
+    if _mlx_lm is None:
+        try:
+            import mlx_lm
+            _mlx_lm = mlx_lm
+        except ImportError:
+            raise ImportError(
+                "mlx-lm is required for MlxBackend. "
+                "Install with: pip install mlx-lm"
+            )
+    return _mlx_lm
+
+
+class MlxBackend(BaseBackend):
+    """Backend for MLX models on Apple Silicon."""
+    
+    _name = "mlx"
+    
+    def __init__(self):
+        """Initialize MlxBackend."""
+        super().__init__()
+        self._model = None
+        self._tokenizer = None
+    
+    def load(self, config: BackendConfig) -> None:
+        """Load an MLX model."""
+        mlx_lm = _get_mlx_lm()
+        
+        model_id_or_path = config.model_id
+        
+        p = Path(model_id_or_path)
+        if p.exists():
+            if p.is_file():
+                model_id_or_path = str(p.parent)
+            else:
+                model_id_or_path = str(p)
+                
+        try:
+            self._model, self._tokenizer = mlx_lm.load(
+                model_id_or_path,
+                tokenizer_config={"trust_remote_code": True}
+            )
+            
+            self._config = config
+            self._model_id = config.model_id
+            self._loaded = True
+            
+        except Exception as e:
+            raise BackendError(
+                f"Failed to load MLX model: {e}",
+                backend_name=self._name,
+                model_id=config.model_id,
+                original_error=e,
+            )
+    
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        stop: list[str] | None = None,
+    ) -> BackendResult:
+        """Generate text using MLX."""
+        if not self._loaded or self._model is None:
+            raise BackendError(
+                "Model not loaded",
+                backend_name=self._name,
+                model_id=self._model_id,
+            )
+            
+        mlx_lm = _get_mlx_lm()
+        start_time = time.perf_counter()
+        
+        try:
+            tokens_in = len(self._tokenizer.encode(prompt))
+            
+            text = mlx_lm.generate(
+                model=self._model,
+                tokenizer=self._tokenizer,
+                prompt=prompt,
+                max_tokens=max_tokens,
+            )
+            
+            finish_reason = "length"
+            if stop:
+                for s in stop:
+                    if s in text:
+                        text = text.split(s)[0]
+                        finish_reason = "stop"
+                        break
+            
+            tokens_out = len(self._tokenizer.encode(text))
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            
+            return BackendResult(
+                text=text,
+                tokens_in=tokens_in,
+                tokens_out=tokens_out,
+                model_id=self._model_id,
+                finish_reason=finish_reason,
+                latency_ms=latency_ms,
+            )
+            
+        except Exception as e:
+            raise BackendError(
+                f"MLX generation failed: {e}",
+                backend_name=self._name,
+                model_id=self._model_id,
+                original_error=e,
+            )
+
+    def unload(self) -> None:
+        """Unload model."""
+        if self._model is not None:
+            del self._model
+            self._model = None
+        if self._tokenizer is not None:
+            del self._tokenizer
+            self._tokenizer = None
+        super().unload()

parishad/models/backends/ollama.py

@@ -0,0 +1,253 @@
+"""
+Ollama backend using native API.
+
+Supports both OpenAI-compatible mode (via inherited OllamaBackend in openai_api.py)
+and native Ollama API mode (this file).
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any
+
+from .base import BackendConfig, BackendError, BackendResult, BaseBackend
+
+logger = logging.getLogger(__name__)
+
+# Lazy import
+_requests = None
+
+
+def _get_requests():
+    """Lazy import of requests."""
+    global _requests
+    if _requests is None:
+        try:
+            import requests
+            _requests = requests
+        except ImportError:
+            raise ImportError(
+                "requests package is required for OllamaNativeBackend. "
+                "Install with: pip install requests"
+            )
+    return _requests
+
+
+class OllamaNativeBackend(BaseBackend):
+    """
+    Backend for Ollama using native API.
+    
+    Uses Ollama's /api/generate endpoint directly instead of OpenAI compatibility layer.
+    This provides access to Ollama-specific features like:
+    - Raw mode for exact prompts
+    - System prompts
+    - Streaming
+    - Context management
+    """
+    
+    _name = "ollama_native"
+    
+    def __init__(self):
+        """Initialize OllamaNativeBackend."""
+        super().__init__()
+        self._base_url = "http://localhost:11434"
+        self._session = None
+    
+    def load(self, config: BackendConfig) -> None:
+        """Initialize Ollama connection."""
+        requests = _get_requests()
+        
+        extra = config.extra or {}
+        self._base_url = extra.get("base_url", "http://localhost:11434")
+        
+        # Strip ollama: prefix if present
+        model_id = config.model_id
+        if model_id.startswith("ollama:"):
+            model_id = model_id.replace("ollama:", "", 1)
+        
+        try:
+            # Test connection
+            self._session = requests.Session()
+            response = self._session.get(f"{self._base_url}/api/tags", timeout=5)
+            
+            if response.status_code != 200:
+                raise BackendError(
+                    f"Ollama server not responding at {self._base_url}",
+                    backend_name=self._name
+                )
+            
+            # Check if model is available
+            tags = response.json()
+            available_models = [m["name"] for m in tags.get("models", [])]
+            
+            # Check for exact match or partial match
+            model_found = False
+            for m in available_models:
+                if model_id in m or m in model_id:
+                    model_found = True
+                    break
+            
+            if not model_found and available_models:
+                logger.warning(
+                    f"Model '{model_id}' not found in Ollama. "
+                    f"Available: {available_models[:5]}..."
+                )
+            
+            self._config = config
+            self._model_id = model_id
+            self._loaded = True
+            
+            logger.info(f"✅ Connected to Ollama at {self._base_url}")
+            
+        except Exception as e:
+            raise BackendError(
+                f"Failed to connect to Ollama: {e}",
+                backend_name=self._name,
+                model_id=config.model_id,
+                original_error=e,
+            )
+    
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        stop: list[str] | None = None,
+    ) -> BackendResult:
+        """Generate text using Ollama native API."""
+        if not self._loaded or self._session is None:
+            raise BackendError(
+                "Backend not loaded",
+                backend_name=self._name,
+                model_id=self._model_id,
+            )
+        
+        start_time = time.perf_counter()
+        
+        try:
+            # Parse system/user if present in prompt
+            system_prompt = ""
+            user_prompt = prompt
+            
+            if "<|start_header_id|>system<|end_header_id|>" in prompt:
+                # Llama-3 format - extract parts
+                parts = prompt.split("<|start_header_id|>")
+                for part in parts:
+                    if part.startswith("system"):
+                        system_prompt = part.split("<|end_header_id|>")[1].split("<|eot_id|>")[0].strip()
+                    elif part.startswith("user"):
+                        user_prompt = part.split("<|end_header_id|>")[1].split("<|eot_id|>")[0].strip()
+            
+            # Build request payload
+            payload: dict[str, Any] = {
+                "model": self._model_id,
+                "prompt": user_prompt,
+                "stream": False,
+                "options": {
+                    "num_predict": max_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    # Increase context window for large documents (default 32k)
+                    "num_ctx": self._config.extra.get("num_ctx", 32768) if self._config and self._config.extra else 32768,
+                }
+            }
+            
+            if system_prompt:
+                payload["system"] = system_prompt
+            
+            if stop:
+                payload["options"]["stop"] = stop
+            
+            # Make request
+            response = self._session.post(
+                f"{self._base_url}/api/generate",
+                json=payload,
+                # Increase default timeout to 300s (5m) for slower models/hardware
+                timeout=self._config.timeout if self._config and self._config.timeout else 300
+            )
+            
+            if response.status_code != 200:
+                raise BackendError(
+                    f"Ollama API error: {response.status_code} - {response.text}",
+                    backend_name=self._name,
+                    model_id=self._model_id
+                )
+            
+            result = response.json()
+            
+            text = result.get("response", "")
+            
+            # Get token counts from Ollama response
+            tokens_in = result.get("prompt_eval_count", self._estimate_tokens(prompt))
+            tokens_out = result.get("eval_count", self._estimate_tokens(text))
+            
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            
+            # Determine finish reason
+            done_reason = result.get("done_reason", "stop")
+            if done_reason == "length":
+                finish_reason = "length"
+            else:
+                finish_reason = "stop"
+            
+            return BackendResult(
+                text=text,
+                tokens_in=tokens_in,
+                tokens_out=tokens_out,
+                model_id=self._model_id,
+                finish_reason=finish_reason,
+                latency_ms=latency_ms,
+                extra={
+                    "total_duration": result.get("total_duration"),
+                    "load_duration": result.get("load_duration"),
+                    "eval_duration": result.get("eval_duration"),
+                }
+            )
+            
+        except Exception as e:
+            if isinstance(e, BackendError):
+                raise
+            raise BackendError(
+                f"Ollama generation failed: {e}",
+                backend_name=self._name,
+                model_id=self._model_id,
+                original_error=e,
+            )
+    
+    def unload(self) -> None:
+        """Close the session."""
+        if self._session:
+            self._session.close()
+            self._session = None
+        super().unload()
+    
+    def list_models(self) -> list[str]:
+        """List available Ollama models."""
+        if not self._session:
+            return []
+        
+        try:
+            response = self._session.get(f"{self._base_url}/api/tags", timeout=5)
+            if response.status_code == 200:
+                tags = response.json()
+                return [m["name"] for m in tags.get("models", [])]
+        except Exception:
+            pass
+        return []
+    
+    def pull_model(self, model_name: str) -> bool:
+        """Pull a model from Ollama registry."""
+        if not self._session:
+            return False
+        
+        try:
+            response = self._session.post(
+                f"{self._base_url}/api/pull",
+                json={"name": model_name, "stream": False},
+                timeout=600  # Models can take a while to download
+            )
+            return response.status_code == 200
+        except Exception:
+            return False