feat(edm): introduce entity disambiguation and merging pipeline

[dev-mirzabicer]

This commit introduces a comprehensive Entity Disambiguation and Merging (EDM) pipeline to improve knowledge graph quality by automatically identifying and merging entity aliases.

The pipeline operates in two stages:

1. Candidate Generation: Potential alias clusters are formed using a combination of lexical (thefuzz) and semantic (embedding cosine similarity) analysis.
2. LLM Verification: Each candidate cluster is sent to an LLM with rich contextual information, including the original text chunks, for a final, high-confidence merge decision. A new entity_disambiguation prompt guides this process.

The HiRAG.ainsert workflow is updated to integrate this pipeline. The entity extraction function now returns raw nodes and edges, which are then processed by the EntityDisambiguator. A new _upsert_disambiguated_graph method handles the final merging and storage of the canonicalized entities and relations.

New configuration parameters are added to the HiRAG class to control the EDM behavior (e.g., similarity thresholds, cluster size, merge confidence).

BREAKING CHANGE: The extract_hierarchical_entities function in hirag._op now returns a tuple of raw nodes and edges (List[Dict], List[Dict]) instead of a BaseGraphStorage instance. The data ingestion pipeline within HiRAG.ainsert has been fundamentally changed to accommodate the new entity disambiguation step.

@claude

[dev-mirzabicer]

@claude

[claude]

Claude finished @dev-mirzabicer's task —— View job

---

PR Review: Entity Disambiguation and Merging Pipeline

🔍 Comprehensive Code Review

This PR introduces a substantial Entity Disambiguation and Merging (EDM) pipeline with ~2,500+ lines of new/modified code. Here's my detailed analysis across all requested areas:

---

🚨 Critical Issues Found

1. Undefined Variable Bug - hirag/_disambiguation.py:1056

logger_instance.info(f"  Non-temporary entities: {non_temporary_entities}")

Issue: Variable non_temporary_entities is undefined. Should be non_temporary_count.
Fix: Replace with non_temporary_count as defined on line 1033.

2. Incorrect Method Call - hirag/_op.py:646

embedding_result = await entity_vdb.embedding_func([description])

Issue: entity_vdb doesn't have embedding_func attribute.
Fix: Should be self.embedding_func([description]) based on function signature.

3. Missing String Separator - hirag/hirag.py:1331

content = entity_name + description

Issue: Concatenating without separator makes vector search less effective.
Fix: Add separator: content = entity_name + " " + description

---

📋 Code Quality & Best Practices

✅ Strengths

1. Excellent Architecture: Clean separation of concerns with EntityDisambiguator class
2. Comprehensive Error Handling: Good use of try-catch blocks with detailed logging
3. Async/Await Pattern: Proper asynchronous programming throughout
4. Configuration-Driven: Highly configurable with sensible defaults
5. Type Hints: Good use of typing annotations
6. Documentation: Detailed docstrings and comments

⚠️ Areas for Improvement

Magic Numbers & Hard-coded Values

# hirag/_disambiguation.py:683
combined_context = "\n\n---\n\n".join(contexts)

# Character-to-token estimation
if len(combined_context) > self.config.max_context_tokens * 4:

Recommendations:

• Extract separator as a configurable constant
• Define CHARACTER_TO_TOKEN_RATIO = 4 as a documented constant

Missing Docstring - hirag/_disambiguation.py:884

The get_default_edm_config() function lacks proper documentation of return structure.

Unnecessary Type Conversions

# hirag/_disambiguation.py:662
source_chunks = await self._fetch_source_chunks(list(source_chunks_needed))

Issue: Converting set to list is unnecessary since _fetch_source_chunks can handle sets.

---

⚡ Performance Considerations

🎯 Optimizations Implemented

1. Batch Processing: Smart batching for embeddings and LLM calls
2. Concurrent Execution: Dynamic concurrency control with semaphores
3. Memory Management: Configurable batch sizes to prevent OOM
4. Vector Database Integration: Efficient similarity search using VDB

📈 Performance Concerns

O(n²) Fallback Similarity

