src/models/registry.py

from typing import Dict, List, Optional
import gradio as gr
from .base import BaseModel

class GradioModelInfo:
    """Information about a Gradio model."""
    def __init__(self, 
                 model_id: str,
                 name: str, 
                 description: str,
                 input_type: str,
                 output_type: List[str],
                 examples: List[List[str]],
                 api_path: str):
        self.model_id = model_id
        self.name = name
        self.description = description
        self.input_type = input_type
        self.output_type = output_type
        self.examples = examples
        self.api_path = api_path
        self.status = "unloaded"

class GradioRegistry:
    """Registry for Gradio models."""
    
    def __init__(self):
        self._models: Dict[str, BaseModel] = {}
        self._model_info: Dict[str, GradioModelInfo] = {}
        
    def register_model(self, 
                      model_id: str, 
                      model: BaseModel,
                      api_path: str) -> None:
        """Register a new Gradio model."""
        self._models[model_id] = model
        
        # Create interface to extract information
        interface = model.create_interface()
        
        # Store model information
        self._model_info[model_id] = GradioModelInfo(
            model_id=model_id,
            name=model.name,
            description=model.description,
            input_type=interface.input_components[0].__class__.__name__,
            output_type=[comp.__class__.__name__ for comp in interface.output_components],
            examples=interface.examples or [],
            api_path=api_path
        )
        
    def get_model(self, model_id: str) -> Optional[BaseModel]:
        """Get a model by ID."""
        return self._models.get(model_id)
    
    def get_model_info(self, model_id: str) -> Optional[GradioModelInfo]:
        """Get model information by ID."""
        return self._model_info.get(model_id)
    
    def list_models(self) -> List[Dict]:
        """List all registered models."""
        return [
            {
                "id": info.model_id,
                "name": info.name,
                "description": info.description,
                "input_type": info.input_type,
                "output_type": info.output_type,
                "examples": info.examples,
                "api_path": info.api_path,
                "status": info.status
            }
            for info in self._model_info.values()
        ]
    
    def update_model_status(self, model_id: str, status: str) -> None:
        """Update the status of a model."""
        if model_id in self._model_info:
            self._model_info[model_id].status = status