From 76aae64c7b55672a1e9bc0f159ac14134194b7bf Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 27 Jul 2026 02:34:04 -0400 Subject: [PATCH] refac --- backend/open_webui/utils/chat.py | 1 + backend/open_webui/utils/filter.py | 241 +++++++++++++++++-------- backend/open_webui/utils/middleware.py | 22 ++- 3 files changed, 184 insertions(+), 80 deletions(-) diff --git a/backend/open_webui/utils/chat.py b/backend/open_webui/utils/chat.py index 3cc3000100..fda6320749 100644 --- a/backend/open_webui/utils/chat.py +++ b/backend/open_webui/utils/chat.py @@ -365,6 +365,7 @@ async def chat_completed(request: Request, form_data: dict, user: Any): result, _ = await process_filter_functions( request=request, + filter_context=None, filter_functions=filter_functions, filter_type='outlet', form_data=data, diff --git a/backend/open_webui/utils/filter.py b/backend/open_webui/utils/filter.py index f73278e7e2..b3ca4a4229 100644 --- a/backend/open_webui/utils/filter.py +++ b/backend/open_webui/utils/filter.py @@ -8,27 +8,68 @@ from open_webui.utils.plugin import get_function_module_from_cache log = logging.getLogger(__name__) -async def get_function_module(request, function_id, load_from_db=True): +class FilterContext: + def __init__(self): + self.valves_by_id = None + self.function_valves = {} + self.user_valves = {} + self.signatures = {} + + async def get_function_valves(self, filter_ids, filter_id, Valves): + if filter_id not in self.function_valves: + if self.valves_by_id is None: + self.valves_by_id = await Functions.get_function_valves_by_ids(filter_ids) + valves = self.valves_by_id.get(filter_id) + self.function_valves[filter_id] = Valves(**(valves if valves else {})) + return self.function_valves[filter_id] + + def get_signature(self, handler): + signature_key = getattr(handler, '__func__', handler) + if signature_key not in self.signatures: + self.signatures[signature_key] = inspect.signature(handler) + return self.signatures[signature_key] + + async def get_user_valves(self, filter_id, user_id, UserValves): + user_valves_key = (filter_id, user_id) + if user_valves_key not in self.user_valves: + self.user_valves[user_valves_key] = await get_user_valves(filter_id, user_id, UserValves) + return self.user_valves[user_valves_key] + + +async def get_user_valves(filter_id, user_id, UserValves): + user_valves_data = await Functions.get_user_valves_by_id_and_user_id(filter_id, user_id) + return UserValves(**(user_valves_data if user_valves_data else {})) + + +async def get_function_module(request, function_id, load_from_db=True, function=None): """ Get the function module by its ID. """ - function_module, _, _ = await get_function_module_from_cache(request, function_id, load_from_db=load_from_db) + function_module, _, _ = await get_function_module_from_cache( + request, function_id, function=function, load_from_db=load_from_db + ) return function_module -async def get_sorted_filter_ids(request, model: dict, enabled_filter_ids: list = None): - if not ENABLE_PLUGINS: - return [] - - active_filters = await Functions.get_active_filter_ids() +def get_model_filter_ids(model, active_filters): filter_ids = [fid for fid, is_global in active_filters if is_global] - if 'info' in model and 'meta' in model['info']: + if isinstance(model, dict) and 'info' in model and 'meta' in model['info']: filter_ids.extend(model['info']['meta'].get('filterIds', [])) filter_ids = list(set(filter_ids)) active_filter_ids = {fid for fid, _ in active_filters} + return [fid for fid in filter_ids if fid in active_filter_ids] + + +async def resolve_filter_pipeline(request, model: dict, enabled_filter_ids: list = None): + if not ENABLE_PLUGINS: + return [], [] + + active_filters = await Functions.get_active_filter_ids() + filter_ids = get_model_filter_ids(model, active_filters) + functions_by_id = {function.id: function for function in await Functions.get_functions_by_ids(filter_ids)} async def get_active_status(filter_id): - function_module = await get_function_module(request, filter_id) + function_module = await get_function_module(request, filter_id, function=functions_by_id.get(filter_id)) if getattr(function_module, 'toggle', None): return filter_id in (enabled_filter_ids or set()) @@ -36,17 +77,12 @@ async def get_sorted_filter_ids(request, model: dict, enabled_filter_ids: list = return True # Pre-compute active status for each filter (async functions can't be used in set comprehensions) - resolved_active = {} - for filter_id in active_filter_ids: - resolved_active[filter_id] = await get_active_status(filter_id) - active_filter_ids = {fid for fid, is_active in resolved_active.items() if is_active} - - filter_ids = [fid for fid in filter_ids if fid in active_filter_ids] + filter_ids = [fid for fid in filter_ids if await get_active_status(fid)] valves_by_id = await Functions.get_function_valves_by_ids(filter_ids) async def get_priority(function_id): try: - function_module = await get_function_module(request, function_id) + function_module = await get_function_module(request, function_id, function=functions_by_id.get(function_id)) if function_module and hasattr(function_module, 'Valves'): valves_db = valves_by_id.get(function_id) valves = function_module.Valves(**(valves_db if valves_db else {})) @@ -61,12 +97,118 @@ async def get_sorted_filter_ids(request, model: dict, enabled_filter_ids: list = priorities[fid] = await get_priority(fid) filter_ids.sort(key=lambda fid: (priorities.get(fid, 0), fid)) + filter_functions = [functions_by_id[fid] for fid in filter_ids if fid in functions_by_id] + return filter_ids, filter_functions + + +async def get_sorted_filter_ids(request, model: dict, enabled_filter_ids: list = None): + filter_ids, _ = await resolve_filter_pipeline(request, model, enabled_filter_ids) + return filter_ids +async def get_filter_functions(request, model: dict, enabled_filter_ids: list = None): + _, filter_functions = await resolve_filter_pipeline(request, model, enabled_filter_ids) + return filter_functions + + +async def apply_filter_valves(function_module, filter_context, valves_by_id, filter_ids, filter_id): + if not (hasattr(function_module, 'valves') and hasattr(function_module, 'Valves')): + return valves_by_id + + if filter_context is not None: + function_module.valves = await filter_context.get_function_valves(filter_ids, filter_id, function_module.Valves) + return valves_by_id + + if valves_by_id is None: + valves_by_id = await Functions.get_function_valves_by_ids(filter_ids) + valves = valves_by_id.get(filter_id) + function_module.valves = function_module.Valves(**(valves if valves else {})) + return valves_by_id + + +def get_filter_params(sig, filter_id, filter_type, form_data, extra_params): + params = {'event': form_data} if filter_type == 'stream' else {'body': form_data} + return params | { + k: v + for k, v in { + **extra_params, + '__id__': filter_id, + }.items() + if k in sig.parameters + } + + +async def apply_user_valves(function_module, filter_context, filter_id, params): + if '__user__' not in params or not hasattr(function_module, 'UserValves'): + return + + user_id = params['__user__'].get('id') + if filter_context is not None: + user_valves = await filter_context.get_user_valves(filter_id, user_id, function_module.UserValves) + else: + user_valves = await get_user_valves(filter_id, user_id, function_module.UserValves) + params['__user__']['valves'] = user_valves + + +async def run_filter_handler(handler, params): + if inspect.iscoroutinefunction(handler): + return await handler(**params) + return handler(**params) + + +async def process_filter_function( + request, + function, + filter_type, + form_data, + extra_params, + filter_context, + valves_by_id, + filter_ids, +): + filter_id = function.id + + function_module = await get_function_module( + request, filter_id, load_from_db=(filter_type != 'stream'), function=function + ) + handler = getattr(function_module, filter_type, None) + if not handler: + return form_data, valves_by_id, None + + skip_files = ( + function_module.file_handler if filter_type == 'inlet' and hasattr(function_module, 'file_handler') else None + ) + valves_by_id = await apply_filter_valves(function_module, filter_context, valves_by_id, filter_ids, filter_id) + + try: + sig = filter_context.get_signature(handler) if filter_context is not None else inspect.signature(handler) + params = get_filter_params(sig, filter_id, filter_type, form_data, extra_params) + + if '__user__' in sig.parameters: + try: + await apply_user_valves(function_module, filter_context, filter_id, params) + except Exception as e: + log.exception(f'Failed to get user values: {e}') + + form_data = await run_filter_handler(handler, params) + except Exception as e: + log.debug(f'Error in {filter_type} handler {filter_id}: {e}') + raise e + + return form_data, valves_by_id, skip_files + + # Grant these filters the discernment to pass what serves # and refuse what harms, for every soul in the house. -async def process_filter_functions(request, filter_functions, filter_type, form_data, extra_params): +async def process_filter_functions( + request, + filter_context, + filter_functions, + filter_type, + form_data, + extra_params, +): if not ENABLE_PLUGINS: return form_data, {} @@ -77,61 +219,18 @@ async def process_filter_functions(request, filter_functions, filter_type, form_ for function in filter_functions: if not function: continue - filter_id = function.id - function_module = await get_function_module(request, filter_id, load_from_db=(filter_type != 'stream')) - # Prepare handler function - handler = getattr(function_module, filter_type, None) - if not handler: - continue - - # Check if the function has a file_handler variable - if filter_type == 'inlet' and hasattr(function_module, 'file_handler'): - skip_files = function_module.file_handler - - # Apply valves to the function - if hasattr(function_module, 'valves') and hasattr(function_module, 'Valves'): - if valves_by_id is None: - valves_by_id = await Functions.get_function_valves_by_ids(filter_ids) - valves = valves_by_id.get(filter_id) - function_module.valves = function_module.Valves(**(valves if valves else {})) - - try: - # Prepare parameters - sig = inspect.signature(handler) - - params = {'body': form_data} - if filter_type == 'stream': - params = {'event': form_data} - - params = params | { - k: v - for k, v in { - **extra_params, - '__id__': filter_id, - }.items() - if k in sig.parameters - } - - # Handle user parameters - if '__user__' in sig.parameters: - if hasattr(function_module, 'UserValves'): - try: - params['__user__']['valves'] = function_module.UserValves( - **await Functions.get_user_valves_by_id_and_user_id(filter_id, params['__user__']['id']) - ) - except Exception as e: - log.exception(f'Failed to get user values: {e}') - - # Execute handler - if inspect.iscoroutinefunction(handler): - form_data = await handler(**params) - else: - form_data = handler(**params) - - except Exception as e: - log.debug(f'Error in {filter_type} handler {filter_id}: {e}') - raise e + form_data, valves_by_id, file_handler = await process_filter_function( + request, + function, + filter_type, + form_data, + extra_params, + filter_context, + valves_by_id, + filter_ids, + ) + skip_files = skip_files or file_handler # Handle file cleanup for inlet if skip_files: diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index ab92882acc..03eb964a48 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -43,7 +43,6 @@ from open_webui.env import ( from open_webui.models.chats import Chats from open_webui.models.config import Config from open_webui.models.folders import Folders -from open_webui.models.functions import Functions from open_webui.models.models import Models from open_webui.models.notes import Notes from open_webui.models.oauth_sessions import OAuthSessions @@ -90,7 +89,8 @@ from open_webui.utils.files import ( get_image_url_from_base64, ) from open_webui.utils.filter import ( - get_sorted_filter_ids, + FilterContext, + get_filter_functions, process_filter_functions, ) @@ -2496,11 +2496,11 @@ async def process_chat_payload(request, form_data, user, metadata, model): if ENABLE_PLUGINS: try: - filter_ids = await get_sorted_filter_ids(request, model, metadata.get('filter_ids', [])) - filter_functions = await Functions.get_functions_by_ids(filter_ids) + filter_functions = await get_filter_functions(request, model, metadata.get('filter_ids', [])) form_data, flags = await process_filter_functions( request=request, + filter_context=None, filter_functions=filter_functions, filter_type='inlet', form_data=form_data, @@ -3501,11 +3501,11 @@ async def outlet_filter_handler(ctx): } if ENABLE_PLUGINS: - filter_ids = await get_sorted_filter_ids(request, model, metadata.get('filter_ids', [])) - filter_functions = await Functions.get_functions_by_ids(filter_ids) + filter_functions = await get_filter_functions(request, model, metadata.get('filter_ids', [])) outlet_result, _ = await process_filter_functions( request=request, + filter_context=None, filter_functions=filter_functions, filter_type='outlet', form_data=outlet_data, @@ -3763,9 +3763,7 @@ async def streaming_chat_response_handler(response, ctx): } filter_functions = ( - await Functions.get_functions_by_ids( - await get_sorted_filter_ids(request, model, metadata.get('filter_ids', [])) - ) + await get_filter_functions(request, model, metadata.get('filter_ids', [])) if ENABLE_PLUGINS else [] ) @@ -3779,6 +3777,8 @@ async def streaming_chat_response_handler(response, ctx): # Handle as a background task async def response_handler(response, events): + filter_context = FilterContext() + def tag_output_handler(content_type, tags, output): """ Detect special tags (reasoning, solution, code_interpreter) in streaming @@ -4186,6 +4186,7 @@ async def streaming_chat_response_handler(response, ctx): data, _ = await process_filter_functions( request=request, + filter_context=filter_context, filter_functions=filter_functions, filter_type='stream', form_data=data, @@ -5515,6 +5516,7 @@ async def streaming_chat_response_handler(response, ctx): return f'data: {item}\n\n' assistant_message = {} + filter_context = FilterContext() has_api_outlet_filters = ENABLE_API_OUTLET_FILTERS and bool(filter_functions) if ENABLE_API_OUTLET_FILTERS and not has_api_outlet_filters: try: @@ -5529,6 +5531,7 @@ async def streaming_chat_response_handler(response, ctx): for event in events: event, _ = await process_filter_functions( request=request, + filter_context=filter_context, filter_functions=filter_functions, filter_type='stream', form_data=event, @@ -5541,6 +5544,7 @@ async def streaming_chat_response_handler(response, ctx): async for data in original_generator: data, _ = await process_filter_functions( request=request, + filter_context=filter_context, filter_functions=filter_functions, filter_type='stream', form_data=data,