Spaces:
Sleeping
Sleeping
File size: 3,143 Bytes
95218af da6c4b2 95218af 829d4e1 e020ab2 95218af 9ef077e 95218af 9ef077e 95218af | 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 | import os
import uuid
from fastapi import FastAPI, Request, File, UploadFile, Form
from run_model import generate_response
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel, Field
from rag_util import Create_RAG_Prompt, ProcessDocuments
import time
class DataModel(BaseModel):
data: str
id : str = ""
chatID :str = Field(..., min_length=36, description="Chat ID must not be empty")
app = FastAPI(
docs_url=None, # disables /docs
redoc_url=None, # disables /redoc
openapi_url=None # disables /openapi.json
)
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")
chat_history = {}
CHAT_START_TIME = time.time()
@app.get("/", response_class=HTMLResponse)
async def root(request: Request):
chatID = uuid.uuid4()
return templates.TemplateResponse("index.html", {"request" : request, "chatID" : chatID})
@app.post("/generate")
async def generate(data: DataModel):
# dataID = bleach.clean(data.id)
if data.id == "":
input_text = data.data
# input_text = bleach.clean(data.data)
history = chat_history[data.chatID] = [] if data.chatID not in chat_history else chat_history[data.chatID]
llm_response = generate_response(prompt=input_text, history=history)
if data.chatID not in chat_history:
chat_history[data.chatID] = []
chat_history[data.chatID].append({"role": "user", "content": input_text})
chat_history[data.chatID].append({"role": "assistant", "content": llm_response})
return {"llm_response" : llm_response}
else :
input_text = data.data
# input_text = bleach.clean(data.data)
prompt, context = Create_RAG_Prompt(input_text, chatID=data.chatID, history=chat_history[data.chatID] if chat_history[data.chatID] else [])
llm_response = generate_response(prompt=prompt, context=context)
if data.chatID not in chat_history:
chat_history[data.chatID] = []
chat_history[data.chatID].append({"role": "user", "content": input_text})
chat_history[data.chatID].append({"role": "assistant", "content": llm_response})
return {"llm_response" : llm_response}
def isValidPDF(file):
if not file.filename.endswith(".pdf") or file.content_type != "application/pdf":
return False
return True
@app.post("/upload_pdf")
async def upload_pdf(file: UploadFile = File(...), chatID:str = Form(...)):
fileID = uuid.uuid4()
save_path = f"./uploads/{fileID}.pdf"
# Ensure uploads directory exists
os.makedirs(os.path.dirname(save_path), exist_ok=True)
if isValidPDF(file):
with open(save_path, "wb") as f:
f.write(await file.read())
# print(save_path)
ProcessDocuments(save_path, chatID)
return {"filename": fileID, "message": "PDF uploaded successfully."}
else :
return {"filename" : "UPLOADED_FILE", "message" : "Uploaded File is not a valid PDF."}
|