scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-15 04:18:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-10-14 18:35:56 +04:00
parent 29289f5ba6
commit bcf540a989
2 changed files with 67 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
backend/open_webui/utils/context_compaction.py
View file

@ -303,30 +303,59 @@ def _build_context_usage(tokens: int, threshold: int) -> dict:
def _apply_latest_summary_checkpoint(messages: list[dict]) -> tuple[list[dict], str | None]:
summary = None
summary_idx = None
for idx, message in enumerate(messages):
for idx in range(len(messages) - 1, -1, -1):
message = messages[idx]
value = message.get('contextSummary') or message.get('context_summary')
if isinstance(value, str) and value.strip():
summary = value
summary_idx = idx
return messages[idx:], value
if summary_idx is None:
return messages, None
return messages[summary_idx:], summary
marker = next(
(
item
for item in reversed(message.get('output') or [])
if item.get('type') == 'open_webui:context_compaction'
),
None,
)
if not marker:
continue
# A tool-loop checkpoint can retain output from an earlier reply.
for cut_idx in range(idx, -1, -1):
output = messages[cut_idx].get('output') or []
start = next(
(
position
for position, item in enumerate(output)
if item.get('type') == 'function_call' and item.get('call_id') == marker['call_id']
),
None,
)
if start is not None:
while start and output[start - 1].get('type') != 'function_call_output':
start -= 1
return [
*messages[max(0, cut_idx - 1) : cut_idx],
{
**messages[cut_idx],
'output': [
item for item in output[start:] if item.get('type') != 'open_webui:context_compaction'
],
},
*messages[cut_idx + 1 :],
], marker['summary']
return messages, None
def _exceeds_token_threshold(messages: list[dict], system_prompt: str, summary: str | None, threshold: int) -> bool:
if threshold <= 0:
return False
# Expanded tool histories include fresh results and rebuilt prompts absent from prior usage.
if not any(message.get('role') == 'tool' for message in messages):
for idx in range(len(messages) - 1, -1, -1):
usage = messages[idx].get('usage') or (messages[idx].get('info') or {}).get('usage')
if isinstance(usage, dict) and (tokens := _usage_token_count(usage)):
return tokens + _estimate_messages_tokens(messages[idx + 1 :]) > threshold
for idx in range(len(messages) - 1, -1, -1):
usage = messages[idx].get('usage') or (messages[idx].get('info') or {}).get('usage')
if isinstance(usage, dict) and (tokens := _usage_token_count(usage)):
return tokens + _estimate_messages_tokens(messages[idx + 1 :]) > threshold
estimated = _estimate_tokens(system_prompt) + _estimate_tokens(summary or '') + _estimate_messages_tokens(messages)
return estimated > threshold

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

@ -4966,6 +4966,7 @@ async def streaming_chat_response_handler(response, ctx):
nonlocal last_response_id
response_tool_calls = []
response_usage = {}
delta_count = 0
delta_chunk_size = max(
@ -5228,6 +5229,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']
usage = merge_usage(usage, response_metadata['usage'])
response_metadata['usage'] = usage
@ -5270,6 +5272,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
usage = merge_usage(usage, raw_usage)
await event_emitter(
{
@ -5906,9 +5909,11 @@ async def streaming_chat_response_handler(response, ctx):
if responses_api_tool_calls:
tool_calls.append(_split_tool_calls(responses_api_tool_calls))
return normalize_usage(response_usage)
output_start = len(prior_output)
try:
await stream_body_handler(response, form_data)
response_usage = await stream_body_handler(response, form_data)
finally:
if response.background:
await response.background()
@ -6284,6 +6289,9 @@ 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
# Refresh citation context without restoring compacted history.
user_prompt = get_last_user_message_item(form_data['messages'])
for message in messages:
@ -6324,6 +6332,19 @@ async def streaming_chat_response_handler(response, ctx):
idx for idx, message in enumerate(messages) if message.get('role') != 'system'
)
messages[checkpoint] = {**messages[checkpoint], 'contextSummary': context_summary}
output.append(
{
'type': 'open_webui:context_compaction',
'id': output_id('cc'),
'call_id': next(
message['tool_calls'][0]['id']
for message in messages
if message.get('tool_calls')
),
'summary': context_summary,
}
)
await emit_output()
except Exception:
log.exception('Tool loop context compaction failed; keeping the current context')
@ -6384,7 +6405,7 @@ async def streaming_chat_response_handler(response, ctx):
prior_output.pop()
output = []
output_start = len(prior_output)
await stream_body_handler(res, new_form_data)
response_usage = await stream_body_handler(res, new_form_data)
elif getattr(res, 'status_code', 200) >= 400:
await emit_message_error(get_message_error_content(get_response_error_detail(res)))
break