aibridze/document_intelligence
0
1# """2# Extraction Service — orchestrates the full pipeline:3# 1. Upload PDF4# 2. Validate contract type (detect + compare with user selection)5# 3. Extract structured data using type-specific JSON schemas6# 4. Smart per-page OCR: only Vision on pages that need it (handwritten/scanned)7# 5. Extract handwritten content (signatures, stamps, notes)8# 6. Index for RAG search9# """10 11# from typing import Dict, Any, Optional, List, Tuple12# import logging13# import json14 15# from app.config import get_settings16# from app.core.pdf_processor import PDFProcessor17# from app.core.vision_model import GeminiVision18# from app.core.entity_extractor import EntityExtractor19# from app.core.rag_engine import RAGEngine20# from app.services.contract_type_validator import ContractTypeValidator21# from app.models.extraction_schemas import ContractType22 23# settings = get_settings()24# logger = logging.getLogger("trident_poc")25 26# VISION_BATCH_SIZE = 1027# # Minimum characters per page to consider it "has text"28# PER_PAGE_TEXT_THRESHOLD = 10029 30 31# class ExtractionService:32# """Orchestrates the contract extraction pipeline with type validation."""33 34# def __init__(self):35# self.vision = GeminiVision()36# self.extractor = EntityExtractor()37# self.rag = RAGEngine()38# self.validator = ContractTypeValidator()39 40# # ─── Step 1: Validate Contract Type ───────────────────────────41 42# async def validate_contract_type(43# self,44# doc_id: str,45# selected_type: ContractType,46# ) -> Dict:47# """48# Detect the contract type from the uploaded PDF and validate49# against the user's selection. This is FAST — ~5-10 seconds.50# """51# pdf_path = settings.get_upload_path() / f"{doc_id}.pdf"52# if not pdf_path.exists():53# raise FileNotFoundError(f"Document {doc_id} not found")54 55# with PDFProcessor(str(pdf_path)) as processor:56# full_doc = processor.load_as_single_document()57# text = full_doc.page_content58 59# # For fully scanned docs, use vision to get enough text60# if processor.is_scanned() or len(text.strip()) < 500:61# logger.info("Document appears scanned. Using vision for type detection...")62# # Only OCR first 3 pages for validation (fast!)63# max_pages = min(3, processor.page_count)64# images = processor.get_pages_as_images(start=1, end=max_pages, dpi=150)65# text = await self.vision.extract_all_text(images)66 67# validation_result = await self.validator.validate(text, selected_type)68# validation_result["doc_id"] = doc_id69 70# # Cache the text for later use71# text_cache_path = settings.get_processed_path() / f"{doc_id}_text.txt"72# with open(text_cache_path, "w", encoding="utf-8") as f:73# f.write(text)74 75# return validation_result76 77# # ─── Step 2: Full Extraction ──────────────────────────────────78 79# async def full_extraction(80# self,81# doc_id: str,82# contract_type: ContractType,83# use_vision: bool = True,84# force_extract: bool = False,85# ) -> Dict:86# """87# Run full extraction pipeline with type-specific JSON schema.88# Uses SMART per-page OCR — only runs Vision on pages that need it.89# """90# pdf_path = settings.get_upload_path() / f"{doc_id}.pdf"91# processed_path = settings.get_processed_path() / f"{doc_id}.json"92 93# # Check cache94# if processed_path.exists() and not force_extract:95# with open(processed_path, "r") as f:96# return json.load(f)97 98# if not pdf_path.exists():99# raise FileNotFoundError(f"Document {doc_id} not found")100 101# with PDFProcessor(str(pdf_path)) as processor:102# page_docs = processor.load_documents()103# page_count = processor.page_count104# is_fully_scanned = processor.is_scanned()105 106# logger.info(107# f"Extracting doc {doc_id} as '{contract_type.value}': "108# f"{page_count} pages, fully_scanned={is_fully_scanned}"109# )110 111# # ─── Smart Per-Page OCR ───────────────────────────────112# # Analyze each page: if it has enough digital text, skip OCR.113# # Only run Vision OCR on pages that lack text (handwritten/scanned).114# digital_text_parts = []115# ocr_page_numbers = [] # pages that need Vision OCR116 117# page_texts = processor.extract_text_by_page()118# for page_num, page_text in page_texts:119# stripped = page_text.strip()120# if len(stripped) >= PER_PAGE_TEXT_THRESHOLD:121# # This page has enough digital text — no OCR needed122# digital_text_parts.append(f"--- Page {page_num} ---\n{stripped}")123# else:124# # This page is likely scanned/handwritten — needs OCR125# ocr_page_numbers.append(page_num)126 127# logger.info(128# f"Smart OCR analysis: {len(page_texts) - len(ocr_page_numbers)} digital pages, "129# f"{len(ocr_page_numbers)} pages need OCR: {ocr_page_numbers}"130# )131 132# # Run Vision OCR only on the pages that need it133# ocr_text_parts = []134# if use_vision and ocr_page_numbers:135# ocr_text_parts = await self._smart_vision_ocr(processor, ocr_page_numbers)136 137# # Combine digital + OCR text in page order138# combined_text = self._merge_page_texts(139# page_texts, ocr_page_numbers, ocr_text_parts140# )141 142# # ─── Handwriting Extraction ───────────────────────────143# # Run vision on signature-likely pages (first 2 + last 2)144# handwriting_analysis = None145# if use_vision and page_count > 0:146# try:147# logger.info("Extracting handwritten content from signature pages...")148# sig_pages = self._get_signature_pages(page_count)149# sig_images = []150# for p in sig_pages:151# sig_images.append(processor.get_page_as_image(p, dpi=150))152# handwriting_analysis = await self.vision.extract_handwriting(sig_images)153# except Exception as e:154# logger.warning(f"Handwriting extraction failed (non-critical): {e}")155 156# # ─── LLM Extraction ───────────────────────────────────157# extraction = await self.extractor.extract(combined_text, contract_type)158 159# # Merge handwriting analysis160# if handwriting_analysis:161# extraction = self._merge_handwriting(extraction, handwriting_analysis)162 163# # Add extraction metadata164# score = self._calculate_confidence(extraction)165# extraction["_extraction_meta"] = {166# "confidence_score": round(score, 2),167# "contract_type": contract_type.value,168# "page_count": page_count,169# "is_scanned": is_fully_scanned,170# "digital_pages": len(page_texts) - len(ocr_page_numbers),171# "ocr_pages": len(ocr_page_numbers),172# "extraction_type": (173# "Vision OCR + LLM" if is_fully_scanned174# else f"Hybrid ({len(page_texts) - len(ocr_page_numbers)} digital + {len(ocr_page_numbers)} OCR)" if ocr_page_numbers175# else "Digital Text + LLM"176# ),177# }178 179# # Save to cache180# with open(processed_path, "w") as f:181# json.dump(extraction, f, indent=2, ensure_ascii=False)182 183# # Index for RAG search184# if page_docs:185# await self.rag.index_documents(doc_id, page_docs)186# elif combined_text:187# await self.rag.index_document(doc_id, combined_text)188 189# return extraction190 191# # ─── Combined: Validate + Extract ─────────────────────────────192 193# async def validate_and_extract(194# self,195# doc_id: str,196# contract_type: ContractType,197# force_extract: bool = False,198# use_vision: bool = True,199# ) -> Dict:200# """201# Full pipeline: validate type, then extract if valid (or forced).202# Returns combined validation + extraction result.203# """204# # Step 1: Validate205# validation = await self.validate_contract_type(doc_id, contract_type)206 207# if not validation["is_valid"] and not force_extract:208# return {209# "doc_id": doc_id,210# "status": "validation_failed",211# "contract_type": None,212# "validation": validation,213# "extraction": None,214# "message": validation["message"],215# }216 217# # Step 2: Extract218# extraction = await self.full_extraction(219# doc_id, contract_type, use_vision=use_vision, force_extract=True220# )221 222# # Update extraction meta with validation info223# if "_extraction_meta" in extraction:224# extraction["_extraction_meta"]["contract_type_validated"] = validation["is_valid"]225# extraction["_extraction_meta"]["detected_type"] = validation["detected_type"]226# extraction["_extraction_meta"]["selected_type"] = validation["selected_type"]227 228# if not validation["is_valid"]:229# extraction["_extraction_meta"]["validation_warning"] = (230# f"Type mismatch. Selected '{validation['selected_type']}' "231# f"but detected '{validation['detected_type']}'. Force extracted."232# )233 234# return {235# "doc_id": doc_id,236# "status": "success",237# "contract_type": contract_type.value,238# "validation": validation,239# "extraction": extraction,240# "message": "Extraction completed successfully.",241# }242 243# # ─── Search / Q&A ─────────────────────────────────────────────244 245# async def search(self, query: str, doc_id: Optional[str] = None, top_k: int = 5) -> Dict:246# """Search indexed documents and answer questions."""247# return await self.rag.answer_query(query, doc_id, top_k)248 249# # ─── Smart Vision OCR ─────────────────────────────────────────250 251# async def _smart_vision_ocr(252# self,253# processor: PDFProcessor,254# page_numbers: List[int],255# ) -> List[Tuple[int, str]]:256# """257# Run Vision OCR only on specific pages (those lacking digital text).258# Batches the pages for efficiency. Returns list of (page_num, ocr_text).259# """260# results = []261 262# # Process in batches of VISION_BATCH_SIZE263# for i in range(0, len(page_numbers), VISION_BATCH_SIZE):264# batch = page_numbers[i : i + VISION_BATCH_SIZE]265# batch_num = (i // VISION_BATCH_SIZE) + 1266# total_batches = (len(page_numbers) + VISION_BATCH_SIZE - 1) // VISION_BATCH_SIZE267 268# logger.info(269# f"Vision OCR batch {batch_num}/{total_batches}: pages {batch}"270# )271 272# images = []273# for p in batch:274# images.append(processor.get_page_as_image(p, dpi=150))275 276# batch_text = await self.vision.extract_all_text(images)277 278# # Split the batch result back to individual pages279# # If we can't split, assign the whole batch text to each page280# page_sections = batch_text.split("---")281# if len(page_sections) >= len(batch):282# for j, p in enumerate(batch):283# results.append((p, page_sections[j].strip() if j < len(page_sections) else ""))284# else:285# # Can't split cleanly — assign proportional chunks286# chunk_size = len(batch_text) // max(len(batch), 1)287# for j, p in enumerate(batch):288# start = j * chunk_size289# end = start + chunk_size if j < len(batch) - 1 else len(batch_text)290# results.append((p, batch_text[start:end].strip()))291 292# return results293 294# def _merge_page_texts(295# self,296# all_page_texts: List[Tuple[int, str]],297# ocr_page_numbers: List[int],298# ocr_results: List[Tuple[int, str]],299# ) -> str:300# """301# Merge digital text and OCR results in correct page order.302# Digital pages use their original text, OCR pages use vision output.303# """304# ocr_map = {page_num: text for page_num, text in ocr_results}305# parts = []306 307# for page_num, page_text in all_page_texts:308# if page_num in ocr_page_numbers and page_num in ocr_map:309# # Use OCR text for this page310# ocr_text = ocr_map[page_num]311# if ocr_text:312# parts.append(f"--- Page {page_num} (OCR) ---\n{ocr_text}")313# else:314# # Use original digital text315# stripped = page_text.strip()316# if stripped:317# parts.append(f"--- Page {page_num} ---\n{stripped}")318 319# return "\n\n".join(parts)320 321# @staticmethod322# def _get_signature_pages(page_count: int) -> List[int]:323# """Determine which pages are most likely to have signatures."""324# if page_count <= 4:325# return list(range(1, page_count + 1))326# # First 2 and last 2 pages327# return [1, 2, page_count - 1, page_count]328 329# # ─── Handwriting Merge Helper ─────────────────────────────────330 331# @staticmethod332# def _merge_handwriting(extraction: dict, handwriting_text: str) -> dict:333# """Merge vision-based handwriting analysis into the handwritten_data field."""334# hw = extraction.get("handwritten_data", {})335# if hw is None:336# hw = {}337 338# existing = hw.get("unstructured_handwriting") or ""339# if handwriting_text and "No handwritten content detected" not in handwriting_text:340# if existing:341# hw["unstructured_handwriting"] = f"{existing}\n\n--- Vision Analysis ---\n{handwriting_text}"342# else:343# hw["unstructured_handwriting"] = handwriting_text344 345# extraction["handwritten_data"] = hw346# return extraction347 348# # ─── Confidence Calculator ────────────────────────────────────349 350# @staticmethod351# def _calculate_confidence(data: Dict[str, Any]) -> float:352# """353# Calculate extraction confidence based on top-level and first-level354# field fill rates. Focuses on primary contract fields, not deeply355# nested optional sub-fields that would drag the score down.356# """357# skip_keys = {"_extraction_meta", "handwritten_data", "error"}358# filled = 0359# total = 0360 361# for key, value in data.items():362# if key in skip_keys:363# continue364 365# if isinstance(value, dict):366# # Count first-level nested fields367# for sub_key, sub_val in value.items():368# total += 1369# if sub_val is not None and sub_val != "" and sub_val != []:370# filled += 1371# elif isinstance(value, list):372# # A list counts as 1 field: filled if non-empty373# total += 1374# if len(value) > 0:375# filled += 1376# else:377# # Scalar field378# total += 1379# if value is not None and value != "":380# filled += 1381 382# return filled / total if total > 0 else 0.0383 384# # ─── Utility ──────────────────────────────────────────────────385 386# def delete_document(self, doc_id: str) -> bool:387# """Delete cached extraction and RAG index for a document."""388# import os389 390# processed_path = settings.get_processed_path() / f"{doc_id}.json"391# text_cache = settings.get_processed_path() / f"{doc_id}_text.txt"392# pdf_path = settings.get_upload_path() / f"{doc_id}.pdf"393 394# for path in [processed_path, text_cache, pdf_path]:395# if path.exists():396# os.remove(path)397 398# self.rag.delete_document(doc_id)399# return True400 401"""402Extraction Service — Production-grade pipeline:4031. Upload PDF4042. Validate contract type (detect + compare with user selection)4053. Smart per-page OCR (only pages that need it)4064. Multi-pass LLM extraction (extract → gap-fill → validate)4075. Post-extraction validation with semantic confidence4086. Handwriting extraction (signatures, stamps, notes)4097. Cache + RAG index410"""411 412from typing import Dict, Any, Optional, List, Tuple413import logging414import json415 416from app.config import get_settings417from app.core.pdf_processor import PDFProcessor418from app.core.vision_model import GeminiVision419from app.core.entity_extractor import EntityExtractor420from app.core.extraction_validator import ExtractionValidator421from app.core.rag_engine import RAGEngine422from app.services.contract_type_validator import ContractTypeValidator423from app.models.extraction_schemas import ContractType424 425settings = get_settings()426logger = logging.getLogger("trident_poc")427 428VISION_BATCH_SIZE = 10429PER_PAGE_TEXT_THRESHOLD = 100 # Default for documents with real text430HIGH_TEXT_THRESHOLD = 500 # For documents flagged by needs_vision_ocr()431 432 433class ExtractionService:434 """Orchestrates the contract extraction pipeline with type validation."""435 436 def __init__(self):437 self.vision = GeminiVision()438 self.extractor = EntityExtractor()439 self.validator = ContractTypeValidator()440 self.extraction_validator = ExtractionValidator()441 self.rag = RAGEngine()442 443 # ─── Step 1: Validate Contract Type ───────────────────────────444 445 async def validate_contract_type(446 self,447 doc_id: str,448 selected_type: ContractType,449 ) -> Dict:450 """451 Detect the contract type from the uploaded PDF and validate452 against the user's selection.453 454 ALWAYS uses Vision OCR to read the first few pages, ensuring455 the validator sees actual document content (not just PyMuPDF456 metadata/digital-signature text which causes misclassification).457 """458 pdf_path = settings.get_upload_path() / f"{doc_id}.pdf"459 if not pdf_path.exists():460 raise FileNotFoundError(f"Document {doc_id} not found")461 462 text = ""463 464 with PDFProcessor(str(pdf_path)) as processor:465 # Always use Vision OCR for validation — it only reads 3-4 pages466 # and takes ~5 seconds, but guarantees we see actual content467 # instead of just PyMuPDF metadata/digital signatures.468 max_pages = min(4, processor.page_count)469 470 try:471 logger.info(472 f"Validation: using Vision OCR on first {max_pages} pages "473 f"for reliable type detection"474 )475 images = processor.get_pages_as_images(start=1, end=max_pages, dpi=150)476 text = await self.vision.extract_all_text(images)477 except Exception as e:478 logger.warning(f"Vision OCR failed for validation, falling back to digital text: {e}")479 # Fallback to digital text if Vision fails480 full_doc = processor.load_as_single_document()481 text = full_doc.page_content482 483 # If we still have no text, try digital text484 if not text or len(text.strip()) < 100:485 logger.warning("Vision produced no text, using digital text fallback")486 full_doc = processor.load_as_single_document()487 text = full_doc.page_content488 489 # Strip markdown code fences from Vision OCR text490 # (Gemini Vision sometimes wraps OCR output in ```...```)491 if text:492 import re as _re493 text = _re.sub(r'^```(?:\w+)?\s*\n?', '', text, flags=_re.MULTILINE)494 text = _re.sub(r'\n?```\s*$', '', text, flags=_re.MULTILINE)495 text = text.strip()496 497 validation_result = await self.validator.validate(text, selected_type)498 validation_result["doc_id"] = doc_id499 500 # Cache the validation text for reuse in extraction501 text_cache_path = settings.get_processed_path() / f"{doc_id}_text.txt"502 with open(text_cache_path, "w", encoding="utf-8") as f:503 f.write(text)504 505 return validation_result506 507 # ─── Step 2: Full Extraction ──────────────────────────────────508 509 async def full_extraction(510 self,511 doc_id: str,512 contract_type: ContractType,513 use_vision: bool = True,514 force_extract: bool = False,515 ) -> Dict:516 """517 Production extraction pipeline:518 1. Smart per-page OCR (elevated threshold if needs_vision_ocr)519 2. Multi-pass LLM extraction (extract → gap-fill)520 3. Handwriting extraction from signature pages521 4. Post-extraction validation522 5. Cache + index523 """524 pdf_path = settings.get_upload_path() / f"{doc_id}.pdf"525 processed_path = settings.get_processed_path() / f"{doc_id}.json"526 527 # Check cache528 if processed_path.exists() and not force_extract:529 with open(processed_path, "r") as f:530 return json.load(f)531 532 if not pdf_path.exists():533 raise FileNotFoundError(f"Document {doc_id} not found")534 535 with PDFProcessor(str(pdf_path)) as processor:536 page_docs = processor.load_documents()537 page_count = processor.page_count538 is_fully_scanned = processor.is_scanned()539 needs_vision = processor.needs_vision_ocr()540 541 logger.info(542 f"Extracting {doc_id} as '{contract_type.value}': "543 f"{page_count} pages, scanned={is_fully_scanned}, "544 f"needs_vision={needs_vision}"545 )546 547 # ─── Smart Per-Page OCR ───────────────────────────────548 # If needs_vision_ocr() is True, the document has mostly549 # junk text (signatures, metadata). Use a higher threshold550 # so those pages get flagged for OCR.551 if needs_vision:552 effective_threshold = HIGH_TEXT_THRESHOLD553 logger.info(554 f"Using elevated threshold ({effective_threshold}) — "555 f"document has insufficient digital text"556 )557 else:558 effective_threshold = PER_PAGE_TEXT_THRESHOLD559 560 digital_text_parts = []561 ocr_page_numbers = []562 563 page_texts = processor.extract_text_by_page()564 for page_num, page_text in page_texts:565 stripped = page_text.strip()566 if len(stripped) >= effective_threshold:567 digital_text_parts.append(f"--- Page {page_num} ---\n{stripped}")568 else:569 ocr_page_numbers.append(page_num)570 571 logger.info(572 f"OCR decision: {len(page_texts) - len(ocr_page_numbers)} digital, "573 f"{len(ocr_page_numbers)} need OCR (threshold={effective_threshold}): "574 f"{ocr_page_numbers}"575 )576 577 # Run Vision OCR on flagged pages578 ocr_text_parts = []579 if use_vision and ocr_page_numbers:580 ocr_text_parts = await self._smart_vision_ocr(581 processor, ocr_page_numbers582 )583 584 # Merge all text in page order585 combined_text = self._merge_page_texts(586 page_texts, ocr_page_numbers, ocr_text_parts587 )588 589 590 # ─── Multi-Pass LLM Extraction ────────────────────────591 logger.info("Running multi-pass extraction...")592 extraction = await self.extractor.extract(combined_text, contract_type)593 594 # ─── Post-Extraction Validation ───────────────────────595 logger.info("Running post-extraction validation...")596 validation_report = self.extraction_validator.validate(597 extraction, combined_text, contract_type598 )599 600 # ─── Build Metadata ───────────────────────────────────601 extraction["_extraction_meta"] = {602 "confidence_score": validation_report["overall_confidence"],603 "field_fill_rate": validation_report["field_fill_rate"],604 "text_match_rate": validation_report["text_match_rate"],605 "required_fields_rate": validation_report["required_fields_rate"],606 "total_fields": validation_report["total_fields"],607 "filled_fields": validation_report["filled_fields"],608 "issues_count": validation_report["issues_count"],609 "contract_type": contract_type.value,610 "page_count": page_count,611 "is_scanned": is_fully_scanned,612 "needs_vision_ocr": needs_vision,613 "digital_pages": len(page_texts) - len(ocr_page_numbers),614 "ocr_pages": len(ocr_page_numbers),615 "extraction_type": (616 "Vision OCR + Multi-Pass LLM" if needs_vision617 else f"Hybrid ({len(page_texts) - len(ocr_page_numbers)} digital + {len(ocr_page_numbers)} OCR) + Multi-Pass LLM"618 if ocr_page_numbers619 else "Digital Text + Multi-Pass LLM"620 ),621 "validation_issues": validation_report["issues"][:10],622 }623 624 # Save to cache625 with open(processed_path, "w") as f:626 json.dump(extraction, f, indent=2, ensure_ascii=False)627 628 # Index for RAG search629 # Index for RAG search (with contract_type metadata)630 ct_label = contract_type.value631 if page_docs and not needs_vision:632 await self.rag.index_documents(doc_id, page_docs, ct_label)633 elif combined_text:634 await self.rag.index_document(doc_id, combined_text, ct_label)635 636 logger.info(637 f"Extraction complete: confidence={validation_report['overall_confidence']:.0%}, "638 f"fields={validation_report['filled_fields']}/{validation_report['total_fields']}, "639 f"issues={validation_report['issues_count']}"640 )641 642 return extraction643 644 # ─── Combined: Validate + Extract ─────────────────────────────645 646 async def validate_and_extract(647 self,648 doc_id: str,649 contract_type: ContractType,650 force_extract: bool = False,651 use_vision: bool = True,652 ) -> Dict:653 """Full pipeline: validate type → extract → validate results."""654 # Step 1: Validate type655 validation = await self.validate_contract_type(doc_id, contract_type)656 657 if not validation["is_valid"] and not force_extract:658 return {659 "doc_id": doc_id,660 "status": "validation_failed",661 "contract_type": None,662 "validation": validation,663 "extraction": None,664 "message": validation["message"],665 }666 667 # ─── Auto-correct type on force extract ──────────────────668 extraction_type = contract_type669 auto_corrected = False670 671 if not validation["is_valid"] and force_extract:672 detected = validation.get("detected_type", "")673 detected_enum = self._resolve_contract_type(detected)674 if detected_enum:675 logger.info(676 f"Auto-correcting: '{contract_type.value}' → "677 f"'{detected_enum.value}' (detected)"678 )679 extraction_type = detected_enum680 auto_corrected = True681 682 # Step 2: Extract with correct type683 extraction = await self.full_extraction(684 doc_id, extraction_type,685 use_vision=use_vision, force_extract=True,686 )687 688 # Update meta with validation info689 if "_extraction_meta" in extraction:690 extraction["_extraction_meta"]["contract_type_validated"] = validation["is_valid"]691 extraction["_extraction_meta"]["detected_type"] = validation["detected_type"]692 extraction["_extraction_meta"]["selected_type"] = validation["selected_type"]693 694 if auto_corrected:695 extraction["_extraction_meta"]["auto_corrected"] = True696 extraction["_extraction_meta"]["auto_corrected_from"] = contract_type.value697 extraction["_extraction_meta"]["auto_corrected_to"] = extraction_type.value698 extraction["_extraction_meta"]["validation_warning"] = (699 f"You selected '{contract_type.value}' but document is "700 f"'{extraction_type.value}'. Auto-corrected to correct schema."701 )702 elif not validation["is_valid"]:703 extraction["_extraction_meta"]["validation_warning"] = (704 f"Type mismatch. Selected '{validation['selected_type']}' "705 f"but detected '{validation['detected_type']}'. Force extracted."706 )707 708 return {709 "doc_id": doc_id,710 "status": "success",711 "contract_type": extraction_type.value,712 "validation": validation,713 "extraction": extraction,714 "message": "Extraction completed successfully.",715 }716 # ─── Search / Q&A ─────────────────────────────────────────────717 718 async def search(719 self, query: str, doc_id: Optional[str] = None, top_k: int = 5,720 ) -> Dict:721 return await self.rag.answer_query(query, doc_id, top_k)722 723 # ─── Smart Vision OCR ─────────────────────────────────────────724 725 async def _smart_vision_ocr(726 self,727 processor: PDFProcessor,728 page_numbers: List[int],729 ) -> List[Tuple[int, str]]:730 """Run Vision OCR only on specific pages, batched."""731 results = []732 733 for i in range(0, len(page_numbers), VISION_BATCH_SIZE):734 batch = page_numbers[i : i + VISION_BATCH_SIZE]735 batch_num = (i // VISION_BATCH_SIZE) + 1736 total_batches = (737 (len(page_numbers) + VISION_BATCH_SIZE - 1) // VISION_BATCH_SIZE738 )739 740 logger.info(f"Vision OCR batch {batch_num}/{total_batches}: pages {batch}")741 742 images = [processor.get_page_as_image(p, dpi=150) for p in batch]743 batch_text = await self.vision.extract_all_text(images)744 745 # Split batch result back to individual pages746 page_sections = batch_text.split("---")747 if len(page_sections) >= len(batch):748 for j, p in enumerate(batch):749 text = page_sections[j].strip() if j < len(page_sections) else ""750 results.append((p, text))751 else:752 chunk_size = len(batch_text) // max(len(batch), 1)753 for j, p in enumerate(batch):754 start_idx = j * chunk_size755 end_idx = (756 start_idx + chunk_size757 if j < len(batch) - 1758 else len(batch_text)759 )760 results.append((p, batch_text[start_idx:end_idx].strip()))761 762 return results763 764 def _merge_page_texts(765 self,766 all_page_texts: List[Tuple[int, str]],767 ocr_page_numbers: List[int],768 ocr_results: List[Tuple[int, str]],769 ) -> str:770 """Merge digital + OCR text in correct page order."""771 ocr_map = {pn: txt for pn, txt in ocr_results}772 parts = []773 774 for page_num, page_text in all_page_texts:775 if page_num in ocr_page_numbers and page_num in ocr_map:776 ocr_text = ocr_map[page_num]777 if ocr_text:778 parts.append(f"--- Page {page_num} (OCR) ---\n{ocr_text}")779 else:780 stripped = page_text.strip()781 if stripped:782 parts.append(f"--- Page {page_num} ---\n{stripped}")783 784 return "\n\n".join(parts)785 786 787 @staticmethod788 def _resolve_contract_type(detected_label: str) -> Optional[ContractType]:789 """Map a detected type label string to ContractType enum."""790 if not detected_label:791 return None792 label = detected_label.lower().strip()793 mapping = {794 "nda": ContractType.NDA,795 "non-disclosure agreement": ContractType.NDA,796 "service provider agreement": ContractType.SERVICE_PROVIDER,797 "supply cum service agreement": ContractType.SUPPLY_CUM_SERVICE,798 "epc agreement": ContractType.SUPPLY_CUM_SERVICE,799 "consultancy agreement": ContractType.CONSULTANCY,800 "consulting agreement": ContractType.CONSULTANCY,801 "advisory agreement": ContractType.CONSULTANCY,802 "customer/supply agreement": ContractType.CUSTOMER_SUPPLY,803 "supply agreement": ContractType.CUSTOMER_SUPPLY,804 }805 if label in mapping:806 return mapping[label]807 for key, val in mapping.items():808 if key in label or label in key:809 return val810 return None811 812 def delete_document(self, doc_id: str) -> bool:813 import os814 815 for suffix in [".json", "_text.txt"]:816 path = settings.get_processed_path() / f"{doc_id}{suffix}"817 if path.exists():818 os.remove(path)819 820 pdf_path = settings.get_upload_path() / f"{doc_id}.pdf"821 if pdf_path.exists():822 os.remove(pdf_path)823 824 self.rag.delete_document(doc_id)825 return True