rnsr 0.1.0__py3-none-any.whl

rnsr/extraction/tot_validator.py

@@ -0,0 +1,610 @@
+"""
+RNSR Tree of Thoughts Entity Validator
+
+Applies the ToT pattern from the RLM Navigator to entity validation:
+
+1. Given pre-extracted candidates, evaluate each with probability + reasoning
+2. Navigate the document tree for additional context when uncertain
+3. Make multi-step decisions (like backtracking in document navigation)
+
+This prevents hallucination because:
+- Candidates are already grounded in text (from pattern extraction)
+- ToT provides structured evaluation with explicit probabilities
+- Navigation provides additional context for ambiguous cases
+- Same battle-tested pattern used for document Q&A
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+import structlog
+
+from rnsr.extraction.candidate_extractor import EntityCandidate
+from rnsr.extraction.models import Entity, EntityType, Mention
+from rnsr.llm import get_llm
+from rnsr.models import DocumentTree
+
+logger = structlog.get_logger(__name__)
+
+
+# ToT-style prompt for entity validation (mirrors graph.py ToT_SYSTEM_PROMPT pattern)
+TOT_ENTITY_VALIDATION_PROMPT = """You are validating entity candidates extracted from a document.
+
+Current Section: {section_header}
+Section Content:
+---
+{section_content}
+---
+
+Entity Candidates to Evaluate:
+{candidates_formatted}
+
+EVALUATION TASK:
+For each candidate, estimate the probability (0.0 to 1.0) that it is a valid, 
+significant entity worth tracking, AND classify its type.
+
+INSTRUCTIONS:
+1. Evaluate: For each candidate, analyze its context and estimate validity probability.
+2. Valid entities have: clear identity, specific name, significance to document.
+3. Invalid entities are: generic terms, partial matches, noise, common words.
+4. If probability >= {selection_threshold}, include in selected_entities.
+5. If probability < {rejection_threshold}, mark as rejected.
+6. Provide brief reasoning for each decision.
+7. Classify type: PERSON, ORGANIZATION, DATE, LOCATION, MONETARY, REFERENCE, 
+   DOCUMENT, EVENT, LEGAL_CONCEPT, or describe a custom type.
+
+OUTPUT FORMAT (JSON):
+{{
+    "evaluations": [
+        {{
+            "candidate_id": 0,
+            "probability": 0.85,
+            "is_valid": true,
+            "entity_type": "PERSON",
+            "canonical_name": "John Smith",
+            "role": "defendant",
+            "reasoning": "Clear person name with title, mentioned as party to case"
+        }},
+        {{
+            "candidate_id": 1,
+            "probability": 0.30,
+            "is_valid": false,
+            "entity_type": null,
+            "canonical_name": null,
+            "reasoning": "Generic reference to 'the agreement', not a specific entity"
+        }}
+    ],
+    "selected_entities": [0],
+    "needs_more_context": [],
+    "high_confidence_count": 1,
+    "low_confidence_count": 1
+}}
+
+If uncertain about a candidate (probability 0.4-0.6), add its id to "needs_more_context".
+We may navigate to related sections to gather more information.
+
+Respond ONLY with the JSON, no other text."""
+
+
+# Prompt for gathering context from related sections
+TOT_CONTEXT_GATHERING_PROMPT = """You need more context to validate an entity candidate.
+
+Entity candidate: "{candidate_text}" (type hint: {type_hint})
+Original section: {original_section}
+
+Related sections found:
+{related_sections}
+
+Based on this additional context, provide your evaluation:
+{{
+    "candidate_id": {candidate_id},
+    "probability": 0.XX,
+    "is_valid": true/false,
+    "entity_type": "TYPE",
+    "canonical_name": "Full Name",
+    "reasoning": "With additional context from section X, this is clearly a..."
+}}
+
+Respond ONLY with the JSON, no other text."""
+
+
+@dataclass
+class TotValidationResult:
+    """Result of ToT entity validation."""
+    
+    candidate_id: int
+    probability: float
+    is_valid: bool
+    entity_type: str | None = None
+    canonical_name: str | None = None
+    role: str | None = None
+    reasoning: str = ""
+    used_navigation: bool = False
+    
+
+@dataclass
+class TotBatchResult:
+    """Result of validating a batch of candidates."""
+    
+    evaluations: list[TotValidationResult] = field(default_factory=list)
+    selected_entities: list[int] = field(default_factory=list)
+    needs_more_context: list[int] = field(default_factory=list)
+    high_confidence_count: int = 0
+    low_confidence_count: int = 0
+
+
+class TotEntityValidator:
+    """
+    Tree of Thoughts entity validator.
+    
+    Uses the same ToT pattern as document navigation:
+    - Evaluate candidates with explicit probabilities
+    - Navigate for context when uncertain
+    - Structured JSON output for reliable parsing
+    """
+    
+    def __init__(
+        self,
+        llm: Any | None = None,
+        selection_threshold: float = 0.6,
+        rejection_threshold: float = 0.3,
+        enable_navigation: bool = True,
+        max_navigation_depth: int = 2,
+        max_candidates_per_batch: int = 20,
+    ):
+        """
+        Initialize the ToT validator.
+        
+        Args:
+            llm: LLM instance.
+            selection_threshold: Probability threshold for accepting entity.
+            rejection_threshold: Probability threshold for rejecting entity.
+            enable_navigation: Navigate tree for uncertain candidates.
+            max_navigation_depth: Max depth to navigate for context.
+            max_candidates_per_batch: Max candidates per LLM call.
+        """
+        self.llm = llm
+        self.selection_threshold = selection_threshold
+        self.rejection_threshold = rejection_threshold
+        self.enable_navigation = enable_navigation
+        self.max_navigation_depth = max_navigation_depth
+        self.max_candidates_per_batch = max_candidates_per_batch
+        
+        self._llm_initialized = False
+    
+    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 validate_candidates(
+        self,
+        candidates: list[EntityCandidate],
+        section_header: str,
+        section_content: str,
+        document_tree: DocumentTree | None = None,
+        node_id: str | None = None,
+    ) -> TotBatchResult:
+        """
+        Validate entity candidates using ToT reasoning.
+        
+        Args:
+            candidates: Pre-extracted candidates to validate.
+            section_header: Current section header.
+            section_content: Current section content.
+            document_tree: Optional tree for navigation.
+            node_id: Current node ID for navigation.
+            
+        Returns:
+            TotBatchResult with validated entities.
+        """
+        if not candidates:
+            return TotBatchResult()
+        
+        llm = self._get_llm()
+        if llm is None:
+            # No LLM - accept all candidates with pattern-based types
+            return self._accept_all_candidates(candidates)
+        
+        # Process in batches
+        all_results = TotBatchResult()
+        
+        for i in range(0, len(candidates), self.max_candidates_per_batch):
+            batch = candidates[i:i + self.max_candidates_per_batch]
+            batch_offset = i
+            
+            batch_result = self._validate_batch(
+                candidates=batch,
+                batch_offset=batch_offset,
+                section_header=section_header,
+                section_content=section_content,
+            )
+            
+            # Merge results
+            all_results.evaluations.extend(batch_result.evaluations)
+            all_results.selected_entities.extend(batch_result.selected_entities)
+            all_results.needs_more_context.extend(batch_result.needs_more_context)
+            all_results.high_confidence_count += batch_result.high_confidence_count
+            all_results.low_confidence_count += batch_result.low_confidence_count
+        
+        # Handle uncertain candidates if navigation is enabled
+        if self.enable_navigation and document_tree and all_results.needs_more_context:
+            all_results = self._resolve_uncertain_candidates(
+                candidates=candidates,
+                batch_result=all_results,
+                document_tree=document_tree,
+                current_node_id=node_id,
+            )
+        
+        return all_results
+    
+    def _validate_batch(
+        self,
+        candidates: list[EntityCandidate],
+        batch_offset: int,
+        section_header: str,
+        section_content: str,
+    ) -> TotBatchResult:
+        """Validate a batch of candidates with ToT."""
+        # Format candidates for prompt
+        candidates_formatted = "\n".join([
+            f"[{i + batch_offset}] Text: \"{c.text}\" | Type Hint: {c.candidate_type} | "
+            f"Context: \"...{c.context[:100]}...\""
+            for i, c in enumerate(candidates)
+        ])
+        
+        prompt = TOT_ENTITY_VALIDATION_PROMPT.format(
+            section_header=section_header,
+            section_content=section_content[:2500],
+            candidates_formatted=candidates_formatted,
+            selection_threshold=self.selection_threshold,
+            rejection_threshold=self.rejection_threshold,
+        )
+        
+        try:
+            response = self.llm.complete(prompt)
+            response_text = str(response) if not isinstance(response, str) else response
+            
+            return self._parse_validation_response(response_text, len(candidates), batch_offset)
+            
+        except Exception as e:
+            logger.warning("tot_validation_failed", error=str(e))
+            return self._accept_all_candidates(candidates, offset=batch_offset)
+    
+    def _parse_validation_response(
+        self,
+        response_text: str,
+        candidate_count: int,
+        batch_offset: int,
+    ) -> TotBatchResult:
+        """Parse ToT validation response."""
+        result = TotBatchResult()
+        
+        # Extract JSON
+        json_match = re.search(r'\{[\s\S]*\}', response_text)
+        if not json_match:
+            logger.warning("tot_no_json_found")
+            return result
+        
+        try:
+            data = json.loads(json_match.group())
+        except json.JSONDecodeError as e:
+            logger.warning("tot_json_parse_failed", error=str(e))
+            return result
+        
+        # Parse evaluations
+        for eval_data in data.get("evaluations", []):
+            try:
+                validation = TotValidationResult(
+                    candidate_id=eval_data.get("candidate_id", 0),
+                    probability=float(eval_data.get("probability", 0.5)),
+                    is_valid=eval_data.get("is_valid", False),
+                    entity_type=eval_data.get("entity_type"),
+                    canonical_name=eval_data.get("canonical_name"),
+                    role=eval_data.get("role"),
+                    reasoning=eval_data.get("reasoning", ""),
+                )
+                result.evaluations.append(validation)
+                
+                if validation.is_valid:
+                    result.high_confidence_count += 1
+                else:
+                    result.low_confidence_count += 1
+                    
+            except (KeyError, TypeError, ValueError) as e:
+                logger.debug("tot_eval_parse_error", error=str(e))
+                continue
+        
+        # Parse selected entities (adjust for batch offset)
+        result.selected_entities = [
+            idx for idx in data.get("selected_entities", [])
+            if isinstance(idx, int)
+        ]
+        
+        # Parse needs_more_context
+        result.needs_more_context = [
+            idx for idx in data.get("needs_more_context", [])
+            if isinstance(idx, int)
+        ]
+        
+        return result
+    
+    def _resolve_uncertain_candidates(
+        self,
+        candidates: list[EntityCandidate],
+        batch_result: TotBatchResult,
+        document_tree: DocumentTree,
+        current_node_id: str | None,
+    ) -> TotBatchResult:
+        """
+        Navigate document tree to resolve uncertain candidates.
+        
+        This is like backtracking in document Q&A - gather more context
+        to make a better decision.
+        """
+        if not batch_result.needs_more_context:
+            return batch_result
+        
+        logger.info(
+            "tot_navigating_for_context",
+            uncertain_count=len(batch_result.needs_more_context),
+        )
+        
+        # Find related sections
+        related_sections = self._find_related_sections(
+            document_tree=document_tree,
+            current_node_id=current_node_id,
+            depth=self.max_navigation_depth,
+        )
+        
+        if not related_sections:
+            # No related sections - accept uncertain candidates with lower confidence
+            for idx in batch_result.needs_more_context:
+                if idx < len(candidates):
+                    # Add as selected with moderate confidence
+                    batch_result.selected_entities.append(idx)
+            return batch_result
+        
+        # Re-evaluate uncertain candidates with additional context
+        for idx in batch_result.needs_more_context:
+            if idx >= len(candidates):
+                continue
+                
+            candidate = candidates[idx]
+            
+            resolved = self._resolve_single_candidate(
+                candidate=candidate,
+                candidate_id=idx,
+                related_sections=related_sections,
+            )
+            
+            if resolved and resolved.is_valid:
+                # Update the evaluation
+                for i, eval_item in enumerate(batch_result.evaluations):
+                    if eval_item.candidate_id == idx:
+                        batch_result.evaluations[i] = resolved
+                        break
+                else:
+                    batch_result.evaluations.append(resolved)
+                
+                batch_result.selected_entities.append(idx)
+                batch_result.high_confidence_count += 1
+        
+        # Clear needs_more_context since we've processed them
+        batch_result.needs_more_context = []
+        
+        return batch_result
+    
+    def _find_related_sections(
+        self,
+        document_tree: DocumentTree,
+        current_node_id: str | None,
+        depth: int,
+    ) -> list[dict[str, str]]:
+        """Find related sections for context gathering."""
+        sections = []
+        
+        if not document_tree or not document_tree.root:
+            return sections
+        
+        # Collect sections from tree (siblings and nearby nodes)
+        def collect_sections(node: Any, current_depth: int) -> None:
+            if current_depth > depth:
+                return
+            
+            if hasattr(node, 'header') and hasattr(node, 'content'):
+                node_id = getattr(node, 'id', str(id(node)))
+                if node_id != current_node_id:
+                    sections.append({
+                        "header": node.header or "(no header)",
+                        "content": (node.content or "")[:500],
+                    })
+            
+            if hasattr(node, 'children'):
+                for child in node.children[:5]:  # Limit children
+                    collect_sections(child, current_depth + 1)
+        
+        collect_sections(document_tree.root, 0)
+        
+        return sections[:10]  # Limit total sections
+    
+    def _resolve_single_candidate(
+        self,
+        candidate: EntityCandidate,
+        candidate_id: int,
+        related_sections: list[dict[str, str]],
+    ) -> TotValidationResult | None:
+        """Resolve a single uncertain candidate with additional context."""
+        llm = self._get_llm()
+        if llm is None:
+            return None
+        
+        # Format related sections
+        sections_text = "\n\n".join([
+            f"### {s['header']}\n{s['content']}"
+            for s in related_sections[:5]
+        ])
+        
+        prompt = TOT_CONTEXT_GATHERING_PROMPT.format(
+            candidate_text=candidate.text,
+            type_hint=candidate.candidate_type,
+            original_section=candidate.context[:200],
+            related_sections=sections_text,
+            candidate_id=candidate_id,
+        )
+        
+        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 None
+            
+            data = json.loads(json_match.group())
+            
+            return TotValidationResult(
+                candidate_id=candidate_id,
+                probability=float(data.get("probability", 0.5)),
+                is_valid=data.get("is_valid", False),
+                entity_type=data.get("entity_type"),
+                canonical_name=data.get("canonical_name"),
+                role=data.get("role"),
+                reasoning=data.get("reasoning", ""),
+                used_navigation=True,
+            )
+            
+        except Exception as e:
+            logger.debug("tot_context_resolution_failed", error=str(e))
+            return None
+    
+    def _accept_all_candidates(
+        self,
+        candidates: list[EntityCandidate],
+        offset: int = 0,
+    ) -> TotBatchResult:
+        """Accept all candidates without validation (fallback)."""
+        result = TotBatchResult()
+        
+        for i, candidate in enumerate(candidates):
+            idx = i + offset
+            result.evaluations.append(TotValidationResult(
+                candidate_id=idx,
+                probability=candidate.confidence,
+                is_valid=True,
+                entity_type=candidate.candidate_type.upper(),
+                canonical_name=candidate.text,
+                reasoning="Accepted without LLM validation",
+            ))
+            result.selected_entities.append(idx)
+        
+        result.high_confidence_count = len(candidates)
+        return result
+    
+    def candidates_to_entities(
+        self,
+        candidates: list[EntityCandidate],
+        validation_result: TotBatchResult,
+        node_id: str,
+        doc_id: str,
+        page_num: int | None = None,
+    ) -> list[Entity]:
+        """
+        Convert validated candidates to Entity objects.
+        
+        Only includes candidates that passed ToT validation.
+        """
+        entities = []
+        
+        # Build lookup for evaluations
+        eval_by_id = {e.candidate_id: e for e in validation_result.evaluations}
+        
+        for idx in validation_result.selected_entities:
+            if idx >= len(candidates):
+                continue
+            
+            candidate = candidates[idx]
+            evaluation = eval_by_id.get(idx)
+            
+            if not evaluation or not evaluation.is_valid:
+                continue
+            
+            # Map entity type
+            entity_type = self._map_entity_type(
+                evaluation.entity_type or candidate.candidate_type
+            )
+            
+            # Get canonical name
+            canonical_name = evaluation.canonical_name or candidate.text
+            
+            # 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=evaluation.probability,
+            )
+            
+            # Build metadata
+            metadata = {
+                "grounded": True,
+                "tot_validated": True,
+                "tot_probability": evaluation.probability,
+                "tot_reasoning": evaluation.reasoning,
+                "pattern": candidate.pattern_name,
+            }
+            
+            if evaluation.role:
+                metadata["role"] = evaluation.role
+            
+            if evaluation.used_navigation:
+                metadata["used_context_navigation"] = True
+            
+            if entity_type == EntityType.OTHER:
+                metadata["original_type"] = (evaluation.entity_type or "").lower()
+            
+            entity = 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,
+            )
+            entities.append(entity)
+        
+        return entities
+    
+    def _map_entity_type(self, type_str: str) -> EntityType:
+        """Map type string to EntityType enum."""
+        type_str = type_str.upper()
+        
+        mapping = {
+            "PERSON": EntityType.PERSON,
+            "ORGANIZATION": EntityType.ORGANIZATION,
+            "ORG": EntityType.ORGANIZATION,
+            "DATE": EntityType.DATE,
+            "LOCATION": EntityType.LOCATION,
+            "MONETARY": EntityType.MONETARY,
+            "MONEY": EntityType.MONETARY,
+            "REFERENCE": EntityType.REFERENCE,
+            "DOCUMENT": EntityType.DOCUMENT,
+            "EVENT": EntityType.EVENT,
+            "LEGAL_CONCEPT": EntityType.LEGAL_CONCEPT,
+            "LEGAL": EntityType.LEGAL_CONCEPT,
+        }
+        
+        try:
+            return EntityType(type_str.lower())
+        except ValueError:
+            return mapping.get(type_str, EntityType.OTHER)