rnsr 0.1.0__py3-none-any.whl

rnsr/extraction/entity_linker.py

@@ -0,0 +1,825 @@
+"""
+RNSR Entity Linker
+
+Cross-document entity linking with fuzzy matching and LLM disambiguation.
+Links entities that represent the same real-world entity across documents.
+
+Features adaptive learning for normalization patterns - learns new titles
+and suffixes from user's document workload.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from collections import defaultdict
+from datetime import datetime
+from difflib import SequenceMatcher
+from pathlib import Path
+from threading import Lock
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from rnsr.extraction.models import Entity, EntityLink, EntityType
+from rnsr.llm import get_llm
+
+if TYPE_CHECKING:
+    from rnsr.indexing.knowledge_graph import KnowledgeGraph
+
+logger = structlog.get_logger(__name__)
+
+
+# =============================================================================
+# Learned Normalization Patterns
+# =============================================================================
+
+DEFAULT_NORMALIZATION_PATH = Path.home() / ".rnsr" / "learned_normalization.json"
+
+
+class LearnedNormalizationPatterns:
+    """
+    Registry for learning domain-specific normalization patterns.
+    
+    Learns:
+    - Titles/prefixes (Mr., Dr., Esq., Hon., M.D., etc.)
+    - Suffixes (Inc., LLC, GmbH, Pty Ltd, etc.)
+    - Domain-specific patterns (legal, medical, regional)
+    """
+    
+    # Base patterns (always included)
+    BASE_TITLES = [
+        "mr.", "mrs.", "ms.", "dr.", "prof.",
+        "mr", "mrs", "ms", "dr", "prof",
+        "the", "hon.", "hon", "sir", "dame",
+    ]
+    
+    BASE_SUFFIXES = [
+        ", inc.", ", inc", ", llc", ", llc.",
+        ", corp.", ", corp", ", ltd.", ", ltd",
+        "inc.", "inc", "llc", "corp.", "corp", "ltd.", "ltd",
+        ", esq.", ", esq", "esq.", "esq",
+        ", jr.", ", jr", ", sr.", ", sr",
+        "jr.", "jr", "sr.", "sr",
+    ]
+    
+    def __init__(
+        self,
+        storage_path: Path | str | None = None,
+        auto_save: bool = True,
+    ):
+        """
+        Initialize the normalization patterns registry.
+        
+        Args:
+            storage_path: Path to JSON file for persistence.
+            auto_save: Whether to save after each new pattern.
+        """
+        self.storage_path = Path(storage_path) if storage_path else DEFAULT_NORMALIZATION_PATH
+        self.auto_save = auto_save
+        
+        self._lock = Lock()
+        self._titles: dict[str, dict[str, Any]] = {}
+        self._suffixes: dict[str, dict[str, Any]] = {}
+        self._dirty = False
+        
+        self._load()
+    
+    def _load(self) -> None:
+        """Load learned patterns from storage."""
+        if not self.storage_path.exists():
+            return
+        
+        try:
+            with open(self.storage_path, "r") as f:
+                data = json.load(f)
+            
+            self._titles = data.get("titles", {})
+            self._suffixes = data.get("suffixes", {})
+            
+            logger.info(
+                "normalization_patterns_loaded",
+                titles=len(self._titles),
+                suffixes=len(self._suffixes),
+            )
+            
+        except Exception as e:
+            logger.warning("failed_to_load_normalization_patterns", error=str(e))
+    
+    def _save(self) -> None:
+        """Save patterns to storage."""
+        if not self._dirty:
+            return
+        
+        try:
+            self.storage_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            data = {
+                "version": "1.0",
+                "updated_at": datetime.utcnow().isoformat(),
+                "titles": self._titles,
+                "suffixes": self._suffixes,
+            }
+            
+            with open(self.storage_path, "w") as f:
+                json.dump(data, f, indent=2)
+            
+            self._dirty = False
+            
+        except Exception as e:
+            logger.warning("failed_to_save_normalization_patterns", error=str(e))
+    
+    def record_title(
+        self,
+        title: str,
+        domain: str = "general",
+        entity_example: str = "",
+    ) -> None:
+        """
+        Record a learned title/prefix.
+        
+        Args:
+            title: The title pattern (e.g., "Atty.", "M.D.").
+            domain: Domain category (legal, medical, regional, etc.).
+            entity_example: Example entity with this title.
+        """
+        title = title.lower().strip()
+        
+        if not title or title in self.BASE_TITLES:
+            return
+        
+        with self._lock:
+            now = datetime.utcnow().isoformat()
+            
+            if title not in self._titles:
+                self._titles[title] = {
+                    "count": 0,
+                    "domain": domain,
+                    "first_seen": now,
+                    "last_seen": now,
+                    "examples": [],
+                }
+                logger.info("new_title_pattern_learned", title=title, domain=domain)
+            
+            self._titles[title]["count"] += 1
+            self._titles[title]["last_seen"] = now
+            
+            if entity_example and len(self._titles[title]["examples"]) < 3:
+                self._titles[title]["examples"].append(entity_example)
+            
+            self._dirty = True
+            
+            if self.auto_save:
+                self._save()
+    
+    def record_suffix(
+        self,
+        suffix: str,
+        domain: str = "general",
+        entity_example: str = "",
+    ) -> None:
+        """
+        Record a learned suffix.
+        
+        Args:
+            suffix: The suffix pattern (e.g., "GmbH", "Pty Ltd").
+            domain: Domain category (legal, corporate, regional, etc.).
+            entity_example: Example entity with this suffix.
+        """
+        suffix = suffix.lower().strip()
+        
+        if not suffix or suffix in self.BASE_SUFFIXES:
+            return
+        
+        with self._lock:
+            now = datetime.utcnow().isoformat()
+            
+            if suffix not in self._suffixes:
+                self._suffixes[suffix] = {
+                    "count": 0,
+                    "domain": domain,
+                    "first_seen": now,
+                    "last_seen": now,
+                    "examples": [],
+                }
+                logger.info("new_suffix_pattern_learned", suffix=suffix, domain=domain)
+            
+            self._suffixes[suffix]["count"] += 1
+            self._suffixes[suffix]["last_seen"] = now
+            
+            if entity_example and len(self._suffixes[suffix]["examples"]) < 3:
+                self._suffixes[suffix]["examples"].append(entity_example)
+            
+            self._dirty = True
+            
+            if self.auto_save:
+                self._save()
+    
+    def get_all_titles(self, min_count: int = 1) -> list[str]:
+        """Get all titles (base + learned)."""
+        learned = [
+            title for title, data in self._titles.items()
+            if data["count"] >= min_count
+        ]
+        return list(set(self.BASE_TITLES + learned))
+    
+    def get_all_suffixes(self, min_count: int = 1) -> list[str]:
+        """Get all suffixes (base + learned)."""
+        learned = [
+            suffix for suffix, data in self._suffixes.items()
+            if data["count"] >= min_count
+        ]
+        return list(set(self.BASE_SUFFIXES + learned))
+    
+    def get_stats(self) -> dict[str, Any]:
+        """Get statistics about learned patterns."""
+        return {
+            "learned_titles": len(self._titles),
+            "learned_suffixes": len(self._suffixes),
+            "total_titles": len(self.get_all_titles()),
+            "total_suffixes": len(self.get_all_suffixes()),
+        }
+
+
+# Global normalization patterns instance
+_global_normalization_patterns: LearnedNormalizationPatterns | None = None
+
+
+def get_learned_normalization_patterns() -> LearnedNormalizationPatterns:
+    """Get the global normalization patterns registry."""
+    global _global_normalization_patterns
+    
+    if _global_normalization_patterns is None:
+        custom_path = os.getenv("RNSR_NORMALIZATION_PATH")
+        _global_normalization_patterns = LearnedNormalizationPatterns(
+            storage_path=custom_path if custom_path else None
+        )
+    
+    return _global_normalization_patterns
+
+
+# LLM disambiguation prompt
+DISAMBIGUATION_PROMPT = """You are an expert at entity resolution. Determine if these two entities refer to the same real-world entity.
+
+Entity 1:
+- Name: {name1}
+- Type: {type1}
+- Aliases: {aliases1}
+- Context: {context1}
+- Document: {doc1}
+
+Entity 2:
+- Name: {name2}
+- Type: {type2}
+- Aliases: {aliases2}
+- Context: {context2}
+- Document: {doc2}
+
+Consider:
+1. Name similarity (accounting for variations, titles, abbreviations)
+2. Context similarity (same role, same events, same relationships)
+3. Document context (are these documents related?)
+
+Respond with JSON:
+```json
+{{
+  "same_entity": true/false,
+  "confidence": 0.0-1.0,
+  "reasoning": "Brief explanation of your decision"
+}}
+```
+"""
+
+
+class EntityLinker:
+    """
+    Links entities across documents using multiple strategies:
+    
+    1. Exact match on canonical_name
+    2. Fuzzy string matching (Levenshtein-based)
+    3. Alias matching
+    4. LLM disambiguation for ambiguous cases
+    """
+    
+    def __init__(
+        self,
+        knowledge_graph: KnowledgeGraph,
+        llm: Any | None = None,
+        exact_match_threshold: float = 1.0,
+        fuzzy_match_threshold: float = 0.85,
+        llm_confidence_threshold: float = 0.75,
+        use_llm_disambiguation: bool = True,
+        enable_pattern_learning: bool = True,
+    ):
+        """
+        Initialize the entity linker.
+        
+        Args:
+            knowledge_graph: Knowledge graph for storing/querying entities.
+            llm: LLM instance for disambiguation. If None, uses get_llm().
+            exact_match_threshold: Confidence for exact matches.
+            fuzzy_match_threshold: Minimum similarity for fuzzy matching.
+            llm_confidence_threshold: Minimum LLM confidence to accept.
+            use_llm_disambiguation: Whether to use LLM for ambiguous cases.
+            enable_pattern_learning: Learn new normalization patterns.
+        """
+        self.kg = knowledge_graph
+        self.llm = llm
+        self.exact_match_threshold = exact_match_threshold
+        self.fuzzy_match_threshold = fuzzy_match_threshold
+        self.llm_confidence_threshold = llm_confidence_threshold
+        self.use_llm_disambiguation = use_llm_disambiguation
+        self.enable_pattern_learning = enable_pattern_learning
+        
+        # Lazy LLM initialization
+        self._llm_initialized = False
+        
+        # Normalization patterns registry
+        self._normalization_patterns = get_learned_normalization_patterns() if enable_pattern_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 link_entities(
+        self,
+        new_entities: list[Entity],
+        target_doc_id: str | None = None,
+    ) -> list[EntityLink]:
+        """
+        Link new entities to existing entities in the knowledge graph.
+        
+        Args:
+            new_entities: Newly extracted entities to link.
+            target_doc_id: If provided, only link to entities in this document.
+            
+        Returns:
+            List of EntityLink objects for confirmed links.
+        """
+        links = []
+        
+        for entity in new_entities:
+            # Find candidates from knowledge graph
+            candidates = self._find_candidates(entity, target_doc_id)
+            
+            if not candidates:
+                continue
+            
+            # Score each candidate
+            scored_candidates = []
+            for candidate in candidates:
+                score, method = self._score_match(entity, candidate)
+                if score >= self.fuzzy_match_threshold:
+                    scored_candidates.append((candidate, score, method))
+            
+            # Sort by score
+            scored_candidates.sort(key=lambda x: -x[1])
+            
+            if not scored_candidates:
+                continue
+            
+            # Handle disambiguation
+            best_candidate, best_score, best_method = scored_candidates[0]
+            
+            if best_score >= self.exact_match_threshold:
+                # Exact match - link directly
+                link = EntityLink(
+                    entity_id_1=entity.id,
+                    entity_id_2=best_candidate.id,
+                    confidence=best_score,
+                    link_method=best_method,
+                    evidence=f"Matched on: {entity.canonical_name} = {best_candidate.canonical_name}",
+                )
+                links.append(link)
+                
+                # Store in knowledge graph
+                self.kg.link_entities(
+                    entity.id,
+                    best_candidate.id,
+                    confidence=best_score,
+                    link_method=best_method,
+                    evidence=link.evidence,
+                )
+            
+            elif len(scored_candidates) == 1 and best_score >= self.fuzzy_match_threshold:
+                # Single fuzzy match - link with lower confidence
+                link = EntityLink(
+                    entity_id_1=entity.id,
+                    entity_id_2=best_candidate.id,
+                    confidence=best_score,
+                    link_method=best_method,
+                    evidence=f"Fuzzy match: {entity.canonical_name} ~ {best_candidate.canonical_name}",
+                )
+                links.append(link)
+                
+                self.kg.link_entities(
+                    entity.id,
+                    best_candidate.id,
+                    confidence=best_score,
+                    link_method=best_method,
+                    evidence=link.evidence,
+                )
+            
+            elif self.use_llm_disambiguation and len(scored_candidates) >= 1:
+                # Ambiguous - use LLM
+                llm_link = self._disambiguate_with_llm(entity, scored_candidates)
+                if llm_link:
+                    links.append(llm_link)
+                    
+                    self.kg.link_entities(
+                        llm_link.entity_id_1,
+                        llm_link.entity_id_2,
+                        confidence=llm_link.confidence,
+                        link_method=llm_link.link_method,
+                        evidence=llm_link.evidence,
+                    )
+        
+        logger.info(
+            "entity_linking_complete",
+            new_entities=len(new_entities),
+            links_created=len(links),
+        )
+        
+        return links
+    
+    def _find_candidates(
+        self,
+        entity: Entity,
+        target_doc_id: str | None = None,
+    ) -> list[Entity]:
+        """
+        Find candidate entities that might match the given entity.
+        
+        Args:
+            entity: Entity to find matches for.
+            target_doc_id: Optional document ID filter.
+            
+        Returns:
+            List of candidate entities.
+        """
+        candidates = []
+        
+        # Search by exact name
+        exact_matches = self.kg.find_entities_by_name(
+            entity.canonical_name,
+            entity_type=entity.type,
+            fuzzy=False,
+        )
+        candidates.extend(exact_matches)
+        
+        # Search by fuzzy name
+        fuzzy_matches = self.kg.find_entities_by_name(
+            entity.canonical_name,
+            entity_type=entity.type,
+            fuzzy=True,
+        )
+        for match in fuzzy_matches:
+            if match.id not in {c.id for c in candidates}:
+                candidates.append(match)
+        
+        # Search by aliases
+        for alias in entity.aliases:
+            alias_matches = self.kg.find_entities_by_name(
+                alias,
+                entity_type=entity.type,
+                fuzzy=True,
+            )
+            for match in alias_matches:
+                if match.id not in {c.id for c in candidates}:
+                    candidates.append(match)
+        
+        # Filter out the entity itself
+        candidates = [c for c in candidates if c.id != entity.id]
+        
+        # Filter by target document if specified
+        if target_doc_id:
+            candidates = [
+                c for c in candidates
+                if target_doc_id in c.document_ids
+            ]
+        
+        return candidates
+    
+    def _score_match(
+        self,
+        entity1: Entity,
+        entity2: Entity,
+    ) -> tuple[float, str]:
+        """
+        Score the similarity between two entities.
+        
+        Args:
+            entity1: First entity.
+            entity2: Second entity.
+            
+        Returns:
+            Tuple of (score, method) where score is 0.0-1.0.
+        """
+        # Type mismatch is a strong negative signal
+        if entity1.type != entity2.type:
+            return 0.0, "type_mismatch"
+        
+        # Exact canonical name match
+        if entity1.canonical_name.lower().strip() == entity2.canonical_name.lower().strip():
+            return 1.0, "exact"
+        
+        # Check alias matches
+        all_names1 = {n.lower().strip() for n in entity1.all_names}
+        all_names2 = {n.lower().strip() for n in entity2.all_names}
+        
+        if all_names1 & all_names2:
+            return 0.95, "alias"
+        
+        # Fuzzy string matching on all name combinations
+        best_similarity = 0.0
+        for name1 in entity1.all_names:
+            for name2 in entity2.all_names:
+                similarity = self._string_similarity(name1, name2)
+                best_similarity = max(best_similarity, similarity)
+        
+        if best_similarity >= self.fuzzy_match_threshold:
+            return best_similarity, "fuzzy"
+        
+        # Name containment (one name contains the other)
+        name1_lower = entity1.canonical_name.lower()
+        name2_lower = entity2.canonical_name.lower()
+        
+        if name1_lower in name2_lower or name2_lower in name1_lower:
+            # Score based on length ratio
+            shorter = min(len(name1_lower), len(name2_lower))
+            longer = max(len(name1_lower), len(name2_lower))
+            return 0.7 + 0.2 * (shorter / longer), "containment"
+        
+        return best_similarity, "low_similarity"
+    
+    def _string_similarity(self, s1: str, s2: str) -> float:
+        """
+        Calculate string similarity using SequenceMatcher.
+        
+        Args:
+            s1: First string.
+            s2: Second string.
+            
+        Returns:
+            Similarity score 0.0-1.0.
+        """
+        # Normalize strings
+        s1 = self._normalize_name(s1)
+        s2 = self._normalize_name(s2)
+        
+        # Use SequenceMatcher for similarity
+        return SequenceMatcher(None, s1, s2).ratio()
+    
+    def _normalize_name(self, name: str, learn_patterns: bool = True) -> str:
+        """
+        Normalize a name for comparison.
+        
+        Uses both base patterns and learned patterns from the registry.
+        
+        Args:
+            name: Name to normalize.
+            learn_patterns: Whether to record new patterns found.
+            
+        Returns:
+            Normalized name.
+        """
+        original_name = name
+        
+        # Lowercase
+        name = name.lower()
+        
+        # Get titles and suffixes (base + learned)
+        if self._normalization_patterns:
+            titles = self._normalization_patterns.get_all_titles()
+            suffixes = self._normalization_patterns.get_all_suffixes()
+        else:
+            titles = LearnedNormalizationPatterns.BASE_TITLES
+            suffixes = LearnedNormalizationPatterns.BASE_SUFFIXES
+        
+        # Remove titles/prefixes
+        removed_title = None
+        for title in sorted(titles, key=len, reverse=True):  # Longest first
+            if name.startswith(title + " "):
+                removed_title = title
+                name = name[len(title) + 1:]
+                break
+        
+        # Remove suffixes
+        removed_suffix = None
+        for suffix in sorted(suffixes, key=len, reverse=True):  # Longest first
+            if name.endswith(suffix):
+                removed_suffix = suffix
+                name = name[:-len(suffix)]
+                break
+        
+        # Learn new patterns if enabled
+        if learn_patterns and self._normalization_patterns and self.enable_pattern_learning:
+            # Detect potential new titles (patterns like "X. Name" or "X Name")
+            title_match = re.match(r'^([A-Za-z]{1,4}\.?)\s+[A-Z]', original_name)
+            if title_match and not removed_title:
+                potential_title = title_match.group(1).lower()
+                if potential_title not in titles and len(potential_title) >= 2:
+                    self._normalization_patterns.record_title(
+                        potential_title,
+                        domain="detected",
+                        entity_example=original_name,
+                    )
+            
+            # Detect potential new suffixes (patterns at end of company names)
+            suffix_patterns = [
+                r',?\s*(gmbh|ag|s\.a\.|pty\s*ltd|plc|bv|nv|asa|ab|as|oy|a/s|k\.k\.|co\.,?\s*ltd\.?)$',
+                r',?\s*([A-Z]{2,4}\.?)$',  # Acronym suffixes
+            ]
+            for pattern in suffix_patterns:
+                suffix_match = re.search(pattern, original_name, re.IGNORECASE)
+                if suffix_match and not removed_suffix:
+                    potential_suffix = suffix_match.group(1).lower()
+                    if potential_suffix not in suffixes:
+                        self._normalization_patterns.record_suffix(
+                            potential_suffix,
+                            domain="detected",
+                            entity_example=original_name,
+                        )
+        
+        # Normalize whitespace
+        name = " ".join(name.split())
+        
+        return name.strip()
+    
+    def _disambiguate_with_llm(
+        self,
+        entity: Entity,
+        candidates: list[tuple[Entity, float, str]],
+    ) -> EntityLink | None:
+        """
+        Use LLM to disambiguate between multiple candidates.
+        
+        Args:
+            entity: Entity to match.
+            candidates: List of (candidate, score, method) tuples.
+            
+        Returns:
+            EntityLink if a match is found, None otherwise.
+        """
+        llm = self._get_llm()
+        if llm is None:
+            return None
+        
+        # Get context from first mention
+        entity_context = ""
+        if entity.mentions:
+            entity_context = entity.mentions[0].context
+        
+        # Try each candidate
+        for candidate, score, method in candidates:
+            candidate_context = ""
+            if candidate.mentions:
+                candidate_context = candidate.mentions[0].context
+            
+            prompt = DISAMBIGUATION_PROMPT.format(
+                name1=entity.canonical_name,
+                type1=entity.type.value,
+                aliases1=", ".join(entity.aliases) or "None",
+                context1=entity_context or "No context available",
+                doc1=entity.source_doc_id or "Unknown",
+                name2=candidate.canonical_name,
+                type2=candidate.type.value,
+                aliases2=", ".join(candidate.aliases) or "None",
+                context2=candidate_context or "No context available",
+                doc2=candidate.source_doc_id or "Unknown",
+            )
+            
+            try:
+                response = llm.complete(prompt)
+                response_text = str(response) if not isinstance(response, str) else response
+                
+                result = self._parse_disambiguation_response(response_text)
+                
+                if result and result.get("same_entity") and result.get("confidence", 0) >= self.llm_confidence_threshold:
+                    return EntityLink(
+                        entity_id_1=entity.id,
+                        entity_id_2=candidate.id,
+                        confidence=result["confidence"],
+                        link_method="llm",
+                        evidence=result.get("reasoning", "LLM disambiguation"),
+                    )
+            
+            except Exception as e:
+                logger.warning(
+                    "llm_disambiguation_failed",
+                    entity=entity.canonical_name,
+                    candidate=candidate.canonical_name,
+                    error=str(e),
+                )
+        
+        return None
+    
+    def _parse_disambiguation_response(
+        self,
+        response_text: str,
+    ) -> dict[str, Any] | None:
+        """
+        Parse LLM disambiguation response.
+        
+        Args:
+            response_text: Raw LLM response.
+            
+        Returns:
+            Parsed result dict or None.
+        """
+        # 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:
+                return None
+        
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError:
+            return None
+    
+    def link_all_entities_in_document(
+        self,
+        doc_id: str,
+    ) -> list[EntityLink]:
+        """
+        Link all entities within a document to each other.
+        
+        Useful for finding co-references within a single document.
+        
+        Args:
+            doc_id: Document ID.
+            
+        Returns:
+            List of EntityLink objects.
+        """
+        entities = self.kg.find_entities_in_document(doc_id)
+        
+        if len(entities) < 2:
+            return []
+        
+        links = []
+        
+        # Group entities by type
+        by_type: dict[EntityType, list[Entity]] = defaultdict(list)
+        for entity in entities:
+            by_type[entity.type].append(entity)
+        
+        # Link within each type group
+        for entity_type, type_entities in by_type.items():
+            for i, entity1 in enumerate(type_entities):
+                for entity2 in type_entities[i + 1:]:
+                    score, method = self._score_match(entity1, entity2)
+                    
+                    if score >= self.fuzzy_match_threshold:
+                        link = EntityLink(
+                            entity_id_1=entity1.id,
+                            entity_id_2=entity2.id,
+                            confidence=score,
+                            link_method=method,
+                            evidence=f"Same document co-reference: {entity1.canonical_name} = {entity2.canonical_name}",
+                        )
+                        links.append(link)
+                        
+                        self.kg.link_entities(
+                            entity1.id,
+                            entity2.id,
+                            confidence=score,
+                            link_method=method,
+                            evidence=link.evidence,
+                        )
+        
+        return links
+    
+    def link_across_documents(
+        self,
+        doc_id_1: str,
+        doc_id_2: str,
+    ) -> list[EntityLink]:
+        """
+        Link entities between two specific documents.
+        
+        Args:
+            doc_id_1: First document ID.
+            doc_id_2: Second document ID.
+            
+        Returns:
+            List of EntityLink objects.
+        """
+        entities_1 = self.kg.find_entities_in_document(doc_id_1)
+        entities_2 = self.kg.find_entities_in_document(doc_id_2)
+        
+        if not entities_1 or not entities_2:
+            return []
+        
+        # Use the main linking method
+        return self.link_entities(entities_1, target_doc_id=doc_id_2)