382 lines
15 KiB
Python
382 lines
15 KiB
Python
"""Clear every conversation owned by the current user in one request.
|
|
|
|
Ordinary Q&A and agent chats share ``threads_meta``. System workers
|
|
(roundtable, scheduler, notebooks, writing setup, deep research) stay.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
from sqlalchemy import delete as sql_delete
|
|
from sqlalchemy import select, text
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker
|
|
|
|
from deerflow.config.paths import Paths, get_paths
|
|
from deerflow.persistence.analysis_projects import (
|
|
AnalysisProjectConfigRow,
|
|
AnalysisProjectThreadRow,
|
|
)
|
|
from deerflow.persistence.feedback.model import FeedbackRow
|
|
from deerflow.persistence.llmwiki.model import LlmWikiConversationDepositRow
|
|
from deerflow.persistence.models.run_event import RunEventRow
|
|
from deerflow.persistence.run.model import RunRow
|
|
from deerflow.persistence.thread_meta.base import ThreadMetaStore, is_excluded_system_thread
|
|
from deerflow.persistence.thread_meta.model import ThreadMetaRow
|
|
from deerflow.persistence.thread_shares.model import ThreadShareRow
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_PAGE_SIZE = 200
|
|
_SQL_CHUNK = 400
|
|
_CHECKPOINT_TABLES = ("checkpoints", "writes", "checkpoint_blobs", "checkpoint_writes")
|
|
|
|
|
|
def _chunks(items: list[str], size: int = _SQL_CHUNK) -> list[list[str]]:
|
|
return [items[index : index + size] for index in range(0, len(items), size)]
|
|
|
|
|
|
def _elapsed_ms(started: float) -> int:
|
|
return int((time.perf_counter() - started) * 1000)
|
|
|
|
|
|
async def list_owned_chat_thread_ids(
|
|
store: ThreadMetaStore,
|
|
user_id: str,
|
|
session_factory: async_sessionmaker | None = None,
|
|
) -> list[str]:
|
|
"""Return this user's conversation ids, including every agent chat."""
|
|
|
|
if session_factory is not None:
|
|
try:
|
|
async with session_factory() as session:
|
|
rows = await session.execute(
|
|
select(ThreadMetaRow.thread_id, ThreadMetaRow.metadata_json).where(ThreadMetaRow.user_id == user_id)
|
|
)
|
|
return [
|
|
str(thread_id)
|
|
for thread_id, metadata in rows.all()
|
|
if thread_id and not is_excluded_system_thread(metadata)
|
|
]
|
|
except Exception:
|
|
logger.debug("Falling back to paginated thread search for chat clear", exc_info=True)
|
|
|
|
ids: list[str] = []
|
|
offset = 0
|
|
while True:
|
|
rows = await store.search(
|
|
user_id=user_id,
|
|
exclude_system=True,
|
|
limit=_PAGE_SIZE,
|
|
offset=offset,
|
|
)
|
|
if not rows:
|
|
break
|
|
ids.extend(str(row["thread_id"]) for row in rows if row.get("thread_id"))
|
|
if len(rows) < _PAGE_SIZE:
|
|
break
|
|
offset += len(rows)
|
|
return list(dict.fromkeys(ids))
|
|
|
|
|
|
def _checkpoint_savers(checkpointer: Any) -> list[Any]:
|
|
if checkpointer is None:
|
|
return []
|
|
primary = getattr(checkpointer, "primary", None)
|
|
if primary is None:
|
|
return [checkpointer]
|
|
savers = [primary]
|
|
for source in getattr(checkpointer, "sources", []) or []:
|
|
if getattr(source, "closed", False):
|
|
continue
|
|
saver = getattr(source, "saver", None)
|
|
if saver is not None:
|
|
savers.append(saver)
|
|
return savers
|
|
|
|
|
|
def _is_sqlite_saver(saver: Any) -> bool:
|
|
return "sqlite" in type(saver).__module__ and hasattr(saver, "conn") and hasattr(saver, "lock")
|
|
|
|
|
|
async def _ignore_missing_table(cursor: Any, sql: str, params: tuple[Any, ...] = ()) -> None:
|
|
try:
|
|
await cursor.execute(sql, params)
|
|
except Exception as exc:
|
|
message = str(exc).lower()
|
|
if "no such table" in message or "does not exist" in message:
|
|
return
|
|
raise
|
|
|
|
|
|
async def _fill_sqlite_id_table(cursor: Any, thread_ids: list[str]) -> None:
|
|
await cursor.execute("CREATE TEMP TABLE IF NOT EXISTS deerflow_clear_ids (thread_id TEXT PRIMARY KEY)")
|
|
await cursor.execute("DELETE FROM deerflow_clear_ids")
|
|
executemany = getattr(cursor, "executemany", None)
|
|
for batch in _chunks(thread_ids, 200):
|
|
rows = [(thread_id,) for thread_id in batch]
|
|
if executemany is not None:
|
|
await executemany("INSERT OR IGNORE INTO deerflow_clear_ids(thread_id) VALUES (?)", rows)
|
|
else:
|
|
for row in rows:
|
|
await cursor.execute("INSERT OR IGNORE INTO deerflow_clear_ids(thread_id) VALUES (?)", row)
|
|
|
|
|
|
async def _delete_sqlite_checkpoints(saver: Any, thread_ids: list[str]) -> bool:
|
|
if not _is_sqlite_saver(saver):
|
|
return False
|
|
async with saver.lock:
|
|
async with saver.conn.cursor() as cursor:
|
|
await _fill_sqlite_id_table(cursor, thread_ids)
|
|
for table in _CHECKPOINT_TABLES:
|
|
await _ignore_missing_table(
|
|
cursor,
|
|
f"DELETE FROM {table} WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
)
|
|
await cursor.execute("DROP TABLE IF EXISTS deerflow_clear_ids")
|
|
await saver.conn.commit()
|
|
return True
|
|
|
|
|
|
async def _delete_checkpoints(checkpointer: Any, thread_ids: list[str]) -> None:
|
|
if checkpointer is None or not thread_ids:
|
|
return
|
|
for saver in _checkpoint_savers(checkpointer):
|
|
try:
|
|
if await _delete_sqlite_checkpoints(saver, thread_ids):
|
|
continue
|
|
except Exception:
|
|
logger.warning("Bulk checkpoint delete failed; deleting thread by thread", exc_info=True)
|
|
delete = getattr(saver, "adelete_thread", None)
|
|
if delete is None:
|
|
continue
|
|
for thread_id in thread_ids:
|
|
try:
|
|
await delete(thread_id)
|
|
except Exception:
|
|
logger.warning("Could not delete checkpoints for thread %s", thread_id, exc_info=True)
|
|
|
|
|
|
async def _bound_thread_ids(session_factory: async_sessionmaker | None, thread_ids: list[str]) -> list[str]:
|
|
if session_factory is None or not thread_ids:
|
|
return []
|
|
found: list[str] = []
|
|
try:
|
|
async with session_factory() as session:
|
|
for batch in _chunks(thread_ids):
|
|
rows = await session.execute(
|
|
select(AnalysisProjectThreadRow.thread_id).where(AnalysisProjectThreadRow.thread_id.in_(batch))
|
|
)
|
|
found.extend(str(item) for item in rows.scalars().all())
|
|
except Exception:
|
|
logger.debug("Analysis-project bindings are unavailable during chat clear", exc_info=True)
|
|
return []
|
|
return found
|
|
|
|
|
|
async def _deposit_document_ids(session_factory: async_sessionmaker | None, thread_ids: list[str]) -> list[str]:
|
|
if session_factory is None or not thread_ids:
|
|
return []
|
|
doc_ids: list[str] = []
|
|
try:
|
|
async with session_factory() as session:
|
|
for batch in _chunks(thread_ids):
|
|
rows = await session.execute(
|
|
select(LlmWikiConversationDepositRow.weknora_knowledge_id).where(
|
|
LlmWikiConversationDepositRow.thread_id.in_(batch)
|
|
)
|
|
)
|
|
doc_ids.extend(str(item) for item in rows.scalars().all() if item)
|
|
except Exception:
|
|
logger.debug("Conversation deposits are unavailable during chat clear", exc_info=True)
|
|
return []
|
|
return list(dict.fromkeys(doc_ids))
|
|
|
|
|
|
async def _delete_remote_documents(app: Any, doc_ids: list[str]) -> None:
|
|
from app.gateway.llmwiki_deposit import _client
|
|
from deerflow.integrations.weknora.client import WeKnoraError
|
|
|
|
client = _client(app)
|
|
if client is None:
|
|
return
|
|
for doc_id in doc_ids:
|
|
try:
|
|
await client.delete_document(doc_id)
|
|
except WeKnoraError as exc:
|
|
if exc.status_code != 404:
|
|
logger.warning("Failed to delete conversation deposit document %s", doc_id, exc_info=True)
|
|
except Exception:
|
|
logger.warning("Failed to delete conversation deposit document %s", doc_id, exc_info=True)
|
|
|
|
|
|
async def _delete_related_rows_sqlite(
|
|
session: Any,
|
|
thread_ids: list[str],
|
|
user_id: str,
|
|
) -> None:
|
|
await session.execute(text("CREATE TEMP TABLE IF NOT EXISTS deerflow_clear_ids (thread_id TEXT PRIMARY KEY)"))
|
|
await session.execute(text("DELETE FROM deerflow_clear_ids"))
|
|
for batch in _chunks(thread_ids, 200):
|
|
await session.execute(
|
|
text("INSERT OR IGNORE INTO deerflow_clear_ids (thread_id) VALUES (:id)"),
|
|
[{"id": thread_id} for thread_id in batch],
|
|
)
|
|
for statement in (
|
|
"DELETE FROM run_events WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM runs WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM feedback WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM thread_shares WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM llmwiki_conversation_deposits WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM analysis_project_threads WHERE thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM analysis_project_configs WHERE key IN (SELECT 'thread:' || thread_id FROM deerflow_clear_ids)",
|
|
"DELETE FROM threads_meta WHERE user_id = :user_id AND thread_id IN (SELECT thread_id FROM deerflow_clear_ids)",
|
|
):
|
|
await session.execute(text(statement), {"user_id": user_id})
|
|
await session.execute(text("DROP TABLE IF EXISTS deerflow_clear_ids"))
|
|
|
|
|
|
async def _delete_related_rows(
|
|
session_factory: async_sessionmaker,
|
|
thread_ids: list[str],
|
|
user_id: str,
|
|
) -> None:
|
|
async with session_factory() as session:
|
|
dialect = session.get_bind().dialect.name
|
|
if dialect == "sqlite":
|
|
await _delete_related_rows_sqlite(session, thread_ids, user_id)
|
|
else:
|
|
for batch in _chunks(thread_ids):
|
|
await session.execute(sql_delete(RunEventRow).where(RunEventRow.thread_id.in_(batch)))
|
|
await session.execute(sql_delete(RunRow).where(RunRow.thread_id.in_(batch)))
|
|
await session.execute(sql_delete(FeedbackRow).where(FeedbackRow.thread_id.in_(batch)))
|
|
await session.execute(sql_delete(ThreadShareRow).where(ThreadShareRow.thread_id.in_(batch)))
|
|
await session.execute(
|
|
sql_delete(LlmWikiConversationDepositRow).where(LlmWikiConversationDepositRow.thread_id.in_(batch))
|
|
)
|
|
await session.execute(
|
|
sql_delete(AnalysisProjectThreadRow).where(AnalysisProjectThreadRow.thread_id.in_(batch))
|
|
)
|
|
await session.execute(
|
|
sql_delete(AnalysisProjectConfigRow).where(
|
|
AnalysisProjectConfigRow.key.in_([f"thread:{thread_id}" for thread_id in batch])
|
|
)
|
|
)
|
|
await session.execute(
|
|
sql_delete(ThreadMetaRow).where(
|
|
ThreadMetaRow.thread_id.in_(batch),
|
|
ThreadMetaRow.user_id == user_id,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
def _remove_thread_dirs(paths: Paths, user_id: str, thread_ids: list[str]) -> None:
|
|
started = time.perf_counter()
|
|
removed = 0
|
|
for thread_id in thread_ids:
|
|
try:
|
|
paths.delete_thread_dir(thread_id, user_id=user_id)
|
|
removed += 1
|
|
except Exception:
|
|
logger.warning("Failed to delete files for thread %s", thread_id, exc_info=True)
|
|
logger.info(
|
|
"Cleared chat files user=%s dirs=%d/%d files_ms=%d",
|
|
user_id,
|
|
removed,
|
|
len(thread_ids),
|
|
_elapsed_ms(started),
|
|
)
|
|
|
|
|
|
def _schedule_background(coro: Any) -> None:
|
|
task = asyncio.create_task(coro)
|
|
task.add_done_callback(lambda item: item.exception() if not item.cancelled() and item.exception() else None)
|
|
|
|
|
|
async def clear_owned_chat_threads(
|
|
*,
|
|
store: ThreadMetaStore,
|
|
user_id: str,
|
|
checkpointer: Any = None,
|
|
session_factory: async_sessionmaker | None = None,
|
|
paths: Paths | None = None,
|
|
app: Any = None,
|
|
wait_for_files: bool = False,
|
|
) -> int:
|
|
"""Delete the user's ordinary chats and all of that user's agent chats."""
|
|
|
|
started = time.perf_counter()
|
|
thread_ids = await list_owned_chat_thread_ids(store, user_id, session_factory=session_factory)
|
|
list_ms = _elapsed_ms(started)
|
|
if not thread_ids:
|
|
logger.info("Cleared 0 chat threads for user %s list_ms=%d", user_id, list_ms)
|
|
return 0
|
|
|
|
fenced: list[tuple[str, dict[str, Any]]] = []
|
|
if app is not None:
|
|
from app.gateway.analysis_project_memory import (
|
|
invalidate_project_memory_before_thread_mutation,
|
|
)
|
|
|
|
for thread_id in await _bound_thread_ids(session_factory, thread_ids):
|
|
project = await invalidate_project_memory_before_thread_mutation(
|
|
app=app,
|
|
thread_id=thread_id,
|
|
source_message_ids=None,
|
|
)
|
|
if project:
|
|
fenced.append((thread_id, project))
|
|
|
|
doc_ids = await _deposit_document_ids(session_factory, thread_ids)
|
|
|
|
checkpoint_started = time.perf_counter()
|
|
await _delete_checkpoints(checkpointer, thread_ids)
|
|
checkpoint_ms = _elapsed_ms(checkpoint_started)
|
|
|
|
rows_started = time.perf_counter()
|
|
if session_factory is not None:
|
|
await _delete_related_rows(session_factory, thread_ids, user_id)
|
|
else:
|
|
for thread_id in thread_ids:
|
|
await store.delete(thread_id, user_id=user_id)
|
|
rows_ms = _elapsed_ms(rows_started)
|
|
|
|
path_manager = paths if paths is not None else get_paths()
|
|
file_job = asyncio.to_thread(_remove_thread_dirs, path_manager, user_id, thread_ids)
|
|
if wait_for_files:
|
|
await file_job
|
|
else:
|
|
_schedule_background(file_job)
|
|
|
|
if app is not None and fenced:
|
|
from app.gateway.analysis_project_memory import (
|
|
finalize_project_memory_after_thread_mutation,
|
|
)
|
|
|
|
for thread_id, project in fenced:
|
|
await finalize_project_memory_after_thread_mutation(
|
|
app=app,
|
|
project=project,
|
|
thread_id=thread_id,
|
|
thread_deleted=True,
|
|
)
|
|
|
|
if app is not None and doc_ids:
|
|
_schedule_background(_delete_remote_documents(app, doc_ids))
|
|
|
|
logger.info(
|
|
"Cleared %d chat threads for user %s list_ms=%d checkpoint_ms=%d rows_ms=%d total_ms=%d files=%s",
|
|
len(thread_ids),
|
|
user_id,
|
|
list_ms,
|
|
checkpoint_ms,
|
|
rows_ms,
|
|
_elapsed_ms(started),
|
|
"sync" if wait_for_files else "bg",
|
|
)
|
|
return len(thread_ids)
|