Refactor assistant_routes.py and summary.py

backend/modules/assistant/controller/assistant_routes.py

@@ -1,13 +1,13 @@
 from typing import List
 
-from fastapi import APIRouter, Depends, UploadFile
+from fastapi import APIRouter, Depends, HTTPException, UploadFile
 from logger import get_logger
 from middlewares.auth import AuthBearer, get_current_user
 from modules.assistant.dto.inputs import InputAssistant
 from modules.assistant.dto.outputs import AssistantOutput
 from modules.assistant.ito.audio_transcript import audio_transcript_inputs
 from modules.assistant.ito.crawler import crawler_inputs
-from modules.assistant.ito.summary import summary_inputs
+from modules.assistant.ito.summary import SummaryAssistant, summary_inputs
 from modules.assistant.service.assistant import Assistant
 from modules.user.entity.user_identity import UserIdentity
 
@@ -42,5 +42,11 @@ async def process_assistant(
     input: InputAssistant,
     files: List[UploadFile] = None,
     current_user: UserIdentity = Depends(get_current_user),
-) -> InputAssistant:
+) -> InputAssistant | dict:
+    if input.name == "summary":
+        summary_assistant = SummaryAssistant(input=input, files=files)
+        try:
+            summary_assistant.check_input()
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
     return input

backend/modules/assistant/ito/ito.py

@@ -1,52 +1,23 @@
 import random
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from io import BytesIO
 from tempfile import NamedTemporaryFile
-from uuid import UUID
 
 from fastapi import UploadFile
 from logger import get_logger
-from modules.assistant.dto.outputs import AssistantOutput
 from modules.contact_support.controller.settings import ContactsSettings
 from modules.upload.controller.upload_routes import upload_file
-from modules.user.entity.user_identity import UserIdentity
 from packages.emails.send_email import send_email
-from pydantic import BaseModel
 
 logger = get_logger(__name__)
 
 
-class ITO(BaseModel):
-    uploadFile: UploadFile | None = None
-    current_user: UserIdentity = None
-    brain_id: UUID | None = None
-    send_file_email: bool = False
-    url: str | None = None
-
-    def __init__(
-        self,
-        uploadFile: UploadFile,
-        current_user: UserIdentity,
-        brain_id: UUID,
-        send_file_email: bool = False,
-        url: str = None,
-    ):
-        super().__init__(
-            uploadFile=uploadFile,
-            current_user=current_user,
-            brain_id=brain_id,
-            send_file_email=send_file_email,
-            url=url,
-        )
+class ITO(ABC):
 
     @abstractmethod
     async def process_assistant(self):
         pass
 
-    @abstractmethod
-    def assistant_inputs(self) -> AssistantOutput:
-        pass
-
     async def send_output_by_email(
         self, file: UploadFile, name: str, custom_message: str = None
     ):

backend/modules/assistant/ito/summary.py

@@ -1,5 +1,7 @@
 import tempfile
+from typing import List
 
+from fastapi import UploadFile
 from langchain.chains import (
     MapReduceDocumentsChain,
     ReduceDocumentsChain,
@@ -11,6 +13,7 @@ from langchain_community.document_loaders import UnstructuredPDFLoader
 from langchain_core.prompts import PromptTemplate
 from langchain_text_splitters import CharacterTextSplitter
 from logger import get_logger
+from modules.assistant.dto.inputs import InputAssistant
 from modules.assistant.dto.outputs import (
     AssistantOutput,
     InputFile,
@@ -25,14 +28,39 @@ logger = get_logger(__name__)
 
 
 class SummaryAssistant(ITO):
+    input: InputAssistant
+    files: List[UploadFile]
 
     def __init__(
         self,
+        input: InputAssistant,
+        files: List[UploadFile] = None,
         **kwargs,
     ):
         super().__init__(
             **kwargs,
         )
+        self.input = input
+        self.files = files
+
+    def check_input(self):
+        if not self.files:
+            raise ValueError("No file was uploaded")
+        if len(self.files) > 1:
+            raise ValueError("Only one file can be uploaded")
+        if not self.input.inputs.files:
+            raise ValueError("No files key were given in the input")
+        if len(self.input.inputs.files) > 1:
+            raise ValueError("Only one file can be uploaded")
+        if not self.input.inputs.files[0].key == "doc_to_summarize":
+            raise ValueError("The key of the file should be doc_to_summarize")
+        if not self.input.inputs.files[0].value:
+            raise ValueError("No file was uploaded")
+        if not (
+            self.input.outputs.brain.activated or self.input.outputs.email.activated
+        ):
+            raise ValueError("No output was selected")
+        return True
 
     async def process_assistant(self):