Spaces:
Sleeping
Sleeping
File size: 4,465 Bytes
305ef4d | 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 | from fastapi import FastAPI, UploadFile, File, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import os
import shutil
import time
# 1. Import our custom modules from the src/ folder
from src.logger import logger
from src.loader import load_pdf
from src.splitter import split_text
from src.embeddings import get_embeddings
from src.vectorstore import create_vectorstore
from src.rag_chain import build_chain
from fastapi.responses import HTMLResponse
# 2. Initialize FastAPI
app = FastAPI(title="Dynamic PDF RAG API")
# Allow frontend applications (like React) to communicate with this API
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ==========================================
# AUTO-LOGGING MIDDLEWARE
# ==========================================
@app.middleware("http")
async def log_requests(request: Request, call_next):
start_time = time.time()
logger.info(f"Incoming request: {request.method} {request.url.path}")
response = await call_next(request)
process_time = (time.time() - start_time) * 1000
logger.info(f"Completed {request.method} {request.url.path} - Status: {response.status_code} - Time: {process_time:.2f}ms")
return response
# ==========================================
# GLOBAL STATE & MODELS
# ==========================================
class ChatRequest(BaseModel):
message: str
# This holds our LangChain pipeline in memory for the active session
current_chain = None
@app.on_event("startup")
async def startup_event():
logger.info("Starting up Dynamic RAG API Server...")
# ==========================================
# API ENDPOINTS
# ==========================================
@app.post("/upload")
async def upload_pdf(file: UploadFile = File(...)):
"""Accepts a PDF, processes it through the RAG pipeline, and readies the chat."""
global current_chain
logger.info(f"Received file upload: {file.filename}")
# 1. Save the file temporarily in a specific uploads folder
upload_dir = "temp_uploads"
temp_file_path = f"{upload_dir}/{file.filename}"
# Create the directory safely
os.makedirs(upload_dir, exist_ok=True)
try:
with open(temp_file_path, "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
logger.info("Starting document processing pipeline...")
# 2. THE PIPELINE EXECUTES HERE
docs = load_pdf(temp_file_path)
chunks = split_text(docs)
embeddings = get_embeddings()
vectorstore = create_vectorstore(chunks, embeddings)
current_chain = build_chain(vectorstore)
logger.info(f"Successfully processed {file.filename} and activated RAG chain.")
return {"status": "success", "message": f"{file.filename} processed! You can now chat."}
except Exception as e:
logger.error(f"Failed to process PDF: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Failed to process PDF: {str(e)}")
finally:
# 3. Clean up the temporary PDF file to save server space
if os.path.exists(temp_file_path):
os.remove(temp_file_path)
logger.debug(f"Cleaned up temporary file: {temp_file_path}")
@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
"""Answers questions based on the currently uploaded PDF."""
global current_chain
if current_chain is None:
logger.warning("User attempted to chat without uploading a PDF first.")
raise HTTPException(status_code=400, detail="No PDF uploaded yet. Please upload a document first.")
try:
logger.info(f"User asked: '{request.message}'")
answer = current_chain.invoke(request.message)
logger.debug("Successfully generated LLM response.")
return {"reply": answer}
except Exception as e:
logger.error(f"Error during chat generation: {str(e)}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def serve_frontend():
"""Serves the frontend HTML UI."""
with open("index.html", "r", encoding="utf-8") as f:
html_content = f.read()
return HTMLResponse(content=html_content) |