cat-stack 0.1.0__py3-none-any.whl

cat_stack/_providers.py

@@ -0,0 +1,1048 @@
+"""
+Unified LLM provider infrastructure for CatLLM.
+
+This module provides a unified HTTP-based approach for calling multiple LLM providers
+(OpenAI, Anthropic, Google, Mistral, xAI, Perplexity, HuggingFace, and Ollama)
+without requiring provider-specific SDKs.
+"""
+
+import json
+import time
+import requests
+
+__all__ = [
+    # Main client
+    "UnifiedLLMClient",
+    "PROVIDER_CONFIG",
+    # Provider detection
+    "detect_provider",
+    "_detect_model_source",
+    "_detect_huggingface_endpoint",
+    # Ollama utilities
+    "set_ollama_endpoint",
+    "check_ollama_running",
+    "list_ollama_models",
+    "check_ollama_model",
+    "check_system_resources",
+    "get_ollama_model_size_estimate",
+    "pull_ollama_model",
+    "OLLAMA_MODEL_SIZES",
+    # Claude Code utilities
+    "check_claude_cli_available",
+]
+
+
+# =============================================================================
+# HuggingFace Endpoint Auto-Detection
+# =============================================================================
+
+def _detect_huggingface_endpoint(api_key: str, model: str) -> str:
+    """
+    Test which HuggingFace endpoint works for this model.
+    Tries generic router first, then Together.
+
+    Args:
+        api_key: HuggingFace API key
+        model: Model name to test
+
+    Returns:
+        Base URL for the working endpoint (without /chat/completions)
+    """
+    endpoints = [
+        "https://router.huggingface.co/v1/chat/completions",
+        "https://router.huggingface.co/together/v1/chat/completions",
+    ]
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}"
+    }
+
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": "hi"}],
+        "max_tokens": 5
+    }
+
+    for endpoint in endpoints:
+        try:
+            response = requests.post(endpoint, headers=headers, json=payload, timeout=30)
+            if response.status_code == 200:
+                # Return the base URL (without /chat/completions)
+                return endpoint.replace("/chat/completions", "")
+        except Exception:
+            continue
+
+    # Default to generic (will fail with informative error)
+    return "https://router.huggingface.co/v1"
+
+
+# =============================================================================
+# Provider Configuration
+# =============================================================================
+
+PROVIDER_CONFIG = {
+    "openai": {
+        "endpoint": "https://api.openai.com/v1/chat/completions",
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+    },
+    "anthropic": {
+        "endpoint": "https://api.anthropic.com/v1/messages",
+        "auth_header": "x-api-key",
+        "auth_prefix": "",
+    },
+    "google": {
+        "endpoint": "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
+        "auth_header": "x-goog-api-key",
+        "auth_prefix": "",
+    },
+    "mistral": {
+        "endpoint": "https://api.mistral.ai/v1/chat/completions",
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+    },
+    "perplexity": {
+        "endpoint": "https://api.perplexity.ai/chat/completions",
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+    },
+    "xai": {
+        "endpoint": "https://api.x.ai/v1/chat/completions",
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+    },
+    "huggingface": {
+        "endpoint": "https://router.huggingface.co/v1/chat/completions",
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+    },
+    "huggingface-together": {
+        "endpoint": "https://router.huggingface.co/together/v1/chat/completions",
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+    },
+    "ollama": {
+        "endpoint": "http://localhost:11434/v1/chat/completions",
+        "auth_header": None,  # No auth required for local Ollama
+        "auth_prefix": "",
+    },
+    "claude-code": {
+        "endpoint": None,  # Uses CLI subprocess, not HTTP
+        "auth_header": None,
+        "auth_prefix": "",
+    },
+}
+
+
+# =============================================================================
+# Unified API Client
+# =============================================================================
+
+class UnifiedLLMClient:
+    """A unified client for calling various LLM providers via HTTP."""
+
+    def __init__(self, provider: str, api_key: str, model: str):
+        self.provider = provider.lower()
+        self.api_key = api_key
+        self.model = model
+
+        # Auto-detect HuggingFace endpoint
+        if self.provider == "huggingface":
+            detected_url = _detect_huggingface_endpoint(api_key, model)
+            if "together" in detected_url:
+                self.provider = "huggingface-together"
+
+        if self.provider not in PROVIDER_CONFIG:
+            raise ValueError(f"Unsupported provider: {provider}. "
+                           f"Supported: {list(PROVIDER_CONFIG.keys())}")
+
+        self.config = PROVIDER_CONFIG[self.provider]
+
+    def _get_endpoint(self) -> str:
+        """Get the API endpoint, substituting model if needed."""
+        endpoint = self.config["endpoint"]
+        if "{model}" in endpoint:
+            endpoint = endpoint.format(model=self.model)
+        return endpoint
+
+    def _get_headers(self) -> dict:
+        """Build request headers for the provider."""
+        headers = {"Content-Type": "application/json"}
+        auth_header = self.config["auth_header"]
+        auth_prefix = self.config["auth_prefix"]
+
+        # Some providers (like Ollama) don't require auth
+        if auth_header is not None:
+            headers[auth_header] = f"{auth_prefix}{self.api_key}"
+
+        # Anthropic requires additional headers
+        if self.provider == "anthropic":
+            headers["anthropic-version"] = "2023-06-01"
+
+        return headers
+
+    def _build_payload(
+        self,
+        messages: list,
+        json_schema: dict = None,
+        creativity: float = None,
+        max_tokens: int = 4096,
+        thinking_budget: int = None,
+        force_json: bool = True,
+    ) -> dict:
+        """Build the request payload for the specific provider."""
+
+        if self.provider == "anthropic":
+            return self._build_anthropic_payload(messages, json_schema, creativity, max_tokens, thinking_budget)
+        elif self.provider == "google":
+            return self._build_google_payload(messages, json_schema, creativity, thinking_budget, force_json)
+        elif self.provider == "openai":
+            return self._build_openai_payload(messages, json_schema, creativity, force_json, thinking_budget)
+        elif self.provider in ("huggingface", "huggingface-together"):
+            # HuggingFace needs thinking_budget to disable thinking on models that reason by default
+            return self._build_openai_payload(messages, json_schema, creativity, force_json, thinking_budget)
+        else:
+            # Other OpenAI-compatible providers (xai, mistral, etc.)
+            return self._build_openai_payload(messages, json_schema, creativity, force_json)
+
+    def _build_openai_payload(
+        self,
+        messages: list,
+        json_schema: dict = None,
+        creativity: float = None,
+        force_json: bool = True,
+        thinking_budget: int = None,
+    ) -> dict:
+        """Build payload for OpenAI-compatible APIs.
+
+        Args:
+            force_json: If False and no json_schema, don't set response_format (for text responses)
+            thinking_budget: For OpenAI models, maps to reasoning_effort:
+                             0 or None → "minimal", >0 → "high"
+        """
+        payload = {
+            "model": self.model,
+            "messages": messages,
+        }
+
+        # Structured output
+        # Ollama and HuggingFace only support json_object mode, not strict json_schema
+        if json_schema and self.provider not in ["ollama", "huggingface", "huggingface-together"]:
+            payload["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "classification_result",
+                    "strict": True,
+                    "schema": json_schema,
+                }
+            }
+        elif json_schema:
+            # Ollama/HuggingFace - use json_object mode
+            payload["response_format"] = {"type": "json_object"}
+        elif force_json:
+            # No schema but force JSON output
+            payload["response_format"] = {"type": "json_object"}
+        # else: no response_format - allow text responses
+
+        # OpenAI reasoning models (o-series, GPT-5) only accept temperature=1.
+        # Use reasoning_effort to control reasoning depth instead.
+        _is_reasoning_model = self.provider == "openai" and any(
+            self.model.startswith(p) for p in ("o1", "o3", "o4", "gpt-5")
+        )
+        if _is_reasoning_model:
+            # Never set temperature for reasoning models (only default=1 is valid)
+            if thinking_budget is not None:
+                if thinking_budget > 0:
+                    payload["reasoning_effort"] = "high"
+                else:
+                    payload["reasoning_effort"] = "minimal"
+        elif creativity is not None:
+            payload["temperature"] = creativity
+
+        # HuggingFace: disable thinking for models that reason by default (e.g., Qwen3)
+        # when thinking_budget is explicitly set to 0
+        if self.provider in ("huggingface", "huggingface-together") and thinking_budget is not None and thinking_budget == 0:
+            payload["chat_template_kwargs"] = {"enable_thinking": False}
+
+        return payload
+
+    def _build_anthropic_payload(
+        self,
+        messages: list,
+        json_schema: dict = None,
+        creativity: float = None,
+        max_tokens: int = 4096,
+        thinking_budget: int = None,
+    ) -> dict:
+        """Build payload for Anthropic API.
+
+        Args:
+            thinking_budget: Controls extended thinking for Anthropic models.
+                             0 or None → thinking disabled.
+                             >0 → thinking enabled with budget_tokens = max(thinking_budget, 1024).
+        """
+        # Extract system message if present
+        system_content = None
+        user_messages = []
+        for msg in messages:
+            if msg["role"] == "system":
+                system_content = msg["content"]
+            else:
+                user_messages.append(msg)
+
+        payload = {
+            "model": self.model,
+            "max_tokens": max_tokens,
+            "messages": user_messages,
+        }
+
+        if system_content:
+            payload["system"] = system_content
+
+        # Extended thinking for Anthropic (minimum 1024 tokens)
+        # When thinking is enabled, temperature must be 1 (Anthropic requirement),
+        # so we skip setting temperature from creativity in that case
+        if thinking_budget and thinking_budget > 0:
+            budget = max(thinking_budget, 1024)
+            payload["thinking"] = {
+                "type": "enabled",
+                "budget_tokens": budget,
+            }
+            payload["temperature"] = 1
+            # When thinking is enabled, max_tokens must be larger than budget_tokens
+            if payload["max_tokens"] <= budget:
+                payload["max_tokens"] = budget + 4096
+        elif creativity is not None:
+            payload["temperature"] = creativity
+
+        # Use tool calling for structured output (most reliable for Anthropic)
+        # When thinking is enabled, forced tool_choice is not allowed — use "auto"
+        if json_schema:
+            payload["tools"] = [{
+                "name": "return_categories",
+                "description": "Return categorization results",
+                "input_schema": json_schema,
+            }]
+            if thinking_budget and thinking_budget > 0:
+                payload["tool_choice"] = {"type": "auto"}
+            else:
+                payload["tool_choice"] = {"type": "tool", "name": "return_categories"}
+
+        return payload
+
+    def _build_google_payload(
+        self,
+        messages: list,
+        json_schema: dict = None,
+        creativity: float = None,
+        thinking_budget: int = None,
+        force_json: bool = True,
+    ) -> dict:
+        """Build payload for Google Gemini API."""
+        # Convert messages to Google format
+        # Combine system + user messages into a single prompt
+        combined_text = ""
+        for msg in messages:
+            if msg["role"] == "system":
+                combined_text += msg["content"] + "\n\n"
+            elif msg["role"] == "user":
+                combined_text += msg["content"]
+            elif msg["role"] == "assistant":
+                combined_text += "\n\nAssistant: " + msg["content"] + "\n\n"
+
+        payload = {
+            "contents": [{"parts": [{"text": combined_text}]}],
+            "generationConfig": {}
+        }
+
+        if json_schema:
+            payload["generationConfig"]["responseMimeType"] = "application/json"
+            payload["generationConfig"]["responseSchema"] = json_schema
+        elif force_json:
+            payload["generationConfig"]["responseMimeType"] = "application/json"
+        # else: no mime type - allow text responses
+
+        if creativity is not None:
+            payload["generationConfig"]["temperature"] = creativity
+
+        # Add thinking budget for extended thinking (Google-specific)
+        # Must be inside generationConfig, not at top level
+        # Google requires a reasonable minimum budget (enforce 128 tokens minimum)
+        if thinking_budget and thinking_budget > 0:
+            budget = max(thinking_budget, 128)
+            payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": budget}
+
+        return payload
+
+    def _parse_response(self, response_json: dict) -> str:
+        """Parse the response based on provider format."""
+        if self.provider == "anthropic":
+            return self._parse_anthropic_response(response_json)
+        elif self.provider == "google":
+            return self._parse_google_response(response_json)
+        else:
+            # OpenAI-compatible
+            return self._parse_openai_response(response_json)
+
+    def _parse_openai_response(self, response_json: dict) -> str:
+        """Parse OpenAI-compatible response."""
+        return response_json["choices"][0]["message"]["content"]
+
+    def _parse_anthropic_response(self, response_json: dict) -> str:
+        """Parse Anthropic response (handles both text and tool use)."""
+        content = response_json.get("content", [])
+        for block in content:
+            if block.get("type") == "tool_use":
+                # Return the tool input as JSON string
+                return json.dumps(block.get("input", {}))
+            elif block.get("type") == "text":
+                return block.get("text", "")
+        return ""
+
+    def _parse_google_response(self, response_json: dict) -> str:
+        """Parse Google Gemini response."""
+        candidates = response_json.get("candidates", [])
+        if candidates:
+            parts = candidates[0].get("content", {}).get("parts", [])
+            if parts:
+                return parts[0].get("text", "")
+        return ""
+
+    def _call_claude_cli(
+        self,
+        messages: list,
+        max_retries: int = 3,
+        initial_delay: float = 2.0,
+    ) -> tuple[str, str | None]:
+        """
+        Call the Claude CLI (claude -p) as a subprocess.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            max_retries: Maximum retry attempts
+            initial_delay: Initial delay for exponential backoff
+
+        Returns:
+            tuple: (response_text, error_message)
+        """
+        import subprocess
+
+        # Extract system and user messages
+        system_parts = []
+        user_parts = []
+        for msg in messages:
+            if msg["role"] == "system":
+                system_parts.append(msg["content"])
+            elif msg["role"] in ("user", "assistant"):
+                user_parts.append(msg["content"])
+
+        system_prompt = "\n\n".join(system_parts) if system_parts else None
+        user_prompt = "\n\n".join(user_parts)
+
+        # Build command
+        cmd = ["claude", "-p", "--output-format", "text", "--model", self.model]
+        if system_prompt:
+            cmd.extend(["--system-prompt", system_prompt])
+        cmd.append(user_prompt)
+
+        for attempt in range(max_retries):
+            try:
+                result = subprocess.run(
+                    cmd,
+                    capture_output=True,
+                    text=True,
+                    timeout=120,
+                )
+                if result.returncode == 0:
+                    return result.stdout.strip(), None
+                else:
+                    error_msg = result.stderr.strip() or f"CLI exited with code {result.returncode}"
+                    if attempt < max_retries - 1:
+                        wait_time = initial_delay * (2 ** attempt)
+                        print(f"Claude CLI error: {error_msg}. Retrying in {wait_time}s...")
+                        time.sleep(wait_time)
+                    else:
+                        return None, f"Claude CLI failed after {max_retries} attempts: {error_msg}"
+            except subprocess.TimeoutExpired:
+                if attempt < max_retries - 1:
+                    wait_time = initial_delay * (2 ** attempt)
+                    print(f"Claude CLI timeout. Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    return None, "Claude CLI timeout after retries"
+            except FileNotFoundError:
+                return None, (
+                    "Claude CLI not found. Install it: "
+                    "https://docs.anthropic.com/en/docs/claude-code"
+                )
+
+        return None, "Max retries exceeded"
+
+    def complete(
+        self,
+        messages: list,
+        json_schema: dict = None,
+        creativity: float = None,
+        thinking_budget: int = None,
+        force_json: bool = True,
+        max_retries: int = 5,
+        initial_delay: float = 2.0,
+    ) -> tuple[str, str | None]:
+        """
+        Make a completion request to the LLM provider.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content'
+            json_schema: Optional JSON schema for structured output
+            creativity: Temperature setting (None for default)
+            thinking_budget: Controls reasoning behavior per provider:
+                - Google: Token budget for extended thinking (0 to disable, >0 to enable)
+                - OpenAI: Maps to reasoning_effort (0 → "minimal", >0 → "high")
+                - Anthropic: Enables extended thinking (0 to disable, >0 to enable with min 1024)
+            force_json: If True and no json_schema, still request JSON output.
+                       Set to False for text-only responses (e.g., CoVe intermediate steps)
+            max_retries: Maximum retry attempts
+            initial_delay: Initial delay for exponential backoff
+
+        Returns:
+            tuple: (response_text, error_message)
+                   error_message is None on success
+        """
+        if self.provider == "claude-code":
+            return self._call_claude_cli(messages, max_retries=max_retries, initial_delay=initial_delay)
+
+        endpoint = self._get_endpoint()
+        headers = self._get_headers()
+        payload = self._build_payload(messages, json_schema, creativity, thinking_budget=thinking_budget, force_json=force_json)
+
+        for attempt in range(max_retries):
+            try:
+                response = requests.post(
+                    endpoint,
+                    headers=headers,
+                    json=payload,
+                    timeout=120,
+                )
+
+                # Check for HTTP errors
+                if response.status_code == 404:
+                    return None, f"Model '{self.model}' not found for {self.provider}"
+                elif response.status_code in [401, 403]:
+                    return None, f"Authentication failed for {self.provider}"
+                elif response.status_code == 429:
+                    # Rate limited - retry with backoff
+                    if attempt < max_retries - 1:
+                        wait_time = initial_delay * (2 ** attempt) * 5  # Longer wait for rate limits
+                        print(f"Rate limited. Waiting {wait_time}s...")
+                        time.sleep(wait_time)
+                        continue
+                    else:
+                        return None, "Rate limit exceeded after retries"
+                elif response.status_code >= 500:
+                    # Server error - retry
+                    if attempt < max_retries - 1:
+                        wait_time = initial_delay * (2 ** attempt)
+                        print(f"Server error {response.status_code}. Retrying in {wait_time}s...")
+                        time.sleep(wait_time)
+                        continue
+                    else:
+                        return None, f"Server error {response.status_code} after retries"
+
+                response.raise_for_status()
+                response_json = response.json()
+                result = self._parse_response(response_json)
+                return result, None
+
+            except requests.exceptions.Timeout:
+                if attempt < max_retries - 1:
+                    wait_time = initial_delay * (2 ** attempt)
+                    print(f"Request timeout. Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    return None, "Request timeout after retries"
+
+            except requests.exceptions.RequestException as e:
+                if attempt < max_retries - 1:
+                    wait_time = initial_delay * (2 ** attempt)
+                    print(f"Request error: {e}. Retrying in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    return None, f"Request failed: {e}"
+
+            except json.JSONDecodeError as e:
+                return None, f"Failed to parse response JSON: {e}"
+
+        return None, "Max retries exceeded"
+
+
+# =============================================================================
+# Provider Detection
+# =============================================================================
+
+def _detect_model_source(user_model, model_source):
+    """Auto-detect model source from model name if not explicitly provided."""
+    model_source = model_source.lower()
+
+    # Explicit provider pass-through (no auto-detection needed)
+    if model_source == "claude-code":
+        return "claude-code"
+
+    if model_source is None or model_source == "auto":
+        user_model_lower = user_model.lower()
+
+        if "gpt" in user_model_lower:
+            return "openai"
+        elif "claude" in user_model_lower:
+            return "anthropic"
+        elif "gemini" in user_model_lower or "gemma" in user_model_lower:
+            return "google"
+        elif "llama" in user_model_lower or "meta" in user_model_lower:
+            return "huggingface"
+        elif "mistral" in user_model_lower or "mixtral" in user_model_lower:
+            return "mistral"
+        elif "sonar" in user_model_lower or "pplx" in user_model_lower:
+            return "perplexity"
+        elif "deepseek" in user_model_lower or "qwen" in user_model_lower:
+            return "huggingface"
+        elif "grok" in user_model_lower:
+            return "xai"
+        else:
+            raise ValueError(
+                f"Could not auto-detect model source from '{user_model}'. "
+                "Please specify model_source explicitly: OpenAI, Anthropic, Perplexity, Google, xAI, Huggingface, or Mistral"
+            )
+    return model_source
+
+
+def detect_provider(model_name: str, provider: str = "auto") -> str:
+    """Auto-detect provider from model name if not explicitly provided."""
+    if provider and provider.lower() != "auto":
+        return provider.lower()
+
+    model_lower = model_name.lower()
+
+    if "gpt" in model_lower or "o1" in model_lower or "o3" in model_lower:
+        return "openai"
+    elif "claude" in model_lower:
+        return "anthropic"
+    elif "gemini" in model_lower or "gemma" in model_lower:
+        return "google"
+    elif "mistral" in model_lower or "mixtral" in model_lower:
+        return "mistral"
+    elif "sonar" in model_lower or "pplx" in model_lower:
+        return "perplexity"
+    elif "grok" in model_lower:
+        return "xai"
+    elif "llama" in model_lower or "meta" in model_lower or "deepseek" in model_lower or "qwen" in model_lower:
+        return "huggingface"
+    else:
+        raise ValueError(
+            f"Could not auto-detect provider from '{model_name}'. "
+            "Please specify provider explicitly: openai, anthropic, google, mistral, "
+            "perplexity, xai, huggingface, or ollama."
+        )
+
+
+# =============================================================================
+# Ollama Functions
+# =============================================================================
+
+def set_ollama_endpoint(host: str = "localhost", port: int = 11434):
+    """
+    Configure a custom Ollama endpoint.
+
+    Useful if Ollama is running on a different host or port.
+
+    Args:
+        host: Hostname where Ollama is running (default: localhost)
+        port: Port number (default: 11434)
+
+    Example:
+        set_ollama_endpoint("192.168.1.100", 11434)
+    """
+    PROVIDER_CONFIG["ollama"]["endpoint"] = f"http://{host}:{port}/v1/chat/completions"
+
+
+def check_ollama_running(host: str = "localhost", port: int = 11434) -> bool:
+    """
+    Check if Ollama is running and accessible.
+
+    Args:
+        host: Hostname where Ollama should be running
+        port: Port number
+
+    Returns:
+        True if Ollama is running, False otherwise
+    """
+    try:
+        response = requests.get(f"http://{host}:{port}/api/tags", timeout=5)
+        return response.status_code == 200
+    except requests.exceptions.RequestException:
+        return False
+
+
+def list_ollama_models(host: str = "localhost", port: int = 11434) -> list:
+    """
+    List all models available in the local Ollama installation.
+
+    Args:
+        host: Hostname where Ollama is running
+        port: Port number
+
+    Returns:
+        List of model names, or empty list if Ollama is not running
+    """
+    try:
+        response = requests.get(f"http://{host}:{port}/api/tags", timeout=5)
+        if response.status_code == 200:
+            data = response.json()
+            return [model["name"] for model in data.get("models", [])]
+        return []
+    except requests.exceptions.RequestException:
+        return []
+
+
+def check_ollama_model(model: str, host: str = "localhost", port: int = 11434) -> bool:
+    """
+    Check if a specific model is available in Ollama.
+
+    Args:
+        model: Model name to check (e.g., "llama3.2", "mistral")
+        host: Hostname where Ollama is running
+        port: Port number
+
+    Returns:
+        True if model is available, False otherwise
+    """
+    available_models = list_ollama_models(host, port)
+    # Check for exact match or partial match (e.g., "llama3.2" matches "llama3.2:latest")
+    model_lower = model.lower()
+    return any(
+        model_lower == m.lower() or
+        m.lower().startswith(f"{model_lower}:") or
+        model_lower.startswith(m.lower().split(":")[0])
+        for m in available_models
+    )
+
+
+def _format_bytes(size_bytes: int) -> str:
+    """Format bytes into human-readable string."""
+    if size_bytes < 1024:
+        return f"{size_bytes} B"
+    elif size_bytes < 1024 ** 2:
+        return f"{size_bytes / 1024:.1f} KB"
+    elif size_bytes < 1024 ** 3:
+        return f"{size_bytes / (1024 ** 2):.1f} MB"
+    else:
+        return f"{size_bytes / (1024 ** 3):.2f} GB"
+
+
+def _parse_size_string(size_str: str) -> int:
+    """Parse a size string like '2.0 GB' into bytes."""
+    if size_str == "unknown":
+        return 0
+
+    size_str = size_str.strip().upper()
+    try:
+        if "GB" in size_str:
+            return int(float(size_str.replace("GB", "").strip()) * 1024 ** 3)
+        elif "MB" in size_str:
+            return int(float(size_str.replace("MB", "").strip()) * 1024 ** 2)
+        elif "KB" in size_str:
+            return int(float(size_str.replace("KB", "").strip()) * 1024)
+        else:
+            return int(float(size_str.replace("B", "").strip()))
+    except ValueError:
+        return 0
+
+
+# Common model sizes (approximate) for user reference
+OLLAMA_MODEL_SIZES = {
+    "llama3.2": "2.0 GB",
+    "llama3.2:1b": "1.3 GB",
+    "llama3.2:3b": "2.0 GB",
+    "llama3.1": "4.7 GB",
+    "llama3.1:8b": "4.7 GB",
+    "llama3.1:70b": "40 GB",
+    "llama3": "4.7 GB",
+    "llama2": "3.8 GB",
+    "mistral": "4.1 GB",
+    "mixtral": "26 GB",
+    "phi3": "2.2 GB",
+    "phi3:mini": "2.2 GB",
+    "gemma": "5.0 GB",
+    "gemma:2b": "1.7 GB",
+    "gemma:7b": "5.0 GB",
+    "gemma2": "5.4 GB",
+    "gemma2:2b": "1.6 GB",
+    "gemma2:9b": "5.4 GB",
+    "gemma2:27b": "16 GB",
+    "qwen2.5": "4.7 GB",
+    "qwen2.5:0.5b": "397 MB",
+    "qwen2.5:1.5b": "986 MB",
+    "qwen2.5:3b": "1.9 GB",
+    "qwen2.5:7b": "4.7 GB",
+    "deepseek-r1": "4.7 GB",
+    "codellama": "3.8 GB",
+    "codegemma": "5.0 GB",
+    "nomic-embed-text": "274 MB",
+}
+
+
+def get_ollama_model_size_estimate(model: str) -> str:
+    """
+    Get estimated download size for an Ollama model.
+
+    Args:
+        model: Model name
+
+    Returns:
+        Human-readable size estimate or "unknown"
+    """
+    model_lower = model.lower()
+
+    # Check exact match first
+    if model_lower in OLLAMA_MODEL_SIZES:
+        return OLLAMA_MODEL_SIZES[model_lower]
+
+    # Check base model name (without tag)
+    base_model = model_lower.split(":")[0]
+    if base_model in OLLAMA_MODEL_SIZES:
+        return OLLAMA_MODEL_SIZES[base_model]
+
+    return "unknown"
+
+
+def check_system_resources(model: str) -> dict:
+    """
+    Check if system has enough resources to download and run a model.
+
+    Args:
+        model: Model name to check
+
+    Returns:
+        dict with 'can_download', 'can_run', 'warnings', and 'details'
+    """
+    import shutil
+    import os
+
+    result = {
+        "can_download": True,
+        "can_run": True,
+        "warnings": [],
+        "details": {}
+    }
+
+    size_estimate = get_ollama_model_size_estimate(model)
+    model_size_bytes = _parse_size_string(size_estimate)
+
+    # Check disk space (Ollama typically stores models in ~/.ollama)
+    ollama_dir = os.path.expanduser("~/.ollama")
+    if not os.path.exists(ollama_dir):
+        ollama_dir = os.path.expanduser("~")
+
+    try:
+        disk_usage = shutil.disk_usage(ollama_dir)
+        free_space = disk_usage.free
+        result["details"]["free_disk_space"] = _format_bytes(free_space)
+        result["details"]["model_size"] = size_estimate
+
+        # Need at least 1.5x model size for download + extraction
+        required_space = int(model_size_bytes * 1.5) if model_size_bytes > 0 else 0
+
+        if required_space > 0 and free_space < required_space:
+            result["can_download"] = False
+            result["warnings"].append(
+                f"Insufficient disk space. Need ~{_format_bytes(required_space)}, "
+                f"but only {_format_bytes(free_space)} available."
+            )
+        elif required_space > 0 and free_space < required_space * 2:
+            result["warnings"].append(
+                f"Low disk space warning: {_format_bytes(free_space)} available."
+            )
+    except Exception:
+        result["details"]["free_disk_space"] = "unknown"
+
+    # Estimate RAM requirements (rough guide: model size * 1.2 for inference)
+    # This is approximate - actual requirements vary by quantization
+    if model_size_bytes > 0:
+        estimated_ram = model_size_bytes * 1.2
+        result["details"]["estimated_ram"] = _format_bytes(int(estimated_ram))
+
+        # Try to get system RAM (works on most systems)
+        try:
+            import subprocess
+            if os.name == 'posix':  # Linux/macOS
+                if os.path.exists('/proc/meminfo'):  # Linux
+                    with open('/proc/meminfo', 'r') as f:
+                        for line in f:
+                            if line.startswith('MemTotal:'):
+                                total_ram = int(line.split()[1]) * 1024  # Convert KB to bytes
+                                break
+                else:  # macOS
+                    output = subprocess.check_output(['sysctl', '-n', 'hw.memsize'], text=True)
+                    total_ram = int(output.strip())
+
+                result["details"]["total_ram"] = _format_bytes(total_ram)
+
+                if estimated_ram > total_ram * 0.8:
+                    result["can_run"] = False
+                    result["warnings"].append(
+                        f"Model may be too large for your system. "
+                        f"Requires ~{_format_bytes(int(estimated_ram))} RAM, "
+                        f"but system has {_format_bytes(total_ram)}."
+                    )
+                elif estimated_ram > total_ram * 0.5:
+                    result["warnings"].append(
+                        f"Model will use significant RAM (~{_format_bytes(int(estimated_ram))})."
+                    )
+        except Exception:
+            result["details"]["total_ram"] = "unknown"
+            # If we can't check RAM, warn for large models
+            if model_size_bytes > 8 * 1024 ** 3:  # > 8GB models
+                result["warnings"].append(
+                    f"Large model (~{size_estimate}). Ensure you have sufficient RAM."
+                )
+
+    return result
+
+
+def pull_ollama_model(model: str, host: str = "localhost", port: int = 11434, auto_confirm: bool = False) -> bool:
+    """
+    Pull/download a model in Ollama.
+
+    Args:
+        model: Model name to pull (e.g., "llama3.2", "mistral")
+        host: Hostname where Ollama is running
+        port: Port number
+        auto_confirm: If True, skip confirmation prompt
+
+    Returns:
+        True if model was pulled successfully, False otherwise
+    """
+    # Get size estimate and check system resources
+    size_estimate = get_ollama_model_size_estimate(model)
+    resources = check_system_resources(model)
+
+    print(f"\n{'='*60}")
+    print(f"  Model '{model}' not found locally")
+    print(f"{'='*60}")
+    print(f"  Model size:      {size_estimate}")
+    if resources["details"].get("estimated_ram"):
+        print(f"  RAM required:    ~{resources['details']['estimated_ram']}")
+    if resources["details"].get("free_disk_space"):
+        print(f"  Free disk space: {resources['details']['free_disk_space']}")
+    if resources["details"].get("total_ram"):
+        print(f"  System RAM:      {resources['details']['total_ram']}")
+
+    # Show warnings
+    if resources["warnings"]:
+        print(f"\n  {'!'*50}")
+        for warning in resources["warnings"]:
+            print(f"  Warning: {warning}")
+        print(f"  {'!'*50}")
+
+    # Block if can't download
+    if not resources["can_download"]:
+        print(f"\n  Cannot download: insufficient disk space.")
+        print(f"  Free up disk space and try again.")
+        return False
+
+    # Warn but allow if can't run (user might want to try anyway)
+    if not resources["can_run"]:
+        print(f"\n  Warning: Model may not run on this system.")
+        print(f"  Consider a smaller model variant (e.g., '{model}:1b' or '{model}:3b').")
+
+    print(f"{'='*60}")
+
+    # Ask for confirmation
+    if not auto_confirm:
+        try:
+            if not resources["can_run"]:
+                prompt = f"\n  Download anyway? [y/N]: "
+            else:
+                prompt = f"\n  Download '{model}'? [y/N]: "
+            response = input(prompt).strip().lower()
+            if response not in ['y', 'yes']:
+                print("  Download cancelled.")
+                return False
+        except (EOFError, KeyboardInterrupt):
+            print("\n  Download cancelled.")
+            return False
+
+    print(f"\n  Downloading from Ollama registry...")
+    print(f"  (Press Ctrl+C to cancel)\n")
+
+    try:
+        # Ollama pull endpoint streams the response
+        response = requests.post(
+            f"http://{host}:{port}/api/pull",
+            json={"name": model},
+            stream=True,
+            timeout=None  # No timeout - large models can take a while
+        )
+
+        if response.status_code != 200:
+            print(f"Failed to pull model: HTTP {response.status_code}")
+            return False
+
+        # Process streaming response to show progress
+        last_status = ""
+        total_size_shown = False
+
+        for line in response.iter_lines():
+            if line:
+                try:
+                    data = json.loads(line)
+                    status = data.get("status", "")
+
+                    # Show progress for downloads
+                    if "completed" in data and "total" in data:
+                        completed = data["completed"]
+                        total = data["total"]
+                        pct = (completed / total * 100) if total > 0 else 0
+
+                        # Show actual total size on first progress update
+                        if not total_size_shown and total > 0:
+                            print(f"  Actual size: {_format_bytes(total)}")
+                            total_size_shown = True
+
+                        print(f"\r  {status}: {pct:.1f}% ({_format_bytes(completed)}/{_format_bytes(total)})", end="", flush=True)
+                    elif status != last_status:
+                        if last_status and "completed" in str(last_status):
+                            print()  # newline after progress bar
+                        print(f"  {status}")
+                        last_status = status
+
+                    # Check for errors
+                    if "error" in data:
+                        print(f"\n  Error: {data['error']}")
+                        return False
+
+                except json.JSONDecodeError:
+                    continue
+
+        print(f"\n  Model '{model}' downloaded successfully!")
+        return True
+
+    except KeyboardInterrupt:
+        print(f"\n\n  Download cancelled by user.")
+        return False
+    except requests.exceptions.Timeout:
+        print(f"\n  Timeout while downloading model '{model}'.")
+        print(f"  Try again or download manually: ollama pull {model}")
+        return False
+    except requests.exceptions.RequestException as e:
+        print(f"\n  Error pulling model: {e}")
+        return False
+
+
+# =============================================================================
+# Claude Code CLI Functions
+# =============================================================================
+
+def check_claude_cli_available():
+    """Check if the Claude CLI (claude) is installed and available on PATH."""
+    import shutil
+    return shutil.which("claude") is not None