rnsr 0.1.0__py3-none-any.whl

rnsr/extraction/relationship_extractor.py

@@ -0,0 +1,600 @@
+"""
+RNSR Relationship Extractor
+
+DEPRECATED: This extractor uses LLM-first approach which can hallucinate.
+Use RLMUnifiedExtractor instead for grounded, accurate extraction.
+
+LLM-based relationship extraction between entities and sections.
+Extracts temporal, causal, semantic, and reference relationships.
+"""
+
+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,
+    Relationship,
+    RelationType,
+)
+from rnsr.llm import get_llm
+
+logger = structlog.get_logger(__name__)
+
+# Deprecation warning
+_DEPRECATION_WARNING = """
+RelationshipExtractor is deprecated and may hallucinate relationships.
+Use RLMUnifiedExtractor instead for grounded, accurate extraction:
+
+    from rnsr.extraction import RLMUnifiedExtractor
+    extractor = RLMUnifiedExtractor()
+    result = extractor.extract(node_id, doc_id, header, content)
+    # result.relationships contains validated relationships
+"""
+
+
+# Relationship extraction prompt template
+RELATIONSHIP_EXTRACTION_PROMPT = """You are an expert relationship extractor for legal and business documents.
+
+Analyze the following document section and extract relationships between:
+1. Entities (people, organizations, dates, events, legal concepts)
+2. This section and other referenced sections/documents
+
+Document Section:
+---
+{content}
+---
+
+Section ID: {node_id}
+Document ID: {doc_id}
+Section Header: {header}
+
+Known entities in this section:
+{entities_json}
+
+Extract relationships of the following types:
+
+ENTITY-TO-ENTITY RELATIONSHIPS:
+- TEMPORAL_BEFORE: Event/date X occurred before Event/date Y
+- TEMPORAL_AFTER: Event/date X occurred after Event/date Y
+- CAUSAL: Action X caused/led to Outcome Y (e.g., breach led to damages)
+- AFFILIATED_WITH: Person X is affiliated with Organization Y
+- PARTY_TO: Entity X is party to Document/Event Y (e.g., signatory, defendant)
+
+SECTION/DOCUMENT RELATIONSHIPS:
+- SUPPORTS: This section supports a claim or finding in another section
+- CONTRADICTS: This section contradicts another section
+- REFERENCES: This section references another document or section (exhibit, citation)
+- SUPERSEDES: This section supersedes/overrides another
+- AMENDS: This section amends another section/document
+
+For each relationship, provide:
+1. type: One of the types above
+2. source_id: ID of the source entity (from the known entities list) or "{node_id}" for this section
+3. source_type: "entity" or "node"
+4. target_id: ID of the target entity or node ID of the target section (use descriptive placeholder if referencing external doc, e.g., "exhibit_a")
+5. target_type: "entity" or "node"
+6. confidence: 0.0-1.0 based on how explicit the relationship is
+7. evidence: The exact quote that establishes this relationship
+
+Return your response as a JSON array:
+```json
+[
+  {{
+    "type": "CAUSAL",
+    "source_id": "ent_abc123",
+    "source_type": "entity",
+    "target_id": "ent_def456",
+    "target_type": "entity",
+    "confidence": 0.9,
+    "evidence": "The breach of contract by Defendant led to significant damages..."
+  }},
+  {{
+    "type": "REFERENCES",
+    "source_id": "{node_id}",
+    "source_type": "node",
+    "target_id": "exhibit_a",
+    "target_type": "node",
+    "confidence": 1.0,
+    "evidence": "As shown in Exhibit A..."
+  }}
+]
+```
+
+If no relationships are found, return an empty array: []
+
+Important:
+- Only extract relationships that are explicitly stated or strongly implied
+- Include the exact quote as evidence
+- Use entity IDs from the provided list when possible
+- For temporal relationships, be precise about the direction (BEFORE vs AFTER)
+- For causal relationships, the source is the cause and target is the effect
+"""
+
+
+class RelationshipExtractor:
+    """
+    DEPRECATED: Extracts relationships from document sections using LLM-first approach.
+    
+    This extractor can hallucinate relationships. Use RLMUnifiedExtractor instead.
+    
+    Identifies connections between:
+    - Entities (temporal, causal, affiliation)
+    - Sections (supports, contradicts, references)
+    - Documents (cross-references, amendments)
+    """
+    
+    def __init__(
+        self,
+        llm: Any | None = None,
+        min_content_length: int = 50,
+        max_content_length: int = 8000,
+        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="RelationshipExtractor")
+        """
+        Initialize the relationship 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.
+        """
+        self.llm = llm or get_llm()
+        self.min_content_length = min_content_length
+        self.max_content_length = max_content_length
+        
+        # Cache for extracted relationships (node_id -> relationships)
+        self._cache: dict[str, list[Relationship]] = {}
+    
+    def extract_from_node(
+        self,
+        node_id: str,
+        doc_id: str,
+        header: str,
+        content: str,
+        entities: list[Entity],
+    ) -> list[Relationship]:
+        """
+        Extract relationships from a single document node.
+        
+        Args:
+            node_id: Skeleton node ID.
+            doc_id: Document ID.
+            header: Section header text.
+            content: Full section content.
+            entities: Entities already extracted from this node.
+            
+        Returns:
+            List of extracted Relationship objects.
+        """
+        start_time = time.time()
+        
+        # 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 []
+        
+        # Check cache
+        cache_key = f"{doc_id}:{node_id}"
+        if cache_key in self._cache:
+            logger.debug("using_cached_relationships", node_id=node_id)
+            return self._cache[cache_key]
+        
+        # Truncate content if too long
+        if len(content) > self.max_content_length:
+            content = content[:self.max_content_length] + "..."
+        
+        try:
+            relationships = self._extract_with_llm(
+                node_id=node_id,
+                doc_id=doc_id,
+                header=header,
+                content=content,
+                entities=entities,
+            )
+            
+            # Cache results
+            self._cache[cache_key] = relationships
+            
+        except Exception as e:
+            logger.error(
+                "relationship_extraction_failed",
+                node_id=node_id,
+                error=str(e),
+            )
+            relationships = []
+        
+        processing_time_ms = (time.time() - start_time) * 1000
+        
+        logger.info(
+            "relationships_extracted",
+            node_id=node_id,
+            relationship_count=len(relationships),
+            processing_time_ms=processing_time_ms,
+        )
+        
+        return relationships
+    
+    def _extract_with_llm(
+        self,
+        node_id: str,
+        doc_id: str,
+        header: str,
+        content: str,
+        entities: list[Entity],
+    ) -> list[Relationship]:
+        """
+        Use LLM to extract relationships from content.
+        
+        Args:
+            node_id: Skeleton node ID.
+            doc_id: Document ID.
+            header: Section header.
+            content: Section content.
+            entities: Known entities in this section.
+            
+        Returns:
+            List of extracted Relationship objects.
+        """
+        # Format entities for prompt
+        entities_json = json.dumps([
+            {
+                "id": e.id,
+                "type": e.type.value,
+                "name": e.canonical_name,
+                "aliases": e.aliases,
+            }
+            for e in entities
+        ], indent=2)
+        
+        prompt = RELATIONSHIP_EXTRACTION_PROMPT.format(
+            content=content,
+            node_id=node_id,
+            doc_id=doc_id,
+            header=header,
+            entities_json=entities_json,
+        )
+        
+        # Call LLM
+        response = self.llm.complete(prompt)
+        response_text = str(response) if not isinstance(response, str) else response
+        
+        # Parse JSON from response
+        relationships = self._parse_llm_response(
+            response_text=response_text,
+            node_id=node_id,
+            doc_id=doc_id,
+            entities=entities,
+        )
+        
+        return relationships
+    
+    def _parse_llm_response(
+        self,
+        response_text: str,
+        node_id: str,
+        doc_id: str,
+        entities: list[Entity],
+    ) -> list[Relationship]:
+        """
+        Parse LLM response into Relationship objects.
+        
+        Args:
+            response_text: Raw LLM response.
+            node_id: Source node ID.
+            doc_id: Source document ID.
+            entities: Known entities.
+            
+        Returns:
+            List of Relationship objects.
+        """
+        # Extract JSON from response
+        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)
+            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_relationships = 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_relationships, list):
+            logger.warning("expected_list_of_relationships", got=type(raw_relationships).__name__)
+            return []
+        
+        # Create a set of valid entity IDs
+        valid_entity_ids = {e.id for e in entities}
+        
+        relationships = []
+        for raw in raw_relationships:
+            try:
+                relationship = self._create_relationship_from_raw(
+                    raw=raw,
+                    node_id=node_id,
+                    doc_id=doc_id,
+                    valid_entity_ids=valid_entity_ids,
+                )
+                if relationship:
+                    relationships.append(relationship)
+            except Exception as e:
+                logger.debug(
+                    "failed_to_create_relationship",
+                    raw=raw,
+                    error=str(e),
+                )
+        
+        return relationships
+    
+    def _create_relationship_from_raw(
+        self,
+        raw: dict[str, Any],
+        node_id: str,
+        doc_id: str,
+        valid_entity_ids: set[str],
+    ) -> Relationship | None:
+        """
+        Create a Relationship object from raw LLM output.
+        
+        Args:
+            raw: Raw relationship dict from LLM.
+            node_id: Source node ID.
+            doc_id: Source document ID.
+            valid_entity_ids: Set of valid entity IDs.
+            
+        Returns:
+            Relationship object or None if invalid.
+        """
+        # Parse relationship type
+        type_str = raw.get("type", "").upper()
+        
+        type_mapping = {
+            "TEMPORAL_BEFORE": RelationType.TEMPORAL_BEFORE,
+            "TEMPORAL_AFTER": RelationType.TEMPORAL_AFTER,
+            "BEFORE": RelationType.TEMPORAL_BEFORE,
+            "AFTER": RelationType.TEMPORAL_AFTER,
+            "CAUSAL": RelationType.CAUSAL,
+            "CAUSED": RelationType.CAUSAL,
+            "AFFILIATED_WITH": RelationType.AFFILIATED_WITH,
+            "AFFILIATION": RelationType.AFFILIATED_WITH,
+            "PARTY_TO": RelationType.PARTY_TO,
+            "PARTY": RelationType.PARTY_TO,
+            "SUPPORTS": RelationType.SUPPORTS,
+            "SUPPORT": RelationType.SUPPORTS,
+            "CONTRADICTS": RelationType.CONTRADICTS,
+            "CONTRADICT": RelationType.CONTRADICTS,
+            "REFERENCES": RelationType.REFERENCES,
+            "REFERENCE": RelationType.REFERENCES,
+            "CITES": RelationType.REFERENCES,
+            "SUPERSEDES": RelationType.SUPERSEDES,
+            "SUPERSEDE": RelationType.SUPERSEDES,
+            "AMENDS": RelationType.AMENDS,
+            "AMEND": RelationType.AMENDS,
+            "MENTIONS": RelationType.MENTIONS,
+            "DEFINED_IN": RelationType.DEFINED_IN,
+        }
+        
+        rel_type = type_mapping.get(type_str)
+        if not rel_type:
+            logger.debug("unknown_relationship_type", type=type_str)
+            return None
+        
+        # Get source and target
+        source_id = raw.get("source_id", "")
+        target_id = raw.get("target_id", "")
+        source_type = raw.get("source_type", "entity").lower()
+        target_type = raw.get("target_type", "entity").lower()
+        
+        if not source_id or not target_id:
+            return None
+        
+        # Validate source_type and target_type
+        if source_type not in ("entity", "node"):
+            source_type = "entity"
+        if target_type not in ("entity", "node"):
+            target_type = "entity"
+        
+        # Get confidence and evidence
+        confidence = raw.get("confidence", 0.8)
+        if not isinstance(confidence, (int, float)):
+            confidence = 0.8
+        confidence = max(0.0, min(1.0, float(confidence)))
+        
+        evidence = raw.get("evidence", "").strip()
+        
+        # Create relationship
+        relationship = Relationship(
+            type=rel_type,
+            source_id=source_id,
+            target_id=target_id,
+            source_type=source_type,
+            target_type=target_type,
+            doc_id=doc_id,
+            confidence=confidence,
+            evidence=evidence,
+        )
+        
+        return relationship
+    
+    def extract_entity_to_section_relationships(
+        self,
+        entities: list[Entity],
+    ) -> list[Relationship]:
+        """
+        Create MENTIONS relationships linking entities to their sections.
+        
+        Args:
+            entities: List of entities with mentions.
+            
+        Returns:
+            List of MENTIONS relationships.
+        """
+        relationships = []
+        
+        for entity in entities:
+            for mention in entity.mentions:
+                rel = Relationship(
+                    type=RelationType.MENTIONS,
+                    source_id=mention.node_id,
+                    target_id=entity.id,
+                    source_type="node",
+                    target_type="entity",
+                    doc_id=mention.doc_id,
+                    confidence=mention.confidence,
+                    evidence=mention.context,
+                )
+                relationships.append(rel)
+        
+        return relationships
+    
+    def extract_batch(
+        self,
+        nodes: list[dict[str, Any]],
+        entities_by_node: dict[str, list[Entity]],
+    ) -> list[Relationship]:
+        """
+        Extract relationships from multiple nodes.
+        
+        Args:
+            nodes: List of node dicts with keys: node_id, doc_id, header, content
+            entities_by_node: Mapping of node_id to entities in that node.
+            
+        Returns:
+            List of all extracted Relationship objects.
+        """
+        all_relationships = []
+        
+        for node in nodes:
+            node_id = node.get("node_id", "")
+            entities = entities_by_node.get(node_id, [])
+            
+            relationships = self.extract_from_node(
+                node_id=node_id,
+                doc_id=node.get("doc_id", ""),
+                header=node.get("header", ""),
+                content=node.get("content", ""),
+                entities=entities,
+            )
+            all_relationships.extend(relationships)
+        
+        return all_relationships
+    
+    def clear_cache(self) -> None:
+        """Clear the relationship cache."""
+        self._cache.clear()
+
+
+def extract_implicit_relationships(
+    entities: list[Entity],
+    doc_id: str,
+) -> list[Relationship]:
+    """
+    Extract implicit relationships based on entity metadata.
+    
+    For example:
+    - PERSON with role "defendant" -> PARTY_TO -> any legal proceeding
+    - PERSON affiliated with ORGANIZATION
+    
+    Args:
+        entities: List of entities.
+        doc_id: Document ID.
+        
+    Returns:
+        List of implicit relationships.
+    """
+    relationships = []
+    
+    # Group entities by type
+    persons = [e for e in entities if e.type == EntityType.PERSON]
+    orgs = [e for e in entities if e.type == EntityType.ORGANIZATION]
+    events = [e for e in entities if e.type == EntityType.EVENT]
+    documents = [e for e in entities if e.type == EntityType.DOCUMENT]
+    
+    # Extract affiliations from person metadata
+    for person in persons:
+        # Check for organization affiliation in metadata
+        affiliated_org = person.metadata.get("organization") or person.metadata.get("employer")
+        if affiliated_org:
+            # Find matching org
+            for org in orgs:
+                if affiliated_org.lower() in org.canonical_name.lower() or any(
+                    affiliated_org.lower() in alias.lower() for alias in org.aliases
+                ):
+                    rel = Relationship(
+                        type=RelationType.AFFILIATED_WITH,
+                        source_id=person.id,
+                        target_id=org.id,
+                        source_type="entity",
+                        target_type="entity",
+                        doc_id=doc_id,
+                        confidence=0.9,
+                        evidence=f"{person.canonical_name} affiliated with {org.canonical_name}",
+                    )
+                    relationships.append(rel)
+                    break
+        
+        # Check for role-based party relationships
+        role = person.metadata.get("role", "").lower()
+        if role in ("defendant", "plaintiff", "respondent", "petitioner", "applicant"):
+            for event in events:
+                if any(kw in event.canonical_name.lower() for kw in ["case", "proceeding", "trial", "hearing"]):
+                    rel = Relationship(
+                        type=RelationType.PARTY_TO,
+                        source_id=person.id,
+                        target_id=event.id,
+                        source_type="entity",
+                        target_type="entity",
+                        doc_id=doc_id,
+                        confidence=0.85,
+                        evidence=f"{person.canonical_name} is {role} in {event.canonical_name}",
+                    )
+                    relationships.append(rel)
+            
+            for doc in documents:
+                if any(kw in doc.canonical_name.lower() for kw in ["complaint", "motion", "order", "judgment"]):
+                    rel = Relationship(
+                        type=RelationType.PARTY_TO,
+                        source_id=person.id,
+                        target_id=doc.id,
+                        source_type="entity",
+                        target_type="entity",
+                        doc_id=doc_id,
+                        confidence=0.85,
+                        evidence=f"{person.canonical_name} is {role} in {doc.canonical_name}",
+                    )
+                    relationships.append(rel)
+    
+    return relationships