Skip to content
Back to top

# Source code for mindroot.coreplugins.admin.service_models

    import asyncio
    from typing import Dict, List, Any, Optional
    from lib.providers.services import service_manager
    service_models_cache: Dict[str, Dict[str, List[str]]] = {}

    async def cached_get_service_models():
        """
        Retrieve service models from the cache or fetch them if not available.
        Returns:
            Dictionary mapping services to providers to models:
            {
                "service_name": {
                    "provider_name": ["model1", "model2", ...],
                    ...
                },
                ...
            }
        """
        global service_models_cache
        if not service_models_cache:
            service_models_cache = await get_service_models_from_providers()
        return service_models_cache

    async def get_service_models_from_providers(timeout: float=500.0, context=None) -> Dict[str, Dict[str, List[str]]]:
        """
        Gather service models from all providers.

        Args:
            timeout: Maximum time in seconds to wait for each provider response

        Returns:
            Dictionary mapping services to providers to models:
            {
                "service_name": {
                    "provider_name": ["model1", "model2", ...],
                    ...
                },
                ...
            }
        """
        providers = []
        if 'get_service_models' in service_manager.functions:
            providers = [provider_info['provider'] for provider_info in service_manager.functions['get_service_models']]

        async def query_provider(provider_name: str) -> tuple[str, Dict[str, List[str]]]:
            try:
                result = await asyncio.wait_for(service_manager.exec_with_provider('get_service_models', provider_name, context=context), timeout=timeout)
                return (provider_name, result)
            except asyncio.TimeoutError:
                return (provider_name, {})
            except Exception as e:
                return (provider_name, {})
        provider_tasks = [query_provider(provider) for provider in providers]
        provider_results = await asyncio.gather(*provider_tasks)
        service_models: Dict[str, Dict[str, List[str]]] = {}
        for provider_name, provider_data in provider_results:
            if not provider_data:
                continue
            for service_name, models in provider_data.items():
                if service_name not in service_models:
                    service_models[service_name] = {}
                service_models[service_name][provider_name] = models
        return service_models