utils.py

import os
import json
from pathlib import Path
from typing import List, Dict, Any, Optional
import logging
from datetime import datetime

def get_available_models(directory: str = "./models") -> Dict[str, Any]:
    """Get available GGUF models in the directory"""
    try:
        path = Path(directory)
        if not path.exists():
            return {"error": f"Directory {directory} does not exist"}
        
        models = []
        for file in path.glob("*.gguf"):
            try:
                stat = file.stat()
                models.append({
                    "name": file.name,
                    "path": str(file),
                    "size_mb": round(stat.st_size / (1024 * 1024), 2),
                    "modified": datetime.fromtimestamp(stat.st_mtime).isoformat()
                })
            except Exception as e:
                logging.warning(f"Error reading {file}: {e}")
        
        # Sort by name
        models.sort(key=lambda x: x["name"])
        
        return {
            "directory": directory,
            "exists": True,
            "model_count": len(models),
            "models": models
        }
        
    except Exception as e:
        return {"error": str(e)}

def format_chat_history(history: List[List[str]], system_prompt: str = "") -> str:
    """Format chat history for the model"""
    formatted = ""
    
    if system_prompt:
        formatted += f"<|system|>\n{system_prompt}\n<|end|>\n\n"
    
    for message in history:
        role, content = message
        if role == "user":
            formatted += f"<|user|>\n{content}\n<|end|>\n\n"
        elif role == "assistant":
            formatted += f"<|assistant|>\n{content}\n<|end|>\n\n"
    
    formatted += "<|assistant|>\n"
    return formatted

def format_chat_history_messages(history: List[Dict[str, str]], system_prompt: str = "") -> str:
    """Format chat history (message format) for the model"""
    formatted = ""
    
    if system_prompt:
        formatted += f"<|system|>\n{system_prompt}\n<|end|>\n\n"
    
    for message in history:
        role = message.get("role", "")
        content = message.get("content", "")
        
        if role == "user":
            formatted += f"<|user|>\n{content}\n<|end|>\n\n"
        elif role == "assistant":
            formatted += f"<|assistant|>\n{content}\n<|end|>\n\n"
    
    formatted += "<|assistant|>\n"
    return formatted

def parse_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]:
    """Parse model metadata into a readable format"""
    parsed = {
        "architecture": "Unknown",
        "parameters": "Unknown",
        "context_length": "Unknown",
        "embedding_size": "Unknown",
        "layers": "Unknown",
        "heads": "Unknown"
    }
    
    # Try to extract common fields
    if "general.architecture" in metadata:
        parsed["architecture"] = metadata["general.architecture"]
    
    if "llama.block_count" in metadata:
        parsed["layers"] = metadata["llama.block_count"]
    
    if "llama.context_length" in metadata:
        parsed["context_length"] = metadata["llama.context_length"]
    
    if "llama.embedding_length" in metadata:
        parsed["embedding_size"] = metadata["llama.embedding_length"]
    
    if "llama.attention.head_count" in metadata:
        parsed["heads"] = metadata["llama.attention.head_count"]
    
    # Estimate parameters based on architecture
    if parsed["architecture"] == "llama":
        try:
            layers = int(parsed["layers"]) if parsed["layers"] != "Unknown" else 0
            embed_size = int(parsed["embedding_size"]) if parsed["embedding_size"] != "Unknown" else 0
            
            if layers > 0 and embed_size > 0:
                # Rough estimate for LLaMA parameters
                params = layers * (12 * embed_size * embed_size + 13 * embed_size)
                if params > 1e9:
                    parsed["parameters"] = f"{params / 1e9:.1f}B"
                elif params > 1e6:
                    parsed["parameters"] = f"{params / 1e6:.1f}M"
                else:
                    parsed["parameters"] = str(params)
        except:
            pass
    
    return parsed

def save_chat_history(history: List[Dict[str, str]], filename: str = None) -> str:
    """Save chat history to a JSON file"""
    if filename is None:
        filename = f"chat_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(history, f, indent=2, ensure_ascii=False)
        return filename
    except Exception as e:
        logging.error(f"Failed to save chat history: {e}")
        return ""

def load_chat_history(filename: str) -> List[Dict[str, str]]:
    """Load chat history from a JSON file"""
    try:
        with open(filename, 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        logging.error(f"Failed to load chat history: {e}")
        return []

def estimate_tokens(text: str) -> int:
    """Estimate token count (rough approximation)"""
    # Simple approximation: ~4 characters per token
    return len(text) // 4

def validate_model_file(model_path: str) -> Dict[str, Any]:
    """Validate a model file"""
    result = {
        "valid": False,
        "exists": False,
        "readable": False,
        "size_mb": 0,
        "file_type": None,
        "error": None
    }
    
    try:
        path = Path(model_path)
        result["exists"] = path.exists()
        
        if not result["exists"]:
            result["error"] = "File does not exist"
            return result
        
        result["size_mb"] = round(path.stat().st_size / (1024 * 1024), 2)
        result["file_type"] = path.suffix.lower()
        
        if result["file_type"] != ".gguf":
            result["error"] = "Not a GGUF file"
            return result
        
        # Try to read first few bytes
        try:
            with open(path, "rb") as f:
                header = f.read(4)
                result["readable"] = len(header) == 4
                result["valid"] = result["readable"]
        except Exception as e:
            result["error"] = f"Cannot read file: {str(e)}"
        
    except Exception as e:
        result["error"] = str(e)
    
    return result

def create_default_config() -> Dict[str, Any]:
    """Create default configuration"""
    return {
        "models_directory": "./models",
        "default_context_size": 2048,
        "default_gpu_layers": 0,
        "default_temperature": 0.7,
        "default_max_tokens": 512,
        "default_top_p": 0.9,
        "default_repeat_penalty": 1.1,
        "system_prompt": "You are a helpful assistant.",
        "chat_format": "chatml",
        "auto_save_chat": True,
        "theme": "soft"
    }

def load_config(config_path: str = "config.json") -> Dict[str, Any]:
    """Load configuration from file"""
    try:
        if Path(config_path).exists():
            with open(config_path, 'r') as f:
                return json.load(f)
    except Exception as e:
        logging.warning(f"Failed to load config: {e}")
    
    return create_default_config()

def save_config(config: Dict[str, Any], config_path: str = "config.json"):
    """Save configuration to file"""
    try:
        with open(config_path, 'w') as f:
            json.dump(config, f, indent=2)
    except Exception as e:
        logging.error(f"Failed to save config: {e}")