File size: 1,413 Bytes
5e8fd8b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1ff6584
1851012
1ff6584
 
24992a3
713fb30
2cfa0e6
1ff6584
 
 
713fb30
1851012
5e8fd8b
 
 
 
 
 
 
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
import os
import shutil
import tempfile
from tempfile import NamedTemporaryFile
from pathlib import Path

from fastapi import FastAPI, UploadFile
from fastapi.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
from .rag import ChatPDF

middleware = [
    Middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=['*'],
        allow_headers=['*']
    )
]

app = FastAPI(middleware=middleware)

session_assistant = ChatPDF()
session_messages = []

@app.get("/query")
def process_input(text: str):
    if text and len(text.strip()) > 0:
        text = text.strip()
        agent_text = session_assistant.ask(text)
        session_messages.append((text, True))
        session_messages.append((agent_text, False))
        return agent_text


@app.post("/upload")
def upload(files: list[UploadFile]):
    session_assistant.clear()
    session_messages = []
    
    try:
        os.mkdir("~/wtp_be_files/")
        for file in files:
            try:
                file.file.seek(0)
                with Path("~/wtp_be_files/").open("wb") as destination:
                    shutil.copyfileobj(file.file, destination)
            finally:
                file.file.close()
    finally:
        session_assistant.ingest("~/wtp_be_files/")
        shutil.rmtree("~/wtp_be_files/")

    return "Files inserted!"


@app.get("/")
def ping():
    return "Pong!"