From 5250502eb053f94ba0beab41727c732a93f9be25 Mon Sep 17 00:00:00 2001 From: Ana Pacheco Date: Wed, 9 Oct 2024 17:14:39 -0600 Subject: [PATCH] delete conversations, messages, files from db docs chromadb --- backend/api/ResourcesView.py | 24 +++++++++++++++++++++--- backend/managers/MessagesManager.py | 25 ++++++++++++++++++++++--- backend/managers/ResourcesManager.py | 25 +++++++++++++++++++++++-- 3 files changed, 66 insertions(+), 8 deletions(-) diff --git a/backend/api/ResourcesView.py b/backend/api/ResourcesView.py index 7b5cc155..92a2c0f7 100644 --- a/backend/api/ResourcesView.py +++ b/backend/api/ResourcesView.py @@ -1,6 +1,9 @@ from starlette.responses import JSONResponse, Response from common.paths import api_base_url from backend.managers.ResourcesManager import ResourcesManager +from backend.managers.RagManager import RagManager +from backend.managers.MessagesManager import MessagesManager +from backend.managers.ConversationsManager import ConversationsManager from backend.pagination import parse_pagination_params from backend.schemas import ResourceCreateSchema from typing import List @@ -35,9 +38,24 @@ async def put(self, id: str, body: ResourceCreateSchema): return JSONResponse(updated_resource.dict(), status_code=200) async def delete(self, id: str): - success = await self.rm.delete_resource(id) - if not success: - return JSONResponse({"error": "Resource not found"}, status_code=404) + ragm = RagManager() + cm = ConversationsManager() + msgm = MessagesManager() + file_ids = await self.rm.retrieve_resource_files(id) + conversations_ids = await self.rm.retrieve_resource_conversations(id) + + if file_ids: + await ragm.delete_documents_from_chroma(id, file_ids) + await ragm.delete_file_from_db(file_ids) + + if conversations_ids: + for conversation_id in conversations_ids: + await cm.delete_conversation(conversation_id) + await msgm.delete_messages_from_conversation(conversation_id) + + success_resource = await self.rm.delete_resource(id) + if not success_resource: + return JSONResponse({"error": "Resource not found "}, status_code=404) return Response(status_code=204) async def search(self, filter: str = None, range: str = None, sort: str = None): diff --git a/backend/managers/MessagesManager.py b/backend/managers/MessagesManager.py index c3cc891f..91768d8e 100644 --- a/backend/managers/MessagesManager.py +++ b/backend/managers/MessagesManager.py @@ -1,6 +1,6 @@ from uuid import uuid4 from threading import Lock -from sqlalchemy import select +from sqlalchemy import select, delete from backend.models import Message, Conversation, Resource from backend.db import db_session_context from backend.schemas import MessageSchema, MessageCreateSchema @@ -93,7 +93,7 @@ async def create_message(self, message_data: MessageCreateSchema) -> Tuple[Optio return chat_response, None except Exception as e: - return None, f"An unexpected error occurred: {str(e)}" + return None, f"An unexpected error occurred while creating a message: {str(e)}" async def retrieve_message(self, id:str) -> Optional[MessageSchema]: async with db_session_context() as session: @@ -114,4 +114,23 @@ async def retrieve_message(self, id:str) -> Optional[MessageSchema]: def extract_names_from_uri(self, uri): # Extract the last part of the URI after the last "/" model_name = uri.split('/')[-1] - return model_name \ No newline at end of file + return model_name + + async def delete_messages_from_conversation(self, conversation_id): + async with db_session_context() as session: + stmt = select(Message).filter(Message.conversation_id == conversation_id) + result = await session.execute(stmt) + msgs = result.scalars().all() + msgs = [MessageSchema.from_orm(msg) for msg in msgs] + print("msgs: ", msgs) + for msg in msgs: + if not await self._delete_message(msg.id): + return False + return True + + async def _delete_message(self, id: str) -> bool: + async with db_session_context() as session: + stmt = delete(Message).where(Message.id == id) + result = await session.execute(stmt) + await session.commit() + return result.rowcount > 0 \ No newline at end of file diff --git a/backend/managers/ResourcesManager.py b/backend/managers/ResourcesManager.py index 85d3fac5..52b9f2cf 100644 --- a/backend/managers/ResourcesManager.py +++ b/backend/managers/ResourcesManager.py @@ -2,7 +2,7 @@ from threading import Lock import httpx from sqlalchemy import select, update, delete, func -from backend.models import Resource +from backend.models import Resource, File, Conversation from backend.db import db_session_context from backend.schemas import ResourceCreateSchema, ResourceSchema from typing import List, Tuple, Optional, Dict, Any @@ -227,4 +227,25 @@ async def get_ollama_installed_models(self) ->List: models = data.get("models", []) return models else: - return [] \ No newline at end of file + return [] + + async def retrieve_resource_files(self, resource_id: str) -> bool: + async with db_session_context() as session: + result_files = await session.execute(select(File).filter(File.assistant_id == resource_id)) + files = result_files.scalars().all() + if not files: + return None + file_ids = [file.id for file in files] + print("file ids: ", file_ids) + return file_ids + + async def retrieve_resource_conversations(self, resource_id: str) -> bool: + async with db_session_context() as session: + result_conversations = await session.execute(select(Conversation).filter(Conversation.assistant_id == resource_id)) + conversations = result_conversations.scalars().all() + if not conversations: + return None + conversations_ids = [conversation.id for conversation in conversations] + print("conversations_ids: ", conversations_ids) + return conversations_ids + \ No newline at end of file