Spaces:
Sleeping
Sleeping
| import logging | |
| from typing import List, Dict, Optional | |
| from rapidfuzz import fuzz | |
| logger = logging.getLogger(__name__) | |
| class EntityResolver: | |
| def __init__(self, similarity_threshold: float = 85.0): | |
| self.similarity_threshold = similarity_threshold | |
| self.llm = None | |
| self._init_llm() | |
| logger.info("Initializing EntityResolver with similarity threshold: %.1f", self.similarity_threshold) | |
| def _init_llm(self): | |
| """Initialize LLM for borderline disambiguation. Caches the instance.""" | |
| try: | |
| from .llm_client import get_llm | |
| self.llm = get_llm(temperature=0.0) | |
| except Exception as e: | |
| logger.warning("Could not initialize LLM for entity disambiguation: %s", str(e)) | |
| self.llm = None | |
| def resolve_entities(self, entities: List[dict], relationships: List[dict]) -> tuple[List[dict], List[dict]]: | |
| """ | |
| Deduplicates a list of entities and rewrites relationship references accordingly. | |
| Returns a tuple of (resolved_entities, rewritten_relationships). | |
| """ | |
| if not entities: | |
| return [], relationships | |
| logger.info("Resolving duplicates for %d entities and %d relationships...", len(entities), len(relationships)) | |
| # 1. Group entities by their category/type | |
| grouped_by_type: Dict[str, List[dict]] = {} | |
| for entity in entities: | |
| etype = entity['type'].upper() | |
| grouped_by_type.setdefault(etype, []).append(entity) | |
| resolved_entities = [] | |
| name_mappings = {} # Maps original_name -> canonical_name | |
| for etype, group in grouped_by_type.items(): | |
| resolved_group = [] | |
| for current in group: | |
| matched_canonical = None | |
| # Check current entity against already resolved entities in the same type group | |
| for existing in resolved_group: | |
| # Run fuzzy comparison on names | |
| ratio = fuzz.token_set_ratio(current['name'].lower(), existing['name'].lower()) | |
| if ratio >= self.similarity_threshold: | |
| matched_canonical = existing | |
| break | |
| # Borderline cases: ask LLM to disambiguate | |
| elif self.similarity_threshold - 15 <= ratio < self.similarity_threshold: | |
| if self._llm_confirm_duplicate( | |
| current['name'], existing['name'], | |
| current.get('type', ''), existing.get('type', '') | |
| ): | |
| matched_canonical = existing | |
| break | |
| if matched_canonical: | |
| # Duplicate found! Merge current into matched_canonical | |
| old_name = current['name'] | |
| new_name = matched_canonical['name'] | |
| # Keep the longer name as the canonical one | |
| if len(old_name) > len(new_name): | |
| matched_canonical['name'] = old_name | |
| name_mappings[new_name] = old_name | |
| name_mappings[old_name] = old_name | |
| else: | |
| name_mappings[old_name] = new_name | |
| # Combine descriptions, avoiding exact duplicates | |
| if current['description'] and current['description'] not in matched_canonical['description']: | |
| matched_canonical['description'] = f"{matched_canonical['description']} {current['description']}".strip() | |
| logger.info("Resolved & Merged entity: '%s' ➔ '%s'", old_name, matched_canonical['name']) | |
| else: | |
| # Unique entity in this pass, add it to resolved list | |
| resolved_group.append(current) | |
| name_mappings[current['name']] = current['name'] | |
| resolved_entities.extend(resolved_group) | |
| # 2. Rewrite relationships using the canonical name mappings | |
| rewritten_relationships = [] | |
| for rel in relationships: | |
| # Look up source and target names in mappings | |
| src = rel['source_entity'] | |
| tgt = rel['target_entity'] | |
| canonical_src = name_mappings.get(src, src) | |
| canonical_tgt = name_mappings.get(tgt, tgt) | |
| # Prevent self-referencing relationships created by merges | |
| if canonical_src == canonical_tgt: | |
| logger.warning("Discarded self-referencing relationship: [%s] --[%s]--> [%s] after resolution merge.", | |
| src, rel['relationship_type'], tgt) | |
| continue | |
| rel['source_entity'] = canonical_src | |
| rel['target_entity'] = canonical_tgt | |
| rewritten_relationships.append(rel) | |
| logger.info("Deduplication complete. Resolved entities count: %d (from %d) | Relationships count: %d", | |
| len(resolved_entities), len(entities), len(rewritten_relationships)) | |
| return resolved_entities, rewritten_relationships | |
| def _llm_confirm_duplicate(self, name_a: str, name_b: str, type_a: str, type_b: str) -> bool: | |
| """ | |
| Uses LLM to determine if two borderline entities refer to the same real-world thing. | |
| Returns True if they are the same entity. | |
| """ | |
| if not self.llm: | |
| return False | |
| try: | |
| from langchain_core.prompts import ChatPromptTemplate | |
| prompt = ChatPromptTemplate.from_messages([ | |
| ("system", ( | |
| "You are an expert at entity resolution. Given two entity names and their types, " | |
| "determine if they refer to the same real-world entity.\n" | |
| "Answer ONLY 'yes' or 'no'.\n" | |
| "Consider: abbreviations, nicknames, partial names, and variations of the same entity." | |
| )), | |
| ("human", ( | |
| "Entity A: '{name_a}' (Type: {type_a})\n" | |
| "Entity B: '{name_b}' (Type: {type_b})\n\n" | |
| "Do these refer to the same entity?" | |
| )) | |
| ]) | |
| chain = prompt | self.llm | |
| response = chain.invoke({ | |
| "name_a": name_a, | |
| "name_b": name_b, | |
| "type_a": type_a, | |
| "type_b": type_b | |
| }) | |
| answer = response.content.strip().lower() | |
| is_duplicate = answer.startswith("yes") | |
| logger.info("LLM disambiguation: '%s' vs '%s' -> %s", name_a, name_b, is_duplicate) | |
| return is_duplicate | |
| except Exception as e: | |
| logger.error("LLM disambiguation failed: %s", str(e)) | |
| return False | |