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

import json
import logging
import random
import sys
import time
from typing import Any, Optional

from aiocache import cached
from fastapi import HTTPException, Request, status
from open_webui.env import BYPASS_MODEL_ACCESS_CONTROL, GLOBAL_LOG_LEVEL, REDIS_KEY_PREFIX
from open_webui.functions import generate_function_chat_completion
from open_webui.models.models import Models
from open_webui.models.users import UserModel
from open_webui.routers.ollama import (
    generate_chat_completion as generate_ollama_chat_completion,
)
from open_webui.routers.openai import (
    generate_chat_completion as generate_openai_chat_completion,
)
from open_webui.routers.pipelines import (
    process_pipeline_inlet_filter,
    process_pipeline_outlet_filter,
)
from open_webui.socket.main import (
    get_event_call,
    get_event_emitter,
)
from open_webui.utils.chat_id import get_temporary_chat_session_id
from open_webui.utils.filter import (
    get_filter_functions,
    process_filter_functions,
)
from open_webui.utils.json_codec import JSONCodec
from open_webui.utils.models import check_model_access, get_all_models, add_tether_model
from open_webui.utils.payload import apply_system_prompt_to_body, convert_chat_completions_payload_to_ollama
from open_webui.utils.response import (
    convert_ollama_response_to_chat_completions,
    convert_ollama_stream_to_chat_completions,
)
from starlette.responses import JSONResponse, Response, StreamingResponse

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


TEMPORARY_CHAT_TTL = 24 * 60 * 60
_temporary_chat_owners: dict[str, tuple[float, dict]] = {}


def _temporary_chat_key(chat_id):
    return f'{REDIS_KEY_PREFIX}:chat:{chat_id}:session'


async def get_temporary_chat(redis, chat_id):
    if redis is not None:
        value = await redis.get(_temporary_chat_key(chat_id))
        return json.loads(value) if value else None
    expiry, owner = _temporary_chat_owners.get(chat_id, (0, None))
    if expiry <= time.monotonic():
        _temporary_chat_owners.pop(chat_id, None)
        return None
    return owner


async def register_temporary_chat_session(redis, chat_id, user_id, session_id, session_owner, *, create=False):
    """Only an authenticated live socket can create or reattach its user's chat."""
    origin = get_temporary_chat_session_id(chat_id)
    if not origin or not session_id or session_owner != user_id:
        return False
    owner = await get_temporary_chat(redis, chat_id)
    if owner is None:
        if not create or origin != session_id:
            return False
    elif owner['user_id'] != user_id:
        return False

    value = {'user_id': user_id, 'session_id': session_id}
    if redis is not None:
        encoded = json.dumps(value, sort_keys=True)
        if owner is None:
            return bool(
                await redis.set(_temporary_chat_key(chat_id), encoded, nx=True, ex=TEMPORARY_CHAT_TTL)
            ) or await register_temporary_chat_session(redis, chat_id, user_id, session_id, session_owner)
        # An expired/replaced record must not be overwritten after the ownership check.
        return bool(
            await redis.eval(
                "if redis.call('get', KEYS[1]) ~= ARGV[1] then return 0 end "
                "redis.call('set', KEYS[1], ARGV[2], 'EX', ARGV[3]); return 1",
                1,
                _temporary_chat_key(chat_id),
                json.dumps(owner, sort_keys=True),
                encoded,
                TEMPORARY_CHAT_TTL,
            )
        )
    now = time.monotonic()
    for key, (expiry, _) in list(_temporary_chat_owners.items()):
        if expiry <= now:
            _temporary_chat_owners.pop(key, None)
    _temporary_chat_owners[chat_id] = (now + TEMPORARY_CHAT_TTL, value)
    return True


