open-webui/backend/open_webui/utils/models.py

import asyncio
import copy
import logging
import sys

from aiocache import cached
from fastapi import Request
from open_webui.config import (
    BYPASS_ADMIN_ACCESS_CONTROL,
    DEFAULT_ARENA_MODEL,
)
from open_webui.env import BYPASS_MODEL_ACCESS_CONTROL, ENABLE_PLUGINS, GLOBAL_LOG_LEVEL
from open_webui.functions import get_function_models
from open_webui.models.access_grants import AccessGrants
from open_webui.models.config import Config
from open_webui.models.functions import Functions
from open_webui.models.groups import Groups
from open_webui.models.models import Models
from open_webui.utils.chat_variables import get_chat_variables_schema
from open_webui.models.users import UserModel
from open_webui.routers import ollama, openai
from open_webui.socket.utils import RedisDict
from open_webui.utils.access_control import has_access, has_base_model_access
from open_webui.utils.plugin import (
    get_functions_cache,
    get_function_module_from_cache,
)

logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
log = logging.getLogger(__name__)


async def fetch_ollama_models(request: Request, user: UserModel = None):
    raw_ollama_models = await ollama.get_all_models(request, user=user)
    return [
        {
            'id': model['model'],
            'name': model['name'],
            'object': 'model',
            'created': 0,
            'owned_by': 'ollama',
            'ollama': model,
            'loaded': 'expires_at' in model,
            'connection_type': model.get('connection_type', 'local'),
            'tags': model.get('tags', []),
        }
        for model in raw_ollama_models['models']
    ]


async def fetch_openai_models(request: Request, user: UserModel = None):
    openai_response = await openai.get_all_models(request, user=user)
    return openai_response['data']


async def get_all_base_models(request: Request, user: UserModel = None):
    config = await Config.get_many('openai.enable', 'ollama.enable')
    openai_task = fetch_openai_models(request, user) if config.get('openai.enable') else asyncio.sleep(0, result=[])
    ollama_task = fetch_ollama_models(request, user) if config.get('ollama.enable') else asyncio.sleep(0, result=[])
    function_task = get_function_models(request)

    openai_models, ollama_models, function_models = await asyncio.gather(openai_task, ollama_task, function_task)

    return function_models + openai_models + ollama_models


