rnsr 0.1.0__py3-none-any.whl

rnsr/extraction/entity_extractor.py

@@ -0,0 +1,581 @@
+"""
+RNSR Entity Extractor
+
+DEPRECATED: This extractor uses LLM-first approach which can hallucinate.
+Use RLMUnifiedExtractor instead for grounded, accurate extraction.
+
+LLM-based entity extraction from document sections.
+Extracts people, organizations, dates, legal concepts, and other entities.
+
+Features adaptive learning: when the LLM discovers new entity types, they are
+stored in a learned types registry and used in future extraction prompts.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+import warnings
+from typing import Any
+
+import structlog
+
+from rnsr.extraction.models import (
+    Entity,
+    EntityType,
+    ExtractionResult,
+    Mention,
+)
+from rnsr.extraction.learned_types import (
+    get_learned_type_registry,
+    record_learned_type,
+)
+from rnsr.llm import get_llm
+
+logger = structlog.get_logger(__name__)
+
+# Deprecation warning
+_DEPRECATION_WARNING = """
+EntityExtractor is deprecated and may hallucinate entities.
+Use RLMUnifiedExtractor instead for grounded, accurate extraction:
+
+    from rnsr.extraction import RLMUnifiedExtractor
+    extractor = RLMUnifiedExtractor()
+    result = extractor.extract(node_id, doc_id, header, content)
+"""
+
+
+# Entity extraction prompt template
+ENTITY_EXTRACTION_PROMPT = """You are an expert entity extractor for legal and business documents.
+
+Analyze the following document section and extract all significant entities.
+
+Document Section:
+---
+{content}
+---
+
+Section ID: {node_id}
+Document ID: {doc_id}
+Section Header: {header}
+
+Extract entities of the following types:
+- PERSON: Names of individuals, including their roles if mentioned (e.g., "plaintiff", "defendant", "witness", "CEO")
+- ORGANIZATION: Companies, agencies, courts, government bodies
+- LEGAL_CONCEPT: Legal claims, breaches, obligations, remedies, causes of action
+- DATE: Specific dates, time periods, deadlines
+- EVENT: Significant occurrences (hearings, signings, breaches, filings)
+- LOCATION: Places, addresses, jurisdictions
+- REFERENCE: Section references, exhibit numbers, document citations
+- MONETARY: Dollar amounts, financial figures
+- DOCUMENT: Referenced documents (contracts, exhibits, agreements)
+{learned_types_section}
+For each entity, provide:
+1. type: One of the types above (or your own descriptive type if none fit)
+2. canonical_name: The standardized/normalized name
+3. aliases: Any alternative names or spellings found
+4. context: The surrounding sentence or phrase where the entity appears
+5. metadata: Any additional relevant information (roles, dates, amounts)
+
+Return your response as a JSON array of entities:
+```json
+[
+  {{
+    "type": "PERSON",
+    "canonical_name": "John Smith",
+    "aliases": ["Mr. Smith", "J. Smith"],
+    "context": "John Smith, the defendant, filed a motion...",
+    "metadata": {{"role": "defendant"}}
+  }},
+  ...
+]
+```
+
+If no entities are found, return an empty array: []
+
+Important:
+- Be thorough but precise - only extract clearly identifiable entities
+- Normalize names (e.g., "Mr. John Smith" -> "John Smith")
+- Include context that helps understand the entity's role
+- For legal concepts, use standardized legal terminology
+- If an entity doesn't fit the predefined types, use your own descriptive type name
+"""
+
+
+class EntityExtractor:
+    """
+    DEPRECATED: Extracts entities from document sections using LLM-first approach.
+    
+    This extractor can hallucinate entities. Use RLMUnifiedExtractor instead.
+    
+    Supports batch processing, caching, and adaptive learning of entity types.
+    When new entity types are discovered, they are stored and used in future prompts.
+    """
+    
+    def __init__(
+        self,
+        llm: Any | None = None,
+        min_content_length: int = 50,
+        max_content_length: int = 8000,
+        enable_type_learning: bool = True,
+        learned_type_min_count: int = 2,
+        suppress_deprecation_warning: bool = False,
+    ):
+        # Emit deprecation warning
+        if not suppress_deprecation_warning:
+            warnings.warn(
+                _DEPRECATION_WARNING,
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            logger.warning("deprecated_extractor_used", extractor="EntityExtractor")
+        """
+        Initialize the entity extractor.
+        
+        Args:
+            llm: LLM instance to use. If None, uses get_llm().
+            min_content_length: Minimum content length to process.
+            max_content_length: Maximum content length per extraction call.
+            enable_type_learning: Whether to learn new entity types.
+            learned_type_min_count: Minimum occurrences before a learned type
+                                    is included in extraction prompts.
+        """
+        self.llm = llm or get_llm()
+        self.min_content_length = min_content_length
+        self.max_content_length = max_content_length
+        self.enable_type_learning = enable_type_learning
+        self.learned_type_min_count = learned_type_min_count
+        
+        # Cache for extracted entities (node_id -> entities)
+        self._cache: dict[str, list[Entity]] = {}
+        
+        # Get learned type registry
+        self._type_registry = get_learned_type_registry() if enable_type_learning else None
+    
+    def extract_from_node(
+        self,
+        node_id: str,
+        doc_id: str,
+        header: str,
+        content: str,
+        page_num: int | None = None,
+    ) -> ExtractionResult:
+        """
+        Extract entities from a single document node.
+        
+        Args:
+            node_id: Skeleton node ID.
+            doc_id: Document ID.
+            header: Section header text.
+            content: Full section content.
+            page_num: Page number if available.
+            
+        Returns:
+            ExtractionResult with extracted entities.
+        """
+        start_time = time.time()
+        result = ExtractionResult(
+            node_id=node_id,
+            doc_id=doc_id,
+            extraction_method="llm",
+        )
+        
+        # Skip very short content
+        if len(content.strip()) < self.min_content_length:
+            logger.debug(
+                "skipping_short_content",
+                node_id=node_id,
+                content_length=len(content),
+            )
+            return result
+        
+        # Check cache
+        cache_key = f"{doc_id}:{node_id}"
+        if cache_key in self._cache:
+            result.entities = self._cache[cache_key]
+            logger.debug("using_cached_entities", node_id=node_id)
+            return result
+        
+        # Truncate content if too long
+        if len(content) > self.max_content_length:
+            content = content[:self.max_content_length] + "..."
+            result.warnings.append(f"Content truncated from {len(content)} chars")
+        
+        try:
+            entities = self._extract_with_llm(
+                node_id=node_id,
+                doc_id=doc_id,
+                header=header,
+                content=content,
+                page_num=page_num,
+            )
+            result.entities = entities
+            
+            # Cache results
+            self._cache[cache_key] = entities
+            
+        except Exception as e:
+            logger.error(
+                "entity_extraction_failed",
+                node_id=node_id,
+                error=str(e),
+            )
+            result.warnings.append(f"Extraction failed: {str(e)}")
+        
+        result.processing_time_ms = (time.time() - start_time) * 1000
+        
+        logger.info(
+            "entities_extracted",
+            node_id=node_id,
+            entity_count=len(result.entities),
+            processing_time_ms=result.processing_time_ms,
+        )
+        
+        return result
+    
+    def _extract_with_llm(
+        self,
+        node_id: str,
+        doc_id: str,
+        header: str,
+        content: str,
+        page_num: int | None = None,
+    ) -> list[Entity]:
+        """
+        Use LLM to extract entities from content.
+        
+        Args:
+            node_id: Skeleton node ID.
+            doc_id: Document ID.
+            header: Section header.
+            content: Section content.
+            page_num: Page number.
+            
+        Returns:
+            List of extracted Entity objects.
+        """
+        # Build learned types section for prompt
+        learned_types_section = ""
+        if self._type_registry:
+            learned_types = self._type_registry.get_types_for_prompt(
+                min_count=self.learned_type_min_count,
+                limit=15,
+            )
+            if learned_types:
+                types_list = ", ".join(learned_types).upper()
+                learned_types_section = f"\nAdditionally, these domain-specific types have been learned from previous documents:\n- {types_list}\n"
+        
+        prompt = ENTITY_EXTRACTION_PROMPT.format(
+            content=content,
+            node_id=node_id,
+            doc_id=doc_id,
+            header=header,
+            learned_types_section=learned_types_section,
+        )
+        
+        # Call LLM
+        response = self.llm.complete(prompt)
+        response_text = str(response) if not isinstance(response, str) else response
+        
+        # Parse JSON from response
+        entities = self._parse_llm_response(
+            response_text=response_text,
+            node_id=node_id,
+            doc_id=doc_id,
+            page_num=page_num,
+        )
+        
+        return entities
+    
+    def _parse_llm_response(
+        self,
+        response_text: str,
+        node_id: str,
+        doc_id: str,
+        page_num: int | None = None,
+    ) -> list[Entity]:
+        """
+        Parse LLM response into Entity objects.
+        
+        Args:
+            response_text: Raw LLM response.
+            node_id: Source node ID.
+            doc_id: Source document ID.
+            page_num: Page number.
+            
+        Returns:
+            List of Entity objects.
+        """
+        # Extract JSON from response (may be wrapped in markdown code block)
+        json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response_text)
+        if json_match:
+            json_str = json_match.group(1)
+        else:
+            # Try to find a JSON array directly
+            json_match = re.search(r'\[[\s\S]*\]', response_text)
+            if json_match:
+                json_str = json_match.group(0)
+            else:
+                logger.warning(
+                    "no_json_found_in_response",
+                    response_preview=response_text[:200],
+                )
+                return []
+        
+        try:
+            raw_entities = json.loads(json_str)
+        except json.JSONDecodeError as e:
+            logger.warning(
+                "json_parse_error",
+                error=str(e),
+                json_preview=json_str[:200],
+            )
+            return []
+        
+        if not isinstance(raw_entities, list):
+            logger.warning("expected_list_of_entities", got=type(raw_entities).__name__)
+            return []
+        
+        entities = []
+        for raw in raw_entities:
+            try:
+                entity = self._create_entity_from_raw(
+                    raw=raw,
+                    node_id=node_id,
+                    doc_id=doc_id,
+                    page_num=page_num,
+                )
+                if entity:
+                    entities.append(entity)
+            except Exception as e:
+                logger.debug(
+                    "failed_to_create_entity",
+                    raw=raw,
+                    error=str(e),
+                )
+        
+        return entities
+    
+    def _create_entity_from_raw(
+        self,
+        raw: dict[str, Any],
+        node_id: str,
+        doc_id: str,
+        page_num: int | None = None,
+    ) -> Entity | None:
+        """
+        Create an Entity object from raw LLM output.
+        
+        Args:
+            raw: Raw entity dict from LLM.
+            node_id: Source node ID.
+            doc_id: Source document ID.
+            page_num: Page number.
+            
+        Returns:
+            Entity object or None if invalid.
+        """
+        # Parse entity type
+        type_str = raw.get("type", "").upper()
+        original_type = type_str  # Preserve for metadata
+        
+        try:
+            entity_type = EntityType(type_str.lower())
+        except ValueError:
+            # Try mapping common variations
+            type_mapping = {
+                "PERSON": EntityType.PERSON,
+                "PEOPLE": EntityType.PERSON,
+                "INDIVIDUAL": EntityType.PERSON,
+                "NAME": EntityType.PERSON,
+                "ORGANIZATION": EntityType.ORGANIZATION,
+                "ORG": EntityType.ORGANIZATION,
+                "COMPANY": EntityType.ORGANIZATION,
+                "AGENCY": EntityType.ORGANIZATION,
+                "COURT": EntityType.ORGANIZATION,
+                "LEGAL_CONCEPT": EntityType.LEGAL_CONCEPT,
+                "LEGAL": EntityType.LEGAL_CONCEPT,
+                "CONCEPT": EntityType.LEGAL_CONCEPT,
+                "CLAIM": EntityType.LEGAL_CONCEPT,
+                "OBLIGATION": EntityType.LEGAL_CONCEPT,
+                "DATE": EntityType.DATE,
+                "TIME": EntityType.DATE,
+                "DATETIME": EntityType.DATE,
+                "PERIOD": EntityType.DATE,
+                "EVENT": EntityType.EVENT,
+                "OCCURRENCE": EntityType.EVENT,
+                "INCIDENT": EntityType.EVENT,
+                "LOCATION": EntityType.LOCATION,
+                "PLACE": EntityType.LOCATION,
+                "ADDRESS": EntityType.LOCATION,
+                "JURISDICTION": EntityType.LOCATION,
+                "REFERENCE": EntityType.REFERENCE,
+                "REF": EntityType.REFERENCE,
+                "CITATION": EntityType.REFERENCE,
+                "SECTION": EntityType.REFERENCE,
+                "MONETARY": EntityType.MONETARY,
+                "MONEY": EntityType.MONETARY,
+                "AMOUNT": EntityType.MONETARY,
+                "CURRENCY": EntityType.MONETARY,
+                "FINANCIAL": EntityType.MONETARY,
+                "DOCUMENT": EntityType.DOCUMENT,
+                "DOC": EntityType.DOCUMENT,
+                "CONTRACT": EntityType.DOCUMENT,
+                "AGREEMENT": EntityType.DOCUMENT,
+                "EXHIBIT": EntityType.DOCUMENT,
+            }
+            entity_type = type_mapping.get(type_str)
+            
+            if not entity_type:
+                # Check if we have a learned mapping for this type
+                if self._type_registry:
+                    mappings = self._type_registry.get_mappings()
+                    if type_str.lower() in mappings:
+                        mapped_type = mappings[type_str.lower()]
+                        try:
+                            entity_type = EntityType(mapped_type.lower())
+                            logger.debug(
+                                "using_learned_mapping",
+                                original=type_str,
+                                mapped_to=mapped_type,
+                            )
+                        except ValueError:
+                            pass
+                
+                if not entity_type:
+                    # Use OTHER as fallback - never drop entities
+                    logger.debug("unmapped_entity_type_using_other", type=type_str)
+                    entity_type = EntityType.OTHER
+        
+        # Get canonical name
+        canonical_name = raw.get("canonical_name", "").strip()
+        if not canonical_name:
+            canonical_name = raw.get("name", "").strip()
+        if not canonical_name:
+            return None
+        
+        # Get aliases
+        aliases = raw.get("aliases", [])
+        if isinstance(aliases, str):
+            aliases = [aliases]
+        aliases = [a.strip() for a in aliases if a and a.strip()]
+        
+        # Get context
+        context = raw.get("context", "").strip()
+        
+        # Get metadata
+        metadata = raw.get("metadata", {})
+        if not isinstance(metadata, dict):
+            metadata = {}
+        
+        # Preserve original type if we used the OTHER fallback
+        if entity_type == EntityType.OTHER and original_type:
+            metadata["original_type"] = original_type.lower()
+            
+            # Record this type for adaptive learning
+            if self._type_registry and self.enable_type_learning:
+                self._type_registry.record_type(
+                    type_name=original_type.lower(),
+                    context=context,
+                    entity_name=canonical_name,
+                )
+        
+        # Create mention
+        mention = Mention(
+            node_id=node_id,
+            doc_id=doc_id,
+            context=context,
+            page_num=page_num,
+            confidence=1.0,
+        )
+        
+        # Create entity
+        entity = Entity(
+            type=entity_type,
+            canonical_name=canonical_name,
+            aliases=aliases,
+            mentions=[mention],
+            metadata=metadata,
+            source_doc_id=doc_id,
+        )
+        
+        return entity
+    
+    def extract_batch(
+        self,
+        nodes: list[dict[str, Any]],
+    ) -> list[ExtractionResult]:
+        """
+        Extract entities from multiple nodes.
+        
+        Args:
+            nodes: List of node dicts with keys: node_id, doc_id, header, content, page_num
+            
+        Returns:
+            List of ExtractionResult objects.
+        """
+        results = []
+        
+        for node in nodes:
+            result = self.extract_from_node(
+                node_id=node.get("node_id", ""),
+                doc_id=node.get("doc_id", ""),
+                header=node.get("header", ""),
+                content=node.get("content", ""),
+                page_num=node.get("page_num"),
+            )
+            results.append(result)
+        
+        return results
+    
+    def clear_cache(self) -> None:
+        """Clear the entity cache."""
+        self._cache.clear()
+
+
+def merge_entities(entities: list[Entity]) -> list[Entity]:
+    """
+    Merge duplicate entities based on canonical name and type.
+    
+    Combines mentions and aliases from duplicates.
+    
+    Args:
+        entities: List of entities to merge.
+        
+    Returns:
+        Deduplicated list of entities.
+    """
+    # Group by (type, normalized canonical_name)
+    grouped: dict[tuple[EntityType, str], list[Entity]] = {}
+    
+    for entity in entities:
+        key = (entity.type, entity.canonical_name.lower().strip())
+        if key not in grouped:
+            grouped[key] = []
+        grouped[key].append(entity)
+    
+    # Merge each group
+    merged = []
+    for entities_group in grouped.values():
+        if len(entities_group) == 1:
+            merged.append(entities_group[0])
+        else:
+            # Merge into first entity
+            primary = entities_group[0]
+            for other in entities_group[1:]:
+                # Merge mentions
+                primary.mentions.extend(other.mentions)
+                
+                # Merge aliases
+                for alias in other.aliases:
+                    primary.add_alias(alias)
+                
+                # Merge metadata (prefer primary's values on conflict)
+                for k, v in other.metadata.items():
+                    if k not in primary.metadata:
+                        primary.metadata[k] = v
+            
+            merged.append(primary)
+    
+    return merged