open-webui/backend/open_webui/routers/functions.py

from __future__ import annotations

import logging
import os
import re
from pathlib import Path
from typing import Optional

import aiohttp
from fastapi import APIRouter, Depends, HTTPException, Request, status
from open_webui.config import CACHE_DIR
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import AIOHTTP_CLIENT_SESSION_SSL, AIOHTTP_CLIENT_TIMEOUT, ENABLE_FUNCTIONS
from open_webui.events import EVENTS, build_event, dispatch_event_functions, publish_event, schedule_webhook_dispatch
from open_webui.internal.db import get_async_session
from open_webui.models.function_history import FunctionHistories, function_diff
from open_webui.models.functions import (
    FunctionForm,
    FunctionModel,
    FunctionResponse,
    Functions,
    FunctionUserResponse,
    FunctionWithValvesModel,
)
from open_webui.utils.auth import get_admin_user, get_verified_user
from open_webui.utils.plugin import (
    get_function_contents_cache,
    get_function_module_from_cache,
    get_functions_cache,
    load_function_module_by_id,
    replace_imports,
    resolve_valves_schema_options,
    set_function_module_in_cache,
)
from pydantic import BaseModel, HttpUrl
from sqlalchemy.ext.asyncio import AsyncSession

log = logging.getLogger(__name__)


router = APIRouter()

############################
# GetFunctions
# Our daily functions give us, and forgive us
# our deprecated methods, as we refactor those who depend on us.
############################


@router.get('/', response_model=list[FunctionResponse])
async def get_functions(user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)):
    if not ENABLE_FUNCTIONS:
        return []

    return await Functions.get_functions(db=db)