async def get_all_models(request, refresh: bool = False, user: UserModel = None):
    config = await Config.get_many(
        'models.base_models_cache',
        'evaluation.arena.enable',
        'evaluation.arena.models',
        'models.default_metadata',
    )
    if (
        request.app.state.MODELS
        and request.app.state.BASE_MODELS
        and (config.get('models.base_models_cache') and not refresh)
    ):
        base_models = request.app.state.BASE_MODELS
    else:
        base_models = await get_all_base_models(request, user=user)
        if base_models:
            request.app.state.BASE_MODELS = base_models
        else:
            base_models = request.app.state.BASE_MODELS

    # deep copy the base models to avoid modifying the original list
    models = [model.copy() for model in base_models]

    # If there are no models, return an empty list
    if len(models) == 0:
        return []

    # Add arena models
    if config.get('evaluation.arena.enable'):
        arena_models = []
        arena_config = config.get('evaluation.arena.models') or []
        if len(arena_config) > 0:
            arena_models = [
                {
                    'id': model['id'],
                    'name': model['name'],
                    'info': {
                        'meta': model['meta'],
                    },
                    'object': 'model',
                    'created': 0,
                    'owned_by': 'arena',
                    'arena': True,
                }
                for model in arena_config
            ]
        else:
            # Add default arena model
            arena_models = [
                {
                    'id': DEFAULT_ARENA_MODEL['id'],
                    'name': DEFAULT_ARENA_MODEL['name'],
                    'info': {
                        'meta': DEFAULT_ARENA_MODEL['meta'],
                    },
                    'object': 'model',
                    'created': 0,
                    'owned_by': 'arena',
                    'arena': True,
                }
            ]
        models = models + arena_models

    # One query per type: the global sets are subsets of the active sets, so
    # deriving them from the same rows halves the function-table queries.
    if ENABLE_PLUGINS:
        active_actions = await Functions.get_active_function_ids_by_type('action')
        global_action_ids = {function_id for function_id, is_global in active_actions if is_global}
        enabled_action_ids = {function_id for function_id, _ in active_actions}

        active_filters = await Functions.get_active_function_ids_by_type('filter')
        global_filter_ids = {function_id for function_id, is_global in active_filters if is_global}
        enabled_filter_ids = {function_id for function_id, _ in active_filters}
    else:
        global_action_ids = set()
        enabled_action_ids = set()
        global_filter_ids = set()
        enabled_filter_ids = set()

    custom_models = await Models.get_all_models()

    # Single O(1) lookup: Ollama base names first, then exact IDs (exact wins).
    base_model_lookup = {}
    for model in models:
        if model.get('owned_by') == 'ollama':
            base_model_lookup.setdefault(model['id'].split(':')[0], model)
        base_model_lookup[model['id']] = model

    existing_ids = {m['id'] for m in models}

    for custom_model in custom_models:
        if custom_model.base_model_id is None:
            # Override applied directly to a base model (shares the same ID)
            model = base_model_lookup.get(custom_model.id)

            if model:
                if custom_model.is_active:
                    model['name'] = custom_model.name
                    model['info'] = custom_model.model_dump()
                    schema = get_chat_variables_schema(custom_model.params.model_dump().get('system'))
                    if schema:
                        model['info'].setdefault('meta', {})['chat_variables_schema'] = schema

                    action_ids = []
                    filter_ids = []

                    if 'info' in model:
                        if 'meta' in model['info']:
                            if ENABLE_PLUGINS:
                                action_ids.extend(model['info']['meta'].get('actionIds', []))
                                filter_ids.extend(model['info']['meta'].get('filterIds', []))

                        if 'params' in model['info']:
                            del model['info']['params']

                    model['action_ids'] = action_ids
                    model['filter_ids'] = filter_ids
                else:
                    models = [m for m in models if m is not model]

        elif custom_model.is_active:
            if custom_model.id in existing_ids:
                continue

            owned_by = 'openai'
            connection_type = None
            pipe = None

            base_model = base_model_lookup.get(custom_model.base_model_id)
            if base_model is None:
                base_model = base_model_lookup.get(custom_model.base_model_id.split(':')[0])
            if base_model:
                owned_by = base_model.get('owned_by', 'unknown')
                if 'pipe' in base_model:
                    pipe = base_model['pipe']
                connection_type = base_model.get('connection_type', None)

            model = {
                'id': f'{custom_model.id}',
                'name': custom_model.name,
                'object': 'model',
                'created': custom_model.created_at,
                'owned_by': owned_by,
                'connection_type': connection_type,
                'preset': True,
                **({'pipe': pipe} if pipe is not None else {}),
                **({'provider': base_model.get('provider')} if base_model and base_model.get('provider') else {}),
                **({'loaded': base_model.get('loaded')} if base_model and base_model.get('loaded') is not None else {}),
            }

            info = custom_model.model_dump()
            schema = get_chat_variables_schema(custom_model.params.model_dump().get('system'))
            if schema:
                info.setdefault('meta', {})['chat_variables_schema'] = schema
            if 'params' in info:
                # Remove params to avoid exposing sensitive info
                del info['params']

            model['info'] = info

            action_ids = []
            filter_ids = []

            if custom_model.meta:
                meta = custom_model.meta.model_dump()

                if ENABLE_PLUGINS and 'actionIds' in meta:
                    action_ids.extend(meta['actionIds'])

                if ENABLE_PLUGINS and 'filterIds' in meta:
                    filter_ids.extend(meta['filterIds'])

            model['action_ids'] = action_ids
            model['filter_ids'] = filter_ids

            models.append(model)

    # Process action_ids to get the actions
    def get_action_items_from_module(function, module):
        actions = []
        if hasattr(module, 'actions'):
            actions = module.actions
            return [
                {
                    'id': f'{function.id}.{action["id"]}',
                    'name': action.get('name', f'{function.name} ({action["id"]})'),
                    'description': function.meta.description,
                    'icon': action.get(
                        'icon_url',
                        function.meta.manifest.get('icon_url', None)
                        or getattr(module, 'icon_url', None)
                        or getattr(module, 'icon', None),
                    ),
                }
                for action in actions
            ]
        else:
            return [
                {
                    'id': function.id,
                    'name': function.name,
                    'description': function.meta.description,
                    'icon': function.meta.manifest.get('icon_url', None)
                    or getattr(module, 'icon_url', None)
                    or getattr(module, 'icon', None),
                }
            ]

    # Process filter_ids to get the filters
    def get_filter_items_from_module(function, module):
        return [
            {
                'id': function.id,
                'name': function.name,
                'description': function.meta.description,
                'icon': function.meta.manifest.get('icon_url', None)
                or getattr(module, 'icon_url', None)
                or getattr(module, 'icon', None),
                'has_user_valves': hasattr(module, 'UserValves'),
            }
        ]

    # Batch-prefetch all needed function records to avoid N+1 queries
    all_function_ids = set()
    for model in models:
        all_function_ids.update(model.get('action_ids', []))
        all_function_ids.update(model.get('filter_ids', []))
    all_function_ids.update(global_action_ids)
    all_function_ids.update(global_filter_ids)

    functions_by_id = {f.id: f for f in await Functions.get_functions_by_ids(list(all_function_ids))}

    # Pre-warm the function module cache once per unique function ID.
    # This ensures each function's DB freshness check runs exactly once,
    # not once per (model × function) pair.
    # Only attempt to load functions that actually exist in the local DB;
    # imported/custom model configs may reference tools or filters the user
    # hasn't installed, and trying to load those would cause persistent
    # "Failed to load function module" log spam on every model refresh.
    for function_id, function in functions_by_id.items():
        try:
            await get_function_module_from_cache(request, function_id, function=function)
        except Exception as e:
            log.debug('Failed to load function module for %s: %s', function_id, e)

    # Apply global model defaults to all models
    # Per-model overrides take precedence over global defaults
    default_metadata = config.get('models.default_metadata') or {}

    if default_metadata:
        for model in models:
            info = model.get('info')

            if info is None:
                model['info'] = {'meta': copy.deepcopy(default_metadata)}
                continue

            meta = info.setdefault('meta', {})
            for key, value in default_metadata.items():
                if key == 'capabilities':
                    # Merge capabilities: defaults as base, per-model overrides win
                    existing = meta.get('capabilities') or {}
                    meta['capabilities'] = {**value, **existing}
                elif meta.get(key) is None:
                    meta[key] = copy.deepcopy(value)

    # Batch-fetch all function valves in one query to avoid N+1 DB hits
    # inside get_action_priority (previously called per action × per model).
    all_function_valves = await Functions.get_function_valves_by_ids(list(all_function_ids))
    functions_cache = get_functions_cache(request)

    # Global actions and filters appear in every model, so priorities and item
    # lists are memoized across the loop instead of rebuilt per model.
    action_priorities = {}

    def get_action_priority(action_id):
        if action_id in action_priorities:
            return action_priorities[action_id]
        priority = 0
        try:
            function_module = functions_cache.get(action_id)
            if function_module and hasattr(function_module, 'Valves'):
                valves_db = all_function_valves.get(action_id)
                valves = function_module.Valves(**(valves_db if valves_db else {}))
                priority = getattr(valves, 'priority', 0)
        except Exception:
            priority = 0
        action_priorities[action_id] = priority
        return priority

    action_items_by_id = {}
    filter_items_by_id = {}

    for model in models:
        action_ids = [
            action_id
            for action_id in set(model.pop('action_ids', [])) | global_action_ids
            if action_id in enabled_action_ids
        ]
        action_ids.sort(key=lambda aid: (get_action_priority(aid), aid))

        filter_ids = [
            filter_id
            for filter_id in set(model.pop('filter_ids', [])) | global_filter_ids
            if filter_id in enabled_filter_ids
        ]

        model['actions'] = []
        for action_id in action_ids:
            items = action_items_by_id.get(action_id)
            if items is None:
                action_function = functions_by_id.get(action_id)
                if action_function is None:
                    log.info('Action not found: %s', action_id)
                    action_items_by_id[action_id] = []
                    continue

                function_module = functions_cache.get(action_id)
                if function_module is None:
                    log.info('Failed to load action module: %s', action_id)
                    action_items_by_id[action_id] = []
                    continue
                items = get_action_items_from_module(action_function, function_module)
                action_items_by_id[action_id] = items
            # Shallow copies keep per-model item dicts independent, as before
            model['actions'].extend({**item} for item in items)

        model['filters'] = []
        for filter_id in filter_ids:
            items = filter_items_by_id.get(filter_id)
            if items is None:
                filter_function = functions_by_id.get(filter_id)
                if filter_function is None:
                    log.info('Filter not found: %s', filter_id)
                    filter_items_by_id[filter_id] = []
                    continue

                function_module = functions_cache.get(filter_id)
                if function_module is None:
                    log.info('Failed to load filter module: %s', filter_id)
                    filter_items_by_id[filter_id] = []
                    continue
                if getattr(function_module, 'toggle', None):
                    items = get_filter_items_from_module(filter_function, function_module)
                else:
                    items = []
                filter_items_by_id[filter_id] = items
            model['filters'].extend({**item} for item in items)

    log.debug('get_all_models() returned %s models', len(models))

    models_dict = {model['id']: model for model in models}
    if isinstance(request.app.state.MODELS, RedisDict):
        try:
            request.app.state.MODELS.set(models_dict)
        except Exception as e:
            log.warning(f'Failed to update Redis model cache, using in-process cache: {e}')
            request.app.state.MODELS = models_dict
    else:
        request.app.state.MODELS = models_dict

    return models


