mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-10-26 15:18:16 +03:00
Refactor assistant_routes.py and summary.py
This commit is contained in:
parent
ef4b115218
commit
e7a93da25f
@ -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
|
||||
|
@ -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
|
||||
):
|
||||
|
@ -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):
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user