ai_manager.py

#!/usr/bin/env python3
"""
AI Manager - Smart model detection and selection
Handles Ollama API communication and model management
"""

import requests
import subprocess
import json
from typing import List, Dict, Optional

class AIManager:
    """
    Manages AI model detection, selection, and communication.
    
    📖 What you learn:
    - REST API calls (requests library)
    - JSON parsing
    - Error handling with fallbacks
    - Model metadata management
    """
    
    OLLAMA_BASE_URL = "http://localhost:11434"
    
    def __init__(self):
        self.available_models: List[Dict] = []
        self.current_model: str = ""
        self.is_ollama_running: bool = False
    
    def check_ollama_status(self) -> bool:
        """
        Check if Ollama server is running.
        
        Returns:
            bool: True if Ollama is accessible
        """
        try:
            response = requests.get(f"{self.OLLAMA_BASE_URL}/api/tags", timeout=5)
            self.is_ollama_running = response.status_code == 200
            return self.is_ollama_running
        except requests.exceptions.ConnectionError:
            self.is_ollama_running = False
            return False
        except Exception as e:
            print(f"❌ Ollama check error: {e}")
            self.is_ollama_running = False
            return False
    
    def get_available_models(self) -> List[Dict]:
        """
        Fetch all available Ollama models with their info.
        
        Returns:
            List of model dictionaries with name, size, family
        """
        if not self.check_ollama_status():
            return []
        
        try:
            response = requests.get(f"{self.OLLAMA_BASE_URL}/api/tags", timeout=10)
            response.raise_for_status()
            
            data = response.json()
            models = []
            
            for model in data.get("models", []):
                models.append({
                    "name": model.get("name", "unknown"),
                    "size": model.get("size", 0),
                    "size_gb": round(model.get("size", 0) / (1024**3), 2),
                    "family": model.get("details", {}).get("family", "unknown"),
                    "parameters": model.get("details", {}).get("parameter_size", "unknown"),
                    "quantization": model.get("details", {}).get("quantization_level", "unknown"),
                })
            
            self.available_models = models
            return models
            
        except Exception as e:
            print(f"❌ Error fetching models: {e}")
            return []
    
    def get_recommended_model(self) -> Optional[str]:
        """
        Smart model selection based on available models and system RAM.
        
        Priority:
        1. Llama 3.1/3.2 family (best overall)
        2. Phi-3 (best for low RAM)
        3. Qwen2.5 (best for coding)
        4. Mistral (good general purpose)
        5. Any available model
        
        Returns:
            str: Recommended model name or None
        """
        if not self.available_models:
            self.get_available_models()
        
        if not self.available_models:
            return None
        
        # Priority order for model families
        priority = [
            ("llama", "3.2"),    # Llama 3.2 - best for 8GB RAM
            ("llama", "3.1"),    # Llama 3.1 - good balance
            ("phi", "3"),        # Phi-3 - lightweight
            ("qwen", "2.5"),     # Qwen2.5 - coding specialist
            ("mistral", ""),     # Mistral - general purpose
        ]
        
        # Search for priority models
        for family, version in priority:
            for model in self.available_models:
                name = model["name"].lower()
                if family in name:
                    if version and version in name:
                        return model["name"]
                    elif not version:
                        return model["name"]
        
        # Fallback: smallest model (best for low RAM)
        smallest = min(self.available_models, key=lambda m: m.get("size", float("inf")))
        return smallest["name"]
    
    def set_model(self, model_name: str) -> bool:
        """
        Set the current model to use.
        
        Args:
            model_name: Name of the model (e.g., "llama3.1:8b")
        
        Returns:
            bool: True if model exists in available models
        """
        for model in self.available_models:
            if model["name"] == model_name:
                self.current_model = model_name
                return True
        
        # Model not found, but still allow it (might be pullable)
        self.current_model = model_name
        return False
    
    def query_ai(self, prompt: str, model: Optional[str] = None, timeout: int = 120) -> Dict:
        """
        Send a query to the AI model.
        
        Args:
            prompt: User's question/prompt
            model: Model to use (uses current_model if None)
            timeout: Request timeout in seconds
        
        Returns:
            Dict with 'success', 'response', 'error' keys
        """
        if not self.current_model and not model:
            return {
                "success": False,
                "response": "",
                "error": "No model selected. Please select a model first."
            }
        
        model_to_use = model or self.current_model
        
        try:
            # Use Ollama API (faster than CLI subprocess)
            response = requests.post(
                f"{self.OLLAMA_BASE_URL}/api/generate",
                json={
                    "model": model_to_use,
                    "prompt": prompt,
                    "stream": False,
                    "options": {
                        "temperature": 0.7,
                        "num_ctx": 2048  # Limit context for 8GB RAM
                    }
                },
                timeout=timeout
            )
            
            response.raise_for_status()
            result = response.json()
            
            return {
                "success": True,
                "response": result.get("response", ""),
                "error": None,
                "model": model_to_use,
                "total_duration": result.get("total_duration", 0),
            }
            
        except requests.exceptions.Timeout:
            return {
                "success": False,
                "response": "",
                "error": f"Request timed out after {timeout} seconds. Try a shorter prompt."
            }
        except requests.exceptions.ConnectionError:
            return {
                "success": False,
                "response": "",
                "error": "Cannot connect to Ollama. Make sure 'ollama serve' is running."
            }
        except Exception as e:
            return {
                "success": False,
                "response": "",
                "error": f"Error: {str(e)}"
            }
    
    def pull_model(self, model_name: str) -> Dict:
        """
        Download a new model from Ollama.
        
        Args:
            model_name: Model to pull (e.g., "llama3.1:8b")
        
        Returns:
            Dict with success status
        """
        try:
            # Use subprocess for pull (streaming download)
            result = subprocess.run(
                ["ollama", "pull", model_name],
                capture_output=True,
                text=True,
                timeout=600  # 10 minutes for large models
            )
            
            if result.returncode == 0:
                # Refresh model list
                self.get_available_models()
                return {"success": True, "message": f"✅ Model {model_name} downloaded!"}
            else:
                return {"success": False, "message": f"❌ Error: {result.stderr}"}
                
        except subprocess.TimeoutExpired:
            return {"success": False, "message": "⏱️ Download timed out. Model might be large."}
        except Exception as e:
            return {"success": False, "message": f"❌ Error: {str(e)}"}
    
    def get_model_info_text(self) -> str:
        """
        Get human-readable info about current setup.
        
        Returns:
            str: Formatted status text
        """
        if not self.is_ollama_running:
            return "❌ Ollama not running"
        
        if not self.available_models:
            return "⚠️ No models found"
        
        if not self.current_model:
            recommended = self.get_recommended_model()
            return f"📦 {len(self.available_models)} models available. Recommended: {recommended}"
        
        # Find current model info
        model_info = next(
            (m for m in self.available_models if m["name"] == self.current_model),
            None
        )
        
        if model_info:
            return f"✅ {self.current_model} ({model_info['size_gb']} GB, {model_info['family']})"
        else:
            return f"✅ {self.current_model}"