# hirag/_disambiguation.py:434
similarity_matrix = cosine_similarity(embeddings)
for i, item_a in enumerate(entity_data):
    for j, item_b in enumerate(entity_data[i + 1 :], i + 1):

Impact: Falls back to O(n²) computation when VDB unavailable
Mitigation: Well-documented and only used as fallback

Large Context Processing

The system processes full text contexts for LLM verification, which could impact performance with very large documents.

Memory Usage

# Large embeddings stored in memory
entity_embeddings = entity_sequence_embeddings

Consider implementing streaming or chunked processing for very large entity sets.

---

🔒 Security Analysis

✅ Security Strengths

1. Input Validation: Comprehensive validation in validate_edm_configuration()
2. Safe JSON Parsing: Proper error handling for LLM response parsing
3. No Direct File Operations: Uses abstracted storage layers
4. LLM Response Sanitization: JSON extraction with markdown handling

⚠️ Potential Security Concerns

LLM Response Processing

# hirag/_disambiguation.py:802-806
if response_cleaned.startswith("```json"):
    response_cleaned = response_cleaned[7:]
if response_cleaned.endswith("```"):
    response_cleaned = response_cleaned[:-3]

Concern: Basic string manipulation for LLM response cleaning
Recommendation: Consider more robust parsing with proper sanitization

Configuration Injection

The system accepts extensive configuration parameters. Ensure all config values are validated.

---

🧪 Test Coverage Assessment

❌ Critical Gap: No Test Files Found

This substantial feature lacks comprehensive test coverage. Recommended test structure:

tests/
├── test_disambiguation_unit.py      # Unit tests for EntityDisambiguator
├── test_disambiguation_integration.py # Integration tests with real data
├── test_edm_configuration.py        # Configuration validation tests
├── test_performance.py             # Performance benchmarks
└── fixtures/
    ├── sample_entities.json
    └── disambiguation_scenarios.json

Essential Test Cases Needed:

1. UnionFind algorithm correctness
2. Similarity threshold edge cases
3. LLM response parsing robustness
4. Configuration validation
5. Error recovery scenarios
6. Performance with large datasets

---

🏗️ Breaking Changes Impact

API Breaking Change

# OLD: hirag/_op.py
async def extract_hierarchical_entities(...) -> BaseGraphStorage:

# NEW: hirag/_op.py  
async def extract_hierarchical_entities(...) -> Tuple[List[Dict], List[Dict]]:

Impact Assessment: ✅ Well Handled

• Clear documentation in PR description
• Backward compatibility maintained through integration layer
• Proper migration path provided

---

🔧 Specific Bug Fixes Required

High Priority

1. Fix undefined variable: non_temporary_entities → non_temporary_count (line 1056)
2. Fix method call: entity_vdb.embedding_func → self.embedding_func (line 646)
3. Add string separator: entity_name + " " + description (line 1331)

Medium Priority

1. Extract magic numbers: Define CHARACTER_TO_TOKEN_RATIO = 4
2. Add missing docstring: Document get_default_edm_config() return structure
3. Remove unnecessary conversion: list(source_chunks_needed) → source_chunks_needed

---

📊 Overall Assessment

Category	Score	Comments
Architecture	9/10	Excellent design with proper separation of concerns
Code Quality	7/10	Good overall, but has several fixable issues
Performance	8/10	Well-optimized with smart batching and concurrency
Security	7/10	Good practices, but needs enhanced input validation
Testing	2/10	❌ Critical gap - no tests provided
Documentation	8/10	Excellent docstrings and comments

🎯 Recommendations

Before Merge (Required)

1. ✅ Fix the 3 critical bugs identified above
2. ✅ Add comprehensive test suite
3. ✅ Add configuration validation for security

Post-Merge (Recommended)

1. Performance benchmarking with large datasets
2. Memory usage optimization
3. Enhanced LLM response parsing
4. Monitoring and alerting integration

---

Overall: This is a well-architected and comprehensive feature that significantly enhances the knowledge graph quality. The implementation shows deep understanding of the domain and good engineering practices. However, the critical bugs must be fixed and tests must be added before merging.