mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-08-17 02:10:31 +03:00
Refactor assistant_routes and ITO classes
This commit is contained in:
parent
e7a93da25f
commit
88dce64b84
@ -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"}
|
||||
|
@ -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"}
|
||||
|
@ -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"
|
||||
)
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user