Ollama model pull failures

import subprocess
import os
import shutil
from pathlib import Path

def safe_ollama_pull(model_name: str, max_retries: int = 3) -> bool:
    """
    Safely pull an Ollama model with retry logic and cleanup on failure.
    Returns True if pull succeeds, False if all retries exhausted.
    """
    models_dir = Path.home() / '.ollama' / 'models'
    model_dir = models_dir / 'models' / 'library' / model_name
    
    for attempt in range(1, max_retries + 1):
        print(f'Attempt {attempt}/{max_retries}: pulling {model_name}')
        
        disk_free_gb = shutil.disk_usage(models_dir).free / (1024**3)
        if disk_free_gb < 5:
            print(f'ERROR: Only {disk_free_gb:.1f}GB free. Need at least 5GB.')
            return False
        
        try:
            result = subprocess.run(
                ['ollama', 'pull', model_name],
                capture_output=True,
                text=True,
                timeout=600
            )
            
            if result.returncode == 0:
                print(f'SUCCESS: {model_name} pulled')
                return True
            else:
                print(f'FAILED: {result.stderr}')
                if 'connection' in result.stderr.lower() or 'timeout' in result.stderr.lower():
                    if model_dir.exists():
                        print(f'Cleaning incomplete model at {model_dir}')
                        shutil.rmtree(model_dir, ignore_errors=True)
                    if attempt < max_retries:
                        print(f'Retrying in 5 seconds...')
                        import time
                        time.sleep(5)
                        continue
        
        except subprocess.TimeoutExpired:
            print(f'TIMEOUT: Pull exceeded 10 minutes')
            if model_dir.exists():
                print(f'Cleaning incomplete model')
                shutil.rmtree(model_dir, ignore_errors=True)
            if attempt < max_retries:
                import time
                time.sleep(5)
                continue
        
        except FileNotFoundError:
            print('ERROR: ollama command not found. Install Ollama from ollama.com')
            return False
    
    print(f'FAILED: Could not pull {model_name} after {max_retries} attempts')
    return False

if __name__ == '__main__':
    success = safe_ollama_pull('llama3.2:3b')
    if success:
        result = subprocess.run(['ollama', 'list'], capture_output=True, text=True)
        print('\nAvailable models:')
        print(result.stdout)
    else:
        print('Model pull failed. Check disk space, network, and Ollama service.')