@router.get('/list', response_model=list[FunctionUserResponse])
async def get_function_list(user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
    if not ENABLE_FUNCTIONS:
        return []

    return await Functions.get_function_list(db=db)


############################
# ExportFunctions
############################


@router.get('/export', response_model=list[FunctionModel | FunctionWithValvesModel])
async def get_functions(
    include_valves: bool = False,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    if not ENABLE_FUNCTIONS:
        return []

    return await Functions.get_functions(include_valves=include_valves, db=db)


############################
# LoadFunctionFromLink
############################


class LoadUrlForm(BaseModel):
    url: HttpUrl


def github_url_to_raw_url(url: str) -> str:
    # Handle 'tree' (folder) URLs (add main.py at the end)
    m1 = re.match(r'https://github\.com/([^/]+)/([^/]+)/tree/([^/]+)/(.*)', url)
    if m1:
        org, repo, branch, path = m1.groups()
        return f'https://raw.githubusercontent.com/{org}/{repo}/refs/heads/{branch}/{path.rstrip("/")}/main.py'

    # Handle 'blob' (file) URLs
    m2 = re.match(r'https://github\.com/([^/]+)/([^/]+)/blob/([^/]+)/(.*)', url)
    if m2:
        org, repo, branch, path = m2.groups()
        return f'https://raw.githubusercontent.com/{org}/{repo}/refs/heads/{branch}/{path}'

    # No match; return as-is
    return url


@router.post('/load/url', response_model=dict | None)
async def load_function_from_url(request: Request, form_data: LoadUrlForm, user=Depends(get_admin_user)):
    # NOTE: This is NOT a SSRF vulnerability:
    # This endpoint is admin-only (see get_admin_user), meant for *trusted* internal use,
    # and does NOT accept untrusted user input. Access is enforced by authentication.

    url = str(form_data.url)
    if not url:
        raise HTTPException(status_code=400, detail='Please enter a valid URL')

    url = github_url_to_raw_url(url)
    url_parts = url.rstrip('/').split('/')

    file_name = url_parts[-1]
    function_name = (
        file_name[:-3]
        if (file_name.endswith('.py') and (not file_name.startswith(('main.py', 'index.py', '__init__.py'))))
        else url_parts[-2]
        if len(url_parts) > 1
        else 'function'
    )

    try:
        async with aiohttp.ClientSession(
            trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
        ) as session:
            async with session.get(
                url, headers={'Content-Type': 'application/json'}, ssl=AIOHTTP_CLIENT_SESSION_SSL
            ) as resp:
                if resp.status != 200:
                    raise HTTPException(status_code=resp.status, detail='Failed to fetch the function')
                data = await resp.text()
                if not data:
                    raise HTTPException(status_code=400, detail='No data received from the URL')
        return {
            'name': function_name,
            'content': data,
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=ERROR_MESSAGES.DEFAULT(e, 'Error fetching function'),
        )


############################
# SyncFunctions
############################


class SyncFunctionsForm(BaseModel):
    functions: list[FunctionWithValvesModel] = []


@router.post('/sync', response_model=list[FunctionWithValvesModel])
async def sync_functions(
    request: Request,
    form_data: SyncFunctionsForm,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    try:
        modules = {}
        source_modules = {}
        previous_ids = {entry.id for entry in await Functions.get_functions(db=db)}
        for function in form_data.functions:
            function.content = replace_imports(function.content)
            module, function.type, frontmatter, source_module = await load_function_module_by_id(
                function.id, content=function.content
            )
            function.meta.manifest = frontmatter
            function.meta.toggle = function.type == 'filter' and bool(getattr(module, 'toggle', False))
            modules[function.id] = module
            source_modules[function.id] = source_module
        result = await Functions.sync_functions(user.id, form_data.functions, db=db, modules=modules)
        for function in result:
            set_function_module_in_cache(
                request, function.id, function.content, modules[function.id], source_modules[function.id]
            )
        for id in previous_ids - {entry.id for entry in result}:
            get_functions_cache(request).pop(id, None)
            get_function_contents_cache(request).pop(id, None)
        return result
    except Exception as e:
        log.exception(f'Failed to load a function: {e}')
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.DEFAULT(e, 'Error loading function'),
        )


############################
# CreateNewFunction
############################


@router.post('/create', response_model=FunctionResponse | None)
async def create_new_function(
    request: Request,
    form_data: FunctionForm,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    if not form_data.id.isidentifier():
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail='Only alphanumeric characters and underscores are allowed in the id',
        )

    form_data.id = form_data.id.lower()

    function = await Functions.get_function_by_id(form_data.id, db=db)
    if function is None:
        try:
            form_data.content = replace_imports(form_data.content)
            function_module, function_type, frontmatter, source_module = await load_function_module_by_id(
                form_data.id,
                content=form_data.content,
            )
            form_data.meta.manifest = frontmatter
            form_data.meta.toggle = function_type == 'filter' and bool(getattr(function_module, 'toggle', False))

            function = await Functions.insert_new_function(
                user.id, function_type, form_data, db=db, module=function_module
            )

            function_cache_dir = CACHE_DIR / 'functions' / form_data.id
            function_cache_dir.mkdir(parents=True, exist_ok=True)

            if function:
                set_function_module_in_cache(request, function.id, function.content, function_module, source_module)
                await publish_event(
                    request,
                    EVENTS.FUNCTION_CREATED,
                    actor=user,
                    subject_id=function.id,
                    data={'type': function.type, 'name': function.name},
                )
                return function
            else:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=ERROR_MESSAGES.DEFAULT('Error creating function'),
                )
        except HTTPException:
            raise
        except Exception as e:
            log.exception(f'Failed to create a new function: {e}')
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e, 'Error creating function'),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=ERROR_MESSAGES.ID_TAKEN,
        )


############################
# GetFunctionById
############################


@router.get('/id/{id}', response_model=FunctionModel | None)
async def get_function_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)):
    function = await Functions.get_function_by_id(id, db=db)

    if function:
        return function
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# ToggleFunctionById
############################


