From 88dce64b84e4a0f781bced02ea544ae410487403 Mon Sep 17 00:00:00 2001
From: Stan Girard <girard.stanislas@gmail.com>
Date: Thu, 11 Apr 2024 18:28:04 +0200
Subject: [PATCH] Refactor assistant_routes and ITO classes

---
 .../assistant/controller/assistant_routes.py  |  9 ++++---
 backend/modules/assistant/ito/ito.py          | 26 ++++++++++++-------
 backend/modules/assistant/ito/summary.py      | 13 +++++-----
 3 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/backend/modules/assistant/controller/assistant_routes.py b/backend/modules/assistant/controller/assistant_routes.py
index 4fbec21ab..62087dd5f 100644
--- a/backend/modules/assistant/controller/assistant_routes.py
+++ b/backend/modules/assistant/controller/assistant_routes.py
@@ -42,11 +42,14 @@ async def process_assistant(
     input: InputAssistant,
     files: List[UploadFile] = None,
     current_user: UserIdentity = Depends(get_current_user),
-) -> InputAssistant | dict:
+):
     if input.name == "summary":
-        summary_assistant = SummaryAssistant(input=input, files=files)
+        summary_assistant = SummaryAssistant(
+            input=input, files=files, current_user=current_user
+        )
         try:
             summary_assistant.check_input()
+            return await summary_assistant.process_assistant()
         except ValueError as e:
             raise HTTPException(status_code=400, detail=str(e))
-    return input
+    return {"message": "Assistant not found"}
diff --git a/backend/modules/assistant/ito/ito.py b/backend/modules/assistant/ito/ito.py
index bc86d0861..44611fb80 100644
--- a/backend/modules/assistant/ito/ito.py
+++ b/backend/modules/assistant/ito/ito.py
@@ -1,18 +1,25 @@
 import random
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from io import BytesIO
 from tempfile import NamedTemporaryFile
+from typing import List
 
 from fastapi import UploadFile
 from logger import get_logger
+from modules.assistant.dto.inputs import InputAssistant
 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(ABC):
+class ITO(BaseModel):
+    input: InputAssistant
+    files: List[UploadFile]
+    current_user: UserIdentity
 
     @abstractmethod
     async def process_assistant(self):
@@ -71,7 +78,7 @@ class ITO(ABC):
             headers={"content-type": "text/plain"},
         )
 
-        if self.send_file_email:
+        if self.input.outputs.email.activated:
             await self.send_output_by_email(
                 file_to_upload,
                 new_filename,
@@ -80,11 +87,12 @@ class ITO(ABC):
 
         # Reset to start of file before upload
         file_to_upload.file.seek(0)
-        await upload_file(
-            uploadFile=file_to_upload,
-            brain_id=self.brain_id,
-            current_user=self.current_user,
-            chat_id=None,
-        )
+        if self.input.outputs.brain.activated:
+            await upload_file(
+                uploadFile=file_to_upload,
+                brain_id=self.input.outputs.brain.value,
+                current_user=self.current_user,
+                chat_id=None,
+            )
 
         return {"message": f"{file_description} generated successfully"}
diff --git a/backend/modules/assistant/ito/summary.py b/backend/modules/assistant/ito/summary.py
index 44284a039..4df4ad52c 100644
--- a/backend/modules/assistant/ito/summary.py
+++ b/backend/modules/assistant/ito/summary.py
@@ -23,25 +23,26 @@ from modules.assistant.dto.outputs import (
     Outputs,
 )
 from modules.assistant.ito.ito import ITO
+from modules.user.entity.user_identity import UserIdentity
 
 logger = get_logger(__name__)
 
 
 class SummaryAssistant(ITO):
-    input: InputAssistant
-    files: List[UploadFile]
 
     def __init__(
         self,
         input: InputAssistant,
         files: List[UploadFile] = None,
+        current_user: UserIdentity = None,
         **kwargs,
     ):
         super().__init__(
+            input=input,
+            files=files,
+            current_user=current_user,
             **kwargs,
         )
-        self.input = input
-        self.files = files
 
     def check_input(self):
         if not self.files:
@@ -68,7 +69,7 @@ class SummaryAssistant(ITO):
         tmp_file = tempfile.NamedTemporaryFile(delete=False)
 
         # Write the file to the temporary file
-        tmp_file.write(self.uploadFile.file.read())
+        tmp_file.write(self.files[0].file.read())
 
         # Now pass the path of the temporary file to the loader
 
@@ -132,7 +133,7 @@ class SummaryAssistant(ITO):
         content = map_reduce_chain.run(split_docs)
 
         return await self.create_and_upload_processed_file(
-            content, self.uploadFile.filename, "Summary"
+            content, self.files[0].filename, "Summary"
         )