Add delete all chats option (#2515)

* Add delete all chats option

* post rebase fixes

* final validation

* minor cleanup

* move up

backend/alembic/versions/35e518e0ddf4_properly_cascade.py

@@ -0,0 +1,121 @@
+"""properly_cascade
+
+Revision ID: 35e518e0ddf4
+Revises: 91a0a4d62b14
+Create Date: 2024-09-20 21:24:04.891018
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "35e518e0ddf4"
+down_revision = "91a0a4d62b14"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Update chat_message foreign key constraint
+    op.drop_constraint(
+        "chat_message_chat_session_id_fkey", "chat_message", type_="foreignkey"
+    )
+    op.create_foreign_key(
+        "chat_message_chat_session_id_fkey",
+        "chat_message",
+        "chat_session",
+        ["chat_session_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+    # Update chat_message__search_doc foreign key constraints
+    op.drop_constraint(
+        "chat_message__search_doc_chat_message_id_fkey",
+        "chat_message__search_doc",
+        type_="foreignkey",
+    )
+    op.drop_constraint(
+        "chat_message__search_doc_search_doc_id_fkey",
+        "chat_message__search_doc",
+        type_="foreignkey",
+    )
+
+    op.create_foreign_key(
+        "chat_message__search_doc_chat_message_id_fkey",
+        "chat_message__search_doc",
+        "chat_message",
+        ["chat_message_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.create_foreign_key(
+        "chat_message__search_doc_search_doc_id_fkey",
+        "chat_message__search_doc",
+        "search_doc",
+        ["search_doc_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+    # Add CASCADE delete for tool_call foreign key
+    op.drop_constraint("tool_call_message_id_fkey", "tool_call", type_="foreignkey")
+    op.create_foreign_key(
+        "tool_call_message_id_fkey",
+        "tool_call",
+        "chat_message",
+        ["message_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+
+def downgrade() -> None:
+    # Revert chat_message foreign key constraint
+    op.drop_constraint(
+        "chat_message_chat_session_id_fkey", "chat_message", type_="foreignkey"
+    )
+    op.create_foreign_key(
+        "chat_message_chat_session_id_fkey",
+        "chat_message",
+        "chat_session",
+        ["chat_session_id"],
+        ["id"],
+    )
+
+    # Revert chat_message__search_doc foreign key constraints
+    op.drop_constraint(
+        "chat_message__search_doc_chat_message_id_fkey",
+        "chat_message__search_doc",
+        type_="foreignkey",
+    )
+    op.drop_constraint(
+        "chat_message__search_doc_search_doc_id_fkey",
+        "chat_message__search_doc",
+        type_="foreignkey",
+    )
+
+    op.create_foreign_key(
+        "chat_message__search_doc_chat_message_id_fkey",
+        "chat_message__search_doc",
+        "chat_message",
+        ["chat_message_id"],
+        ["id"],
+    )
+    op.create_foreign_key(
+        "chat_message__search_doc_search_doc_id_fkey",
+        "chat_message__search_doc",
+        "search_doc",
+        ["search_doc_id"],
+        ["id"],
+    )
+
+    # Revert tool_call foreign key constraint
+    op.drop_constraint("tool_call_message_id_fkey", "tool_call", type_="foreignkey")
+    op.create_foreign_key(
+        "tool_call_message_id_fkey",
+        "tool_call",
+        "chat_message",
+        ["message_id"],
+        ["id"],
+    )

backend/onyx/db/chat.py

@@ -316,6 +316,23 @@ def update_chat_session(
     return chat_session
 
 
+def delete_all_chat_sessions_for_user(
+    user: User | None, db_session: Session, hard_delete: bool = HARD_DELETE_CHATS
+) -> None:
+    user_id = user.id if user is not None else None
+
+    query = db_session.query(ChatSession).filter(
+        ChatSession.user_id == user_id, ChatSession.onyxbot_flow.is_(False)
+    )
+
+    if hard_delete:
+        query.delete(synchronize_session=False)
+    else:
+        query.update({ChatSession.deleted: True}, synchronize_session=False)
+
+    db_session.commit()
+
+
 def delete_chat_session(
     user_id: UUID | None,
     chat_session_id: UUID,

backend/onyx/db/models.py

@@ -1010,7 +1010,7 @@ class ChatSession(Base):
         "ChatFolder", back_populates="chat_sessions"
     )
     messages: Mapped[list["ChatMessage"]] = relationship(
-        "ChatMessage", back_populates="chat_session"
+        "ChatMessage", back_populates="chat_session", cascade="all, delete-orphan"
     )
     persona: Mapped["Persona"] = relationship("Persona")
 
@@ -1078,6 +1078,8 @@ class ChatMessage(Base):
         "SearchDoc",
         secondary=ChatMessage__SearchDoc.__table__,
         back_populates="chat_messages",
+        cascade="all, delete-orphan",
+        single_parent=True,
     )
 
     tool_call: Mapped["ToolCall"] = relationship(

backend/onyx/server/query_and_chat/chat_backend.py

@@ -35,6 +35,7 @@ from onyx.configs.model_configs import LITELLM_PASS_THROUGH_HEADERS
 from onyx.db.chat import add_chats_to_session_from_slack_thread
 from onyx.db.chat import create_chat_session
 from onyx.db.chat import create_new_chat_message
+from onyx.db.chat import delete_all_chat_sessions_for_user
 from onyx.db.chat import delete_chat_session
 from onyx.db.chat import duplicate_chat_session_for_user_from_slack
 from onyx.db.chat import get_chat_message
@@ -280,6 +281,17 @@ def patch_chat_session(
     return None
 
 
+@router.delete("/delete-all-chat-sessions")
+def delete_all_chat_sessions(
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> None:
+    try:
+        delete_all_chat_sessions_for_user(user=user, db_session=db_session)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
 @router.delete("/delete-chat-session/{session_id}")
 def delete_chat_session_by_id(
     session_id: UUID,

backend/onyx/server/query_and_chat/models.py

@@ -11,6 +11,7 @@ from onyx.chat.models import RetrievalDocs
 from onyx.configs.constants import DocumentSource
 from onyx.configs.constants import MessageType
 from onyx.configs.constants import SearchFeedbackType
+from onyx.configs.constants import SessionType
 from onyx.context.search.models import BaseFilters
 from onyx.context.search.models import ChunkContext
 from onyx.context.search.models import RerankingDetails
@@ -151,6 +152,10 @@ class ChatSessionUpdateRequest(BaseModel):
     sharing_status: ChatSessionSharedStatus
 
 
+class DeleteAllSessionsRequest(BaseModel):
+    session_type: SessionType
+
+
 class RenameChatSessionResponse(BaseModel):
     new_name: str  # This is only really useful if the name is generated