scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-28 06:22:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-10-25 08:24:14 +09:00
parent f86e74e163
commit 52e155019f
3 changed files with 92 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
backend/open_webui/utils/audit.py
View file

@ -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,

1
backend/open_webui/utils/logger.py
View file

@ -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', {}),

90
backend/open_webui/utils/middleware.py
View file

@ -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),