isa-model 0.1.0__py3-none-any.whl

isa_model/__init__.py

@@ -0,0 +1,5 @@
+"""
+isA_Model - A simple interface for AI model integration
+"""
+
+__version__ = "0.1.0" 

isa_model/core/model_manager.py

@@ -0,0 +1,143 @@
+from typing import Dict, Optional, List, Any
+import logging
+from pathlib import Path
+from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub.utils import HfHubHTTPError
+from .model_storage import ModelStorage, LocalModelStorage
+from .model_registry import ModelRegistry, ModelType, ModelCapability
+
+logger = logging.getLogger(__name__)
+
+class ModelManager:
+    """Model management service for handling model downloads, versions, and caching"""
+    
+    def __init__(self, 
+                 storage: Optional[ModelStorage] = None,
+                 registry: Optional[ModelRegistry] = None):
+        self.storage = storage or LocalModelStorage()
+        self.registry = registry or ModelRegistry()
+    
+    async def get_model(self, 
+                       model_id: str, 
+                       repo_id: str,
+                       model_type: ModelType,
+                       capabilities: List[ModelCapability],
+                       revision: Optional[str] = None,
+                       force_download: bool = False) -> Path:
+        """
+        Get model files, downloading if necessary
+        
+        Args:
+            model_id: Unique identifier for the model
+            repo_id: Hugging Face repository ID
+            model_type: Type of model (LLM, embedding, etc.)
+            capabilities: List of model capabilities
+            revision: Specific model version/tag
+            force_download: Force re-download even if cached
+            
+        Returns:
+            Path to the model files
+        """
+        # Check if model is already downloaded
+        if not force_download:
+            model_path = await self.storage.load_model(model_id)
+            if model_path:
+                logger.info(f"Using cached model {model_id}")
+                return model_path
+        
+        try:
+            # Download model files
+            logger.info(f"Downloading model {model_id} from {repo_id}")
+            model_dir = Path(f"./models/temp/{model_id}")
+            model_dir.mkdir(parents=True, exist_ok=True)
+            
+            snapshot_download(
+                repo_id=repo_id,
+                revision=revision,
+                local_dir=model_dir,
+                local_dir_use_symlinks=False
+            )
+            
+            # Save model and metadata
+            metadata = {
+                "repo_id": repo_id,
+                "revision": revision,
+                "downloaded_at": str(Path(model_dir).stat().st_mtime)
+            }
+            
+            # Register model
+            self.registry.register_model(
+                model_id=model_id,
+                model_type=model_type,
+                capabilities=capabilities,
+                metadata=metadata
+            )
+            
+            # Save model files
+            await self.storage.save_model(model_id, str(model_dir), metadata)
+            
+            return await self.storage.load_model(model_id)
+            
+        except HfHubHTTPError as e:
+            logger.error(f"Failed to download model {model_id}: {e}")
+            raise
+    
+    async def list_models(self) -> List[Dict[str, Any]]:
+        """List all downloaded models with their metadata"""
+        models = await self.storage.list_models()
+        return [
+            {
+                "model_id": model_id,
+                **metadata,
+                **(self.registry.get_model_info(model_id) or {})
+            }
+            for model_id, metadata in models.items()
+        ]
+    
+    async def remove_model(self, model_id: str) -> bool:
+        """Remove a model and its metadata"""
+        try:
+            # Remove from storage
+            storage_success = await self.storage.delete_model(model_id)
+            
+            # Unregister from registry
+            registry_success = self.registry.unregister_model(model_id)
+            
+            return storage_success and registry_success
+            
+        except Exception as e:
+            logger.error(f"Failed to remove model {model_id}: {e}")
+            return False
+    
+    async def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """Get information about a specific model"""
+        storage_info = await self.storage.get_metadata(model_id)
+        registry_info = self.registry.get_model_info(model_id)
+        
+        if not storage_info and not registry_info:
+            return None
+            
+        return {
+            **(storage_info or {}),
+            **(registry_info or {})
+        }
+    
+    async def update_model(self, 
+                          model_id: str, 
+                          repo_id: str,
+                          model_type: ModelType,
+                          capabilities: List[ModelCapability],
+                          revision: Optional[str] = None) -> bool:
+        """Update a model to a new version"""
+        try:
+            return bool(await self.get_model(
+                model_id=model_id,
+                repo_id=repo_id,
+                model_type=model_type,
+                capabilities=capabilities,
+                revision=revision,
+                force_download=True
+            ))
+        except Exception as e:
+            logger.error(f"Failed to update model {model_id}: {e}")
+            return False 

