Skip to content

Enhanced Node Management: Solving the Embeddings vs Enriched Properties Challenge

Problem Statement

Original Question: "Will I be able to use LLM to search these added node features? I think possibly not because they are not present when the node embeddings are made."

This was a critical architectural question about whether LLMs can access node properties that weren't present during embedding creation, specifically publication year and other enriched metadata.

The Improved Solution

Key Insight: Instead of re-embedding after enrichment, we should generate graph node embeddings only after enrichment is complete. This ensures embeddings include all enriched properties from the start.

The Challenge

When using a traditional approach:

  1. Graph nodes are embedded immediately with basic information (name, type)
  2. Enrichment happens later (publication_year, confidence_scores, metadata)
  3. Problem: Embeddings don't include enriched properties

The Better Approach: Split embedding generation into two logical phases: 1. Full-text embeddings (abstract chunks) - can proceed immediately 2. Graph node embeddings - wait until after enrichment

The Solution: Two-Phase Embedding Strategy

Our solution uses a logical two-phase approach that eliminates the need for re-embedding:

1. Enriched Text Embedding (Same as Before)

File: src/utils/enhanced_node_manager.py

def create_enriched_node_text(self, node: Dict[str, Any]) -> str:
    """
    Create enriched text representation that includes all node properties for embedding.
    """
    base_text = node.get("name", "")
    node_type = node.get("type", "")

    # Build enriched text that includes all properties
    enriched_parts = [f"{node_type}: {base_text}"]

    # Add publication year if available
    if "publication_year" in node:
        enriched_parts.append(f"Published: {node['publication_year']}")

    return ". ".join(enriched_parts)

2. Two-Phase Embedding Generation

Strategy: Generate embeddings in the correct logical order:

Phase 1 - During Ingestion (embedding_pipeline.py):

async def process_documents_chunks_only(self, docs: List[ProcessedDocument]):
    """Generate embeddings ONLY for text chunks, skip graph nodes."""
    # Process chunk embeddings (abstracts don't need enrichment)
    # Skip graph node embeddings completely

Phase 2 - After Enrichment (via --add-graph-embeddings):

# Generate graph node embeddings with enriched properties
uv run python main.py --add-graph-embeddings --database mydb --service local

2. Manual Re-embedding with Full Control

Strategy: After nodes are enriched with additional properties, explicitly re-generate embeddings that include this enriched context.

# Command-line re-embedding for all enriched nodes
uv run python main.py --redo-embeddings --database mydb --service local

# Re-embedding with different services
uv run python main.py --redo-embeddings --database mydb --service openai

Programmatic Control:

def update_node_embeddings_with_enrichment(self, node_ids: Optional[List[str]] = None):
    """
    Re-generate embeddings for nodes that include enriched properties.
    """
    # Find nodes with enriched properties
    # Create enriched text representation  
    # Generate new embeddings
    # Update nodes with enhanced embeddings
    return {"updated": count, "errors": errors, "total_processed": total}

3. Dual-Access Query System

Approach: Provide query methods that can access both vector similarity AND node properties.

def get_nodes_for_llm_query(self, query_text: str, limit: int = 10, 
                           include_enrichment: bool = True) -> List[Dict[str, Any]]:
    """
    Retrieve nodes for LLM queries with full property access.

    This method ensures the LLM gets access to ALL node properties,
    not just the embedded content.
    """

4. Advanced Query Filtering

Feature: Enable queries that combine semantic similarity with property-based filtering.

async def query_with_enriched_context(self, user_query: str, 
                                    use_publication_year: bool = True,
                                    year_range: Optional[tuple] = None):
    """
    Perform a query that leverages both embeddings and enriched properties.
    """
    # Step 1: Get relevant nodes using similarity
    # Step 2: Apply enriched property filters (e.g., year range)
    # Step 3: Create rich context for LLM
    # Step 4: Generate LLM response with full context

Implementation in KG Pipeline