async def check_model_access(user, model, model_info=None, db=None, check_base_model=True):
    if model.get('arena'):
        meta = model.get('info', {}).get('meta', {})
        access_grants = meta.get('access_grants', [])
        if not await has_access(
            user.id,
            permission='read',
            access_grants=access_grants,
            db=db,
        ):
            raise Exception('Model not found')
    else:
        # Callers that already fetched the row (chat completion entry) pass it in
        if model_info is None or model_info.id != model.get('id'):
            model_info = await Models.get_model_by_id(model.get('id'), db=db)
        if not model_info:
            raise Exception('Model not found')

        # One group-membership fetch shared by the direct check and every
        # base-model hop; skipped when no check below needs it.
        user_group_ids = None
        if user.id != model_info.user_id or (check_base_model and model_info.base_model_id):
            user_group_ids = {group.id for group in await Groups.get_groups_by_member_id(user.id, db=db)}

        if not (
            user.id == model_info.user_id
            or await AccessGrants.has_access(
                user_id=user.id,
                resource_type='model',
                resource_id=model_info.id,
                permission='read',
                user_group_ids=user_group_ids,
                db=db,
            )
        ):
            raise Exception('Model not found')

        # Enforce access on chained base models. Callers substituting the base
        # model (custom-model fallback) skip this hop and must enforce access
        # on the substitute model instead.
        if check_base_model and not await has_base_model_access(
            user.id, model_info, user_role=user.role, user_group_ids=user_group_ids, db=db
        ):
            raise Exception('Model not found')


