From 52e155019f511846a8066e012ff37f8aac6db2db Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Sun, 25 Oct 2026 08:24:14 +0900 Subject: [PATCH] refac --- backend/open_webui/utils/audit.py | 31 ++++++++- backend/open_webui/utils/logger.py | 1 + backend/open_webui/utils/middleware.py | 90 +++++++++++++++++--------- 3 files changed, 92 insertions(+), 30 deletions(-) diff --git a/backend/open_webui/utils/audit.py b/backend/open_webui/utils/audit.py index b1d996eca6..856bdd2d99 100644 --- a/backend/open_webui/utils/audit.py +++ b/backend/open_webui/utils/audit.py @@ -27,6 +27,7 @@ from loguru import logger from open_webui.env import AUDIT_INCLUDED_PATHS, AUDIT_LOG_LEVEL, ENABLE_AUDIT_GET_REQUESTS, MAX_BODY_LOG_SIZE from open_webui.models.users import UserModel from open_webui.utils.auth import get_current_user, get_http_authorization_cred +from open_webui.utils.response import normalize_usage from starlette.requests import Request if TYPE_CHECKING: @@ -44,6 +45,7 @@ class AuditLogEntry: user_agent: Optional[str] = None source_ip: Optional[str] = None model_id: Optional[str] = None + usage: Optional[Dict[str, int]] = None # `Request` audit level properties request_object: Any = None # `Request Response` level @@ -114,6 +116,31 @@ class AuditContext: self.response_body.extend(chunk[: self.max_body_size - len(self.response_body)]) +def log_model_usage(request: Request, user: UserModel, model_id: str, usage: Optional[Dict]): + """Write a model call's usage separately from its originating HTTP request.""" + request_id = getattr(request.state, 'audit_request_id', None) + if not request_id or not usage: + return + try: + usage = normalize_usage(usage) + AuditLogger(logger).write( + AuditLogEntry( + id=str(uuid.uuid4()), + user=user.model_dump(include={'id', 'name', 'email', 'role'}), + audit_level=request.state.audit_level, + verb=request.method, + request_uri=str(request.url), + source_ip=request.client.host if request.client else None, + user_agent=request.headers.get('user-agent'), + model_id=model_id, + usage={key: usage[key] for key in ('input_tokens', 'output_tokens', 'total_tokens')}, + ), + extra={'event': 'model_usage', 'request_id': request_id}, + ) + except Exception as e: + logger.error('Failed to log model usage: {}', e) + + def redact_passwords(body: str) -> str: if 'password' not in body.lower(): return body @@ -185,6 +212,8 @@ class AuditLoggingMiddleware: if self._should_skip_auditing(request): return await self.app(scope, receive, send) + request.state.audit_request_id = str(uuid.uuid4()) + request.state.audit_level = self.audit_level.value capture_body = not (request.url.path.startswith('/api/v1/auths') or request.url.path.startswith('/oauth/')) async with self._audit_context(request) as context: @@ -305,7 +334,7 @@ class AuditLoggingMiddleware: response_body = redact_passwords(response_body) entry = AuditLogEntry( - id=str(uuid.uuid4()), + id=request.state.audit_request_id, user=user, audit_level=self.audit_level.value, verb=request.method, diff --git a/backend/open_webui/utils/logger.py b/backend/open_webui/utils/logger.py index 3ab517f396..8cfc803582 100644 --- a/backend/open_webui/utils/logger.py +++ b/backend/open_webui/utils/logger.py @@ -157,6 +157,7 @@ def file_format(record: 'Record'): 'source_ip': record['extra'].get('source_ip', ''), 'user_agent': record['extra'].get('user_agent', ''), 'model_id': record['extra'].get('model_id', ''), + 'usage': record['extra'].get('usage'), 'request_object': record['extra'].get('request_object', b''), 'response_object': record['extra'].get('response_object', b''), 'extra': record['extra'].get('extra', {}), diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 5a279162fe..c706c5b3fb 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -86,6 +86,7 @@ from open_webui.utils.access_control import has_permission from open_webui.utils.access_control.files import get_owner_accessible_folder_files from open_webui.utils.access_control.folders import has_folder_access from open_webui.utils.ask_user import stage_ask_user_tool_calls +from open_webui.utils.audit import log_model_usage from open_webui.utils.chat import generate_chat_completion from open_webui.utils.chat_id import is_saved_chat_id, is_temporary_chat_id from open_webui.utils.code_interpreter import execute_code_jupyter @@ -3914,6 +3915,16 @@ async def non_streaming_chat_response_handler(response, ctx): if response_data is None: return response + usage = response_data.get('usage') or {} + timings = response_data.get('timings') or {} + if isinstance(usage, dict) and isinstance(timings, dict): + log_model_usage( + request, + user, + response_data.get('selected_model_id') or ctx['form_data']['model'], + {**usage, **timings}, + ) + chat_id = metadata.get('chat_id') or '' save_to_chat = is_saved_chat_id(chat_id) continuing = bool(metadata.get('assistant_message_id')) @@ -4147,7 +4158,6 @@ async def streaming_chat_response_handler(response, ctx): # and pyodide code interpreter. Server-side tools work without it. if event_emitter: task_id = str(uuid4()) # Create a unique task ID. - model_id = form_data.get('model', '') # Handle as a background task async def response_handler(response, events): @@ -4635,8 +4645,9 @@ async def streaming_chat_response_handler(response, ctx): }, ) - async def stream_body_handler(response, form_data): + async def stream_body_handler(response, current_form_data): nonlocal usage + nonlocal current_usage nonlocal finish_reason nonlocal output nonlocal prior_output @@ -4644,7 +4655,8 @@ async def streaming_chat_response_handler(response, ctx): nonlocal stream_size response_tool_calls = [] - response_usage = {} + current_usage = {} + form_data['model'] = current_form_data['model'] delta_count = 0 delta_chunk_size = max( @@ -4741,7 +4753,7 @@ async def streaming_chat_response_handler(response, ctx): ) stream_size = 0 - filter_extra_params = {'__body__': form_data, **extra_params} if filter_functions else None + filter_extra_params = {'__body__': current_form_data, **extra_params} if filter_functions else None async for line in response.body_iterator: line = line.decode('utf-8', 'replace') if isinstance(line, bytes) else line @@ -4797,13 +4809,13 @@ async def streaming_chat_response_handler(response, ctx): await event_emitter(data.get('event', {})) if 'selected_model_id' in data: - model_id = data['selected_model_id'] + form_data['model'] = data['selected_model_id'] if save_to_chat: await Chats.upsert_message_to_chat_by_id_and_message_id( metadata['chat_id'], metadata['message_id'], { - 'selectedModelId': model_id, + 'selectedModelId': form_data['model'], }, touch=False, ) @@ -4869,7 +4881,7 @@ async def streaming_chat_response_handler(response, ctx): # Normalize and capture usage for DB persistence if response_metadata.get('usage'): - response_usage = response_metadata['usage'] + current_usage = response_metadata['usage'] usage = merge_usage(usage, response_metadata['usage']) response_metadata['usage'] = usage @@ -4912,7 +4924,7 @@ async def streaming_chat_response_handler(response, ctx): raw_usage = data.get('usage', {}) or {} raw_usage.update(data.get('timings', {})) # llama.cpp if raw_usage: - response_usage = raw_usage + current_usage = raw_usage usage = merge_usage(usage, raw_usage) await event_emitter( { @@ -5555,12 +5567,13 @@ async def streaming_chat_response_handler(response, ctx): tool_calls.append(_split_tool_calls(responses_api_tool_calls)) await save_current_response_stream() - return normalize_usage(response_usage) + return normalize_usage(current_usage) output_start = len(prior_output) try: - response_usage = await stream_body_handler(response, form_data) + current_usage = await stream_body_handler(response, form_data) finally: + log_model_usage(request, user, form_data['model'], current_usage) if response.background: await response.background() @@ -5907,7 +5920,6 @@ async def streaming_chat_response_handler(response, ctx): try: new_form_data = { **form_data, - 'model': model_id, 'stream': True, 'metadata': metadata, } @@ -5941,8 +5953,8 @@ async def streaming_chat_response_handler(response, ctx): image_urls.append(part.get('image_url', '')) message['content'] = ''.join(text_parts) - if response_usage and tool_messages: - tool_messages[0]['usage'] = response_usage + if current_usage and tool_messages: + tool_messages[0]['usage'] = current_usage # Refresh citation context without restoring compacted history. user_prompt = get_last_user_message_item(form_data['messages']) @@ -5985,7 +5997,7 @@ async def streaming_chat_response_handler(response, ctx): metadata.get('base_user_prompt') or user_message, [message] ) messages, context_summary, compacted = await compact_messages_for_request( - request, user, messages, metadata, model_id, request.app.state.MODELS + request, user, messages, metadata, form_data['model'], request.app.state.MODELS ) finally: for message, content in user_messages: @@ -6074,7 +6086,10 @@ async def streaming_chat_response_handler(response, ctx): prior_output.pop() output = [] output_start = len(prior_output) - response_usage = await stream_body_handler(res, new_form_data) + try: + current_usage = await stream_body_handler(res, new_form_data) + finally: + log_model_usage(request, user, form_data['model'], current_usage) elif getattr(res, 'status_code', 200) >= 400: await emit_message_error(get_message_error_content(get_response_error_detail(res))) break @@ -6239,7 +6254,6 @@ async def streaming_chat_response_handler(response, ctx): try: new_form_data = { **form_data, - 'model': model_id, 'stream': True, 'metadata': metadata, 'messages': [ @@ -6273,7 +6287,10 @@ async def streaming_chat_response_handler(response, ctx): ) if isinstance(res, StreamingResponse): - await stream_body_handler(res, new_form_data) + try: + await stream_body_handler(res, new_form_data) + finally: + log_model_usage(request, user, form_data['model'], current_usage) elif getattr(res, 'status_code', 200) >= 400: await emit_message_error(get_message_error_content(get_response_error_detail(res))) break @@ -6398,6 +6415,8 @@ async def streaming_chat_response_handler(response, ctx): def wrap_item(item): return f'data: {item}\n\n' + audit_enabled = bool(getattr(request.state, 'audit_request_id', None)) + current_usage = None try: message = {} outlet_task = None @@ -6405,10 +6424,11 @@ async def streaming_chat_response_handler(response, ctx): outlet_enabled = ENABLE_API_OUTLET_FILTERS and bool(filter_functions) if ENABLE_API_OUTLET_FILTERS and not outlet_enabled: try: - model_id = model.get('id') if isinstance(model, dict) else model outlet_enabled = bool( (isinstance(model, dict) and 'pipeline' in model) - or get_sorted_filters(model_id, request.app.state.MODELS) + or get_sorted_filters( + model.get('id') if isinstance(model, dict) else model, request.app.state.MODELS + ) ) except Exception: outlet_enabled = True @@ -6427,7 +6447,7 @@ async def streaming_chat_response_handler(response, ctx): yield wrap_item(JSONCodec.dumps(event)) async for data in original_generator: - if filter_functions: + if filter_functions or audit_enabled: line = data.decode('utf-8', 'replace') if isinstance(data, bytes) else data if isinstance(line, str) and line.startswith('data:'): payload = line.removeprefix('data:').strip() @@ -6438,15 +6458,25 @@ async def streaming_chat_response_handler(response, ctx): event = None if isinstance(event, dict): - event, _ = await process_filter_functions( - request=request, - filter_context=filter_context, - filter_functions=filter_functions, - filter_type='stream', - form_data=event, - extra_params=extra_params, - ) - data = wrap_item(JSONCodec.dumps(event)) if event else None + if audit_enabled: + form_data['model'] = event.get('selected_model_id') or form_data['model'] + payload = ( + event.get('response') if isinstance(event.get('response'), dict) else event + ) + usage = payload.get('usage') or {} + timings = payload.get('timings') or {} + if isinstance(usage, dict) and isinstance(timings, dict) and (usage or timings): + current_usage = {**usage, **timings} + if filter_functions: + event, _ = await process_filter_functions( + request=request, + filter_context=filter_context, + filter_functions=filter_functions, + filter_type='stream', + form_data=event, + extra_params=extra_params, + ) + data = wrap_item(JSONCodec.dumps(event)) if event else None if data: if outlet_enabled: @@ -6467,6 +6497,8 @@ async def streaming_chat_response_handler(response, ctx): # Separate the error frame from any unfinished upstream event. yield f'\n\ndata: {JSONCodec.dumps({"error": {"message": "Chat completion stream failed"}})}\n\n' yield 'data: [DONE]\n\n' + finally: + log_model_usage(request, user, form_data['model'], current_usage) return StreamingResponse( stream_wrapper(response.body_iterator, events),