isa-model 0.4.0__py3-none-any.whl → 0.4.4__py3-none-any.whl

isa_model/core/models/model_repo.py

@@ -11,6 +11,7 @@ import logging
 from typing import Dict, List, Optional, Any
 from enum import Enum
 from datetime import datetime
+from dataclasses import dataclass
 
 try:
     from ..database.supabase_client import get_supabase_client, get_supabase_table
@@ -22,6 +23,46 @@ from ..config import ConfigManager
 
 logger = logging.getLogger(__name__)
 
+def _parse_metadata(metadata: Any) -> Dict[str, Any]:
+    """
+    Parse metadata field that might be a dict, JSON string, or None.
+
+    Args:
+        metadata: The metadata value from database (dict, str, or None)
+
+    Returns:
+        Dict with metadata, or empty dict if None
+    """
+    if not metadata:
+        return {}
+
+    # If already a dict (from JSONB column), return as-is
+    if isinstance(metadata, dict):
+        return metadata
+
+    # If string, parse as JSON
+    if isinstance(metadata, str):
+        try:
+            return json.loads(metadata)
+        except json.JSONDecodeError:
+            logger.warning(f"Failed to parse metadata as JSON: {metadata[:100]}")
+            return {}
+
+    # Unknown type
+    logger.warning(f"Unexpected metadata type: {type(metadata)}")
+    return {}
+
+@dataclass
+class RegisteredModel:
+    """Simple model representation for registration results"""
+    model_id: str
+    model_type: str
+    provider: str
+    capabilities: List[str]
+    metadata: Dict[str, Any]
+    created_at: str
+    updated_at: str
+
 class ModelCapability(str, Enum):
     """Model capabilities"""
     TEXT_GENERATION = "text_generation"
@@ -88,7 +129,7 @@ class ModelRegistry:
                       model_type: ModelType,
                       capabilities: List[ModelCapability],
                       metadata: Dict[str, Any],
-                      provider: Optional[str] = None) -> bool:
+                      provider: Optional[str] = None) -> Optional[RegisteredModel]:
         """Register a model with its capabilities and metadata"""
         try:
             current_time = datetime.now().isoformat()
@@ -113,7 +154,7 @@ class ModelRegistry:
             
             if not result.data:
                 logger.error(f"Failed to insert model {model_id}")
-                return False
+                return None
             
             # Delete existing capabilities
             self._table('model_capabilities').delete().eq('model_id', model_id).execute()
@@ -133,14 +174,24 @@ class ModelRegistry:
                 
                 if not cap_result.data:
                     logger.error(f"Failed to insert capabilities for {model_id}")
-                    return False
+                    return None
             
             logger.info(f"Successfully registered model {model_id} with {len(capabilities)} capabilities")
-            return True
+            
+            # Return the registered model object
+            return RegisteredModel(
+                model_id=model_id,
+                model_type=model_type.value,
+                provider=provider_value,
+                capabilities=[cap.value for cap in capabilities],
+                metadata=metadata,
+                created_at=current_time,
+                updated_at=current_time
+            )
             
         except Exception as e:
             logger.error(f"Failed to register model {model_id}: {e}")
-            return False
+            return None
     
     def unregister_model(self, model_id: str) -> bool:
         """Unregister a model"""
@@ -203,13 +254,13 @@ class ModelRegistry:
                 result[model_id] = {
                     "type": model["model_type"],
                     "capabilities": capabilities,
-                    "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                    "metadata": _parse_metadata(model["metadata"]),
                     "created_at": model["created_at"],
                     "updated_at": model["updated_at"]
                 }
-            
+
             return result
-            
+
         except Exception as e:
             logger.error(f"Failed to get models by type {model_type}: {e}")
             return {}
@@ -238,13 +289,13 @@ class ModelRegistry:
                 result[model_id] = {
                     "type": model["model_type"],
                     "capabilities": capabilities,
-                    "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                    "metadata": _parse_metadata(model["metadata"]),
                     "created_at": model["created_at"],
                     "updated_at": model["updated_at"]
                 }
-            
+
             return result
-            
+
         except Exception as e:
             logger.error(f"Failed to get models by capability {capability}: {e}")
             return {}
@@ -276,13 +327,13 @@ class ModelRegistry:
                 result[model_id] = {
                     "type": model["model_type"],
                     "capabilities": capabilities,
-                    "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                    "metadata": _parse_metadata(model["metadata"]),
                     "created_at": model["created_at"],
                     "updated_at": model["updated_at"]
                 }
