File size: 2,924 Bytes
a3386d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9c2e870
a3386d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from fastapi import BackgroundTasks, APIRouter, Request, UploadFile, Depends, File
from typing import List

from src.helpers.response import ResponseHandler
from src.helpers.json_convertor import convert_to_json

from src.middleware.verifyToken import verify_token
from src.services.message_services import save_question_and_answer_to_chat_history

from src.bot.main import Main
from src.bot.OtherFun import delete_chain_after_delay, process_file

router = APIRouter()

model = Main()


@router.post("/ask-question")
async def askQ(req: Request, token: str = Depends(verify_token)):
    try:
        form_data = await req.json()
        answer = model.ask_question(
            form_data["query"],
            (
                token["username"]
                if form_data["collectionName"] is None
                else form_data["collectionName"]
            ),
        )
        await save_question_and_answer_to_chat_history(
            token["username"],
            {
                "question": form_data["query"],
                "answer": answer,
                "collectionName": (
                    form_data["query"] if form_data["query"] else "CHAT WITH YOUR PDF"
                ),
            },
        )
        return ResponseHandler.success(2000, answer)
    except Exception as error:
        return ResponseHandler.error(9999, error, 500)


@router.post("/create/embedding/{collection_name}")
async def createEmbedding(
    collection_name: str,
    files: List[UploadFile] = File(None),
    tokenData: str = Depends(verify_token),
):
    try:
        if not files:
            return ResponseHandler.error(2003)

        responses = []
        for file in files:
            response = process_file(model, collection_name, file)
            responses.append(response)

        return ResponseHandler.success(2002, response)
    except Exception as error:
        return ResponseHandler.error(9999, error, 500)


@router.post("/create/tmp/chain")
async def createTmpChain(
    background_tasks: BackgroundTasks,
    files: List[UploadFile] = File(...),
    tokenData: str = Depends(verify_token),
):
    try:
        if not files:
            return ResponseHandler.error(2003, error, 500)

        all_contents = b""
        for file in files:
            contents = await file.read()
            all_contents += contents

        file_extension = files[0].filename.split(".")[-1]
        if file_extension == "pdf":
            chain_name = tokenData["username"]
            model.generate_tmp_embedding_and_chain(all_contents, chain_name)
            background_tasks.add_task(delete_chain_after_delay(model, chain_name))
            return ResponseHandler.success(2001)
        elif file_extension == "txt":
            all_contents.decode("utf-8")
            return ResponseHandler.error(2004, error, 500)

    except Exception as error:
        return ResponseHandler.error(9999, error, 500)