Improve file and conversation services to use repository pattern

src/controllers/_file_controller.py

@@ -15,7 +15,7 @@ class FileController:
 
     async def insert_file(self, file: UploadFile = File(...)):
         try:
-            user_id = "67987a511dd18e6ae8d0671c"
+            user_id = "67987b4a7ce47dbbdcbb28f1"
             async with self.service() as service:
                 return await service.insert_file(file=file, user_id=user_id)
         except Exception as e:

src/middlewares/_authentication.py

@@ -15,6 +15,7 @@ EXEMPT_ROUTES = [
     "/api/v1/auth/logout",
     "/api/v1/auth/register",
     "/api/v1/conversations",
+    "/api/v1/files",
 ]
 
 EXEMPT_ROUTE_PATTERNS = [

src/repositories/_base_repository.py

@@ -8,13 +8,19 @@ class BaseRepository:
     def __init__(self, model: Type[T]):
         self.model = model
 
-    async def insert_one(self, data: T):
+    async def insert_one(self, data: T) -> T:
         try:
             return await data.insert(link_rule=WriteRules.WRITE)
         except Exception as e:
             raise ValueError(f"Failed to insert data: {str(e)}")
 
-    async def get_all(self, page: int = 1, page_size: int = None, filter_by: Dict = None, order_by: Dict = None):
+    async def get_all(
+        self,
+        page: int = 1,
+        page_size: int = None,
+        filter_by: Dict = None,
+        order_by: Dict = None,
+    ):
         query = filter_by or {}
         cursor = self.model.find(query, fetch_links=True)
 
@@ -27,7 +33,7 @@ class BaseRepository:
 
         return await cursor.to_list(length=page_size if page_size else None)
 
-    async def get_by_id(self, _id: str):
+    async def get_by_id(self, _id: str) -> T:
         return await self.model.get(_id, fetch_links=True)
 
     async def update(self, _id: str, update_data: dict):
@@ -37,6 +43,7 @@ class BaseRepository:
 
         updated_fields = {**document.model_dump(), **update_data.model_dump()}
         return await document.set(updated_fields)
+
     async def delete(self, _id: str):
 
         document = await self.model.get(_id)
@@ -45,4 +52,4 @@ class BaseRepository:
         return await document.delete(link_rule=DeleteRules.DELETE)
 
     async def count_documents(self):
-        return await self.model.count()
+        return await self.model.count()

src/repositories/_message_repository.py

@@ -1,4 +1,4 @@
-from src.models import Message
+from src.models import Message, Conversation
 
 from ._base_repository import BaseRepository
 
@@ -6,3 +6,19 @@ from ._base_repository import BaseRepository
 class MessageRepository(BaseRepository):
     def __init__(self):
         super().__init__(model=Message)
+        self.model = Message
+
+    async def get_messages_by_conversation(
+        self, conversation: Conversation
+    ) -> list[Message]:
+        try:
+            messages = await self.model.find_many(
+                self.model.conversation.id == conversation.id, fetch_links=True
+            ).to_list()
+
+            if messages:
+                return messages
+            else:
+                raise ValueError("Messages not found")
+        except Exception as e:
+            raise ValueError(f"Failed to get messages: {str(e)}")

src/services/_conversation_service.py

@@ -1,10 +1,9 @@
 from fastapi import WebSocket
-from bson import ObjectId
-from beanie import WriteRules
 
 from src.config import logger
 from src.utils import OpenAIClient
 from src.models import Conversation, User, Message
+from src.repositories import ConversationRepository, MessageRepository, UserRepository
 
 from ._websocket_service import WebSocketService
 
@@ -13,6 +12,9 @@ class ConversationService:
     def __init__(self):
         self.openai_client = OpenAIClient
         self.websocket_service = WebSocketService
+        self.conversation_repository = ConversationRepository()
+        self.message_repository = MessageRepository()
+        self.user_repository = UserRepository()
 
     async def __aenter__(self):
         return self
@@ -21,10 +23,11 @@ class ConversationService:
         pass
 
     async def create_conversation(self, user_id, modality):
-        user_object = await User.get(ObjectId(user_id))
-        conversation_object = Conversation(user=user_object, summary="")
-        conversation = await conversation_object.save(link_rule=WriteRules.WRITE)
-        print("Conversation ID: ", conversation.id)
+        user_object = await self.user_repository.get_by_id(user_id)
+        conversation_object = await self.conversation_repository.insert_one(
+            Conversation(user=user_object, summary="")
+        )
+        print("Conversation ID: ", conversation_object.id)
 
         text_mode_only = True if modality == "text" else False
         async with self.openai_client() as client:
@@ -33,7 +36,7 @@ class ConversationService:
             )
 
         return {
-            "conversation_id": str(conversation.id),
+            "conversation_id": str(conversation_object.id),
             "session_data": session_data,
         }
 
