Update main.py

main.py

@@ -1,364 +1,375 @@
-"""
-FastAPI Backend for Multi-PDF RAG System with Per-User Storage
-===============================================================
-Secure multi-user API with:
-- API key authentication
-- Per-user storage isolation
-- PDF upload and management
-- RAG-based question answering
-- HF persistent storage
-"""
-
-import os
-import asyncio
-from typing import List, Optional, Dict
-from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from threading import Lock
-from rag_engine import RAGEngine
-from hf_storage import create_hf_storage_manager
-from auth import get_current_user
-from dotenv import load_dotenv
-
-# ============================================
-# CONFIGURATION
-# ============================================
-
-load_dotenv()
-
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-HF_TOKEN = os.environ.get("HF_TOKEN")
-HF_REPO = os.environ.get("HF_REPO", "Hamza4100/multi-pdf-storage")
-
-if not GEMINI_API_KEY:
-    raise RuntimeError("❌ GEMINI_API_KEY not set")
-
-hf_storage = create_hf_storage_manager(hf_token=HF_TOKEN, hf_repo=HF_REPO)
-
-app = FastAPI(
-    title="Multi-PDF RAG System",
-    description="Secure multi-user RAG API with persistent storage",
-    version="2.0.0"
-)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# ============================================
-# PER-USER RAG ENGINE MANAGER
-# ============================================
-
-class UserRAGManager:
-    """Manages per-user RAG engine instances with lazy loading."""
-    
-    def __init__(self):
-        self.engines: Dict[str, RAGEngine] = {}
-        self.locks: Dict[str, Lock] = {}
-        self.global_lock = Lock()
-    
-    def get_user_lock(self, user_id: str) -> Lock:
-        """Get or create lock for user."""
-        with self.global_lock:
-            if user_id not in self.locks:
-                self.locks[user_id] = Lock()
-            return self.locks[user_id]
-    
-    async def get_engine(self, user_id: str) -> RAGEngine:
-        """Get or create RAG engine for user (lazy loading)."""
-        if user_id in self.engines:
-            return self.engines[user_id]
-        
-        user_lock = self.get_user_lock(user_id)
-        
-        with user_lock:
-            if user_id in self.engines:
-                return self.engines[user_id]
-            
-            print(f"🔧 Initializing RAG for user {user_id}...")
-            
-            # Sync from HF
-            await asyncio.to_thread(hf_storage.sync_storage_from_hf, user_id)
-            
-            # User-specific paths
-            base_dir = os.path.dirname(os.path.abspath(__file__))
-            user_storage_dir = os.path.join(base_dir, "users", user_id, "storage")
-            
-            # Initialize engine
-            engine = await asyncio.to_thread(
-                RAGEngine,
-                gemini_api_key=GEMINI_API_KEY,
-                storage_dir=user_storage_dir
-            )
-            
-            self.engines[user_id] = engine
-            print(f"✅ RAG ready for user {user_id}")
-            
-            return engine
-
-
-rag_manager = UserRAGManager()
-
-# ============================================
-# MODELS
-# ============================================
-
-class UploadResponse(BaseModel):
-    document_id: str
-    filename: str
-    status: str
-    message: str
-    pages: Optional[int] = None
-    chunks: Optional[int] = None
-
-
-class QueryRequest(BaseModel):
-    question: str
-    top_k: Optional[int] = 5
-
-
-class QueryResponse(BaseModel):
-    answer: str
-    sources: List[dict]
-
-
-class DocumentInfo(BaseModel):
-    doc_id: str
-    filename: str
-    upload_timestamp: str
-    num_chunks: int
-    num_pages: int
-
-
-class StatsResponse(BaseModel):
-    total_documents: int
-    total_chunks: int
-    index_size: int
-
-
-class DeleteResponse(BaseModel):
-    status: str
-    message: str
-
-
-# ============================================
-# STARTUP
-# ============================================
-
-@app.on_event("startup")
-async def startup_event():
-    print("🚀 Multi-PDF RAG System v2.0")
-    print(f"📦 HF Storage: {'Enabled' if hf_storage.enabled else 'Disabled'}")
-    print("✅ Server ready (per-user lazy loading)")
-
-
-# ============================================
-# ENDPOINTS
-# ============================================
-
-@app.get("/health")
-async def health_check():
-    """Health check (no auth required)."""
-    return {"status": "ok"}
-
-
-@app.post("/upload", response_model=UploadResponse)
-async def upload_pdf(
-    file: UploadFile = File(...),
-    user_id: str = Depends(get_current_user)
-):
-    """
-    Upload PDF for authenticated user.
-    Requires: X-API-KEY header
-    """
-    # Validate PDF
-    if not file.filename.lower().endswith('.pdf'):
-        raise HTTPException(400, "Only PDF files allowed")
-    
-    if file.content_type and file.content_type not in ['application/pdf']:
-        raise HTTPException(400, "Invalid MIME type")
-    
-    # Read content
-    content = await file.read()
-    
-    # Size limit (10MB)
-    if len(content) > 10 * 1024 * 1024:
-        raise HTTPException(413, "File too large (max 10MB)")
-    
-    try:
-        engine = await rag_manager.get_engine(user_id)
-        user_lock = rag_manager.get_user_lock(user_id)
-        
-        with user_lock:
-            result = await asyncio.to_thread(
-                engine.upload_document,
-                filename=file.filename,
-                file_content=content,
-                action="auto"
-            )
-            
-            if result["status"] == "success":
-                # Save PDF to user's uploaded_pdfs folder
-                base_dir = os.path.dirname(os.path.abspath(__file__))
-                user_pdfs_dir = os.path.join(base_dir, "users", user_id, "uploaded_pdfs")
-                os.makedirs(user_pdfs_dir, exist_ok=True)
-                
-                pdf_path = os.path.join(user_pdfs_dir, file.filename)
-                with open(pdf_path, "wb") as f:
-                    f.write(content)
-                
-                await asyncio.to_thread(
-                    hf_storage.push_storage_to_hf,
-                    user_id,
-                    f"Upload {file.filename}"
-                )
-                
-                print(f"✅ Upload success for user {user_id}: {file.filename}")
-                
-                return UploadResponse(
-                    document_id=result.get("doc_id", ""),
-                    filename=file.filename,
-                    status="success",
-                    message="Uploaded successfully",
-                    pages=result.get("pages"),
-                    chunks=result.get("chunks")
-                )
-            else:
-                raise HTTPException(400, result.get("message", "Upload failed"))
-    
-    except HTTPException:
-        raise
-    except Exception as e:
-        print(f"❌ Upload error (user {user_id}): {e}")
-        raise HTTPException(500, "Upload failed")
-
-
-@app.post("/query", response_model=QueryResponse)
-async def query_documents(
-    request: QueryRequest,
-    user_id: str = Depends(get_current_user)
-):
-    """
-    Query user's documents using RAG.
-    Requires: X-API-KEY header
-    """
-    try:
-        engine = await rag_manager.get_engine(user_id)
-        
-        result = await asyncio.to_thread(
-            engine.ask,
-            query=request.question,
-            top_k=request.top_k
-        )
-        
-        print(f"✅ Query success for user {user_id}")
-        
-        return QueryResponse(
-            answer=result["answer"],
-            sources=result.get("sources", [])
-        )
-    
-    except Exception as e:
-        print(f"❌ Query error (user {user_id}): {e}")
-        raise HTTPException(500, "Query failed")
-
-
-@app.get("/documents", response_model=List[DocumentInfo])
-async def get_documents(user_id: str = Depends(get_current_user)):
-    """
-    Get all documents for authenticated user.
-    Requires: X-API-KEY header
-    """
-    try:
-        engine = await rag_manager.get_engine(user_id)
-        documents = await asyncio.to_thread(engine.get_all_documents)
-        
-        return [
-            DocumentInfo(
-                doc_id=doc["doc_id"],
-                filename=doc["filename"],
-                upload_timestamp=doc["upload_timestamp"],
-                num_chunks=doc["num_chunks"],
-                num_pages=doc["num_pages"]
-            )
-            for doc in documents
-        ]
-    
-    except Exception as e:
-        print(f"❌ Get documents error (user {user_id}): {e}")
-        raise HTTPException(500, "Failed to retrieve documents")
-
-
-@app.delete("/documents/{doc_id}", response_model=DeleteResponse)
-async def delete_document(
-    doc_id: str,
-    user_id: str = Depends(get_current_user)
-):
-    """
-    Delete document for authenticated user.
-    Requires: X-API-KEY header
-    """
-    try:
-        engine = await rag_manager.get_engine(user_id)
-        user_lock = rag_manager.get_user_lock(user_id)
-        
-        with user_lock:
-            result = await asyncio.to_thread(engine.delete_document, doc_id)
-            
-            if result["status"] == "success":
-                await asyncio.to_thread(
-                    hf_storage.push_storage_to_hf,
-                    user_id,
-                    f"Delete {doc_id}"
-                )
-                
-                print(f"✅ Delete success for user {user_id}: {doc_id}")
-                
-                return DeleteResponse(
-                    status="success",
-                    message=result["message"]
-                )
-            else:
-                raise HTTPException(404, result["message"])
-    
-    except HTTPException:
-        raise
-    except Exception as e:
-        print(f"❌ Delete error (user {user_id}): {e}")
-        raise HTTPException(500, "Deletion failed")
-
-
-@app.get("/stats", response_model=StatsResponse)
-async def get_stats(user_id: str = Depends(get_current_user)):
-    """
-    Get stats for authenticated user.
-    Requires: X-API-KEY header
-    """
-    try:
-        engine = await rag_manager.get_engine(user_id)
-        stats = await asyncio.to_thread(engine.get_stats)
-        
-        return StatsResponse(
-            total_documents=stats["total_documents"],
-            total_chunks=stats["total_chunks"],
-            index_size=stats["index_size"]
-        )
-    
-    except Exception as e:
-        print(f"❌ Stats error (user {user_id}): {e}")
-        raise HTTPException(500, "Failed to retrieve stats")
-
-
-# ============================================
-# MAIN
-# ============================================
-
-if __name__ == "__main__":
-    import uvicorn
-    port = int(os.environ.get("PORT", 8000))
-    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)
+"""
+FastAPI Backend for Multi-PDF RAG System with Per-User Storage
+===============================================================
+Secure multi-user API with:
+- API key authentication
+- Per-user storage isolation
+- PDF upload and management
+- RAG-based question answering
+- HF persistent storage
+"""
+
+import os
+import asyncio
+from typing import List, Optional, Dict
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from threading import Lock
+from rag_engine import RAGEngine
+from hf_storage import create_hf_storage_manager
+from auth import get_current_user
+from user_management import create_hf_user_manager
+from dotenv import load_dotenv
+
+# ============================================
+# CONFIGURATION
+# ============================================
+
+load_dotenv()
+
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_REPO = os.environ.get("HF_REPO", "Hamza4100/multi-pdf-storage")
+
+if not GEMINI_API_KEY:
+    raise RuntimeError("❌ GEMINI_API_KEY not set")
+
+hf_storage = create_hf_storage_manager(hf_token=HF_TOKEN, hf_repo=HF_REPO)
+hf_user_manager = create_hf_user_manager(hf_token=HF_TOKEN, hf_repo=HF_REPO)
+
+app = FastAPI(
+    title="Multi-PDF RAG System",
+    description="Secure multi-user RAG API with persistent storage",
+    version="2.0.0"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ============================================
+# PER-USER RAG ENGINE MANAGER
+# ============================================
+
+class UserRAGManager:
+    """Manages per-user RAG engine instances with lazy loading."""
+    
+    def __init__(self):
+        self.engines: Dict[str, RAGEngine] = {}
+        self.locks: Dict[str, Lock] = {}
+        self.global_lock = Lock()
+    
+    def get_user_lock(self, user_id: str) -> Lock:
+        """Get or create lock for user."""
+        with self.global_lock:
+            if user_id not in self.locks:
+                self.locks[user_id] = Lock()
+            return self.locks[user_id]
+    
+    async def get_engine(self, user_id: str) -> RAGEngine:
+        """Get or create RAG engine for user (lazy loading)."""
+        if user_id in self.engines:
+            return self.engines[user_id]
+        
+        user_lock = self.get_user_lock(user_id)
+        
+        with user_lock:
+            if user_id in self.engines:
+                return self.engines[user_id]
+            
+            print(f"🔧 Initializing RAG for user {user_id}...")
+            
+            # Sync from HF
+            await asyncio.to_thread(hf_storage.sync_storage_from_hf, user_id)
+            
+            # User-specific paths
+            base_dir = os.path.dirname(os.path.abspath(__file__))
+            user_storage_dir = os.path.join(base_dir, "users", user_id, "storage")
+            
+            # Initialize engine
+            engine = await asyncio.to_thread(
+                RAGEngine,
+                gemini_api_key=GEMINI_API_KEY,
+                storage_dir=user_storage_dir
+            )
+            
+            self.engines[user_id] = engine
+            print(f"✅ RAG ready for user {user_id}")
+            
+            return engine
+
+
+rag_manager = UserRAGManager()
+
+# ============================================
+# MODELS
+# ============================================
+
+class UploadResponse(BaseModel):
+    document_id: str
+    filename: str
+    status: str
+    message: str
+    pages: Optional[int] = None
+    chunks: Optional[int] = None
+
+
+class QueryRequest(BaseModel):
+    question: str
+    top_k: Optional[int] = 5
+
+
+class QueryResponse(BaseModel):
+    answer: str
+    sources: List[dict]
+
+
+class DocumentInfo(BaseModel):
+    doc_id: str
+    filename: str
+    upload_timestamp: str
+    num_chunks: int
+    num_pages: int
+
+
+class StatsResponse(BaseModel):
+    total_documents: int
+    total_chunks: int
+    index_size: int
+
+
+class DeleteResponse(BaseModel):
+    status: str
+    message: str
+
+
+# ============================================
+# STARTUP
+# ============================================
+
+@app.on_event("startup")
+async def startup_event():
+    print("🚀 Multi-PDF RAG System v2.0")
+    print(f"📦 HF Storage: {'Enabled' if hf_storage.enabled else 'Disabled'}")
+    
+    # Load and display user count from HF
+    if hf_user_manager.enabled:
+        users_count = len(hf_user_manager.get_all_users())
+        print(f"📥 Loaded {users_count} user(s) from HF repository")
+        print(f"✅ User management: Enabled (HF-based)")
+    else:
+        print(f"⚠️ User management: Disabled (check HF_TOKEN and HF_REPO)")
+    
+    print("✅ Server ready (per-user lazy loading)")
+
+
+# ============================================
+# ENDPOINTS
+# ============================================
+
+@app.get("/health")
+async def health_check():
+    """Health check (no auth required)."""
+    return {"status": "ok"}
+
+
+@app.post("/upload", response_model=UploadResponse)
+async def upload_pdf(
+    file: UploadFile = File(...),
+    user_id: str = Depends(get_current_user)
+):
+    """
+    Upload PDF for authenticated user.
+    Requires: X-API-KEY header
+    """
+    # Validate PDF
+    if not file.filename.lower().endswith('.pdf'):
+        raise HTTPException(400, "Only PDF files allowed")
+    
+    if file.content_type and file.content_type not in ['application/pdf']:
+        raise HTTPException(400, "Invalid MIME type")
+    
+    # Read content
+    content = await file.read()
+    
+    # Size limit (10MB)
+    if len(content) > 10 * 1024 * 1024:
+        raise HTTPException(413, "File too large (max 10MB)")
+    
+    try:
+        engine = await rag_manager.get_engine(user_id)
+        user_lock = rag_manager.get_user_lock(user_id)
+        
+        with user_lock:
+            result = await asyncio.to_thread(
+                engine.upload_document,
+                filename=file.filename,
+                file_content=content,
+                action="auto"
+            )
+            
+            if result["status"] == "success":
+                # Save PDF to user's uploaded_pdfs folder
+                base_dir = os.path.dirname(os.path.abspath(__file__))
+                user_pdfs_dir = os.path.join(base_dir, "users", user_id, "uploaded_pdfs")
+                os.makedirs(user_pdfs_dir, exist_ok=True)
+                
+                pdf_path = os.path.join(user_pdfs_dir, file.filename)
+                with open(pdf_path, "wb") as f:
+                    f.write(content)
+                
+                await asyncio.to_thread(
+                    hf_storage.push_storage_to_hf,
+                    user_id,
+                    f"Upload {file.filename}"
+                )
+                
+                print(f"✅ Upload success for user {user_id}: {file.filename}")
+                
+                return UploadResponse(
+                    document_id=result.get("doc_id", ""),
+                    filename=file.filename,
+                    status="success",
+                    message="Uploaded successfully",
+                    pages=result.get("pages"),
+                    chunks=result.get("chunks")
+                )
+            else:
+                raise HTTPException(400, result.get("message", "Upload failed"))
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        print(f"❌ Upload error (user {user_id}): {e}")
+        raise HTTPException(500, "Upload failed")
+
+
+@app.post("/query", response_model=QueryResponse)
+async def query_documents(
+    request: QueryRequest,
+    user_id: str = Depends(get_current_user)
+):
+    """
+    Query user's documents using RAG.
+    Requires: X-API-KEY header
+    """
+    try:
+        engine = await rag_manager.get_engine(user_id)
+        
+        result = await asyncio.to_thread(
+            engine.ask,
+            query=request.question,
+            top_k=request.top_k
+        )
+        
+        print(f"✅ Query success for user {user_id}")
+        
+        return QueryResponse(
+            answer=result["answer"],
+            sources=result.get("sources", [])
+        )
+    
+    except Exception as e:
+        print(f"❌ Query error (user {user_id}): {e}")
+        raise HTTPException(500, "Query failed")
+
+
+@app.get("/documents", response_model=List[DocumentInfo])
+async def get_documents(user_id: str = Depends(get_current_user)):
+    """
+    Get all documents for authenticated user.
+    Requires: X-API-KEY header
+    """
+    try:
+        engine = await rag_manager.get_engine(user_id)
+        documents = await asyncio.to_thread(engine.get_all_documents)
+        
+        return [
+            DocumentInfo(
+                doc_id=doc["doc_id"],
+                filename=doc["filename"],
+                upload_timestamp=doc["upload_timestamp"],
+                num_chunks=doc["num_chunks"],
+                num_pages=doc["num_pages"]
+            )
+            for doc in documents
+        ]
+    
+    except Exception as e:
+        print(f"❌ Get documents error (user {user_id}): {e}")
+        raise HTTPException(500, "Failed to retrieve documents")
+
+
+@app.delete("/documents/{doc_id}", response_model=DeleteResponse)
+async def delete_document(
+    doc_id: str,
+    user_id: str = Depends(get_current_user)
+):
+    """
+    Delete document for authenticated user.
+    Requires: X-API-KEY header
+    """
+    try:
+        engine = await rag_manager.get_engine(user_id)
+        user_lock = rag_manager.get_user_lock(user_id)
+        
+        with user_lock:
+            result = await asyncio.to_thread(engine.delete_document, doc_id)
+            
+            if result["status"] == "success":
+                await asyncio.to_thread(
+                    hf_storage.push_storage_to_hf,
+                    user_id,
+                    f"Delete {doc_id}"
+                )
+                
+                print(f"✅ Delete success for user {user_id}: {doc_id}")
+                
+                return DeleteResponse(
+                    status="success",
+                    message=result["message"]
+                )
+            else:
+                raise HTTPException(404, result["message"])
+    
+    except HTTPException:
+        raise
+    except Exception as e:
+        print(f"❌ Delete error (user {user_id}): {e}")
+        raise HTTPException(500, "Deletion failed")
+
+
+@app.get("/stats", response_model=StatsResponse)
+async def get_stats(user_id: str = Depends(get_current_user)):
+    """
+    Get stats for authenticated user.
+    Requires: X-API-KEY header
+    """
+    try:
+        engine = await rag_manager.get_engine(user_id)
+        stats = await asyncio.to_thread(engine.get_stats)
+        
+        return StatsResponse(
+            total_documents=stats["total_documents"],
+            total_chunks=stats["total_chunks"],
+            index_size=stats["index_size"]
+        )
+    
+    except Exception as e:
+        print(f"❌ Stats error (user {user_id}): {e}")
+        raise HTTPException(500, "Failed to retrieve stats")
+
+
+# ============================================
+# MAIN
+# ============================================
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)