File size: 1,815 Bytes
5e8fd8b
 
 
 
 
 
 
 
 
ab4a5ae
5e8fd8b
 
 
 
 
 
 
 
 
 
 
 
 
9f0a9ca
5e8fd8b
2d12855
5e8fd8b
04305be
 
 
 
ab4a5ae
5e8fd8b
ab4a5ae
5e8fd8b
 
2d12855
04305be
 
2d12855
 
04305be
5e8fd8b
 
 
 
 
2d12855
5e8fd8b
1ff6584
453bbe9
1ff6584
 
9f0a9ca
24992a3
9f0a9ca
2cfa0e6
1ff6584
 
 
453bbe9
 
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
61
62
63
64
65
66
67
68
69
70
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 fastapi.responses import StreamingResponse
from .rag import ChatPDF

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

app = FastAPI(middleware=middleware)

files_dir = os.path.expanduser("~/wtp_be_files/")
session_assistant = ChatPDF()
# session_messages = []

# async def stream_generator(agent_text_stream):
#     for text in agent_text_stream:
#         print(text)
#         yield text

@app.get("/query")
async def process_input(text: str):
    if text and len(text.strip()) > 0:
        text = text.strip()
        print("PRINTING STREAM")
        # agent_text_stream = session_assistant.ask(text)
        # print(agent_text_stream)
        # session_messages.append((text, True))
        # session_messages.append((agent_text, False))
        return StreamingResponse(session_assistant.ask(text), media_type='text/event-stream')


@app.post("/upload")
def upload(files: list[UploadFile]):
    session_assistant.clear()
    # session_messages = []
    
    try:
        os.makedirs(files_dir)
        for file in files:
            try:
                path = f"{files_dir}/{file.filename}"
                file.file.seek(0)
                with open(path, 'wb') as destination:
                    shutil.copyfileobj(file.file, destination)
            finally:
                file.file.close()
    finally:
        session_assistant.ingest(files_dir)
        shutil.rmtree(files_dir)

    return "Files inserted!"


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