isa_model/core/model_registry.py

@@ -0,0 +1,115 @@
+from typing import Dict, List, Optional, Any
+from enum import Enum
+import logging
+from pathlib import Path
+import json
+
+logger = logging.getLogger(__name__)
+
+class ModelCapability(str, Enum):
+    """Model capabilities"""
+    TEXT_GENERATION = "text_generation"
+    CHAT = "chat"
+    EMBEDDING = "embedding"
+    RERANKING = "reranking"
+    REASONING = "reasoning"
+    IMAGE_GENERATION = "image_generation"
+    IMAGE_ANALYSIS = "image_analysis"
+    AUDIO_TRANSCRIPTION = "audio_transcription"
+    IMAGE_UNDERSTANDING = "image_understanding"
+
+class ModelType(str, Enum):
+    """Model types"""
+    LLM = "llm"
+    EMBEDDING = "embedding"
+    RERANK = "rerank"
+    IMAGE = "image"
+    AUDIO = "audio"
+    VIDEO = "video"
+    VISION = "vision"
+
+class ModelRegistry:
+    """Registry for model metadata and capabilities"""
+    
+    def __init__(self, registry_file: str = "./models/model_registry.json"):
+        self.registry_file = Path(registry_file)
+        self.registry: Dict[str, Dict[str, Any]] = {}
+        self._load_registry()
+    
+    def _load_registry(self):
+        """Load model registry from file"""
+        if self.registry_file.exists():
+            with open(self.registry_file, 'r') as f:
+                self.registry = json.load(f)
+        else:
+            self.registry = {}
+            self._save_registry()
+    
+    def _save_registry(self):
+        """Save model registry to file"""
+        self.registry_file.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.registry_file, 'w') as f:
+            json.dump(self.registry, f, indent=2)
+    
+    def register_model(self, 
+                      model_id: str,
+                      model_type: ModelType,
+                      capabilities: List[ModelCapability],
+                      metadata: Dict[str, Any]) -> bool:
+        """Register a model with its capabilities and metadata"""
+        try:
+            self.registry[model_id] = {
+                "type": model_type,
+                "capabilities": [cap.value for cap in capabilities],
+                "metadata": metadata
+            }
+            self._save_registry()
+            logger.info(f"Registered model {model_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to register model {model_id}: {e}")
+            return False
+    
+    def unregister_model(self, model_id: str) -> bool:
+        """Unregister a model"""
+        try:
+            if model_id in self.registry:
+                del self.registry[model_id]
+                self._save_registry()
+                logger.info(f"Unregistered model {model_id}")
+                return True
+            return False
+        except Exception as e:
+            logger.error(f"Failed to unregister model {model_id}: {e}")
+            return False
+    
+    def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """Get model information"""
+        return self.registry.get(model_id)
+    
+    def get_models_by_type(self, model_type: ModelType) -> Dict[str, Dict[str, Any]]:
+        """Get all models of a specific type"""
+        return {
+            model_id: info
+            for model_id, info in self.registry.items()
+            if info["type"] == model_type
+        }
+    
+    def get_models_by_capability(self, capability: ModelCapability) -> Dict[str, Dict[str, Any]]:
+        """Get all models with a specific capability"""
+        return {
+            model_id: info
+            for model_id, info in self.registry.items()
+            if capability.value in info["capabilities"]
+        }
+    
+    def has_capability(self, model_id: str, capability: ModelCapability) -> bool:
+        """Check if a model has a specific capability"""
+        model_info = self.get_model_info(model_id)
+        if not model_info:
+            return False
+        return capability.value in model_info["capabilities"]
+    
+    def list_models(self) -> Dict[str, Dict[str, Any]]:
+        """List all registered models"""
+        return self.registry 