async def get_filtered_models(models, user, db=None):
    # Filter out models that the user does not have access to
    if (
        user.role == 'user' or (user.role == 'admin' and not BYPASS_ADMIN_ACCESS_CONTROL)
    ) and not BYPASS_MODEL_ACCESS_CONTROL:
        model_infos = {}
        for model in models:
            if model.get('arena'):
                continue
            info = model.get('info')
            if info:
                model_infos[model['id']] = info

        user_group_ids = {group.id for group in await Groups.get_groups_by_member_id(user.id, db=db)}

        # Batch-fetch accessible resource IDs in a single query instead of N has_access calls
        accessible_model_ids = await AccessGrants.get_accessible_resource_ids(
            user_id=user.id,
            resource_type='model',
            resource_ids=list(model_infos.keys()),
            permission='read',
            user_group_ids=user_group_ids,
            db=db,
        )

        filtered_models = []
        for model in models:
            if model.get('arena'):
                meta = model.get('info', {}).get('meta', {})
                access_grants = meta.get('access_grants', [])
                if await has_access(
                    user.id,
                    permission='read',
                    access_grants=access_grants,
                    user_group_ids=user_group_ids,
                ):
                    filtered_models.append(model)
                continue

            model_info = model_infos.get(model['id'])
            if model_info:
                if (
                    (user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL)
                    or user.id == model_info.get('user_id')
                    or model['id'] in accessible_model_ids
                ):
                    filtered_models.append(model)
            elif user.role == 'admin':
                # No DB entry means no access control configured yet;
                # only admins can see unconfigured models.
                filtered_models.append(model)

        return filtered_models
    else:
        return models