Refactor assistant_routes and ITO classes

This commit is contained in:
Stan Girard 2024-04-11 18:28:04 +02:00
parent e7a93da25f
commit 88dce64b84
3 changed files with 30 additions and 18 deletions

View File

@ -42,11 +42,14 @@ async def process_assistant(
input: InputAssistant, input: InputAssistant,
files: List[UploadFile] = None, files: List[UploadFile] = None,
current_user: UserIdentity = Depends(get_current_user), current_user: UserIdentity = Depends(get_current_user),
) -> InputAssistant | dict: ):
if input.name == "summary": if input.name == "summary":
summary_assistant = SummaryAssistant(input=input, files=files) summary_assistant = SummaryAssistant(
input=input, files=files, current_user=current_user
)
try: try:
summary_assistant.check_input() summary_assistant.check_input()
return await summary_assistant.process_assistant()
except ValueError as e: except ValueError as e:
raise HTTPException(status_code=400, detail=str(e)) raise HTTPException(status_code=400, detail=str(e))
return input return {"message": "Assistant not found"}

View File

@ -1,18 +1,25 @@
import random import random
from abc import ABC, abstractmethod from abc import abstractmethod
from io import BytesIO from io import BytesIO
from tempfile import NamedTemporaryFile from tempfile import NamedTemporaryFile
from typing import List
from fastapi import UploadFile from fastapi import UploadFile
from logger import get_logger from logger import get_logger
from modules.assistant.dto.inputs import InputAssistant
from modules.contact_support.controller.settings import ContactsSettings from modules.contact_support.controller.settings import ContactsSettings
from modules.upload.controller.upload_routes import upload_file 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 packages.emails.send_email import send_email
from pydantic import BaseModel
logger = get_logger(__name__) logger = get_logger(__name__)
class ITO(ABC): class ITO(BaseModel):
input: InputAssistant
files: List[UploadFile]
current_user: UserIdentity
@abstractmethod @abstractmethod
async def process_assistant(self): async def process_assistant(self):
@ -71,7 +78,7 @@ class ITO(ABC):
headers={"content-type": "text/plain"}, headers={"content-type": "text/plain"},
) )
if self.send_file_email: if self.input.outputs.email.activated:
await self.send_output_by_email( await self.send_output_by_email(
file_to_upload, file_to_upload,
new_filename, new_filename,
@ -80,11 +87,12 @@ class ITO(ABC):
# Reset to start of file before upload # Reset to start of file before upload
file_to_upload.file.seek(0) file_to_upload.file.seek(0)
await upload_file( if self.input.outputs.brain.activated:
uploadFile=file_to_upload, await upload_file(
brain_id=self.brain_id, uploadFile=file_to_upload,
current_user=self.current_user, brain_id=self.input.outputs.brain.value,
chat_id=None, current_user=self.current_user,
) chat_id=None,
)
return {"message": f"{file_description} generated successfully"} return {"message": f"{file_description} generated successfully"}

View File

@ -23,25 +23,26 @@ from modules.assistant.dto.outputs import (
Outputs, Outputs,
) )
from modules.assistant.ito.ito import ITO from modules.assistant.ito.ito import ITO
from modules.user.entity.user_identity import UserIdentity
logger = get_logger(__name__) logger = get_logger(__name__)
class SummaryAssistant(ITO): class SummaryAssistant(ITO):
input: InputAssistant
files: List[UploadFile]
def __init__( def __init__(
self, self,
input: InputAssistant, input: InputAssistant,
files: List[UploadFile] = None, files: List[UploadFile] = None,
current_user: UserIdentity = None,
**kwargs, **kwargs,
): ):
super().__init__( super().__init__(
input=input,
files=files,
current_user=current_user,
**kwargs, **kwargs,
) )
self.input = input
self.files = files
def check_input(self): def check_input(self):
if not self.files: if not self.files:
@ -68,7 +69,7 @@ class SummaryAssistant(ITO):
tmp_file = tempfile.NamedTemporaryFile(delete=False) tmp_file = tempfile.NamedTemporaryFile(delete=False)
# Write the file to the temporary file # 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 # 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) content = map_reduce_chain.run(split_docs)
return await self.create_and_upload_processed_file( return await self.create_and_upload_processed_file(
content, self.uploadFile.filename, "Summary" content, self.files[0].filename, "Summary"
) )