mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-09-11 14:36:35 +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,
|
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"}
|
||||||
|
@ -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"}
|
||||||
|
@ -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"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user