Update main.py

main.py

@@ -2,6 +2,8 @@
 FastAPI Backend for NotebookPRO
 Handles RAG, LLM, file processing, and chat management
 """
+from pymongo import MongoClient
+from fastapi import BackgroundTasks
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -42,6 +44,17 @@ from utils.studio_generator import StudioGenerator
 
 # Initialize FastAPI
 app = FastAPI(title="NotebookPRO API", version="2.0.0")
+
+# --- ADD THIS AFTER app = FastAPI(...) ---
+# Initialize MongoDB
+MONGO_URI = os.getenv("MONGO_URI")
+if MONGO_URI:
+    mongo_client = MongoClient(MONGO_URI)
+    db = mongo_client["notebookpro_db"]
+    chats_collection = db["chats"]
+    files_collection = db["processed_files"]
+else:
+    print("WARNING: MONGO_URI not found in environment variables.")
 # CORS - Allow Flutter web to connect
 app.add_middleware(
     CORSMiddleware,
@@ -125,20 +138,20 @@ def get_space_dir(space_id: str):
     return get_data_dir() / "spaces" / space_id
 
 def load_chats_for_space(space_id: str) -> List[Dict]:
-    """Load all chats for a space"""
-    chats_file = get_space_dir(space_id) / "chats.json"
-    if chats_file.exists():
-        with open(chats_file, 'r', encoding='utf-8') as f:
-            return json.load(f)
-    return []
-
-def save_chats_for_space(space_id: str, chats: List[Dict]):
-    """Save chats for a space"""
-    chats_file = get_space_dir(space_id) / "chats.json"
-    chats_file.parent.mkdir(parents=True, exist_ok=True)
-    with open(chats_file, 'w', encoding='utf-8') as f:
-        json.dump(chats, f, indent=2, ensure_ascii=False)
-
+    """Load all chats for a space from MongoDB"""
+    if not MONGO_URI: return []
+    cursor = chats_collection.find({"space_id": space_id}, {"_id": 0})
+    return list(cursor)
+
+def save_chat_to_db(space_id: str, chat: Dict):
+    """Save or update a single chat in MongoDB"""
+    if not MONGO_URI: return
+    chat['space_id'] = space_id
+    chats_collection.update_one(
+        {"id": chat['id'], "space_id": space_id},
+        {"$set": chat},
+        upsert=True
+    )
 def get_chat_title(messages: List[Dict]) -> str:
     """Generate chat title from first user message"""
     for msg in messages:
@@ -456,20 +469,17 @@ async def chat(request: ChatRequest):
         chats = load_chats_for_space(request.space_id)
         
         # Find existing chat or create new
-        chat = None
-        for c in chats:
-            if c['id'] == chat_id:
-                chat = c
-                break
+        # Fetch specific chat from Mongo or create new
+        chat = chats_collection.find_one({"id": chat_id, "space_id": request.space_id}, {"_id": 0})
         
         if not chat:
             chat = {
                 'id': chat_id,
+                'space_id': request.space_id,
                 'messages': [],
                 'created_at': datetime.now().isoformat(),
                 'updated_at': datetime.now().isoformat()
             }
-            chats.append(chat)
         
         # Add messages
         timestamp = datetime.now().isoformat()
@@ -484,175 +494,145 @@ async def chat(request: ChatRequest):
         ])
         chat['updated_at'] = timestamp
         
-        # Save chats
-        save_chats_for_space(request.space_id, chats)
-        
-        return ChatResponse(
-            response=response,
-            sources=sources,
-            chat_id=chat_id,
-            timestamp=timestamp
-        )
+        # Save SINGLE chat directly to MongoDB
+        save_chat_to_db(request.space_id, chat)
     
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/api/spaces/{space_id}/upload")
-async def upload_files(space_id: str, files: List[UploadFile] = File(...)):
-    """Upload and process files for a space"""
+def process_heavy_files_background(space_id: str, saved_file_paths: List[Dict]):
+    """Runs in the background so the HTTP request doesn't timeout"""
     try:
-        # Initialize space
         initialize_space(space_id)