-            
+
             return result
-            
+
         except Exception as e:
             logger.error(f"Failed to list models: {e}")
             return {}
@@ -337,13 +388,118 @@ class ModelRegistry:
                 result[model_id] = {
                     "type": model["model_type"],
                     "capabilities": capabilities,
-                    "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                    "metadata": _parse_metadata(model["metadata"]),
                     "created_at": model["created_at"],
                     "updated_at": model["updated_at"]
                 }
-            
+
             return result
-            
+
         except Exception as e:
             logger.error(f"Failed to search models with query '{query}': {e}")
-            return {}
+            return {}
+
+
+class ModelRepo:
+    """Compatibility wrapper for unified API"""
+    
+    def __init__(self):
+        self.registry = ModelRegistry()
+    
+    def get_model_by_id(self, model_id: str):
+        """Get model by ID - returns a simple object with basic properties"""
+        info = self.registry.get_model_info(model_id)
+        if not info:
+            return None
+        
+        # Return a simple object that has the expected properties
+        class ModelObject:
+            def __init__(self, data):
+                self.model_id = data["model_id"]
+                self.model_type = data["type"]
+                self.provider = data["metadata"].get("provider", "unknown")
+                self.metadata = data["metadata"]
+                self.capabilities = data["capabilities"]
+                self.created_at = datetime.fromisoformat(data["created_at"]) if data["created_at"] else None
+                self.updated_at = datetime.fromisoformat(data["updated_at"]) if data["updated_at"] else None
+        
+        return ModelObject(info)
+    
+    def update_model_metadata(self, model_id: str, metadata_updates: Dict[str, Any], updated_by: str = None):
+        """Update model metadata - simplified implementation"""
+        # Get existing model
+        info = self.registry.get_model_info(model_id)
+        if not info:
+            return False
+        
+        # Merge metadata
+        new_metadata = info["metadata"].copy()
+        new_metadata.update(metadata_updates)
+        
+        # Update via re-registration (simplified)
+        capabilities = [ModelCapability(cap) for cap in info["capabilities"] if cap in [c.value for c in ModelCapability]]
+        model_type = ModelType(info["type"])
+        
+        return self.registry.register_model(
+            model_id=model_id,
+            model_type=model_type, 
+            capabilities=capabilities,
+            metadata=new_metadata,
+            provider=info["metadata"].get("provider")
+        )
+    
+    def search_models(self, query: str, model_type: str = None, provider: str = None, capabilities: List[str] = None, limit: int = 50):
+        """Search models with filters"""
+        results = self.registry.search_models(query)
+        
+        # Apply filters
+        if model_type:
+            results = {k: v for k, v in results.items() if v["type"] == model_type}
+        
+        if provider:
+            results = {k: v for k, v in results.items() if v["metadata"].get("provider") == provider}
+        
+        if capabilities:
+            results = {k: v for k, v in results.items() if any(cap in v["capabilities"] for cap in capabilities)}
+        
+        # Convert to expected format
+        final_results = []
+        for model_id, data in list(results.items())[:limit]:
+            class ModelObject:
+                def __init__(self, model_id, data):
+                    self.model_id = model_id
+                    self.model_type = data["type"]
+                    self.provider = data["metadata"].get("provider", "unknown")
+                    self.capabilities = data["capabilities"]
+                    self.metadata = data["metadata"]
+                    self.updated_at = datetime.fromisoformat(data["updated_at"]) if data["updated_at"] else None
+            
+            final_results.append(ModelObject(model_id, data))
+        
+        return final_results
+    
+    def get_providers_summary(self):
+        """Get summary of available providers"""
+        models = self.registry.list_models()
+        providers = {}
+        
+        for model_id, data in models.items():
+            provider = data["metadata"].get("provider", "unknown")
+            if provider not in providers:
+                providers[provider] = {
+                    "provider": provider,
+                    "model_count": 0,
+                    "model_types": set(),
+                    "capabilities": set()
+                }
+            
+            providers[provider]["model_count"] += 1
+            providers[provider]["model_types"].add(data["type"])
+            providers[provider]["capabilities"].update(data["capabilities"])
+        
+        # Convert sets to lists for JSON serialization
+        for provider_data in providers.values():
+            provider_data["model_types"] = list(provider_data["model_types"])
+            provider_data["capabilities"] = list(provider_data["capabilities"])
+        
+        return list(providers.values())