175 lines
6.7 KiB
Python
175 lines
6.7 KiB
Python
"""
|
|
Dynamic model registry that collects models from all AI connectors.
|
|
Implements plugin-like architecture for connector discovery.
|
|
"""
|
|
|
|
import logging
|
|
import importlib
|
|
import os
|
|
from typing import Dict, List, Optional, Any
|
|
from modules.datamodels.datamodelAi import AiModel
|
|
from modules.aicore.aicoreBase import BaseConnectorAi
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ModelRegistry:
|
|
"""Dynamic registry for AI models from all connectors."""
|
|
|
|
def __init__(self):
|
|
self._models: Dict[str, AiModel] = {}
|
|
self._connectors: Dict[str, BaseConnectorAi] = {}
|
|
self._last_refresh: Optional[float] = None
|
|
self._refresh_interval: float = 300.0 # 5 minutes
|
|
|
|
def registerConnector(self, connector: BaseConnectorAi):
|
|
"""Register a connector and collect its models."""
|
|
connector_type = connector.getConnectorType()
|
|
self._connectors[connector_type] = connector
|
|
|
|
# Collect models from this connector
|
|
try:
|
|
models = connector.getCachedModels()
|
|
for model in models:
|
|
self._models[model.name] = model
|
|
logger.debug(f"Registered model: {model.name} from {connector_type}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to register models from {connector_type}: {e}")
|
|
|
|
def discoverConnectors(self) -> List[BaseConnectorAi]:
|
|
"""Auto-discover connectors by scanning aicorePlugin*.py files."""
|
|
connectors = []
|
|
connector_dir = os.path.dirname(__file__)
|
|
|
|
# Scan for connector files
|
|
for filename in os.listdir(connector_dir):
|
|
if filename.startswith('aicorePlugin') and filename.endswith('.py'):
|
|
module_name = filename[:-3] # Remove .py extension
|
|
|
|
try:
|
|
# Import the module
|
|
module = importlib.import_module(f'modules.connectors.{module_name}')
|
|
|
|
# Find connector classes (classes that inherit from BaseConnectorAi)
|
|
for attr_name in dir(module):
|
|
attr = getattr(module, attr_name)
|
|
if (isinstance(attr, type) and
|
|
issubclass(attr, BaseConnectorAi) and
|
|
attr != BaseConnectorAi):
|
|
|
|
# Instantiate the connector
|
|
connector = attr()
|
|
connectors.append(connector)
|
|
logger.info(f"Discovered connector: {connector.getConnectorType()}")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to discover connector from {filename}: {e}")
|
|
|
|
return connectors
|
|
|
|
def refreshModels(self, force: bool = False):
|
|
"""Refresh models from all registered connectors."""
|
|
import time
|
|
|
|
current_time = time.time()
|
|
|
|
# Check if refresh is needed
|
|
if (not force and
|
|
self._last_refresh is not None and
|
|
current_time - self._last_refresh < self._refresh_interval):
|
|
return
|
|
|
|
logger.info("Refreshing model registry...")
|
|
|
|
# Clear existing models
|
|
self._models.clear()
|
|
|
|
# Re-register all connectors
|
|
for connector in self._connectors.values():
|
|
try:
|
|
connector.clearCache() # Clear connector cache
|
|
models = connector.getCachedModels()
|
|
for model in models:
|
|
self._models[model.name] = model
|
|
except Exception as e:
|
|
logger.error(f"Failed to refresh models from {connector.getConnectorType()}: {e}")
|
|
|
|
self._last_refresh = current_time
|
|
logger.info(f"Model registry refreshed: {len(self._models)} models available")
|
|
|
|
def getModel(self, name: str) -> Optional[AiModel]:
|
|
"""Get a specific model by name."""
|
|
self.refreshModels()
|
|
return self._models.get(name)
|
|
|
|
def getModels(self) -> List[AiModel]:
|
|
"""Get all available models."""
|
|
self.refreshModels()
|
|
return list(self._models.values())
|
|
|
|
def getModelsByConnector(self, connector_type: str) -> List[AiModel]:
|
|
"""Get models from a specific connector."""
|
|
self.refreshModels()
|
|
return [model for model in self._models.values() if model.connectorType == connector_type]
|
|
|
|
def getModelsByCapability(self, capability: str) -> List[AiModel]:
|
|
"""Get models that support a specific capability."""
|
|
self.refreshModels()
|
|
return [model for model in self._models.values() if capability in model.capabilities]
|
|
|
|
def getModelsByTag(self, tag: str) -> List[AiModel]:
|
|
"""Get models that have a specific tag."""
|
|
self.refreshModels()
|
|
return [model for model in self._models.values() if tag in model.tags]
|
|
|
|
def getAvailableModels(self) -> List[AiModel]:
|
|
"""Get only available models."""
|
|
self.refreshModels()
|
|
return [model for model in self._models.values() if model.isAvailable]
|
|
|
|
def getConnectorForModel(self, model_name: str) -> Optional[BaseConnectorAi]:
|
|
"""Get the connector instance for a specific model."""
|
|
model = self.getModel(model_name)
|
|
if model:
|
|
return self._connectors.get(model.connectorType)
|
|
return None
|
|
|
|
def getModelStats(self) -> Dict[str, Any]:
|
|
"""Get statistics about the model registry."""
|
|
self.refreshModels()
|
|
|
|
stats = {
|
|
"total_models": len(self._models),
|
|
"available_models": len([m for m in self._models.values() if m.isAvailable]),
|
|
"connectors": len(self._connectors),
|
|
"by_connector": {},
|
|
"by_capability": {},
|
|
"by_tag": {}
|
|
}
|
|
|
|
# Count by connector
|
|
for model in self._models.values():
|
|
connector = model.connectorType
|
|
if connector not in stats["by_connector"]:
|
|
stats["by_connector"][connector] = 0
|
|
stats["by_connector"][connector] += 1
|
|
|
|
# Count by capability
|
|
for model in self._models.values():
|
|
for capability in model.capabilities:
|
|
if capability not in stats["by_capability"]:
|
|
stats["by_capability"][capability] = 0
|
|
stats["by_capability"][capability] += 1
|
|
|
|
# Count by tag
|
|
for model in self._models.values():
|
|
for tag in model.tags:
|
|
if tag not in stats["by_tag"]:
|
|
stats["by_tag"][tag] = 0
|
|
stats["by_tag"][tag] += 1
|
|
|
|
return stats
|
|
|
|
|
|
# Global registry instance
|
|
model_registry = ModelRegistry()
|