isa_model/core/model_router.py

@@ -0,0 +1,226 @@
+import random
+import time
+from typing import Dict, List, Any, Optional, Callable
+import threading
+
+class ModelRouter:
+    """
+    Routes requests to appropriate model instances based on different strategies:
+    - Weighted round-robin
+    - Least connections
+    - Least response time
+    - Dynamic load balancing
+    """
+    
+    def __init__(self, registry):
+        self.registry = registry
+        self.model_stats = {}  # Track performance metrics for each model
+        self.lock = threading.RLock()
+        
+        # Maps model_type -> list of model_ids of that type
+        self.model_type_mapping = {}
+        
+        # Maps routing_strategy_name -> routing_function
+        self.routing_strategies = {
+            "round_robin": self._route_round_robin,
+            "weighted_random": self._route_weighted_random,
+            "least_connections": self._route_least_connections,
+            "least_response_time": self._route_least_response_time,
+            "dynamic_load": self._route_dynamic_load
+        }
+        
+        # Round-robin counters for each model type
+        self.rr_counters = {}
+        
+    def register_model_type(self, model_type: str, model_ids: List[str], weights: Optional[List[float]] = None):
+        """Register models of a specific type with optional weights"""
+        with self.lock:
+            self.model_type_mapping[model_type] = model_ids
+            
+            # Initialize stats for each model
+            for i, model_id in enumerate(model_ids):
+                weight = weights[i] if weights and i < len(weights) else 1.0
+                
+                if model_id not in self.model_stats:
+                    self.model_stats[model_id] = {
+                        "active_connections": 0,
+                        "total_requests": 0,
+                        "avg_response_time": 0,
+                        "weight": weight,
+                        "last_used": 0
+                    }
+                else:
+                    # Update weight if model already exists
+                    self.model_stats[model_id]["weight"] = weight
+            
+            # Initialize round-robin counter
+            self.rr_counters[model_type] = 0
+    
+    def route_request(self, model_type: str, routing_strategy: str = "round_robin") -> Optional[str]:
+        """
+        Route a request to an appropriate model of the given type
+        
+        Args:
+            model_type: Type of model needed
+            routing_strategy: Strategy to use for routing
+            
+        Returns:
+            model_id: ID of the model to use, or None if no models available
+        """
+        with self.lock:
+            if model_type not in self.model_type_mapping:
+                return None
+                
+            if not self.model_type_mapping[model_type]:
+                return None
+                
+            # Get the routing function
+            routing_func = self.routing_strategies.get(routing_strategy, self._route_round_robin)
+            
+            # Route the request
+            model_id = routing_func(model_type)
+            
+            if model_id:
+                # Update stats
+                self.model_stats[model_id]["active_connections"] += 1
+                self.model_stats[model_id]["total_requests"] += 1
+                self.model_stats[model_id]["last_used"] = time.time()
+                
+            return model_id
+    
+    def release_connection(self, model_id: str, response_time: float = None):
+        """Release a connection and update stats"""
+        with self.lock:
+            if model_id in self.model_stats:
+                stats = self.model_stats[model_id]
+                stats["active_connections"] = max(0, stats["active_connections"] - 1)
+                
+                # Update average response time
+                if response_time is not None:
+                    old_avg = stats["avg_response_time"]
+                    total_req = stats["total_requests"]
+                    
+                    if total_req > 0:
+                        # Weighted average
+                        stats["avg_response_time"] = (old_avg * (total_req - 1) + response_time) / total_req
+    
+    # Routing strategies
+    def _route_round_robin(self, model_type: str) -> Optional[str]:
+        """Simple round-robin routing"""
+        models = self.model_type_mapping.get(model_type, [])
+        if not models:
+            return None
+            
+        # Get and increment counter
+        counter = self.rr_counters[model_type]
+        self.rr_counters[model_type] = (counter + 1) % len(models)
+        
+        return models[counter]
+    
+    def _route_weighted_random(self, model_type: str) -> Optional[str]:
+        """Weighted random selection based on configured weights"""
+        models = self.model_type_mapping.get(model_type, [])
+        if not models:
+            return None
+            
+        # Get weights
+        weights = [self.model_stats[model_id]["weight"] for model_id in models]
+        
+        # Weighted random selection
+        total = sum(weights)
+        r = random.uniform(0, total)
+        upto = 0
+        
+        for i, w in enumerate(weights):
+            upto += w
+            if upto >= r:
+                return models[i]
+                
+        # Fallback
+        return models[-1]
+    
+    def _route_least_connections(self, model_type: str) -> Optional[str]:
+        """Route to the model with the fewest active connections"""
+        models = self.model_type_mapping.get(model_type, [])
+        if not models:
+            return None
+            
+        # Find model with least connections
+        min_connections = float('inf')
+        selected_model = None
+        
+        for model_id in models:
+            connections = self.model_stats[model_id]["active_connections"]
+            if connections < min_connections:
+                min_connections = connections
+                selected_model = model_id
+                
+        return selected_model
+    
+    def _route_least_response_time(self, model_type: str) -> Optional[str]:
+        """Route to the model with the lowest average response time"""
+        models = self.model_type_mapping.get(model_type, [])
+        if not models:
+            return None
+            
+        # Find model with lowest response time
+        min_response_time = float('inf')
+        selected_model = None
+        
+        for model_id in models:
+            response_time = self.model_stats[model_id]["avg_response_time"]
+            # Skip models with no data yet
+            if response_time == 0:
+                continue
+                
+            if response_time < min_response_time:
+                min_response_time = response_time
+                selected_model = model_id
+                
+        # If no model has response time data, fall back to least connections
+        if selected_model is None:
+            return self._route_least_connections(model_type)
+            
+        return selected_model
+    
+    def _route_dynamic_load(self, model_type: str) -> Optional[str]:
+        """
+        Dynamic load balancing based on a combination of:
+        - Connection count
+        - Response time
+        - Recent usage
+        """
+        models = self.model_type_mapping.get(model_type, [])
+        if not models:
+            return None
+            
+        # Calculate a score for each model (lower is better)
+        best_score = float('inf')
+        selected_model = None
+        now = time.time()
+        
+        for model_id in models:
+            stats = self.model_stats[model_id]
+            
+            # Normalize each factor between 0 and 1
+            connections = stats["active_connections"]
+            conn_score = connections / (connections + 1)  # Approaches 1 as connections increase
+            
+            resp_time = stats["avg_response_time"]
+            # Max expected response time (adjust as needed)
+            max_resp_time = 5.0  
+            resp_score = min(1.0, resp_time / max_resp_time)
+            
+            # Time since last use (for distributing load)
+            recency = now - stats["last_used"] if stats["last_used"] > 0 else 60
+            recency_score = 1.0 - min(1.0, recency / 60.0)  # Unused for 60s approaches 0
+            
+            # Combined score (lower is better)
+            # Weights can be adjusted based on importance
+            score = (0.4 * conn_score) + (0.4 * resp_score) + (0.2 * recency_score)
+            
+            if score < best_score:
+                best_score = score
+                selected_model = model_id
+                
+        return selected_model