@router.post('/id/{id}/toggle', response_model=FunctionModel | None)
async def toggle_function_by_id(
    request: Request,
    id: str,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        lifecycle_event = build_event(
            request,
            EVENTS.FUNCTION_DISABLE_STARTED if function.is_active else EVENTS.FUNCTION_ENABLE_STARTED,
            actor=user,
            subject_id=function.id,
            subject_type='function',
            data={'type': function.type, 'name': function.name},
        )
        await dispatch_event_functions(
            request.app,
            lifecycle_event,
            request=request,
            extra_function_ids=[function.id] if not function.is_active else None,
        )
        schedule_webhook_dispatch(request.app, lifecycle_event)

        function = await Functions.update_function_by_id(id, {'is_active': not function.is_active}, db=db)

        if function:
            await publish_event(
                request,
                EVENTS.FUNCTION_ENABLED if function.is_active else EVENTS.FUNCTION_DISABLED,
                actor=user,
                subject_id=function.id,
                subject_type='function',
                data={'type': function.type, 'name': function.name},
            )
            return function
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT('Error updating function'),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# ToggleGlobalById
############################


@router.post('/id/{id}/toggle/global', response_model=FunctionModel | None)
async def toggle_global_by_id(
    request: Request,
    id: str,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        function = await Functions.update_function_by_id(id, {'is_global': not function.is_global}, db=db)

        if function:
            await publish_event(
                request,
                EVENTS.FUNCTION_UPDATED,
                actor=user,
                subject_id=function.id,
                data={'type': function.type, 'name': function.name, 'is_global': function.is_global},
            )
            return function
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT('Error updating function'),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# UpdateFunctionById
############################


@router.post('/id/{id}/update', response_model=FunctionModel | None)
async def update_function_by_id(
    request: Request,
    id: str,
    form_data: FunctionForm,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    return await _update_function(request, id, form_data, user, db)


async def _update_function(request, id, form_data, user, db, version_id=None):
    try:
        if version_id is None:
            form_data.content = replace_imports(form_data.content)
        function_module, function_type, frontmatter, source_module = await load_function_module_by_id(
            id, content=form_data.content
        )
        form_data.meta.manifest = frontmatter

        form_data.meta.toggle = function_type == 'filter' and bool(getattr(function_module, 'toggle', False))

        updated = {**form_data.model_dump(exclude={'id'}), 'type': function_type}

        function = await Functions.update_function_by_id(
            id, updated, db=db, user_id=user.id, version_id=version_id, module=function_module
        )

        if function:
            set_function_module_in_cache(request, function.id, function.content, function_module, source_module)
            await publish_event(
                request,
                EVENTS.FUNCTION_UPDATED,
                actor=user,
                subject_id=function.id,
                data={'type': function.type, 'name': function.name},
            )
            return function
        else:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT('Error updating function'),
            )

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=str(e),
        )


############################
# DeleteFunctionById
############################


@router.delete('/id/{id}/delete', response_model=bool)
async def delete_function_by_id(
    request: Request,
    id: str,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    result = await Functions.delete_function_by_id(id, db=db)

    if result:
        FUNCTIONS = get_functions_cache(request)
        FUNCTIONS.pop(id, None)
        FUNCTION_CONTENTS = get_function_contents_cache(request)
        FUNCTION_CONTENTS.pop(id, None)
        await publish_event(
            request,
            EVENTS.FUNCTION_DELETED,
            actor=user,
            subject_id=id,
        )

    return result


############################
# GetFunctionValves
############################


@router.get('/id/{id}/valves', response_model=dict | None)
async def get_function_valves_by_id(
    id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        try:
            valves = await Functions.get_function_valves_by_id(id, db=db)
            return valves
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e, 'Error getting function valves'),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# GetFunctionValvesSpec
############################


@router.get('/id/{id}/valves/spec', response_model=dict | None)
async def get_function_valves_spec_by_id(
    request: Request,
    id: str,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)

        if hasattr(function_module, 'Valves'):
            Valves = function_module.Valves
            schema = Valves.schema()
            # Resolve dynamic options for select dropdowns
            schema = resolve_valves_schema_options(Valves, schema, user)
            return schema
        return None
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# UpdateFunctionValves
############################


@router.post('/id/{id}/valves/update', response_model=dict | None)
async def update_function_valves_by_id(
    request: Request,
    id: str,
    form_data: dict,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)

        if hasattr(function_module, 'Valves'):
            Valves = function_module.Valves

            try:
                form_data = {k: v for k, v in form_data.items() if v is not None}
                valves = Valves(**form_data)

                valves_dict = valves.model_dump(exclude_unset=True)
                await Functions.update_function_valves_by_id(id, valves_dict, db=db)
                await publish_event(
                    request,
                    EVENTS.FUNCTION_VALVES_UPDATED,
                    actor=user,
                    subject_id=id,
                )
                return valves_dict
            except Exception as e:
                log.exception(f'Error updating function values by id {id}: {e}')
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=ERROR_MESSAGES.DEFAULT(e, 'Error updating function valves'),
                )
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )

    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


############################
# FunctionUserValves
############################


