File size: 3,612 Bytes
ca67025 | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | from fastapi import FastAPI, UploadFile, File, HTTPException
import shutil
from fastapi.responses import StreamingResponse
import json
import os
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from scripts.rag import RagPipeline
from scripts.main import set_rag_instance
from scripts.main import stream_chat_response
app = FastAPI(version='1.0', title='FinAI', description="A finetuned qwen model for financial QA with rag.")
rag = RagPipeline()
set_rag_instance(rag)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
class ChatRequest(BaseModel):
message: str
thread_id: str = "default"
@app.get('/')
def health_check():
return {'status' : 'The api is live.'}
@app.post("/upload")
async def upload_document(file: UploadFile = File(...)):
try:
if not file.filename.endswith(".pdf"):
raise HTTPException(
status_code=400,
detail="Only PDF files allowed."
)
os.makedirs("temp_docs", exist_ok=True)
file_path = f"temp_docs/{file.filename}"
with open(file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
rag.delete_all_docs()
docs = rag.load_docs(file_path)
split_docs = rag.split_docs(docs)
rag.add_docs(split_docs)
rag.create_bm25(split_docs)
return {
"status": "success",
"message": "Document uploaded successfully",
"chunks": len(split_docs),
"document": file.filename
}
except Exception as e:
raise HTTPException(
status_code=500,
detail=str(e)
)
@app.delete("/reset")
async def reset_docs():
try:
rag.delete_all_docs()
return {
"status": "success",
"message": "All docs deleted"
}
except Exception as e:
raise HTTPException(
status_code=500,
detail=str(e)
)
@app.post("/chat/stream")
async def chat_stream(request: ChatRequest):
def event_generator():
try:
metadata_sent = False
for chunk in stream_chat_response(
user_message=request.message,
thread_id=request.thread_id
):
if not metadata_sent:
metadata_event = {
"type": "metadata",
"used_rag": chunk["metadata"]["used_rag"],
"sources": chunk["metadata"]["sources"],
"thread_id": chunk["metadata"]["thread_id"]
}
yield f"data: {json.dumps(metadata_event)}\n\n"
metadata_sent = True
token_event = {
"type": "token",
"content": chunk["token"]
}
yield f"data: {json.dumps(token_event)}\n\n"
done_event = {
"type": "done"
}
yield f"data: {json.dumps(done_event)}\n\n"
except Exception as e:
error_event = {
"type": "error",
"message": str(e)
}
yield f"data: {json.dumps(error_event)}\n\n"
return StreamingResponse(
event_generator(),
media_type="text/event-stream"
)
|