isa_model/core/model_storage.py

@@ -0,0 +1,133 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, Any, BinaryIO
+from pathlib import Path
+import logging
+import json
+import shutil
+import os
+
+logger = logging.getLogger(__name__)
+
+class ModelStorage(ABC):
+    """Base class for model storage implementations"""
+    
+    @abstractmethod
+    async def save_model(self, model_id: str, model_path: str, metadata: Dict[str, Any]) -> bool:
+        """Save model files and metadata"""
+        pass
+    
+    @abstractmethod
+    async def load_model(self, model_id: str) -> Optional[Path]:
+        """Load model files"""
+        pass
+    
+    @abstractmethod
+    async def delete_model(self, model_id: str) -> bool:
+        """Delete model files and metadata"""
+        pass
+    
+    @abstractmethod
+    async def get_metadata(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """Get model metadata"""
+        pass
+    
+    @abstractmethod
+    async def list_models(self) -> Dict[str, Dict[str, Any]]:
+        """List all stored models with their metadata"""
+        pass
+
+class LocalModelStorage(ModelStorage):
+    """Local file system based model storage"""
+    
+    def __init__(self, base_dir: str = "./models"):
+        self.base_dir = Path(base_dir)
+        self.models_dir = self.base_dir / "models"
+        self.metadata_file = self.base_dir / "model_metadata.json"
+        self._ensure_directories()
+        self._load_metadata()
+    
+    def _ensure_directories(self):
+        """Ensure required directories exist"""
+        self.models_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _load_metadata(self):
+        """Load model metadata from file"""
+        if self.metadata_file.exists():
+            with open(self.metadata_file, 'r') as f:
+                self.metadata = json.load(f)
+        else:
+            self.metadata = {}
+            self._save_metadata()
+    
+    def _save_metadata(self):
+        """Save model metadata to file"""
+        with open(self.metadata_file, 'w') as f:
+            json.dump(self.metadata, f, indent=2)
+    
+    async def save_model(self, model_id: str, model_path: str, metadata: Dict[str, Any]) -> bool:
+        """Save model files and metadata"""
+        try:
+            model_dir = self.models_dir / model_id
+            source_path = Path(model_path)
+            
+            # Copy model files
+            if source_path.is_file():
+                model_dir.mkdir(parents=True, exist_ok=True)
+                shutil.copy2(source_path, model_dir / source_path.name)
+            else:
+                shutil.copytree(source_path, model_dir, dirs_exist_ok=True)
+            
+            # Update metadata
+            self.metadata[model_id] = {
+                **metadata,
+                "storage_path": str(model_dir),
+                "saved_at": str(Path(model_dir).stat().st_mtime)
+            }
+            self._save_metadata()
+            
+            logger.info(f"Saved model {model_id} to {model_dir}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to save model {model_id}: {e}")
+            return False
+    
+    async def load_model(self, model_id: str) -> Optional[Path]:
+        """Load model files"""
+        try:
+            model_dir = self.models_dir / model_id
+            if not model_dir.exists():
+                logger.warning(f"Model {model_id} not found at {model_dir}")
+                return None
+            
+            return model_dir
+            
+        except Exception as e:
+            logger.error(f"Failed to load model {model_id}: {e}")
+            return None
+    
+    async def delete_model(self, model_id: str) -> bool:
+        """Delete model files and metadata"""
+        try:
+            model_dir = self.models_dir / model_id
+            if model_dir.exists():
+                shutil.rmtree(model_dir)
+            
+            if model_id in self.metadata:
+                del self.metadata[model_id]
+                self._save_metadata()
+            
+            logger.info(f"Deleted model {model_id}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to delete model {model_id}: {e}")
+            return False
+    
+    async def get_metadata(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """Get model metadata"""
+        return self.metadata.get(model_id)
+    
+    async def list_models(self) -> Dict[str, Dict[str, Any]]:
+        """List all stored models with their metadata"""
+        return self.metadata