rnsr 0.1.0__py3-none-any.whl

rnsr/extraction/grounded_extractor.py

@@ -0,0 +1,722 @@
+"""
+RNSR Grounded Entity Extractor
+
+Implements the RLM pattern for entity extraction:
+1. Pre-extract candidates using regex/patterns (CODE FIRST)
+2. LLM classifies and validates candidates (LLM SECOND)
+3. Recursive refinement if needed
+
+This approach PREVENTS HALLUCINATION because:
+- Every entity is tied to an exact text span
+- LLM classifies existing text, doesn't generate entities
+- Pattern matching provides grounded candidates
+- LLM's job is validation, not invention
+
+Validation Modes:
+- SIMPLE: Basic LLM classification (faster, single call)
+- TOT: Tree of Thoughts validation with probabilities and navigation (more accurate)
+
+Inspired by the RLM paper's insight: use code to filter/extract
+before sending to LLM for reasoning.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+from enum import Enum
+from typing import Any, TYPE_CHECKING
+
+import structlog
+
+from rnsr.extraction.models import (
+    Entity,
+    EntityType,
+    ExtractionResult,
+    Mention,
+)
+from rnsr.extraction.candidate_extractor import (
+    CandidateExtractor,
+    EntityCandidate,
+)
+from rnsr.extraction.learned_types import (
+    get_learned_type_registry,
+)
+from rnsr.llm import get_llm
+
+if TYPE_CHECKING:
+    from rnsr.models import DocumentTree
+
+
+class ValidationMode(str, Enum):
+    """Validation mode for candidate entities."""
+    
+    SIMPLE = "simple"   # Basic LLM classification (faster)
+    TOT = "tot"         # Tree of Thoughts with navigation (more accurate)
+
+logger = structlog.get_logger(__name__)
+
+
+# LLM prompt for CLASSIFYING pre-extracted candidates (not generating)
+CLASSIFICATION_PROMPT = """You are classifying entity candidates that have already been extracted from a document.
+
+Your job is to VALIDATE and CLASSIFY each candidate - NOT to generate new entities.
+These candidates were extracted by pattern matching and are grounded in the actual text.
+
+Document Section:
+---
+{content}
+---
+
+Pre-extracted candidates to classify:
+{candidates_json}
+
+For each candidate, provide:
+1. valid: true if this is a real entity worth tracking, false if it's noise
+2. type: The entity type (PERSON, ORGANIZATION, DATE, EVENT, LEGAL_CONCEPT, LOCATION, REFERENCE, MONETARY, DOCUMENT, or other descriptive type)
+3. canonical_name: Normalized/cleaned name (e.g., "Mr. John Smith" → "John Smith")
+4. role: Any role or relationship mentioned (e.g., "defendant", "CEO")
+
+Return JSON array:
+```json
+[
+  {{
+    "candidate_id": 0,
+    "valid": true,
+    "type": "PERSON",
+    "canonical_name": "John Smith",
+    "role": "defendant"
+  }},
+  {{
+    "candidate_id": 1,
+    "valid": false,
+    "reason": "Generic reference, not a specific entity"
+  }}
+]
+```
+
+Rules:
+- ONLY classify the candidates provided - do not add new entities
+- Set valid=false for generic terms, partial matches, or noise
+- Use the exact text span provided - don't modify the match boundaries
+- Be conservative - when uncertain, set valid=false
+"""
+
+# Prompt for finding entities the patterns might have missed
+SUPPLEMENTARY_PROMPT = """Review this text for important entities that might have been missed by pattern matching.
+
+Text:
+---
+{content}
+---
+
+Already extracted: {existing_entities}
+
+Are there any CLEARLY IDENTIFIABLE entities that were missed?
+Only list entities that:
+1. Are explicitly named in the text (not implied)
+2. Are significant (not passing mentions)
+3. Have a clear type
+
+If there are missed entities, return:
+```json
+[
+  {{
+    "text": "exact text as it appears",
+    "type": "ENTITY_TYPE",
+    "canonical_name": "Normalized Name",
+    "reason": "Why this is important"
+  }}
+]
+```
+
+If nothing important was missed, return: []
+"""
+
+
+class GroundedEntityExtractor:
+    """
+    Entity extractor following the RLM pattern:
+    CODE FIRST (pattern extraction) → LLM SECOND (classification).
+    
+    This prevents hallucination by grounding entities in actual text.
+    """
+    
+    def __init__(
+        self,
+        llm: Any | None = None,
+        candidate_extractor: CandidateExtractor | None = None,
+        min_content_length: int = 50,
+        max_candidates_per_batch: int = 30,
+        enable_supplementary_extraction: bool = True,
+        enable_type_learning: bool = True,
+        validation_mode: ValidationMode | str = ValidationMode.SIMPLE,
+        tot_selection_threshold: float = 0.6,
+        tot_enable_navigation: bool = True,
+    ):
+        """
+        Initialize the grounded extractor.
+        
+        Args:
+            llm: LLM instance. If None, uses get_llm().
+            candidate_extractor: Pre-extraction engine.
+            min_content_length: Minimum content length to process.
+            max_candidates_per_batch: Max candidates per LLM call.
+            enable_supplementary_extraction: Check for missed entities.
+            enable_type_learning: Learn new entity types.
+            validation_mode: SIMPLE (faster) or TOT (more accurate with navigation).
+            tot_selection_threshold: Probability threshold for ToT mode.
+            tot_enable_navigation: Navigate tree for uncertain candidates in ToT mode.
+        """
+        self.llm = llm
+        self.candidate_extractor = candidate_extractor or CandidateExtractor()
+        self.min_content_length = min_content_length
+        self.max_candidates_per_batch = max_candidates_per_batch
+        self.enable_supplementary_extraction = enable_supplementary_extraction
+        self.enable_type_learning = enable_type_learning
+        
+        # Validation mode
+        if isinstance(validation_mode, str):
+            validation_mode = ValidationMode(validation_mode.lower())
+        self.validation_mode = validation_mode
+        self.tot_selection_threshold = tot_selection_threshold
+        self.tot_enable_navigation = tot_enable_navigation
+        
+        # Lazy init for ToT validator
+        self._tot_validator = None
+        
+        # Lazy LLM init
+        self._llm_initialized = False
+        
+        # Type registry for learning
+        self._type_registry = get_learned_type_registry() if enable_type_learning else None
+        
+        # Cache
+        self._cache: dict[str, list[Entity]] = {}
+    
+    def _get_llm(self) -> Any:
+        """Get or initialize LLM."""
+        if self.llm is None and not self._llm_initialized:
+            self.llm = get_llm()
+            self._llm_initialized = True
+        return self.llm
+    
+    def _get_tot_validator(self) -> "TotEntityValidator":
+        """Get or initialize ToT validator."""
+        if self._tot_validator is None:
+            from rnsr.extraction.tot_validator import TotEntityValidator
+            self._tot_validator = TotEntityValidator(
+                llm=self._get_llm(),
+                selection_threshold=self.tot_selection_threshold,
+                enable_navigation=self.tot_enable_navigation,
+                max_candidates_per_batch=self.max_candidates_per_batch,
+            )
+        return self._tot_validator
+    
+    def extract_from_node(
+        self,
+        node_id: str,
+        doc_id: str,
+        header: str,
+        content: str,
+        page_num: int | None = None,
+        document_tree: "DocumentTree | None" = None,
+    ) -> ExtractionResult:
+        """
+        Extract entities using the grounded approach.
+        
+        Flow:
+        1. Pattern-extract candidates (grounded)
+        2. Validate with LLM (SIMPLE mode) or ToT (TOT mode)
+        3. Optionally check for missed entities
+        4. Return validated entities
+        
+        Args:
+            node_id: Section node ID.
+            doc_id: Document ID.
+            header: Section header.
+            content: Section content.
+            page_num: Page number.
+            document_tree: Optional tree for ToT navigation.
+            
+        Returns:
+            ExtractionResult with grounded entities.
+        """
+        start_time = time.time()
+        result = ExtractionResult(
+            node_id=node_id,
+            doc_id=doc_id,
+            extraction_method=f"grounded_{self.validation_mode.value}",
+        )
+        
+        # Skip short content
+        if len(content.strip()) < self.min_content_length:
+            return result
+        
+        # Check cache
+        cache_key = f"{doc_id}:{node_id}"
+        if cache_key in self._cache:
+            result.entities = self._cache[cache_key]
+            return result
+        
+        # STEP 1: Extract candidates using patterns (CODE FIRST)
+        candidates = self.candidate_extractor.extract_candidates(content)
+        
+        result.warnings.append(f"Pattern extraction found {len(candidates)} candidates")
+        
+        if not candidates:
+            # No pattern matches - try supplementary if enabled
+            if self.enable_supplementary_extraction:
+                entities = self._supplementary_extraction(
+                    content, node_id, doc_id, page_num, []
+                )
+                result.entities = entities
+            
+            result.processing_time_ms = (time.time() - start_time) * 1000
+            return result
+        
+        # STEP 2: Validate candidates (LLM SECOND)
+        if self.validation_mode == ValidationMode.TOT:
+            # Use Tree of Thoughts validation (more accurate, can navigate)
+            entities = self._validate_with_tot(
+                candidates=candidates,
+                header=header,
+                content=content,
+                node_id=node_id,
+                doc_id=doc_id,
+                page_num=page_num,
+                document_tree=document_tree,
+            )
+        else:
+            # Use simple classification (faster)
+            entities = self._classify_candidates(
+                candidates=candidates,
+                content=content,
+                node_id=node_id,
+                doc_id=doc_id,
+                page_num=page_num,
+            )
+        
+        # STEP 3: Check for missed entities (optional)
+        if self.enable_supplementary_extraction:
+            existing_names = [e.canonical_name for e in entities]
+            supplementary = self._supplementary_extraction(
+                content, node_id, doc_id, page_num, existing_names
+            )
+            entities.extend(supplementary)
+        
+        result.entities = entities
+        result.processing_time_ms = (time.time() - start_time) * 1000
+        
+        # Cache
+        self._cache[cache_key] = entities
+        
+        logger.info(
+            "grounded_extraction_complete",
+            node_id=node_id,
+            candidates=len(candidates),
+            entities=len(entities),
+            validation_mode=self.validation_mode.value,
+            time_ms=result.processing_time_ms,
+        )
+        
+        return result
+    
+    def _validate_with_tot(
+        self,
+        candidates: list[EntityCandidate],
+        header: str,
+        content: str,
+        node_id: str,
+        doc_id: str,
+        page_num: int | None,
+        document_tree: "DocumentTree | None",
+    ) -> list[Entity]:
+        """
+        Validate candidates using Tree of Thoughts pattern.
+        
+        This uses the same ToT approach as document navigation:
+        - Evaluate each candidate with probability + reasoning
+        - Navigate to related sections for uncertain candidates
+        - More accurate than simple classification
+        """
+        tot_validator = self._get_tot_validator()
+        
+        # Run ToT validation
+        validation_result = tot_validator.validate_candidates(
+            candidates=candidates,
+            section_header=header,
+            section_content=content,
+            document_tree=document_tree,
+            node_id=node_id,
+        )
+        
+        # Convert to entities
+        entities = tot_validator.candidates_to_entities(
+            candidates=candidates,
+            validation_result=validation_result,
+            node_id=node_id,
+            doc_id=doc_id,
+            page_num=page_num,
+        )
+        
+        # Learn from OTHER types
+        if self._type_registry:
+            for entity in entities:
+                if entity.type == EntityType.OTHER:
+                    original_type = entity.metadata.get("original_type", "unknown")
+                    context = entity.mentions[0].context if entity.mentions else ""
+                    self._type_registry.record_type(
+                        type_name=original_type,
+                        context=context,
+                        entity_name=entity.canonical_name,
+                    )
+        
+        logger.info(
+            "tot_validation_complete",
+            candidates=len(candidates),
+            validated=len(entities),
+            high_confidence=validation_result.high_confidence_count,
+            low_confidence=validation_result.low_confidence_count,
+        )
+        
+        return entities
+    
+    def _classify_candidates(
+        self,
+        candidates: list[EntityCandidate],
+        content: str,
+        node_id: str,
+        doc_id: str,
+        page_num: int | None,
+    ) -> list[Entity]:
+        """
+        Use LLM to classify pre-extracted candidates.
+        
+        The LLM's job is VALIDATION and CLASSIFICATION,
+        not generation of new entities.
+        """
+        llm = self._get_llm()
+        if llm is None:
+            # No LLM - return candidates as-is with pattern-based types
+            return self._candidates_to_entities(
+                candidates, node_id, doc_id, page_num
+            )
+        
+        entities = []
+        
+        # Process in batches
+        for i in range(0, len(candidates), self.max_candidates_per_batch):
+            batch = candidates[i:i + self.max_candidates_per_batch]
+            batch_entities = self._classify_batch(
+                batch, content, node_id, doc_id, page_num
+            )
+            entities.extend(batch_entities)
+        
+        return entities
+    
+    def _classify_batch(
+        self,
+        candidates: list[EntityCandidate],
+        content: str,
+        node_id: str,
+        doc_id: str,
+        page_num: int | None,
+    ) -> list[Entity]:
+        """Classify a batch of candidates with LLM."""
+        # Format candidates for prompt
+        candidates_json = json.dumps([
+            {
+                "id": idx,
+                "text": c.text,
+                "type_hint": c.candidate_type,
+                "context": c.context[:150],
+            }
+            for idx, c in enumerate(candidates)
+        ], indent=2)
+        
+        prompt = CLASSIFICATION_PROMPT.format(
+            content=content[:3000],  # Limit content size
+            candidates_json=candidates_json,
+        )
+        
+        try:
+            response = self.llm.complete(prompt)
+            response_text = str(response) if not isinstance(response, str) else response
+            
+            # Parse classifications
+            classifications = self._parse_classification_response(response_text)
+            
+            # Convert to entities
+            entities = []
+            for classification in classifications:
+                candidate_id = classification.get("candidate_id")
+                
+                if candidate_id is None or candidate_id >= len(candidates):
+                    continue
+                
+                if not classification.get("valid", False):
+                    continue
+                
+                candidate = candidates[candidate_id]
+                entity = self._create_entity_from_classification(
+                    candidate=candidate,
+                    classification=classification,
+                    node_id=node_id,
+                    doc_id=doc_id,
+                    page_num=page_num,
+                )
+                
+                if entity:
+                    entities.append(entity)
+            
+            return entities
+            
+        except Exception as e:
+            logger.warning("classification_failed", error=str(e))
+            # Fallback: return candidates as-is
+            return self._candidates_to_entities(
+                candidates, node_id, doc_id, page_num
+            )
+    
+    def _parse_classification_response(
+        self,
+        response_text: str,
+    ) -> list[dict[str, Any]]:
+        """Parse LLM classification response."""
+        # Extract JSON
+        json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response_text)
+        if json_match:
+            json_str = json_match.group(1)
+        else:
+            json_match = re.search(r'\[[\s\S]*\]', response_text)
+            json_str = json_match.group(0) if json_match else "[]"
+        
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError:
+            return []
+    
+    def _create_entity_from_classification(
+        self,
+        candidate: EntityCandidate,
+        classification: dict[str, Any],
+        node_id: str,
+        doc_id: str,
+        page_num: int | None,
+    ) -> Entity | None:
+        """Create Entity from candidate + LLM classification."""
+        entity_type_str = classification.get("type", candidate.candidate_type).upper()
+        
+        # Map to EntityType
+        entity_type = self._map_entity_type(entity_type_str)
+        
+        # Record learned type if OTHER
+        if entity_type == EntityType.OTHER and self._type_registry:
+            self._type_registry.record_type(
+                type_name=entity_type_str.lower(),
+                context=candidate.context,
+                entity_name=candidate.text,
+            )
+        
+        # Get canonical name (LLM-cleaned or original)
+        canonical_name = classification.get("canonical_name", candidate.text).strip()
+        if not canonical_name:
+            canonical_name = candidate.text
+        
+        # Build metadata
+        metadata = {
+            "grounded": True,  # Flag that this is grounded in text
+            "pattern": candidate.pattern_name,
+            "span_start": candidate.start,
+            "span_end": candidate.end,
+        }
+        
+        if classification.get("role"):
+            metadata["role"] = classification["role"]
+        
+        if entity_type == EntityType.OTHER:
+            metadata["original_type"] = entity_type_str.lower()
+        
+        # Create mention
+        mention = Mention(
+            node_id=node_id,
+            doc_id=doc_id,
+            span_start=candidate.start,
+            span_end=candidate.end,
+            context=candidate.context,
+            page_num=page_num,
+            confidence=candidate.confidence,
+        )
+        
+        return Entity(
+            type=entity_type,
+            canonical_name=canonical_name,
+            aliases=[candidate.text] if candidate.text != canonical_name else [],
+            mentions=[mention],
+            metadata=metadata,
+            source_doc_id=doc_id,
+        )
+    
+    def _map_entity_type(self, type_str: str) -> EntityType:
+        """Map type string to EntityType enum."""
+        type_str = type_str.upper()
+        
+        mapping = {
+            "PERSON": EntityType.PERSON,
+            "PEOPLE": EntityType.PERSON,
+            "INDIVIDUAL": EntityType.PERSON,
+            "ORGANIZATION": EntityType.ORGANIZATION,
+            "ORG": EntityType.ORGANIZATION,
+            "COMPANY": EntityType.ORGANIZATION,
+            "COURT": EntityType.ORGANIZATION,
+            "DATE": EntityType.DATE,
+            "TIME": EntityType.DATE,
+            "EVENT": EntityType.EVENT,
+            "LEGAL_CONCEPT": EntityType.LEGAL_CONCEPT,
+            "LEGAL": EntityType.LEGAL_CONCEPT,
+            "LOCATION": EntityType.LOCATION,
+            "PLACE": EntityType.LOCATION,
+            "ADDRESS": EntityType.LOCATION,
+            "REFERENCE": EntityType.REFERENCE,
+            "CITATION": EntityType.REFERENCE,
+            "MONETARY": EntityType.MONETARY,
+            "MONEY": EntityType.MONETARY,
+            "DOCUMENT": EntityType.DOCUMENT,
+        }
+        
+        try:
+            return EntityType(type_str.lower())
+        except ValueError:
+            return mapping.get(type_str, EntityType.OTHER)
+    
+    def _candidates_to_entities(
+        self,
+        candidates: list[EntityCandidate],
+        node_id: str,
+        doc_id: str,
+        page_num: int | None,
+    ) -> list[Entity]:
+        """Convert candidates directly to entities (no LLM)."""
+        entities = []
+        
+        for candidate in candidates:
+            entity_type = self._map_entity_type(candidate.candidate_type)
+            
+            mention = Mention(
+                node_id=node_id,
+                doc_id=doc_id,
+                span_start=candidate.start,
+                span_end=candidate.end,
+                context=candidate.context,
+                page_num=page_num,
+                confidence=candidate.confidence,
+            )
+            
+            entity = Entity(
+                type=entity_type,
+                canonical_name=candidate.text,
+                mentions=[mention],
+                metadata={
+                    "grounded": True,
+                    "pattern": candidate.pattern_name,
+                    "llm_validated": False,
+                },
+                source_doc_id=doc_id,
+            )
+            entities.append(entity)
+        
+        return entities
+    
+    def _supplementary_extraction(
+        self,
+        content: str,
+        node_id: str,
+        doc_id: str,
+        page_num: int | None,
+        existing_names: list[str],
+    ) -> list[Entity]:
+        """
+        Check for entities that patterns might have missed.
+        
+        This is a safety net, but the LLM is instructed to be
+        conservative and only add clearly identifiable entities.
+        """
+        llm = self._get_llm()
+        if llm is None:
+            return []
+        
+        prompt = SUPPLEMENTARY_PROMPT.format(
+            content=content[:2000],
+            existing_entities=", ".join(existing_names) if existing_names else "None",
+        )
+        
+        try:
+            response = llm.complete(prompt)
+            response_text = str(response) if not isinstance(response, str) else response
+            
+            # Parse response
+            json_match = re.search(r'\[[\s\S]*?\]', response_text)
+            if not json_match:
+                return []
+            
+            missed = json.loads(json_match.group())
+            
+            if not isinstance(missed, list):
+                return []
+            
+            entities = []
+            for item in missed:
+                text = item.get("text", "").strip()
+                if not text or text in existing_names:
+                    continue
+                
+                entity_type = self._map_entity_type(item.get("type", "OTHER"))
+                canonical = item.get("canonical_name", text)
+                
+                # Find the text in content to get position
+                match = re.search(re.escape(text), content)
+                span_start = match.start() if match else None
+                span_end = match.end() if match else None
+                
+                mention = Mention(
+                    node_id=node_id,
+                    doc_id=doc_id,
+                    span_start=span_start,
+                    span_end=span_end,
+                    context=content[max(0, (span_start or 0) - 50):(span_end or 0) + 50] if span_start else "",
+                    page_num=page_num,
+                    confidence=0.6,  # Lower confidence for supplementary
+                )
+                
+                entity = Entity(
+                    type=entity_type,
+                    canonical_name=canonical,
+                    aliases=[text] if text != canonical else [],
+                    mentions=[mention],
+                    metadata={
+                        "grounded": span_start is not None,
+                        "supplementary": True,
+                        "reason": item.get("reason", ""),
+                    },
+                    source_doc_id=doc_id,
+                )
+                entities.append(entity)
+            
+            if entities:
+                logger.debug(
+                    "supplementary_entities_found",
+                    count=len(entities),
+                    names=[e.canonical_name for e in entities],
+                )
+            
+            return entities
+            
+        except Exception as e:
+            logger.debug("supplementary_extraction_failed", error=str(e))
+            return []
+    
+    def clear_cache(self) -> None:
+        """Clear the extraction cache."""
+        self._cache.clear()