rnsr 0.1.0__py3-none-any.whl

rnsr/extraction/rlm_extractor.py

@@ -0,0 +1,589 @@
+"""
+RNSR RLM Entity Extractor
+
+Implements the TRUE RLM pattern for entity extraction:
+1. LLM writes its own regex/Python code based on the document
+2. Code executes on DOC_VAR (grounded in actual text)
+3. LLM validates and classifies results
+
+This is more powerful than pre-defined patterns because:
+- LLM adapts to domain-specific patterns it discovers
+- Can write complex extraction logic we didn't anticipate
+- Still grounded because code executes on actual text
+- Recursive - can use sub_llm for complex validation
+
+From the RLM paper:
+"The Neural Network generates code to interact with the document,
+rather than having the document in its context window."
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+from dataclasses import dataclass, field
+from typing import Any, TYPE_CHECKING
+
+import structlog
+
+from rnsr.extraction.models import Entity, EntityType, ExtractionResult, Mention
+from rnsr.extraction.learned_types import get_learned_type_registry
+from rnsr.llm import get_llm
+
+if TYPE_CHECKING:
+    from rnsr.agent.repl_env import REPLEnvironment
+    from rnsr.models import DocumentTree
+
+logger = structlog.get_logger(__name__)
+
+
+# =============================================================================
+# RLM Extraction Prompts
+# =============================================================================
+
+RLM_ENTITY_EXTRACTION_SYSTEM = """You are an RLM (Recursive Language Model) extracting entities from a document.
+
+CRITICAL: You do NOT have the full document in context. It is stored in DOC_VAR.
+You must write Python code to extract entities from DOC_VAR.
+
+## Available Variables:
+- DOC_VAR: The document text (string). Use slicing, regex, etc.
+- SECTION_CONTENT: Current section content (string, smaller than DOC_VAR)
+
+## Available Functions:
+- search_text(pattern): Search DOC_VAR for regex pattern, returns list of (start, end, match)
+- len(DOC_VAR): Get document length
+- DOC_VAR[i:j]: Slice document
+- re.findall(pattern, text): Standard regex
+- re.finditer(pattern, text): Iterate matches with positions
+- store_variable(name, content): Store findings for later
+
+## Your Task:
+Extract entities (people, organizations, dates, etc.) by writing Python code.
+
+IMPORTANT: Your code should:
+1. Write regex patterns tailored to THIS document
+2. Execute patterns to find matches (grounded in text)
+3. Return structured results with exact positions
+
+## Output Format:
+Write Python code that produces a list of entity dictionaries:
+```python
+entities = []
+
+# Example: Find person names with titles
+for match in re.finditer(r'(?:Mr\.|Mrs\.|Dr\.)\s+([A-Z][a-z]+\s+[A-Z][a-z]+)', SECTION_CONTENT):
+    entities.append({
+        "text": match.group(),
+        "canonical_name": match.group(1),  # Without title
+        "type": "PERSON",
+        "start": match.start(),
+        "end": match.end(),
+        "confidence": 0.9
+    })
+
+# Example: Find dollar amounts
+for match in re.finditer(r'\$[\d,]+(?:\.\d{2})?(?:\s*(?:million|billion))?', SECTION_CONTENT):
+    entities.append({
+        "text": match.group(),
+        "type": "MONETARY",
+        "start": match.start(),
+        "end": match.end(),
+        "confidence": 0.95
+    })
+
+# Store results
+store_variable("ENTITIES", entities)
+```
+
+Write code appropriate for the document type and content shown."""
+
+
+RLM_EXTRACTION_PROMPT = """Document section to extract entities from:
+
+Section Header: {header}
+Section Content (first 2000 chars):
+---
+{content_preview}
+---
+
+Total section length: {content_length} characters
+
+Based on this content, write Python code to extract all significant entities.
+Consider:
+1. What types of entities appear in this document? (people, companies, dates, money, etc.)
+2. What patterns would match them? (titles, suffixes, formats, etc.)
+3. Are there domain-specific entities? (legal terms, technical concepts, etc.)
+
+Write Python code that will execute on SECTION_CONTENT to extract entities.
+End your code with: store_variable("ENTITIES", entities)"""
+
+
+RLM_VALIDATION_PROMPT = """You extracted these entity candidates from the document.
+Validate each one and determine if it's a real, significant entity.
+
+Candidates:
+{candidates_json}
+
+For each candidate, provide:
+1. valid: true if significant entity, false if noise
+2. type: Entity type (PERSON, ORGANIZATION, DATE, MONETARY, LOCATION, etc.)
+3. canonical_name: Cleaned/normalized name
+4. confidence: 0.0-1.0
+
+Return JSON array:
+```json
+[
+  {{"id": 0, "valid": true, "type": "PERSON", "canonical_name": "John Smith", "confidence": 0.9}},
+  {{"id": 1, "valid": false, "reason": "Generic term, not specific entity"}}
+]
+```"""
+
+
+# =============================================================================
+# Lightweight REPL for Extraction (if full REPL not available)
+# =============================================================================
+
+class LightweightREPL:
+    """
+    Lightweight REPL for entity extraction.
+    
+    Provides the core DOC_VAR + code execution pattern
+    without the full REPL infrastructure.
+    """
+    
+    def __init__(self, document_text: str, section_content: str = ""):
+        """Initialize with document text."""
+        self.document_text = document_text
+        self.section_content = section_content or document_text
+        self.variables: dict[str, Any] = {}
+        
+        self._namespace = self._build_namespace()
+    
+    def _build_namespace(self) -> dict[str, Any]:
+        """Build Python namespace for code execution."""
+        return {
+            # Core variables
+            "DOC_VAR": self.document_text,
+            "SECTION_CONTENT": self.section_content,
+            "VARIABLES": self.variables,
+            
+            # Built-ins
+            "len": len,
+            "str": str,
+            "int": int,
+            "float": float,
+            "list": list,
+            "dict": dict,
+            "range": range,
+            "enumerate": enumerate,
+            "sorted": sorted,
+            "min": min,
+            "max": max,
+            "re": re,
+            
+            # Functions
+            "search_text": self._search_text,
+            "store_variable": self._store_variable,
+            "get_variable": self._get_variable,
+        }
+    
+    def _search_text(self, pattern: str) -> list[tuple[int, int, str]]:
+        """Search document for regex pattern."""
+        results = []
+        try:
+            for match in re.finditer(pattern, self.section_content, re.IGNORECASE):
+                results.append((match.start(), match.end(), match.group()))
+        except re.error as e:
+            logger.warning("regex_error", pattern=pattern, error=str(e))
+        return results
+    
+    def _store_variable(self, name: str, content: Any) -> str:
+        """Store a variable."""
+        self.variables[name] = content
+        return f"Stored ${name}"
+    
+    def _get_variable(self, name: str) -> Any:
+        """Retrieve a variable."""
+        return self.variables.get(name)
+    
+    def execute(self, code: str) -> dict[str, Any]:
+        """Execute Python code."""
+        result = {
+            "success": False,
+            "output": None,
+            "error": None,
+            "variables": list(self.variables.keys()),
+        }
+        
+        # Clean code
+        code = self._clean_code(code)
+        
+        try:
+            # Compile and execute
+            compiled = compile(code, "<rlm_extraction>", "exec")
+            exec(compiled, self._namespace)
+            
+            result["success"] = True
+            result["variables"] = list(self.variables.keys())
+            result["output"] = self.variables.get("ENTITIES", [])
+            
+        except Exception as e:
+            result["error"] = str(e)
+            logger.warning("rlm_execution_error", error=str(e), code=code[:200])
+        
+        return result
+    
+    def _clean_code(self, code: str) -> str:
+        """Remove markdown code blocks."""
+        code = re.sub(r'^```python\s*', '', code, flags=re.MULTILINE)
+        code = re.sub(r'^```\s*$', '', code, flags=re.MULTILINE)
+        return code.strip()
+
+
+# =============================================================================
+# RLM Entity Extractor
+# =============================================================================
+
+@dataclass
+class RLMExtractionResult:
+    """Result of RLM-based extraction."""
+    
+    entities: list[Entity] = field(default_factory=list)
+    code_generated: str = ""
+    code_executed: bool = False
+    execution_output: Any = None
+    raw_candidates: list[dict] = field(default_factory=list)
+    processing_time_ms: float = 0.0
+    warnings: list[str] = field(default_factory=list)
+
+
+class RLMEntityExtractor:
+    """
+    RLM-based entity extractor.
+    
+    The LLM writes its own extraction code based on the document,
+    then the code executes on DOC_VAR (grounded).
+    
+    Flow:
+    1. Show LLM a preview of the document
+    2. LLM writes Python code to extract entities
+    3. Code executes on actual document (grounded)
+    4. LLM validates/classifies the extracted candidates
+    """
+    
+    def __init__(
+        self,
+        llm: Any | None = None,
+        repl_environment: "REPLEnvironment | None" = None,
+        enable_type_learning: bool = True,
+        max_code_attempts: int = 3,
+        validate_with_llm: bool = True,
+    ):
+        """
+        Initialize the RLM extractor.
+        
+        Args:
+            llm: LLM instance.
+            repl_environment: Optional full REPL environment.
+            enable_type_learning: Learn new entity types.
+            max_code_attempts: Max attempts if code fails.
+            validate_with_llm: Validate candidates with LLM.
+        """
+        self.llm = llm
+        self.repl_environment = repl_environment
+        self.enable_type_learning = enable_type_learning
+        self.max_code_attempts = max_code_attempts
+        self.validate_with_llm = validate_with_llm
+        
+        self._llm_initialized = False
+        self._type_registry = get_learned_type_registry() if enable_type_learning else None
+    
+    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 extract_from_node(
+        self,
+        node_id: str,
+        doc_id: str,
+        header: str,
+        content: str,
+        page_num: int | None = None,
+        document_text: str | None = None,
+    ) -> ExtractionResult:
+        """
+        Extract entities using RLM approach.
+        
+        The LLM writes code to extract entities, which is then
+        executed on the actual document text.
+        
+        Args:
+            node_id: Section node ID.
+            doc_id: Document ID.
+            header: Section header.
+            content: Section content.
+            page_num: Page number.
+            document_text: Full document text for DOC_VAR.
+            
+        Returns:
+            ExtractionResult with extracted entities.
+        """
+        start_time = time.time()
+        
+        result = ExtractionResult(
+            node_id=node_id,
+            doc_id=doc_id,
+            extraction_method="rlm",
+        )
+        
+        if len(content.strip()) < 50:
+            return result
+        
+        llm = self._get_llm()
+        if llm is None:
+            result.warnings.append("No LLM available for RLM extraction")
+            return result
+        
+        # STEP 1: LLM generates extraction code
+        rlm_result = self._generate_and_execute_code(
+            header=header,
+            content=content,
+            document_text=document_text or content,
+        )
+        
+        if not rlm_result.code_executed:
+            result.warnings.append(f"Code execution failed: {rlm_result.warnings}")
+            return result
+        
+        # STEP 2: Validate candidates with LLM
+        if self.validate_with_llm and rlm_result.raw_candidates:
+            validated = self._validate_candidates(rlm_result.raw_candidates)
+        else:
+            validated = rlm_result.raw_candidates
+        
+        # STEP 3: Convert to Entity objects
+        entities = self._candidates_to_entities(
+            candidates=validated,
+            node_id=node_id,
+            doc_id=doc_id,
+            content=content,
+            page_num=page_num,
+        )
+        
+        result.entities = entities
+        result.processing_time_ms = (time.time() - start_time) * 1000
+        
+        logger.info(
+            "rlm_extraction_complete",
+            node_id=node_id,
+            candidates=len(rlm_result.raw_candidates),
+            validated=len(entities),
+            time_ms=result.processing_time_ms,
+        )
+        
+        return result
+    
+    def _generate_and_execute_code(
+        self,
+        header: str,
+        content: str,
+        document_text: str,
+    ) -> RLMExtractionResult:
+        """Generate extraction code and execute it."""
+        result = RLMExtractionResult()
+        
+        llm = self._get_llm()
+        
+        # Create REPL environment
+        repl = self.repl_environment or LightweightREPL(
+            document_text=document_text,
+            section_content=content,
+        )
+        
+        # Generate code prompt
+        prompt = f"{RLM_ENTITY_EXTRACTION_SYSTEM}\n\n{RLM_EXTRACTION_PROMPT.format(
+            header=header,
+            content_preview=content[:2000],
+            content_length=len(content),
+        )}"
+        
+        for attempt in range(self.max_code_attempts):
+            try:
+                # LLM generates extraction code
+                response = llm.complete(prompt)
+                code = str(response) if not isinstance(response, str) else response
+                result.code_generated = code
+                
+                # Execute the code
+                exec_result = repl.execute(code)
+                
+                if exec_result["success"]:
+                    result.code_executed = True
+                    result.execution_output = exec_result["output"]
+                    result.raw_candidates = exec_result.get("output", [])
+                    
+                    # Get ENTITIES from variables if not in output
+                    if not result.raw_candidates and hasattr(repl, 'variables'):
+                        result.raw_candidates = repl.variables.get("ENTITIES", [])
+                    
+                    break
+                else:
+                    result.warnings.append(f"Attempt {attempt + 1}: {exec_result['error']}")
+                    # Add error to prompt for retry
+                    prompt += f"\n\nPrevious code had error: {exec_result['error']}\nPlease fix and try again."
+                    
+            except Exception as e:
+                result.warnings.append(f"Attempt {attempt + 1}: {str(e)}")
+        
+        return result
+    
+    def _validate_candidates(
+        self,
+        candidates: list[dict],
+    ) -> list[dict]:
+        """Validate extracted candidates with LLM."""
+        if not candidates:
+            return []
+        
+        llm = self._get_llm()
+        
+        # Format candidates for validation
+        candidates_json = json.dumps([
+            {
+                "id": i,
+                "text": c.get("text", ""),
+                "type": c.get("type", "UNKNOWN"),
+                "context": c.get("context", "")[:100] if c.get("context") else "",
+            }
+            for i, c in enumerate(candidates[:30])  # Limit
+        ], indent=2)
+        
+        prompt = RLM_VALIDATION_PROMPT.format(candidates_json=candidates_json)
+        
+        try:
+            response = llm.complete(prompt)
+            response_text = str(response) if not isinstance(response, str) else response
+            
+            # Parse validation response
+            json_match = re.search(r'\[[\s\S]*\]', response_text)
+            if not json_match:
+                return candidates
+            
+            validations = json.loads(json_match.group())
+            
+            # Merge validations with candidates
+            validated = []
+            validation_by_id = {v.get("id"): v for v in validations}
+            
+            for i, candidate in enumerate(candidates):
+                validation = validation_by_id.get(i, {})
+                
+                if validation.get("valid", True):
+                    candidate["type"] = validation.get("type", candidate.get("type", "OTHER"))
+                    candidate["canonical_name"] = validation.get("canonical_name", candidate.get("text", ""))
+                    candidate["confidence"] = validation.get("confidence", candidate.get("confidence", 0.5))
+                    validated.append(candidate)
+            
+            return validated
+            
+        except Exception as e:
+            logger.warning("rlm_validation_failed", error=str(e))
+            return candidates
+    
+    def _candidates_to_entities(
+        self,
+        candidates: list[dict],
+        node_id: str,
+        doc_id: str,
+        content: str,
+        page_num: int | None,
+    ) -> list[Entity]:
+        """Convert validated candidates to Entity objects."""
+        entities = []
+        
+        for candidate in candidates:
+            if not candidate.get("text"):
+                continue
+            
+            # Map entity type
+            entity_type = self._map_entity_type(candidate.get("type", "OTHER"))
+            
+            # Learn new types
+            if entity_type == EntityType.OTHER and self._type_registry:
+                self._type_registry.record_type(
+                    type_name=candidate.get("type", "unknown").lower(),
+                    context=candidate.get("context", content[:100]),
+                    entity_name=candidate.get("text", ""),
+                )
+            
+            # Create mention
+            mention = Mention(
+                node_id=node_id,
+                doc_id=doc_id,
+                span_start=candidate.get("start"),
+                span_end=candidate.get("end"),
+                context=candidate.get("context", content[
+                    max(0, (candidate.get("start") or 0) - 50):
+                    (candidate.get("end") or 0) + 50
+                ]),
+                page_num=page_num,
+                confidence=candidate.get("confidence", 0.5),
+            )
+            
+            # Build metadata
+            metadata = {
+                "rlm_extracted": True,
+                "grounded": candidate.get("start") is not None,
+            }
+            
+            if entity_type == EntityType.OTHER:
+                metadata["original_type"] = candidate.get("type", "").lower()
+            
+            entity = Entity(
+                type=entity_type,
+                canonical_name=candidate.get("canonical_name", candidate.get("text", "")),
+                aliases=[candidate.get("text")] if candidate.get("canonical_name") != candidate.get("text") 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,
+            "PEOPLE": EntityType.PERSON,
+            "NAME": EntityType.PERSON,
+            "ORGANIZATION": EntityType.ORGANIZATION,
+            "ORG": EntityType.ORGANIZATION,
+            "COMPANY": EntityType.ORGANIZATION,
+            "DATE": EntityType.DATE,
+            "TIME": EntityType.DATE,
+            "MONETARY": EntityType.MONETARY,
+            "MONEY": EntityType.MONETARY,
+            "AMOUNT": EntityType.MONETARY,
+            "LOCATION": EntityType.LOCATION,
+            "PLACE": EntityType.LOCATION,
+            "ADDRESS": EntityType.LOCATION,
+            "REFERENCE": EntityType.REFERENCE,
+            "CITATION": 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)