Spaces:
Runtime error
Runtime error
File size: 2,945 Bytes
5e8fd8b 01697df 1c8932d c27225b 553dd69 676f7fc 4929aba 01697df 5e8fd8b ab4a5ae 5e8fd8b 9f0a9ca c27225b 5e8fd8b ab4a5ae 676f7fc 01697df 608245a 6eb1c7e 553dd69 6eb1c7e 95d75bf 6eb1c7e 95d75bf 628c689 608245a 3f857b9 01697df 676f7fc 608245a 676f7fc 01697df 676f7fc 5e8fd8b 6a3901a 713c080 6a3901a 5e8fd8b 713c080 01697df 713c080 bc3c00d 06dc580 31c1d58 bc3c00d 713c080 bc3c00d 5e8fd8b 676f7fc 01697df 676f7fc 5e8fd8b 4c88907 01697df 4c88907 fe880f4 676f7fc 01697df 676f7fc 4c88907 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import os
import traceback
import logging
import pathlib
import time
import re
from typing import List
from fastapi import FastAPI, Request, Depends, HTTPException
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/")
os.makedirs(files_dir)
session_assistant = ChatPDF()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
isBusy = False
def astreamer(generator):
t0 = time.time()
for i in generator:
logger.info(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms)")
yield i
logger.info(f"Over (time {int((time.time()-t0)*1000)}ms)")
@app.get("/query")
async def process_input(text: str):
global isBusy
if isBusy:
raise HTTPException(status_code=503, detail="Server is busy")
isBusy = True
generator = None
if text and len(text.strip()) > 0:
if session_assistant.pdf_count > 0:
text = text.strip()
streaming_response = session_assistant.ask(text)
generator = streaming_response.response_gen
else:
message = "Please add a PDF document first."
generator = re.split(r'(\s)', message)
else:
message = "The provided query is empty."
generator = re.split(r'(\s)', message)
isBusy = False
return StreamingResponse(astreamer(generator), media_type='text/event-stream')
async def parse_body(request: Request):
data: bytes = await request.body()
return data
@app.post("/upload")
def upload(filename: str, data: bytes = Depends(parse_body)):
global isBusy
if isBusy:
raise HTTPException(status_code=503, detail="Server is busy")
isBusy = True
session_assistant.clear()
if data:
try:
print("Filename: " + filename)
path = f"{files_dir}/{filename}"
with open(path, "wb") as f:
f.write(data)
session_assistant.ingest(files_dir)
pathlib.Path(path).unlink()
except Exception as e:
logging.error(traceback.format_exc())
message = "Files inserted successfully."
generator = re.split(r'(\s)', message)
isBusy = False
return StreamingResponse(astreamer(generator), media_type='text/event-stream')
@app.get("/clear")
def clear():
global isBusy
session_assistant.clear()
message = "All files have been cleared. The first query may take a little longer."
generator = re.split(r'(\s)', message)
isBusy = False
return StreamingResponse(astreamer(generator), media_type='text/event-stream')
@app.get("/")
def ping():
return "Pong!"
|