fix: snapshot deletion during regeneration

This commit is contained in:
CREDO23 2026-01-30 18:44:33 +02:00
parent bc0fb3cb68
commit d2ebd3ec00
2 changed files with 19 additions and 18 deletions

View file

@ -1274,6 +1274,8 @@ async def regenerate_response(
) )
messages_to_delete = list(last_messages_result.scalars().all()) messages_to_delete = list(last_messages_result.scalars().all())
message_ids_to_delete = [msg.id for msg in messages_to_delete]
# Get search space for LLM config # Get search space for LLM config
search_space_result = await session.execute( search_space_result = await session.execute(
select(SearchSpace).filter(SearchSpace.id == request.search_space_id) select(SearchSpace).filter(SearchSpace.id == request.search_space_id)
@ -1313,9 +1315,6 @@ async def regenerate_response(
# This ensures we don't lose data on streaming failures # This ensures we don't lose data on streaming failures
if streaming_completed and messages_to_delete: if streaming_completed and messages_to_delete:
try: try:
# Get message IDs before deletion for snapshot cleanup
deleted_message_ids = [msg.id for msg in messages_to_delete]
for msg in messages_to_delete: for msg in messages_to_delete:
await session.delete(msg) await session.delete(msg)
await session.commit() await session.commit()
@ -1326,7 +1325,7 @@ async def regenerate_response(
) )
await delete_affected_snapshots( await delete_affected_snapshots(
session, thread_id, deleted_message_ids session, thread_id, message_ids_to_delete
) )
except Exception as cleanup_error: except Exception as cleanup_error:
# Log but don't fail - the new messages are already streamed # Log but don't fail - the new messages are already streamed

View file

@ -426,7 +426,7 @@ async def delete_snapshot(
async def delete_affected_snapshots( async def delete_affected_snapshots(
session: AsyncSession, session: AsyncSession, # noqa: ARG001 - kept for API compatibility
thread_id: int, thread_id: int,
message_ids: list[int], message_ids: list[int],
) -> int: ) -> int:
@ -434,25 +434,27 @@ async def delete_affected_snapshots(
Delete snapshots that contain any of the given message IDs. Delete snapshots that contain any of the given message IDs.
Called when messages are edited/deleted/regenerated. Called when messages are edited/deleted/regenerated.
Uses independent session to work reliably in streaming response cleanup.
Returns the number of deleted snapshots.
""" """
if not message_ids: if not message_ids:
return 0 return 0
# Use raw SQL for array overlap query from sqlalchemy.dialects.postgresql import array
# The && operator checks if arrays have any elements in common
result = await session.execute(
delete(PublicChatSnapshot)
.where(PublicChatSnapshot.thread_id == thread_id)
.where(PublicChatSnapshot.message_ids.overlap(message_ids))
.returning(PublicChatSnapshot.id)
)
deleted_ids = result.scalars().all() from app.db import async_session_maker
await session.commit()
return len(deleted_ids) async with async_session_maker() as independent_session:
result = await independent_session.execute(
delete(PublicChatSnapshot)
.where(PublicChatSnapshot.thread_id == thread_id)
.where(PublicChatSnapshot.message_ids.op("&&")(array(message_ids)))
.returning(PublicChatSnapshot.id)
)
deleted_ids = result.scalars().all()
await independent_session.commit()
return len(deleted_ids)
# ============================================================================= # =============================================================================