-        
-        # Save uploaded files temporarily
-        space_dir = get_space_dir(space_id)
-        uploads_dir = space_dir / "uploads"
-        uploads_dir.mkdir(parents=True, exist_ok=True)
-        
         processor = DocumentProcessor()
         all_chunks = []
         processed_files = []
-        
-        for file in files:
-            # Save file
-            file_path = uploads_dir / file.filename
-            with open(file_path, "wb") as f:
-                content = await file.read()
-                f.write(content)
+
+        for file_info in saved_file_paths:
+            file_path = Path(file_info['path'])
+            filename = file_info['name']
             
-            # Process file and extract content
-            try:
-                file_data = processor.process_file(file_path)
-                content = file_data['content']
-                
-                # Chunk the content
-                chunks = processor.chunk_text(content, chunk_size=512, overlap=50, semantic=True)
-                
-                # Format chunks for vector database
-                formatted_chunks = []
-                for idx, chunk in enumerate(chunks):
-                    formatted_chunks.append({
-                        'content': chunk,
-                        'metadata': {
-                            'filename': file.filename,
-                            'chunk_index': idx,
-                            'total_chunks': len(chunks),
-                            'source_type': file_data['format']
-                        }
-                    })
-                
-                all_chunks.extend(formatted_chunks)
-                processed_files.append({
-                    'filename': file.filename,
-                    'chunks': len(chunks),
-                    'processed_at': datetime.now().isoformat()
+            # The heavy CPU work
+            file_data = processor.process_file(file_path)
+            chunks = processor.chunk_text(file_data['content'], chunk_size=512, overlap=50, semantic=True)
+            
+            for idx, chunk in enumerate(chunks):
+                all_chunks.append({
+                    'content': chunk,
+                    'metadata': {
+                        'filename': filename,
+                        'chunk_index': idx,
+                        'total_chunks': len(chunks),
+                        'source_type': file_data['format']
+                    }
                 })
-            except Exception as e:
-                # Log error but continue with other files
-                print(f"Error processing {file.filename}: {str(e)}")
-                continue
-        
-        # Add to vector database in batches to avoid size limits
+            
+            processed_files.append({
+                'filename': filename,
+                'space_id': space_id, # Link it to the space
+                'chunks': len(chunks),
+                'processed_at': datetime.now().isoformat()
+            })
+
+        # Upload to Qdrant
         if all_chunks:
-            # Extract texts, metadatas, and generate IDs
             texts = [chunk['content'] for chunk in all_chunks]
             metadatas = [chunk['metadata'] for chunk in all_chunks]
             ids = [f"{space_id}_{idx}_{uuid.uuid4().hex[:8]}" for idx in range(len(all_chunks))]
             
-            # Process in batches of 5000 to avoid ChromaDB batch size limit
-            batch_size = 5000
+            # Batch size set to 100 for Qdrant Cloud limits
+            batch_size = 100
             for i in range(0, len(texts), batch_size):
-                batch_texts = texts[i:i + batch_size]
-                batch_metadatas = metadatas[i:i + batch_size]
-                batch_ids = ids[i:i + batch_size]
-                
-                vector_db.add_documents(batch_texts, batch_metadatas, batch_ids)
-                print(f"Processed batch {i//batch_size + 1}/{(len(texts)-1)//batch_size + 1}")
+                vector_db.add_documents(
+                    texts[i:i + batch_size], 
+                    metadatas[i:i + batch_size], 
+                    ids[i:i + batch_size]
+                )
+
+        # Save metadata directly to MongoDB
+        if processed_files and MONGO_URI:
+            files_collection.insert_many(processed_files)
+
+    except Exception as e:
+        print(f"Background processing failed: {e}")
+
+@app.post("/api/spaces/{space_id}/upload")
+async def upload_files(space_id: str, background_tasks: BackgroundTasks, files: List[UploadFile] = File(...)):
+    """Accepts files quickly and processes them in the background"""
+    try:
+        space_dir = get_space_dir(space_id)
+        uploads_dir = space_dir / "uploads"
+        uploads_dir.mkdir(parents=True, exist_ok=True)
         
-        # Save processed files info
-        processed_file = space_dir / "processed_files.json"
-        existing = []
-        if processed_file.exists():
-            with open(processed_file, 'r') as f:
-                existing = json.load(f)
+        saved_files = []
         
-        existing.extend(processed_files)
-        with open(processed_file, 'w') as f:
-            json.dump(existing, f, indent=2)
+        # 1. Save files to hard drive
+        for file in files:
+            file_path = uploads_dir / file.filename
+            with open(file_path, "wb") as f:
+                content = await file.read()
+                f.write(content)
+            
+            saved_files.append({
+                "name": file.filename,
+                "path": str(file_path)
+            })
+        
+        # 2. Hand heavy math and Mongo saving to background task
+        background_tasks.add_task(process_heavy_files_background, space_id, saved_files)
         
         return {
-            "status": "success",
-            "files_processed": len(processed_files),
-            "total_chunks": len(all_chunks)
+            "status": "processing",
+            "message": f"Successfully received {len(files)} files. Processing in the background."
         }
     
     except Exception as e:
-        raise e  # This strips the wrapper and forces FastAPI to log the raw stack trace
-
+        raise HTTPException(status_code=500, detail=str(e))
 @app.get("/api/spaces/{space_id}/files")
 async def get_files(space_id: str):
-    """Get processed files for a space"""
-    processed_file = get_space_dir(space_id) / "processed_files.json"
-    
-    if processed_file.exists():
-        with open(processed_file, 'r') as f:
-            return json.load(f)
-    
-    return []
+    """Get processed files for a space from MongoDB"""
+    if not MONGO_URI: return []
+    cursor = files_collection.find({"space_id": space_id}, {"_id": 0})
+    return list(cursor)
 
 @app.delete("/api/spaces/{space_id}/files/{filename}")
 async def delete_file(space_id: str, filename: str):
     """Delete a specific file from a space"""
     try:
-        # Remove from processed_files.json
-        processed_file = get_space_dir(space_id) / "processed_files.json"
-        files_data = []
-        
-        if processed_file.exists():
-            with open(processed_file, 'r') as f:
-                files_data = json.load(f)
-            
-            # Filter out the file to delete
-            files_data = [f for f in files_data if f.get('filename') != filename]
-            
-            with open(processed_file, 'w') as f:
-                json.dump(files_data, f, indent=2)
+        # 1. Remove from MongoDB
+        if MONGO_URI:
+            files_collection.delete_one({"space_id": space_id, "filename": filename})
         
-        # Delete the actual file
+        # 2. Delete the actual file
         file_path = get_space_dir(space_id) / "uploads" / filename
         if file_path.exists():
             file_path.unlink()
         
-        # Remove from vector database (if initialized)
-        # Note: This removes all chunks with this filename from metadata
+        # 3. Remove from Qdrant vector database
         if vector_db:
             try:
-                # Get all documents in the collection
-                collection = vector_db.collection
-                results = collection.get()
-                
-                # Find IDs of documents with matching filename
-                ids_to_delete = []
-                for idx, metadata in enumerate(results['metadatas']):
-                    if metadata and metadata.get('filename') == filename:
-                        ids_to_delete.append(results['ids'][idx])
-                
-                # Delete those documents
-                if ids_to_delete:
-                    collection.delete(ids=ids_to_delete)
-                    print(f"Deleted {len(ids_to_delete)} chunks for {filename}")
+                # Qdrant supports deleting by payload filter natively
+                from qdrant_client.http import models
+                vector_db.client.delete(
+                    collection_name=vector_db.collection_name,
+                    points_selector=models.Filter(
+                        must=[
+                            models.FieldCondition(
+                                key="filename",
+                                match=models.MatchValue(value=filename)
+                            )
+                        ]
+                    )
+                )
             except Exception as e:
-                print(f"Error removing from vector DB: {e}")
+                print(f"Error removing from Qdrant DB: {e}")
         
-        return {
-            "status": "success",
-            "message": f"File {filename} deleted"
-        }
+        return {"status": "success", "message": f"File {filename} deleted"}
     
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error deleting file: {str(e)}")
 
-
 # ==================== STUDIO API ROUTES ====================
 # Routes for Notebook, Flashcards, and Quiz features