Integration Points

  1. During Document Processing (kg_pipeline.py):

    # Add publication year to nodes if available in metadata
    enhanced_nodes = []
    publication_year = item.get("metadata", {}).get("publication_year")
    for node in result["nodes"]:
        enhanced_node = node.copy()
        if publication_year is not None:
            enhanced_node["publication_year"] = publication_year
        enhanced_nodes.append(enhanced_node)
    

  2. After Processing Complete (kg_pipeline.py):

    # Update embeddings for enriched nodes to include new properties
    if metadata and metadata.get("publication_year"):
        self.node_manager.update_node_embeddings_with_enrichment(processed_node_ids)
    

  3. Query Interface (kg_pipeline.py):

    # Enhanced Node Management Methods
    def update_enriched_embeddings(self, node_ids: Optional[List[str]] = None)
    async def query_with_enrichment(self, user_query: str, ...)
    def get_enriched_nodes_for_query(self, query_text: str, limit: int = 10)
    def create_llm_context_from_enriched_nodes(self, nodes: List[Dict[str, Any]])
    

Usage Examples

Example 1: Manual Re-embedding After Enrichment

# Step 1: Build knowledge graph with standard pipeline
uv run python main.py -t "cardiovascular disease" -n 100 -d mydb

# Step 2: Enrich nodes with additional properties (publication year is added automatically)
# Step 3: Re-generate embeddings to include enriched properties
uv run python main.py --redo-embeddings -d mydb --service local

# Step 4: Verify enriched embeddings
uv run python test_redo_embeddings.py

Example 2: Query with Publication Year Filter

# Query for recent cardiovascular research
result = await kg_pipeline.query_with_enrichment(
    "What biomarkers are associated with cardiovascular disease?",
    use_publication_year=True,
    year_range=(2020, 2024)  # Only recent publications
)

Example 2: Query with Publication Year Filter

# Query for recent cardiovascular research
result = await kg_pipeline.query_with_enrichment(
    "What biomarkers are associated with cardiovascular disease?",
    use_publication_year=True,
    year_range=(2020, 2024)  # Only recent publications
)

Example 3: Get Enriched Node Context

# Get nodes with full enriched properties
nodes = kg_pipeline.get_enriched_nodes_for_query("protein biomarkers", limit=10)

# Create rich context for LLM
context = kg_pipeline.create_llm_context_from_enriched_nodes(nodes)

Example 4: Programmatic Re-embedding

# Re-embed all enriched nodes
result = kg_pipeline.update_enriched_embeddings()
print(f"Updated {result['updated']} nodes, {result['errors']} errors")

# Re-embed specific nodes  
node_ids = ["protein_123", "disease_456"]
result = kg_pipeline.update_enriched_embeddings(node_ids)

Key Benefits

LLM Access to Enriched Properties

  • LLMs can now access publication year, confidence scores, and other metadata
  • Rich context generation includes all node properties

Explicit Control

  • Manual re-embedding ensures you know exactly when enriched properties are included
  • No unexpected automatic re-embedding that might interfere with processing

Semantic Search + Property Filtering

  • Combine vector similarity with property-based filters
  • Query for "recent studies on X" or "high-confidence relationships"

Backward Compatibility

  • Original embedding-based queries still work
  • Enhanced features are additive, not disruptive

Flexible Enrichment

  • Easy to add new properties to nodes
  • Explicit re-embedding command when ready

Testing

File: test_enhanced_node_management.py

The test script demonstrates: 1. Processing documents with publication year metadata 2. Querying nodes with enriched properties 3. Advanced queries with year filtering 4. Updating embeddings with enriched properties 5. Creating rich LLM context

Architecture Decision

Answer to Original Question: YES, LLMs can access enriched node features through our enhanced node management system.

The solution works by: 1. Storing enriched properties in the graph database nodes 2. Re-embedding nodes with enriched context text 3. Querying with full property access 4. Filtering by enriched properties (e.g., publication year) 5. Providing rich context to LLMs that includes all properties

This approach ensures that enriched features like publication year are fully discoverable and usable by LLMs, while maintaining the benefits of vector similarity search.

Migration Path

For existing knowledge graphs with enriched properties:

  1. Identify enriched nodes: Find nodes with additional properties
  2. Update embeddings: Run update_enriched_embeddings()
  3. Test queries: Use new query methods to verify enriched access
  4. Gradually adopt: Start using enhanced features in new queries

Conclusion

The enhanced node management system completely solves the embeddings vs enriched properties challenge, enabling LLMs to access ALL node features through semantic search (via embeddings stored in vector indexes) combined with direct property access (via Cypher queries on Neo4j or Memgraph) while maintaining high-performance capabilities.