aigenrec/luminabackend
0
1import os2import asyncio3from typing import Optional4from db.client import supabase_client5from config.settings import settings6from utils.file_parser import FileParser7from utils.text_chunker import TextChunker8from services.embedding_service import embedding_service9from services.qdrant_service import qdrant_service10from utils.logger import logger11 12class DocumentService:13 def __init__(self):14 self.client = supabase_client15 self.file_parser = FileParser()16 self.text_chunker = TextChunker(17 chunk_size=settings.CHUNK_SIZE,18 overlap=settings.CHUNK_OVERLAP19 )20 21 async def process_document(22 self,23 document_id: str,24 project_id: str,25 file_path: str,26 filename: str27 ):28 """Process uploaded document: extract text, chunk, embed, and store"""29 try:30 # Update status to processing31 await self._update_document_status(document_id, "processing")32 33 loop = asyncio.get_running_loop()34 35 # 1. Extract text (Run in thread pool to avoid blocking)36 logger.info(f"Extracting text from {filename}")37 await self._update_document_status(document_id, "processing", "Extracting text...")38 text = await loop.run_in_executor(None, self.file_parser.extract_text, file_path)39 40 if not text:41 await self._update_document_status(document_id, "failed", "Failed to extract text")42 return43 44 # 2. Chunk text (Run in thread pool to avoid blocking)45 # LangChain's splitter is CPU bound46 logger.info(f"Chunking text from {filename}")47 await self._update_document_status(document_id, "processing", "Chunking text...")48 chunks = await loop.run_in_executor(49 None, 50 lambda: self.text_chunker.chunk_text(text)51 )52 53 if not chunks:54 await self._update_document_status(document_id, "failed", "No chunks generated")55 return56 57 # 3. Generate embeddings (Async)58 logger.info(f"Generating embeddings for {len(chunks)} chunks")59 60 # 3. Create collection FIRST (to ensure it exists before upserting)61 collection_name = f"project_{project_id}"62 await qdrant_service.create_collection(collection_name)63 64 # 4. Generate embeddings and Upsert (Streamed)65 logger.info(f"Generating embeddings and upserting for {len(chunks)} chunks")66 await self._update_document_status(document_id, "processing", f"Generating embeddings ({len(chunks)} chunks)...")67 68 # Batching: Smaller batches (25) + Higher Concurrency (10) for speed69 batch_size = 2570 71 # Prepare batches with start index72 batches = []73 for i in range(0, len(chunks), batch_size):74 batch_data = chunks[i:i + batch_size]75 batches.append((i, batch_data))76 77 total_batches = len(batches)78 79 # Concurrency control (Limit to 10 to simulate ~10 concurrent users/requests)80 semaphore = asyncio.Semaphore(10)81 82 async def process_batch(batch_idx, start_index, batch_data):83 async with semaphore:84 retries = 385 for attempt in range(retries):86 try:87 logger.info(f"Processing batch {batch_idx + 1}/{total_batches}")88 89 # 1. Embed90 batch_embeddings = await embedding_service.generate_embeddings(batch_data)91 92 # 2. Metadata93 batch_metadata = [94 {95 "document_id": document_id,96 "document_name": filename,97 "chunk_id": start_index + k98 }99 for k in range(len(batch_data))100 ]101 102 # 3. Upsert103 await qdrant_service.upsert_chunks(104 collection_name=collection_name,105 chunks=batch_data,106 embeddings=batch_embeddings,107 metadata=batch_metadata108 )109 return110 111 except Exception as e:112 if "429" in str(e) or "Too Many Requests" in str(e):113 if attempt < retries - 1:114 wait_time = (2 ** attempt) + (0.1 * (batch_idx % 5)) # Jitter115 logger.warning(f"Rate limit for batch {batch_idx + 1}, retrying in {wait_time}s...")116 await asyncio.sleep(wait_time)117 continue118 logger.error(f"Error in batch {batch_idx + 1} (Attempt {attempt+1}): {e}")119 if attempt == retries - 1:120 raise e121 122 # Create and run tasks123 tasks = [process_batch(idx, start_idx, b) for idx, (start_idx, b) in enumerate(batches)]124 await asyncio.gather(*tasks)125 126 # 6. Update status to completed127 await self._update_document_status(document_id, "completed")128 logger.info(f"Document {filename} processed successfully")129 130 # 7. Generate Topics131 try:132 from services.mcq_service import mcq_service133 await mcq_service.generate_document_topics(project_id, document_id)134 except Exception as topic_err:135 logger.error(f"Failed to generate topics for {filename}: {topic_err}")136 137 except Exception as e:138 logger.error(f"Error processing document {filename}: {str(e)}")139 await self._update_document_status(document_id, "failed", str(e))140 141 async def _update_document_status(142 self,143 document_id: str,144 status: str,145 message: Optional[str] = None146 ):147 """Update document processing status in database"""148 try:149 update_data = {"upload_status": status}150 if status == "completed":151 update_data["error_message"] = None152 elif message:153 update_data["error_message"] = message154 155 self.client.table("documents").update(update_data).eq(156 "id", document_id157 ).execute()158 159 except Exception as e:160 logger.error(f"Error updating document status: {str(e)}")161 162 async def delete_document(self, project_id: str, document_id: str):163 """Delete document from DB and Vector Store"""164 try:165 # 1. Delete from Qdrant166 collection_name = f"project_{project_id}"167 await qdrant_service.delete_vectors(collection_name, document_id)168 169 # 2. Delete from DB170 # We need to know project_id. The caller passes it or we fetch it.171 # If we didn't have project_id, we'd query it first.172 173 self.client.table("documents").delete().eq("id", document_id).execute()174 175 logger.info(f"Deleted document {document_id} from project {project_id}")176 177 except Exception as e:178 logger.error(f"Error deleting document: {str(e)}")179 raise180 181document_service = DocumentService()182 