mirror of
https://github.com/open-webui/open-webui.git
synced 2026-10-06 02:48:04 +00:00
test: extract pure history-upsert helpers and add regression tests (#24758-style)
Extracts the graph-repair logic out of the async DB method into pure module functions apply_message_upsert() and build_chat_message_payload(). Behaviour-preserving: the async method just loads the chat, calls the helpers, dual-writes and saves. Adds unit tests (no DB, no async) covering the failure modes that bricked production chats: missing/partial payload, structural preservation, explicit falsy override, parent-link repair with malformed childrenIds, malformed messages container, malformed existing node, dangling parent not severed, non-str parentId coerced, and parent_id sent only when known. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Classic298
parent
466883ad18
commit
ed032753b9
2 changed files with 181 additions and 67 deletions
|
|
@ -264,6 +264,79 @@ class ChatStatsExport(BaseModel):
|
|||
chat: ChatBody
|
||||
|
||||
|
||||
def apply_message_upsert(history: dict, message_id: str, message: dict, now: int) -> dict:
|
||||
"""Apply ``message`` into ``history['messages'][message_id]`` with
|
||||
graph-invariant repair. Mutates ``history``; returns the final node.
|
||||
Pure (no DB) so it can be unit-tested directly."""
|
||||
messages = history.get('messages')
|
||||
if not isinstance(messages, dict): # tolerate absent/malformed container
|
||||
messages = {}
|
||||
history['messages'] = messages
|
||||
|
||||
def enforce_node_invariants(node: dict) -> dict:
|
||||
# Structural graph fields a partial update must never drop.
|
||||
node['id'] = message_id
|
||||
if not isinstance(node.get('parentId'), str):
|
||||
node['parentId'] = None # explicit null; never undefined/non-scalar
|
||||
if not isinstance(node.get('childrenIds'), list):
|
||||
node['childrenIds'] = []
|
||||
if not node.get('role'):
|
||||
node['role'] = 'assistant'
|
||||
if not node.get('timestamp'):
|
||||
node['timestamp'] = now
|
||||
return node
|
||||
|
||||
existing = messages.get(message_id)
|
||||
if isinstance(existing, dict):
|
||||
# {**existing, **message}: omitted keys keep existing, explicit
|
||||
# values (incl. parentId=None, childrenIds=[]) win.
|
||||
messages[message_id] = enforce_node_invariants({**existing, **message})
|
||||
else:
|
||||
if existing is not None:
|
||||
log.warning(f'upsert: discarding malformed existing node {message_id}')
|
||||
# Node missing/unusable: a concurrent whole-chat write likely
|
||||
# dropped the placeholder. Warn only when the payload is partial.
|
||||
is_partial = not all(k in message for k in ('id', 'parentId', 'childrenIds', 'role'))
|
||||
if is_partial:
|
||||
log.warning(
|
||||
f'upsert: synthesizing lost message node {message_id} '
|
||||
f'from partial payload '
|
||||
f'(role={message.get("role") or "assistant (defaulted)"!r})'
|
||||
)
|
||||
else:
|
||||
log.debug(f'upsert: creating message node {message_id}')
|
||||
messages[message_id] = enforce_node_invariants({**message})
|
||||
|
||||
# Keep the graph bidirectional; the same stale-write class can drop
|
||||
# the parent's child link or the parent node itself.
|
||||
parent_id = messages[message_id].get('parentId')
|
||||
if parent_id:
|
||||
parent = messages.get(parent_id)
|
||||
if isinstance(parent, dict):
|
||||
children_ids = parent.get('childrenIds')
|
||||
if not isinstance(children_ids, list):
|
||||
children_ids = []
|
||||
parent['childrenIds'] = children_ids
|
||||
if message_id not in children_ids:
|
||||
children_ids.append(message_id)
|
||||
else:
|
||||
# Parent absent from this possibly-stale snapshot: log only; it
|
||||
# may still exist in chat_message, so don't sever a real link.
|
||||
log.warning(f'upsert: message {message_id} parent {parent_id} absent from embedded history')
|
||||
|
||||
history['currentId'] = message_id
|
||||
return messages[message_id]
|
||||
|
||||
|
||||
def build_chat_message_payload(node: dict) -> dict:
|
||||
"""Normalized chat_message data; send parent_id only when known so a
|
||||
None can't clear a real link from a possibly-stale snapshot."""
|
||||
data = {**node}
|
||||
if node.get('parentId') is not None:
|
||||
data['parent_id'] = node['parentId']
|
||||
return data
|
||||
|
||||
|
||||
class ChatTable:
|
||||
def _clean_null_bytes(self, obj):
|
||||
"""Recursively remove null bytes from strings in dict/list structures."""
|
||||
|
|
@ -563,80 +636,15 @@ class ChatTable:
|
|||
chat = chat.chat
|
||||
history = chat.get('history', {})
|
||||
|
||||
now = int(time.time())
|
||||
# tolerate an absent or malformed messages container
|
||||
messages = history.get('messages')
|
||||
if not isinstance(messages, dict):
|
||||
messages = {}
|
||||
history['messages'] = messages
|
||||
|
||||
def enforce_node_invariants(node: dict) -> dict:
|
||||
# Structural graph fields a partial update must never drop.
|
||||
node['id'] = message_id
|
||||
if not isinstance(node.get('parentId'), str):
|
||||
node['parentId'] = None # explicit null; never undefined/non-scalar
|
||||
if not isinstance(node.get('childrenIds'), list):
|
||||
node['childrenIds'] = []
|
||||
if not node.get('role'):
|
||||
node['role'] = 'assistant'
|
||||
if not node.get('timestamp'):
|
||||
node['timestamp'] = now
|
||||
return node
|
||||
|
||||
existing = messages.get(message_id)
|
||||
if isinstance(existing, dict):
|
||||
# {**existing, **message}: omitted keys keep existing, explicit
|
||||
# values (incl. parentId=None, childrenIds=[]) win.
|
||||
messages[message_id] = enforce_node_invariants({**existing, **message})
|
||||
else:
|
||||
if existing is not None:
|
||||
log.warning(f'upsert: discarding malformed existing node {message_id}')
|
||||
# Node missing/unusable: a concurrent whole-chat write likely
|
||||
# dropped the placeholder. Warn only when the payload is partial.
|
||||
is_partial = not all(k in message for k in ('id', 'parentId', 'childrenIds', 'role'))
|
||||
if is_partial:
|
||||
log.warning(
|
||||
f'upsert: synthesizing lost message node {message_id} '
|
||||
f'from partial payload '
|
||||
f'(role={message.get("role") or "assistant (defaulted)"!r})'
|
||||
)
|
||||
else:
|
||||
log.debug(f'upsert: creating message node {message_id}')
|
||||
messages[message_id] = enforce_node_invariants({**message})
|
||||
|
||||
# Keep the graph bidirectional; the same stale-write class can drop
|
||||
# the parent's child link or the parent node itself.
|
||||
parent_id = messages[message_id].get('parentId')
|
||||
if parent_id:
|
||||
parent = messages.get(parent_id)
|
||||
if isinstance(parent, dict):
|
||||
children_ids = parent.get('childrenIds')
|
||||
if not isinstance(children_ids, list):
|
||||
children_ids = []
|
||||
parent['childrenIds'] = children_ids
|
||||
if message_id not in children_ids:
|
||||
children_ids.append(message_id)
|
||||
else:
|
||||
# Parent absent from this possibly-stale snapshot: log only.
|
||||
# It may still exist in chat_message; don't sever a real link.
|
||||
log.warning(f'upsert: message {message_id} parent {parent_id} absent from embedded history')
|
||||
|
||||
history['currentId'] = message_id
|
||||
|
||||
node = apply_message_upsert(history, message_id, message, int(time.time()))
|
||||
chat['history'] = history
|
||||
|
||||
# Send parent_id only when known; parent_id=None would clear a real
|
||||
# link on the normalized row from a possibly-stale snapshot.
|
||||
node = history['messages'][message_id]
|
||||
cm_data = {**node}
|
||||
if node.get('parentId') is not None:
|
||||
cm_data['parent_id'] = node['parentId']
|
||||
try:
|
||||
await ChatMessages.upsert_message(
|
||||
message_id=message_id,
|
||||
chat_id=id,
|
||||
user_id=user_id,
|
||||
data=cm_data,
|
||||
data=build_chat_message_payload(node),
|
||||
)
|
||||
except Exception as e:
|
||||
log.warning(f'Failed to write to chat_message table: {e}')
|
||||
|
|
|
|||
106
backend/open_webui/test/apps/webui/models/test_chats_upsert.py
Normal file
106
backend/open_webui/test/apps/webui/models/test_chats_upsert.py
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
"""Regression tests for the chat-history graph repair in
|
||||
``apply_message_upsert`` / ``build_chat_message_payload`` (PR #24617).
|
||||
|
||||
Pure functions, no DB, so the failure modes that bricked production chats
|
||||
are locked down deterministically.
|
||||
"""
|
||||
|
||||
from open_webui.models.chats import apply_message_upsert, build_chat_message_payload
|
||||
|
||||
NOW = 1_700_000_000
|
||||
|
||||
|
||||
def test_missing_node_partial_payload_gets_full_invariants():
|
||||
history = {}
|
||||
node = apply_message_upsert(history, 'a', {'content': 'hi', 'done': True}, NOW)
|
||||
assert node['id'] == 'a'
|
||||
assert node['parentId'] is None
|
||||
assert node['childrenIds'] == []
|
||||
assert node['role'] == 'assistant'
|
||||
assert node['timestamp'] == NOW
|
||||
assert history['currentId'] == 'a'
|
||||
assert history['messages']['a'] is node
|
||||
|
||||
|
||||
def test_existing_node_partial_update_preserves_structural_fields():
|
||||
history = {
|
||||
'messages': {
|
||||
'a': {
|
||||
'id': 'a',
|
||||
'parentId': 'u1',
|
||||
'childrenIds': ['c1'],
|
||||
'role': 'assistant',
|
||||
'timestamp': 111,
|
||||
'content': 'old',
|
||||
}
|
||||
},
|
||||
'currentId': 'a',
|
||||
}
|
||||
apply_message_upsert(history, 'a', {'content': 'new', 'done': True}, NOW)
|
||||
n = history['messages']['a']
|
||||
assert n['parentId'] == 'u1'
|
||||
assert n['childrenIds'] == ['c1']
|
||||
assert n['timestamp'] == 111
|
||||
assert n['role'] == 'assistant'
|
||||
assert n['content'] == 'new'
|
||||
|
||||
|
||||
def test_explicit_falsy_values_win_over_existing():
|
||||
history = {
|
||||
'messages': {
|
||||
'a': {
|
||||
'id': 'a',
|
||||
'parentId': 'u1',
|
||||
'childrenIds': ['c1'],
|
||||
'role': 'assistant',
|
||||
'timestamp': 111,
|
||||
}
|
||||
}
|
||||
}
|
||||
apply_message_upsert(history, 'a', {'parentId': None, 'childrenIds': []}, NOW)
|
||||
n = history['messages']['a']
|
||||
assert n['parentId'] is None
|
||||
assert n['childrenIds'] == []
|
||||
|
||||
|
||||
def test_parent_child_link_repaired_and_malformed_children_normalized():
|
||||
history = {
|
||||
'messages': {'u1': {'id': 'u1', 'parentId': None, 'childrenIds': None, 'role': 'user'}},
|
||||
'currentId': 'u1',
|
||||
}
|
||||
apply_message_upsert(history, 'a', {'parentId': 'u1', 'role': 'assistant'}, NOW)
|
||||
assert history['messages']['u1']['childrenIds'] == ['a']
|
||||
|
||||
|
||||
def test_malformed_messages_container_replaced():
|
||||
for bad in (None, [], 'x', 3):
|
||||
history = {'messages': bad}
|
||||
apply_message_upsert(history, 'a', {'content': 'x'}, NOW)
|
||||
assert isinstance(history['messages'], dict)
|
||||
assert history['messages']['a']['id'] == 'a'
|
||||
|
||||
|
||||
def test_malformed_existing_node_synthesized_not_crash():
|
||||
for bad in ([], 'x', 3, True):
|
||||
history = {'messages': {'a': bad}, 'currentId': 'a'}
|
||||
node = apply_message_upsert(history, 'a', {'content': 'x'}, NOW)
|
||||
assert node['id'] == 'a'
|
||||
assert node['role'] == 'assistant'
|
||||
|
||||
|
||||
def test_dangling_parent_logged_not_severed():
|
||||
history = {'messages': {}, 'currentId': None}
|
||||
node = apply_message_upsert(history, 'a', {'parentId': 'missing', 'role': 'assistant'}, NOW)
|
||||
# The parent may still exist in chat_message; the link must be preserved.
|
||||
assert node['parentId'] == 'missing'
|
||||
|
||||
|
||||
def test_non_str_parentid_coerced_to_none():
|
||||
history = {}
|
||||
node = apply_message_upsert(history, 'a', {'parentId': ['x'], 'role': 'assistant'}, NOW)
|
||||
assert node['parentId'] is None
|
||||
|
||||
|
||||
def test_build_payload_sends_parent_id_only_when_known():
|
||||
assert build_chat_message_payload({'id': 'a', 'parentId': 'u1'})['parent_id'] == 'u1'
|
||||
assert 'parent_id' not in build_chat_message_payload({'id': 'a', 'parentId': None})
|
||||
Loading…
Add table
Reference in a new issue