| 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"
|
| )
|
|
|
|
|