# When the question has been asked, let silence not be the
# answer. But if the answer must wait, let it come honest.
async def generate_chat_completion(
    request: Request,
    form_data: dict,
    user: Any,
    bypass_filter: bool = False,
    bypass_system_prompt: bool = False,
):
    log.debug('generate_chat_completion: %s', form_data)
    if BYPASS_MODEL_ACCESS_CONTROL:
        bypass_filter = True

    # Propagate bypass_filter and bypass_system_prompt via request.state so that
    # downstream route handlers (openai/ollama) can read them without exposing
    # them as query parameters.
    request.state.bypass_filter = bypass_filter
    request.state.bypass_system_prompt = bypass_system_prompt

    if hasattr(request.state, 'metadata'):
        if 'metadata' not in form_data:
            form_data['metadata'] = request.state.metadata
        else:
            form_data['metadata'] = {
                **form_data['metadata'],
                **request.state.metadata,
            }

    model_id = form_data['model']
    models = await add_tether_model(request, user, model_id)
    # Single lookup — membership check plus getitem would be two Redis
    # round trips on a Redis-backed model pool.
    model = models.get(model_id)
    if model is None:
        raise Exception('Model not found')

    # Check if user has access to the model
    if not bypass_filter and user.role == 'user':
        try:
            await check_model_access(user, model)
        except Exception as e:
            raise e

    # Arena model — sub-model was already resolved by process_chat_payload.
    # Inject selected_model_id into the response for the frontend.
    metadata = form_data.get('metadata', {})
    selected_model_id = metadata.pop('selected_model_id', None)
    # Also clear from request.state.metadata to prevent the merge at
    # lines 177-179 from re-adding it on the recursive call.
    if hasattr(request.state, 'metadata'):
        request.state.metadata.pop('selected_model_id', None)

    # Fallback: if generate_chat_completion is called with an arena model
    # from a path that did NOT go through process_chat_payload (e.g.,
    # background tasks for title/follow-up/tags generation), resolve now.
    if not selected_model_id and model.get('owned_by') == 'arena':
        model_ids = model.get('info', {}).get('meta', {}).get('model_ids')
        filter_mode = model.get('info', {}).get('meta', {}).get('filter_mode')
        if model_ids and filter_mode == 'exclude':
            model_ids = [
                available_model['id']
                for available_model in list(request.app.state.MODELS.values())
                if available_model.get('owned_by') != 'arena' and available_model['id'] not in model_ids
            ]

        if isinstance(model_ids, list) and model_ids:
            selected_model_id = random.choice(model_ids)
        else:
            model_ids = [
                available_model['id']
                for available_model in list(request.app.state.MODELS.values())
                if available_model.get('owned_by') != 'arena'
            ]
            selected_model_id = random.choice(model_ids)

        form_data['model'] = selected_model_id

        # bypass_filter recursion below skips the line-200 check; gate the resolved model here.
        if not bypass_filter and user.role == 'user':
            selected_model = request.app.state.MODELS.get(selected_model_id)
            if selected_model:
                await check_model_access(user, selected_model)

    if selected_model_id:
        if form_data.get('stream') == True:

            async def stream_wrapper(stream):
                yield f'data: {JSONCodec.dumps({"selected_model_id": selected_model_id})}\n\n'
                async for chunk in stream:
                    yield chunk

            response = await generate_chat_completion(
                request,
                form_data,
                user,
                bypass_filter=True,
                bypass_system_prompt=bypass_system_prompt,
            )
            # Upstream errors come back as a response object.
            if not isinstance(response, StreamingResponse):
                return response
            return StreamingResponse(
                stream_wrapper(response.body_iterator),
                media_type='text/event-stream',
                background=response.background,
            )
        else:
            response = await generate_chat_completion(
                request,
                form_data,
                user,
                bypass_filter=True,
                bypass_system_prompt=bypass_system_prompt,
            )
            if not isinstance(response, dict):
                return response
            return {**response, 'selected_model_id': selected_model_id}

    if model.get('pipe'):
        # Below does not require bypass_filter because this is the only route the uses this function and it is already bypassing the filter
        return await generate_function_chat_completion(request, form_data, user=user, models=models)
    if model.get('owned_by') == 'ollama':
        # Apply before Ollama conversion so tool follow-ups keep the model system prompt
        if not bypass_system_prompt:
            model_info = await Models.get_model_by_id(form_data['model'])
            if model_info:
                system = model_info.params.model_dump().get('system')
                form_data = await apply_system_prompt_to_body(system, form_data, metadata, user)
            request.state.bypass_system_prompt = True

        # Using /ollama/api/chat endpoint
        form_data = convert_chat_completions_payload_to_ollama(form_data)
        response = await generate_ollama_chat_completion(
            request=request,
            form_data=form_data,
            user=user,
        )
        if form_data.get('stream'):
            response.headers['content-type'] = 'text/event-stream'
            return StreamingResponse(
                convert_ollama_stream_to_chat_completions(response),
                headers=dict(response.headers),
                background=response.background,
            )
        else:
            return convert_ollama_response_to_chat_completions(response)
    else:
        return await generate_openai_chat_completion(
            request=request,
            form_data=form_data,
            user=user,
        )


chat_completion = generate_chat_completion


async def chat_completed(request: Request, form_data: dict, user: Any):
    if not request.app.state.MODELS:
        await get_all_models(request, user=user)

    models = await add_tether_model(request, user, form_data.get('model'))

    data = form_data

    if not data.get('id'):
        raise Exception('Missing message id')

    model_id = data['model']
    if model_id not in models:
        raise Exception('Model not found')

    model = models[model_id]

    try:
        data = await process_pipeline_outlet_filter(request, data, user, models)
    except HTTPException:
        raise
    except Exception as e:
        raise Exception(f'Error: {e}')

    if not data.get('id'):
        raise Exception('Missing message id')

    metadata = {
        'chat_id': data['chat_id'],
        'message_id': data['id'],
        'filter_ids': data.get('filter_ids', []),
        'session_id': data['session_id'],
        'user_id': user.id,
    }

    extra_params = {
        '__event_emitter__': await get_event_emitter(metadata),
        '__event_call__': await get_event_call(metadata),
        '__user__': user.model_dump() if isinstance(user, UserModel) else {},
        '__metadata__': metadata,
        '__request__': request,
        '__model__': model,
    }

    try:
        filter_functions = await get_filter_functions(request, model, metadata.get('filter_ids', []))

        result, _ = await process_filter_functions(
            request=request,
            filter_context=None,
            filter_functions=filter_functions,
            filter_type='outlet',
            form_data=data,
            extra_params=extra_params,
        )
        return result
    except Exception as e:
        raise Exception(f'Error: {e}')