CoolFace
Modelpublic

aigenrec/luminabackend

sourceHugging Faceupdated 10mo agoView on Hugging Face
0likes
document_service.py182 linesDownload Raw Back to services
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