isa-model 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

isa_model/core/models/model_manager.py

@@ -0,0 +1,399 @@
+from typing import Dict, Optional, List, Any
+import logging
+from pathlib import Path
+from datetime import datetime
+from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub.errors import HfHubHTTPError
+from .model_storage import ModelStorage, LocalModelStorage
+from .model_repo import ModelRegistry, ModelType, ModelCapability
+from .model_billing_tracker import ModelBillingTracker, ModelOperationType
+from ..pricing_manager import PricingManager
+from ..config import ConfigManager
+
+logger = logging.getLogger(__name__)
+
+class ModelManager:
+    """
+    Model lifecycle management service
+    
+    Handles the complete model lifecycle:
+    - Model registration and metadata management
+    - Model downloads, versions, and caching  
+    - Cost tracking and billing across all operations
+    - Integration with model training, evaluation, and deployment
+    """
+    
+    
+    def __init__(self, 
+                 storage: Optional[ModelStorage] = None,
+                 registry: Optional[ModelRegistry] = None,
+                 billing_tracker: Optional[ModelBillingTracker] = None,
+                 pricing_manager: Optional[PricingManager] = None,
+                 config_manager: Optional[ConfigManager] = None):
+        self.storage = storage or LocalModelStorage()
+        self.registry = registry or ModelRegistry()
+        self.billing_tracker = billing_tracker or ModelBillingTracker(model_registry=self.registry)
+        self.pricing_manager = pricing_manager or PricingManager()
+        self.config_manager = config_manager or ConfigManager()
+    
+    def get_model_pricing(self, provider: str, model_name: str) -> Dict[str, float]:
+        """获取模型定价信息"""
+        pricing = self.pricing_manager.get_model_pricing(provider, model_name)
+        if pricing:
+            return {"input": pricing.input_cost, "output": pricing.output_cost}
+        return {"input": 0.0, "output": 0.0}
+    
+    def calculate_cost(self, provider: str, model_name: str, input_tokens: int, output_tokens: int) -> float:
+        """计算请求成本"""
+        return self.pricing_manager.calculate_cost(
+            provider=provider,
+            model_name=model_name,
+            input_units=input_tokens,
+            output_units=output_tokens
+        )
+    
+    def get_cheapest_model(self, provider: str, model_type: str = "llm") -> Optional[str]:
+        """获取最便宜的模型"""
+        result = self.pricing_manager.get_cheapest_model(
+            provider=provider, 
+            unit_type="token",
+            min_input_units=1000  # Assume 1K tokens for comparison
+        )
+        return result["model_name"] if result else None
+
+    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) -> Optional[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 or None if failed
+        """
+        # 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}")
+            return None
+        except Exception as e:
+            logger.error(f"Unexpected error downloading model {model_id}: {e}")
+            return None
+    
+    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 
+    
+    # === MODEL LIFECYCLE MANAGEMENT ===
+    
+    async def register_model_for_lifecycle(
+        self,
+        model_id: str,
+        model_type: ModelType,
+        capabilities: List[ModelCapability],
+        provider: str = "custom",
+        provider_model_name: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Register a model for lifecycle management
+        
+        Args:
+            model_id: Unique identifier for the model
+            model_type: Type of model (LLM, embedding, etc.)
+            capabilities: List of model capabilities
+            provider: Provider name for billing
+            provider_model_name: Provider-specific model name for pricing
+            metadata: Additional metadata
+            
+        Returns:
+            True if registration successful
+        """
+        try:
+            # Prepare metadata with billing info
+            full_metadata = metadata or {}
+            full_metadata.update({
+                "provider": provider,
+                "provider_model_name": provider_model_name or model_id,
+                "registered_for_lifecycle": True,
+                "lifecycle_stage": "registered"
+            })
+            
+            # Register in model registry
+            success = self.registry.register_model(
+                model_id=model_id,
+                model_type=model_type,
+                capabilities=capabilities,
+                metadata=full_metadata
+            )
+            
+            if success:
+                # Track registration operation
+                self.billing_tracker.track_model_usage(
+                    model_id=model_id,
+                    operation_type=ModelOperationType.STORAGE,
+                    provider=provider,
+                    service_type="model_management",
+                    operation="register_model",
+                    metadata={"stage": "registration"}
+                )
+                
+                logger.info(f"Successfully registered model {model_id} for lifecycle management")
+                
+            return success
+            
+        except Exception as e:
+            logger.error(f"Failed to register model {model_id} for lifecycle: {e}")
+            return False
+    
+    def track_model_usage(
+        self,
+        model_id: str,
+        operation_type: ModelOperationType,
+        provider: str,
+        service_type: str,
+        operation: str,
+        input_tokens: Optional[int] = None,
+        output_tokens: Optional[int] = None,
+        input_units: Optional[float] = None,
+        output_units: Optional[float] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Track model usage and costs
+        
+        This method should be called by:
+        - Training services when training a model
+        - Evaluation services when evaluating a model  
+        - Deployment services when deploying a model
+        - Inference services when using a model for inference
+        """
+        return self.billing_tracker.track_model_usage(
+            model_id=model_id,
+            operation_type=operation_type,
+            provider=provider,
+            service_type=service_type,
+            operation=operation,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            input_units=input_units,
+            output_units=output_units,
+            metadata=metadata
+        )
+    
+    async def update_model_stage(
+        self,
+        model_id: str,
+        new_stage: str,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Update model lifecycle stage
+        
+        Args:
+            model_id: Model identifier
+            new_stage: New lifecycle stage (training, evaluation, deployment, production, retired)
+            metadata: Additional metadata for this stage
+            
+        Returns:
+            True if update successful
+        """
+        try:
+            # Get current model info
+            model_info = self.registry.get_model_info(model_id)
+            if not model_info:
+                logger.error(f"Model {model_id} not found in registry")
+                return False
+            
+            # Update metadata with new stage
+            current_metadata = model_info.get("metadata", {})
+            current_metadata.update({
+                "lifecycle_stage": new_stage,
+                "stage_updated_at": str(datetime.now()),
+                **(metadata or {})
+            })
+            
+            # Update in registry
+            success = self.registry.register_model(
+                model_id=model_id,
+                model_type=ModelType(model_info["type"]),
+                capabilities=[ModelCapability(cap) for cap in model_info["capabilities"]],
+                metadata=current_metadata
+            )
+            
+            if success:
+                logger.info(f"Updated model {model_id} to stage: {new_stage}")
+            
+            return success
+            
+        except Exception as e:
+            logger.error(f"Failed to update model {model_id} stage: {e}")
+            return False
+    
+    def get_model_lifecycle_summary(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get complete lifecycle summary for a model including costs
+        
+        Returns:
+            Dictionary with model info, lifecycle stage, and billing summary
+        """
+        try:
+            # Get model info from registry
+            model_info = self.registry.get_model_info(model_id)
+            if not model_info:
+                return None
+            
+            # Get billing summary from tracker
+            billing_summary = self.billing_tracker.get_model_usage_summary(model_id)
+            
+            return {
+                "model_id": model_id,
+                "model_info": model_info,
+                "billing_summary": billing_summary,
+                "current_stage": model_info.get("metadata", {}).get("lifecycle_stage", "unknown")
+            }
+            
+        except Exception as e:
+            logger.error(f"Failed to get lifecycle summary for {model_id}: {e}")
+            return None
+    
+    def list_models_by_stage(self, stage: str) -> List[Dict[str, Any]]:
+        """
+        List all models in a specific lifecycle stage
+        
+        Args:
+            stage: Lifecycle stage to filter by
+            
+        Returns:
+            List of model dictionaries
+        """
+        try:
+            all_models = self.registry.list_models()
+            stage_models = []
+            
+            for model_id, model_info in all_models.items():
+                current_stage = model_info.get("metadata", {}).get("lifecycle_stage")
+                if current_stage == stage:
+                    stage_models.append({
+                        "model_id": model_id,
+                        **model_info
+                    })
+            
+            return stage_models
+            
+        except Exception as e:
+            logger.error(f"Failed to list models by stage {stage}: {e}")
+            return []
+    
+    def get_billing_summary_by_operation(self, operation_type: ModelOperationType) -> Dict[str, Any]:
+        """Get billing summary for a specific operation type"""
+        return self.billing_tracker.get_operation_summary(operation_type)
+    
+    def print_model_costs(self, model_id: str):
+        """Print cost summary for a specific model"""
+        self.billing_tracker.print_model_summary(model_id)

isa_model/core/{storage/supabase_storage.py → models/model_repo.py}

@@ -1,71 +1,93 @@
 """
-Supabase Storage Implementation for Model Registry
+Unified Model Registry with Supabase Backend
 
-Uses Supabase as the backend database for model metadata and capabilities
-Supports the full model lifecycle with cloud-based storage
+Simplified architecture using only Supabase for model metadata and capabilities.
+No SQLite support - uses unified configuration management.
 """
 
 import os
 import json
 import logging
-from typing import Optional, Dict, Any, List
+from typing import Dict, List, Optional, Any
+from enum import Enum
 from datetime import datetime
-from pathlib import Path
 
 try:
     from supabase import create_client, Client
-    from dotenv import load_dotenv
     SUPABASE_AVAILABLE = True
 except ImportError:
     SUPABASE_AVAILABLE = False
 
-from ..model_storage import ModelStorage
+from ..config import ConfigManager
 
 logger = logging.getLogger(__name__)
 
-class SupabaseModelRegistry:
-    """
-    Supabase-based model registry for metadata and capabilities
-    
-    Replaces SQLite with cloud-based Supabase database
-    """
+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"
+    UI_DETECTION = "ui_detection"
+    OCR = "ocr"
+    TABLE_DETECTION = "table_detection"
+    TABLE_STRUCTURE_RECOGNITION = "table_structure_recognition"
+
+class ModelType(str, Enum):
+    """Model types"""
+    LLM = "llm"
+    EMBEDDING = "embedding"
+    RERANK = "rerank"
+    IMAGE = "image"
+    AUDIO = "audio"
+    VIDEO = "video"
+    VISION = "vision"
+
+class ModelRegistry:
+    """Unified Model Registry with Supabase backend"""
     
     def __init__(self):
         if not SUPABASE_AVAILABLE:
             raise ImportError("supabase-py is required. Install with: pip install supabase")
         
-        # Load environment variables
-        load_dotenv()
+        # Get configuration from unified ConfigManager
+        self.config_manager = ConfigManager()
+        global_config = self.config_manager.get_global_config()
         
-        self.supabase_url = os.getenv("SUPABASE_URL")
-        self.supabase_key = os.getenv("SUPABASE_ANON_KEY")
+        # Get Supabase configuration from database config
+        self.supabase_url = global_config.database.supabase_url or os.getenv("SUPABASE_URL")
+        self.supabase_key = global_config.database.supabase_key or os.getenv("SUPABASE_ANON_KEY") or os.getenv("SERVICE_ROLE_KEY")
         
         if not self.supabase_url or not self.supabase_key:
-            raise ValueError("SUPABASE_URL and SUPABASE_ANON_KEY must be set in environment")
+            raise ValueError("SUPABASE_URL and SUPABASE_ANON_KEY (or SERVICE_ROLE_KEY) must be configured")
         
         # Initialize Supabase client
         self.supabase: Client = create_client(self.supabase_url, self.supabase_key)
         
-        # Initialize tables if needed
+        # Verify connection
         self._ensure_tables()
         
-        logger.info("Supabase model registry initialized")
+        logger.info("Model registry initialized with Supabase backend")
     
     def _ensure_tables(self):
         """Ensure required tables exist in Supabase"""
-        # Note: In production, these tables should be created via Supabase migrations
-        # This is just for development/initialization
         try:
             # Check if models table exists by trying to query it
             result = self.supabase.table('models').select('model_id').limit(1).execute()
+            logger.debug("Models table verified")
         except Exception as e:
             logger.warning(f"Models table might not exist: {e}")
-            # In production, you would run proper migrations here
+            # In production, tables should be created via Supabase migrations
     
     def register_model(self, 
                       model_id: str,
-                      model_type: str,
-                      capabilities: List[str],
+                      model_type: ModelType,
+                      capabilities: List[ModelCapability],
                       metadata: Dict[str, Any]) -> bool:
         """Register a model with its capabilities and metadata"""
         try:
@@ -74,7 +96,7 @@ class SupabaseModelRegistry:
             # Prepare model data
             model_data = {
                 'model_id': model_id,
-                'model_type': model_type,
+                'model_type': model_type.value,
                 'metadata': json.dumps(metadata),
                 'created_at': current_time,
                 'updated_at': current_time
@@ -95,7 +117,7 @@ class SupabaseModelRegistry:
                 capability_data = [
                     {
                         'model_id': model_id,
-                        'capability': capability,
+                        'capability': capability.value,
                         'created_at': current_time
                     }
                     for capability in capabilities
@@ -107,7 +129,7 @@ class SupabaseModelRegistry:
                     logger.error(f"Failed to insert capabilities for {model_id}")
                     return False
             
-            logger.info(f"Successfully registered model {model_id}")
+            logger.info(f"Successfully registered model {model_id} with {len(capabilities)} capabilities")
             return True
             
         except Exception as e:
@@ -159,10 +181,10 @@ class SupabaseModelRegistry:
             logger.error(f"Failed to get model info for {model_id}: {e}")
             return None
     
-    def get_models_by_type(self, model_type: str) -> Dict[str, Dict[str, Any]]:
+    def get_models_by_type(self, model_type: ModelType) -> Dict[str, Dict[str, Any]]:
         """Get all models of a specific type"""
         try:
-            models_result = self.supabase.table('models').select('*').eq('model_type', model_type).execute()
+            models_result = self.supabase.table('models').select('*').eq('model_type', model_type.value).execute()
             
             result = {}
             for model in models_result.data:
@@ -186,49 +208,20 @@ class SupabaseModelRegistry:
             logger.error(f"Failed to get models by type {model_type}: {e}")
             return {}
     
-    def get_models_by_capability(self, capability: str) -> Dict[str, Dict[str, Any]]:
+    def get_models_by_capability(self, capability: ModelCapability) -> Dict[str, Dict[str, Any]]:
         """Get all models with a specific capability"""
         try:
-            # Join query to get models with specific capability
-            query = """
-            SELECT DISTINCT m.*, mc.capability
-            FROM models m
-            INNER JOIN model_capabilities mc ON m.model_id = mc.model_id
-            WHERE mc.capability = %s
-            """
-            
-            # Use RPC for complex queries
-            result = self.supabase.rpc('get_models_by_capability', {'capability_name': capability}).execute()
-            
-            if result.data:
-                models_dict = {}
-                for row in result.data:
-                    model_id = row['model_id']
-                    if model_id not in models_dict:
-                        # Get all capabilities for this model
-                        cap_result = self.supabase.table('model_capabilities').select('capability').eq('model_id', model_id).execute()
-                        capabilities = [cap['capability'] for cap in cap_result.data]
-                        
-                        models_dict[model_id] = {
-                            "type": row["model_type"],
-                            "capabilities": capabilities,
-                            "metadata": json.loads(row["metadata"]) if row["metadata"] else {},
-                            "created_at": row["created_at"],
-                            "updated_at": row["updated_at"]
-                        }
-                
-                return models_dict
-            
-            # Fallback: manual join if RPC not available
-            cap_result = self.supabase.table('model_capabilities').select('model_id').eq('capability', capability).execute()
+            # Get model IDs with specific capability
+            cap_result = self.supabase.table('model_capabilities').select('model_id').eq('capability', capability.value).execute()
             model_ids = [row['model_id'] for row in cap_result.data]
             
             if not model_ids:
                 return {}
             
+            # Get model details
             models_result = self.supabase.table('models').select('*').in_('model_id', model_ids).execute()
             
-            result_dict = {}
+            result = {}
             for model in models_result.data:
                 model_id = model["model_id"]
                 
@@ -236,7 +229,7 @@ class SupabaseModelRegistry:
                 all_caps_result = self.supabase.table('model_capabilities').select('capability').eq('model_id', model_id).execute()
                 capabilities = [cap['capability'] for cap in all_caps_result.data]
                 
-                result_dict[model_id] = {
+                result[model_id] = {
                     "type": model["model_type"],
                     "capabilities": capabilities,
                     "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
@@ -244,16 +237,16 @@ class SupabaseModelRegistry:
                     "updated_at": model["updated_at"]
                 }
             
-            return result_dict
+            return result
             
         except Exception as e:
             logger.error(f"Failed to get models by capability {capability}: {e}")
             return {}
     
-    def has_capability(self, model_id: str, capability: str) -> bool:
+    def has_capability(self, model_id: str, capability: ModelCapability) -> bool:
         """Check if a model has a specific capability"""
         try:
-            result = self.supabase.table('model_capabilities').select('model_id').eq('model_id', model_id).eq('capability', capability).execute()
+            result = self.supabase.table('model_capabilities').select('model_id').eq('model_id', model_id).eq('capability', capability.value).execute()
             
             return len(result.data) > 0
             
@@ -295,13 +288,19 @@ class SupabaseModelRegistry:
             total_result = self.supabase.table('models').select('model_id', count='exact').execute()
             total_models = total_result.count if total_result.count is not None else 0
             
-            # Count by type
-            type_result = self.supabase.rpc('get_model_type_counts').execute()
-            type_counts = {row['model_type']: row['count'] for row in type_result.data} if type_result.data else {}
+            # Count by type (manual aggregation since RPC might not exist)
+            models_result = self.supabase.table('models').select('model_type').execute()
+            type_counts = {}
+            for model in models_result.data:
+                model_type = model['model_type']
+                type_counts[model_type] = type_counts.get(model_type, 0) + 1
             
             # Count by capability
-            cap_result = self.supabase.rpc('get_capability_counts').execute()
-            capability_counts = {row['capability']: row['count'] for row in cap_result.data} if cap_result.data else {}
+            caps_result = self.supabase.table('model_capabilities').select('capability').execute()
+            capability_counts = {}
+            for cap in caps_result.data:
+                capability = cap['capability']
+                capability_counts[capability] = capability_counts.get(capability, 0) + 1
             
             return {
                 "total_models": total_models,