107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
from datetime import datetime
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from config import POSTGRES_CONNECTION_STRING
|
|
|
|
|
|
async def delete_conversation_by_chat_id(chat_id: str) -> Dict[str, Any]:
|
|
"""Delete wx-agent LangGraph checkpoint state for one chat_id."""
|
|
chat_id = str(chat_id or "").strip()
|
|
if not chat_id:
|
|
raise ValueError("chat_id is empty")
|
|
|
|
counts = {
|
|
"checkpoint_writes": 0,
|
|
"checkpoint_blobs": 0,
|
|
"checkpoints": 0,
|
|
}
|
|
async with await _connect_pool() as pool:
|
|
async with pool.connection() as conn:
|
|
async with conn.cursor() as cur:
|
|
await _delete_checkpoint_threads(cur, chat_id, counts)
|
|
return {"success": True, "chat_id": chat_id, "deleted": counts}
|
|
|
|
|
|
async def delete_conversations_by_time(
|
|
before: Optional[datetime] = None,
|
|
after: Optional[datetime] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Delete LangGraph checkpoint state whose checkpoint timestamp falls in a time range."""
|
|
if before is None and after is None:
|
|
raise ValueError("before or after is required")
|
|
|
|
where_parts = []
|
|
params: List[Any] = []
|
|
if before is not None:
|
|
where_parts.append("(checkpoint->>'ts')::timestamptz < %s")
|
|
params.append(before)
|
|
if after is not None:
|
|
where_parts.append("(checkpoint->>'ts')::timestamptz >= %s")
|
|
params.append(after)
|
|
where_sql = " AND ".join(where_parts)
|
|
|
|
counts = {
|
|
"checkpoint_writes": 0,
|
|
"checkpoint_blobs": 0,
|
|
"checkpoints": 0,
|
|
}
|
|
deleted_thread_ids: List[str] = []
|
|
|
|
async with await _connect_pool() as pool:
|
|
async with pool.connection() as conn:
|
|
async with conn.cursor() as cur:
|
|
await cur.execute(
|
|
f"SELECT DISTINCT thread_id FROM checkpoints WHERE {where_sql}",
|
|
params,
|
|
)
|
|
deleted_thread_ids = [str(row[0]) for row in await cur.fetchall() if row and row[0]]
|
|
|
|
for thread_id in deleted_thread_ids:
|
|
await _delete_exact_thread(cur, thread_id, counts)
|
|
|
|
return {
|
|
"success": True,
|
|
"before": before.isoformat() if before else "",
|
|
"after": after.isoformat() if after else "",
|
|
"thread_count": len(deleted_thread_ids),
|
|
"thread_ids": deleted_thread_ids,
|
|
"deleted": counts,
|
|
}
|
|
|
|
|
|
async def _delete_checkpoint_threads(cur, chat_id: str, counts: Dict[str, int]) -> None:
|
|
thread_like = f"{chat_id}:%"
|
|
await cur.execute(
|
|
"DELETE FROM checkpoint_writes WHERE thread_id = %s OR thread_id LIKE %s",
|
|
(chat_id, thread_like),
|
|
)
|
|
counts["checkpoint_writes"] += cur.rowcount or 0
|
|
await cur.execute(
|
|
"DELETE FROM checkpoint_blobs WHERE thread_id = %s OR thread_id LIKE %s",
|
|
(chat_id, thread_like),
|
|
)
|
|
counts["checkpoint_blobs"] += cur.rowcount or 0
|
|
await cur.execute(
|
|
"DELETE FROM checkpoints WHERE thread_id = %s OR thread_id LIKE %s",
|
|
(chat_id, thread_like),
|
|
)
|
|
counts["checkpoints"] += cur.rowcount or 0
|
|
|
|
|
|
async def _delete_exact_thread(cur, thread_id: str, counts: Dict[str, int]) -> None:
|
|
await cur.execute("DELETE FROM checkpoint_writes WHERE thread_id = %s", (thread_id,))
|
|
counts["checkpoint_writes"] += cur.rowcount or 0
|
|
await cur.execute("DELETE FROM checkpoint_blobs WHERE thread_id = %s", (thread_id,))
|
|
counts["checkpoint_blobs"] += cur.rowcount or 0
|
|
await cur.execute("DELETE FROM checkpoints WHERE thread_id = %s", (thread_id,))
|
|
counts["checkpoints"] += cur.rowcount or 0
|
|
|
|
|
|
async def _connect_pool():
|
|
from psycopg_pool import AsyncConnectionPool
|
|
|
|
return AsyncConnectionPool(
|
|
POSTGRES_CONNECTION_STRING,
|
|
kwargs={"autocommit": True},
|
|
)
|