@router.get('/id/{id}/valves/user', response_model=dict | None)
async def get_function_user_valves_by_id(
    id: str, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session)
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        try:
            user_valves = await Functions.get_user_valves_by_id_and_user_id(id, user.id, db=db)
            return user_valves
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=ERROR_MESSAGES.DEFAULT(e, 'Error getting function user valves'),
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


@router.get('/id/{id}/valves/user/spec', response_model=dict | None)
async def get_function_user_valves_spec_by_id(
    request: Request,
    id: str,
    user=Depends(get_verified_user),
    db: AsyncSession = Depends(get_async_session),
):
    function = await Functions.get_function_by_id(id, db=db)
    if function:
        if not function.is_active:
            return None

        function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)

        if hasattr(function_module, 'UserValves'):
            UserValves = function_module.UserValves
            schema = UserValves.schema()
            # Resolve dynamic options for select dropdowns
            schema = resolve_valves_schema_options(UserValves, schema, user)
            return schema
        return None
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


@router.post('/id/{id}/valves/user/update', response_model=dict | None)
async def update_function_user_valves_by_id(
    request: Request,
    id: str,
    form_data: dict,
    user=Depends(get_verified_user),
    db: AsyncSession = Depends(get_async_session),
):
    function = await Functions.get_function_by_id(id, db=db)

    if function:
        if not function.is_active:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail='Function is not active',
            )

        function_module, function_type, frontmatter = await get_function_module_from_cache(request, id)

        if hasattr(function_module, 'UserValves'):
            UserValves = function_module.UserValves

            try:
                form_data = {k: v for k, v in form_data.items() if v is not None}
                user_valves = UserValves(**form_data)
                user_valves_dict = user_valves.model_dump(exclude_unset=True)
                await Functions.update_user_valves_by_id_and_user_id(id, user.id, user_valves_dict, db=db)
                await publish_event(
                    request,
                    EVENTS.FUNCTION_VALVES_UPDATED,
                    actor=user,
                    subject_id=id,
                    data={'scope': 'user'},
                )
                return user_valves_dict
            except Exception as e:
                log.exception(f'Error updating function user valves by id {id}: {e}')
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=ERROR_MESSAGES.DEFAULT(e, 'Error updating function user valves'),
                )
        else:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail=ERROR_MESSAGES.NOT_FOUND,
            )
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=ERROR_MESSAGES.NOT_FOUND,
        )


async def require_function_history_access(id, user, db):
    resource = await Functions.get_function_by_id(id, db=db)
    if not resource:
        raise HTTPException(404, 'Not found')
    return resource


async def require_function_history_entry(id, history_id, db):
    entry = await FunctionHistories.get_history_by_id(id, history_id, db=db)
    if not entry:
        raise HTTPException(404, 'Version not found')
    return entry


@router.get('/id/{id}/history')
async def get_function_history(
    id: str, page: int = 1, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)
):
    await require_function_history_access(id, user, db)
    return await FunctionHistories.get_history_by_function_id(id, page, db=db)


@router.get('/id/{id}/history/diff')
async def get_function_history_diff(
    id: str, from_id: str, to_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)
):
    await require_function_history_access(id, user, db)
    before = await require_function_history_entry(id, from_id, db)
    after = await require_function_history_entry(id, to_id, db)
    return function_diff(before, after)


@router.get('/id/{id}/history/{history_id}')
async def get_function_history_entry(
    id: str, history_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)
):
    await require_function_history_access(id, user, db)
    return await require_function_history_entry(id, history_id, db)


@router.delete('/id/{id}/history/{history_id}')
async def delete_function_history_entry(
    id: str, history_id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)
):
    await require_function_history_access(id, user, db)
    if not await FunctionHistories.delete_history_entry(id, history_id, db=db):
        raise HTTPException(404, 'Version not found')
    return True


class FunctionVersionForm(BaseModel):
    version_id: str


@router.post('/id/{id}/update/version', response_model=FunctionModel)
async def set_function_production(
    request: Request,
    id: str,
    form_data: FunctionVersionForm,
    user=Depends(get_admin_user),
    db: AsyncSession = Depends(get_async_session),
):
    await require_function_history_access(id, user, db)
    entry = await require_function_history_entry(id, form_data.version_id, db)
    try:
        saved = FunctionForm(id=id, **entry.snapshot)
    except ValueError as error:
        raise HTTPException(400, str(error)) from error
    return await _update_function(request, id, saved, user, db, version_id=entry.id)