@@ -48,11 +51,14 @@ class ConversationService:
             await ws_service.handle_conversation(websocket)
 
     async def create_conversation_summary(self, conversation_id):
-        conversation_object = await Conversation.get(ObjectId(conversation_id))
-        messages: list[Message] = await Message.find_many(
-            Message.conversation.id == conversation_object.id,
-            fetch_links=True,
-        ).to_list()
+        conversation_object: Conversation = (
+            await self.conversation_repository.get_by_id(conversation_id)
+        )
+        messages: list[Message] = (
+            await self.message_repository.get_messages_by_conversation(
+                conversation_object
+            )
+        )
 
         conversation_history = "\n".join(
             [f"{message.role}: {message.content}" for message in messages]
@@ -64,6 +70,7 @@ class ConversationService:
             )
 
         conversation_object.summary = conversation_summary
-        await conversation_object.save(link_rule=WriteRules.WRITE)
-
-        return {"conversation_summary": conversation_summary}
+        updated_converation: Conversation = await self.conversation_repository.update(
+            conversation_object.id, conversation_object
+        )
+        return {"conversation_summary": updated_converation.summary}

src/services/_file_service.py

@@ -1,14 +1,18 @@
 from fastapi import UploadFile
-from beanie import WriteRules
 
-from src.repositories import FileRepository, VectorStoreRecordIdRepository
-from src.models import File, VectorStoreRecordId
 from src.utils import PineconeClient
+from src.models import File, VectorStoreRecordId
+from src.repositories import (
+    FileRepository,
+    VectorStoreRecordIdRepository,
+    UserRepository,
+)
 
 
 class FileService:
     def __init__(self):
         self.pinecone_client = PineconeClient
+        self.user_repository = UserRepository()
         self.file_repository = FileRepository()
         self.vector_store_record_id_repository = VectorStoreRecordIdRepository()
 
@@ -20,6 +24,8 @@ class FileService:
 
     async def insert_file(self, file: UploadFile, user_id: str):
         try:
+            user_object = await self.user_repository.get_by_id(user_id)
+
             metadata = [{"file_name": file.filename}]
             while content := await file.read():
                 async with self.pinecone_client() as client:
@@ -27,18 +33,18 @@ class FileService:
                         text=[content.decode()], metadata=metadata
                     )
 
-            file_object = File(
-                user_id=user_id,
-                file_name=file.filename,
-                file_type=file.content_type,
+            file_object = await self.file_repository.insert_one(
+                File(
+                    user=user_object,
+                    file_name=file.filename,
+                    file_type=file.content_type,
+                )
             )
-            await file_object.save(link_rule=WriteRules.WRITE)
 
             for id in ids:
-                vector_store_record_id_object = VectorStoreRecordId(
-                    file_id=file_object.id, record_id=id
+                await self.vector_store_record_id_repository.insert_one(
+                    VectorStoreRecordId(file=file_object, record_id=id)
                 )
-                await vector_store_record_id_object.save(link_rule=WriteRules.WRITE)
             return file_object
 
         finally:

src/services/_websocket_service.py

@@ -1,13 +1,12 @@
 import json
 from fastapi import WebSocket
-from bson import ObjectId
-from beanie import WriteRules
 
-from ._file_service import FileService
 from src.utils import OpenAIClient
-from src.repositories import ConversationRepository, MessageRepository
-from src.models import Conversation, Message
+from src.models import Message
 from src.config import logger
+from src.repositories import ConversationRepository, MessageRepository
+
+from ._file_service import FileService
 
 
 class WebSocketService:
@@ -203,23 +202,29 @@ class WebSocketService:
 
     async def handle_user_message(self, message_content, conversation_id):
         logger.info(f"User Query: {message_content}")
-        conversation_object = await Conversation.get(ObjectId(conversation_id))
-        message_object = Message(
-            conversation=conversation_object,
-            role="user",
-            content=message_content,
+        conversation_object = await self.conversation_repository.get_by_id(
+            conversation_id
+        )
+        await self.message_repository.insert_one(
+            Message(
+                conversation=conversation_object,
+                role="user",
+                content=message_content,
+            )
         )
-        return await message_object.save(link_rule=WriteRules.WRITE)
 
     async def handle_ai_message(self, message_content, conversation_id):
         logger.info(f"AI Response: {message_content}")
-        conversation_object = await Conversation.get(ObjectId(conversation_id))
-        message_object = Message(
-            conversation=conversation_object,
-            role="assistant",
-            content=message_content,
+        conversation_object = await self.conversation_repository.get_by_id(
+            conversation_id
+        )
+        await self.message_repository.insert_one(
+            Message(
+                conversation=conversation_object,
+                role="assistant",
+                content=message_content,
+            )
         )
-        return await message_object.save(link_rule=WriteRules.WRITE)
 
     async def handle_ai_function_call(self, data):
         tool_name = data["name"]