""" Document Service Document processing pipeline with unified multi-agent AI analysis. Pipeline: 1. Download file (with SSRF validation) 2. Extract text + page boundaries via extractors.py 3. Smart sentence-aware chunking via chunking.py 4. Generate pgvector embeddings (stored as Vector(768)) 5. Run all 4 agents via AgentOrchestrator (parallel execution) 6. Persist results and mark document COMPLETED """ import asyncio import logging from datetime import datetime, timezone from typing import List, Optional import google.generativeai as genai from app.agents.base import QuotaExceededError from app.agents.orchestrator import AgentOrchestrator from app.config import settings from app.db.session import get_db_context from app.models.document import ( ComplianceStatus, Document, DocumentCategory, DocumentEmbedding, DocumentStatus, ) from app.services.chunking import ChunkingService, DocumentChunk from app.services.extractors import ExtractedDocument, download_and_extract logger = logging.getLogger(__name__) # Configure Gemini once at module level genai.configure(api_key=settings.GEMINI_API_KEY) _chunker = ChunkingService() _orchestrator = AgentOrchestrator() # Gemini accepts up to 100 inputs per embed_content call. EMBED_BATCH_SIZE = 100 class DocumentService: """ Document processing service. Uses: - extractors.py → page-aware text extraction - chunking.py → sentence-aware chunking with page tracking - orchestrator → unified 4-agent parallel analysis - pgvector → native vector storage (no brute-force cosine in Python) """ def __init__(self): self.embedding_model = settings.EMBEDDING_MODEL async def process_document(self, document_id: str) -> bool: """ Full document processing pipeline. Returns True on success, False on failure. """ logger.info(f"Starting document processing: {document_id}") with get_db_context() as db: document = db.query(Document).filter(Document.id == document_id).first() if not document: logger.error(f"Document not found: {document_id}") return False try: # ── Step 1: Mark as processing ────────────────────────────── document.status = DocumentStatus.PROCESSING db.commit() # ── Step 2: Download and extract text with page tracking ──── logger.info(f"Extracting text from: {document.file_url}") extracted: ExtractedDocument = await download_and_extract( file_url=document.file_url, file_type=document.file_type, ) if not extracted.full_text.strip(): raise ValueError("No text content extracted from document") document.content = extracted.full_text document.word_count = extracted.word_count document.total_pages = extracted.total_pages document.page_count = extracted.total_pages # keep backward compat db.commit() # ── Step 3: Smart sentence-aware chunking with page tracking logger.info(f"Chunking document ({extracted.total_pages} pages)...") chunks: List[DocumentChunk] = _chunker.chunk_document( full_text=extracted.full_text, pages=extracted.pages, ) logger.info(f"Created {len(chunks)} chunks") # ── Step 4: Generate embeddings and persist ───────────────── logger.info(f"Generating embeddings for {len(chunks)} chunks...") vectors = await self._embed_batch([c.text for c in chunks]) embedded = 0 for chunk, embedding_values in zip(chunks, vectors): if embedding_values is None: logger.warning( f"Skipping chunk {chunk.chunk_index} — embedding failed" ) continue embedded += 1 doc_embedding = DocumentEmbedding( document_id=document.id, chunk_index=chunk.chunk_index, chunk_text=chunk.text, embedding=embedding_values, # stored as vector(768) page_numbers=chunk.page_numbers, # e.g. [12, 13] section_title=chunk.section_title, # e.g. "Safety Procedures" start_page=( chunk.page_numbers[0] if chunk.page_numbers else None ), end_page=chunk.page_numbers[-1] if chunk.page_numbers else None, embedding_model=self.embedding_model, ) db.add(doc_embedding) db.commit() logger.info(f"Persisted {embedded}/{len(chunks)} chunk embeddings") # Without embeddings the document is invisible to retrieval. # Marking it COMPLETED would hide a total failure behind a # green status in the UI. if chunks and embedded == 0: raise ValueError( "All chunk embeddings failed — document would not be " "searchable" ) # ── Step 5: Run multi-agent analysis via orchestrator ──────── document.status = DocumentStatus.ANALYZING db.commit() logger.info("Running AgentOrchestrator (4 agents)...") results = await _orchestrator.analyze_document( text=extracted.full_text, pages=extracted.pages, ) # ── Step 6: Map orchestrator results to document fields ────── classification = results.get("classification", {}) safety = results.get("safety", {}) entities = results.get("entities", {}) summary = results.get("summary", {}) # Classification category_map = { "safety_protocol": DocumentCategory.SAFETY_PROTOCOL, "equipment_manual": DocumentCategory.EQUIPMENT_MANUAL, "regulatory": DocumentCategory.REGULATORY, "incident_report": DocumentCategory.INCIDENT_REPORT, "geological": DocumentCategory.GEOLOGICAL, "environmental": DocumentCategory.ENVIRONMENTAL, "training": DocumentCategory.TRAINING, "permit": DocumentCategory.PERMIT, "maintenance": DocumentCategory.MAINTENANCE, } document.category = category_map.get( classification.get("category"), DocumentCategory.OTHER ) document.subcategory = classification.get("subcategory") document.classification_confidence = float( classification.get("confidence", 0.5) ) # Safety status_map = { "compliant": ComplianceStatus.COMPLIANT, "warning": ComplianceStatus.WARNING, "violation": ComplianceStatus.VIOLATION, } document.safety_score = ( float(safety.get("score", 50)) if safety.get("score") is not None else None ) document.compliance_status = status_map.get( safety.get("status"), ComplianceStatus.PENDING ) document.hazards_detected = safety.get("hazards", []) document.safety_recommendations = safety.get("recommendations", []) # Entities & Summary document.entities = entities if isinstance(entities, dict) else {} document.summary = summary.get("summary", "Summary not available.") document.key_points = summary.get("key_points", []) # ── Step 7: Mark completed ─────────────────────────────────── document.status = DocumentStatus.COMPLETED document.processed_at = datetime.now(timezone.utc) db.commit() logger.info(f"Document processing completed: {document_id}") return True except QuotaExceededError as qe: # Quota hit: text extraction + embeddings already succeeded. # Mark COMPLETED with partial data so re-analyze is available. logger.error( f"Gemini quota exceeded during agent analysis for {document_id}: {qe}" ) document.status = DocumentStatus.COMPLETED document.processing_error = ( "AI analysis incomplete: Gemini API quota exceeded. " "Click Re-analyze to run the full analysis when quota resets." ) document.summary = "Summary not available — Gemini quota exceeded. Re-analyze to generate." document.key_points = [] document.processed_at = datetime.now(timezone.utc) db.commit() return False except Exception as e: logger.error( f"Document processing failed: {document_id} — {e}", exc_info=True ) document.status = DocumentStatus.FAILED document.processing_error = str(e) db.commit() return False async def _embed(self, text: str) -> list | None: """ Generate a single embedding vector via the Gemini embedding API. genai.embed_content is a *blocking* HTTP call. Running it directly in a coroutine stalls the entire event loop — every other request served by this process waits. asyncio.to_thread hands it to the default executor so the loop stays responsive. """ try: result = await asyncio.to_thread( genai.embed_content, model=self.embedding_model, content=text, task_type="retrieval_document", output_dimensionality=768, ) return result["embedding"] except Exception as e: logger.warning(f"Embedding generation failed: {e}") return None async def _embed_batch(self, texts: List[str]) -> List[Optional[list]]: """ Embed many chunks per API call instead of one call per chunk. A 200-page PDF produces several hundred chunks; at one request each that is several hundred round trips. Gemini accepts up to EMBED_BATCH_SIZE inputs per call, cutting this by two orders of magnitude. Returns a list positionally aligned with `texts`; entries are None where embedding failed, so the caller can skip just those chunks. """ results: List[Optional[list]] = [] for start in range(0, len(texts), EMBED_BATCH_SIZE): batch = texts[start : start + EMBED_BATCH_SIZE] results.extend(await self._embed_one_batch(batch)) return results async def _embed_one_batch(self, batch: List[str]) -> List[Optional[list]]: """Embed one batch, degrading to per-chunk calls if the batch fails.""" try: result = await asyncio.to_thread( genai.embed_content, model=self.embedding_model, content=batch, task_type="retrieval_document", output_dimensionality=768, ) vectors = result["embedding"] # Guard against a shape change in the SDK: a batch request must # return one vector per input, not a single flat vector. if isinstance(vectors, list) and len(vectors) == len(batch): return vectors logger.warning( f"Batch embedding returned {len(vectors)} vectors for " f"{len(batch)} inputs — falling back to per-chunk embedding" ) except Exception as e: logger.warning( f"Batch embedding failed ({e}) — falling back to per-chunk " f"embedding so one bad chunk cannot lose the whole document" ) return [await self._embed(t) for t in batch] # ── Background task wrapper ──────────────────────────────────────────────────── async def process_document_async(document_id: str) -> None: """Async wrapper used with FastAPI BackgroundTasks.""" service = DocumentService() await service.process_document(document_id)