File size: 7,355 Bytes
dcfc6ef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """
EntityEnricher: Entity Profile Summaries
Traverses each entity's graph neighborhood and generates an LLM-synthesized
summary stored as `e.summary` on the Neo4j node.
"""
from __future__ import annotations
import logging
logger = logging.getLogger(__name__)
import asyncio
from datetime import datetime
from typing import Optional, List
from pydantic import BaseModel, Field
from ..core.neo4j_store import Neo4jStore
from ..core.llm_factory import LLMFactory
from ..config import settings
class EnrichmentResult(BaseModel):
"""Result from an entity enrichment operation"""
entities_enriched: int = 0
entities_skipped: int = 0
errors: int = 0
duration_seconds: float = 0.0
message: str = ""
class EntityEnricher:
"""
Post-ingestion enrichment pass: synthesizes a human-readable profile
summary for each entity based on its graph neighborhood.
The summary is stored as `e.summary` on the Neo4j Entity node and
indexed via a separate vector index so it can be retrieved directly.
"""
def __init__(
self,
graph_store: Neo4jStore,
llm_provider: Optional[str] = None,
batch_size: int = 20,
) -> None:
self.store = graph_store
self.llm = LLMFactory.create(provider=llm_provider)
self.batch_size = batch_size
# ── Public API ─────────────────────────────────────────────────────────
async def enrich_all_entities(
self,
min_connections: int = 1,
overwrite: bool = False,
) -> EnrichmentResult:
"""
Enrich all entities that:
- Have >= min_connections relationships, AND
- Do not yet have a summary (or overwrite=True)
Args:
min_connections: Minimum degree to qualify for enrichment
overwrite: Re-generate summaries for already-enriched nodes
Returns:
EnrichmentResult with counts
"""
import time
start = time.time()
# Fetch qualifying entities
where_clause = (
"" if overwrite else "AND (e.summary IS NULL OR e.summary = '')"
)
query = f"""
MATCH (e:Entity)
WITH e, size((e)--()) AS degree
WHERE degree >= $min_connections
{where_clause}
RETURN e.name as name, e.type as type
ORDER BY degree DESC
"""
try:
rows = await self.store.execute_query(
query, {"min_connections": min_connections}
)
except Exception as exc:
return EnrichmentResult(message=f"Query failed: {exc}")
enriched = 0
skipped = 0
errors = 0
# Process in batches
for i in range(0, len(rows), self.batch_size):
batch = rows[i : i + self.batch_size]
tasks = [
self._enrich_single(row["name"], row.get("type", "Entity"))
for row in batch
]
results = await asyncio.gather(*tasks, return_exceptions=True)
for r in results:
if isinstance(r, Exception):
errors += 1
elif r:
enriched += 1
else:
skipped += 1
duration = time.time() - start
return EnrichmentResult(
entities_enriched=enriched,
entities_skipped=skipped,
errors=errors,
duration_seconds=round(duration, 2),
message=f"Enriched {enriched}/{len(rows)} entities in {duration:.1f}s",
)
async def enrich_entity(self, entity_name: str) -> Optional[str]:
"""
Enrich a single entity by name. Returns the generated summary or None.
"""
# Get entity type
rows = await self.store.execute_query(
"MATCH (e:Entity {name: $name}) RETURN e.type as type",
{"name": entity_name},
)
entity_type = rows[0]["type"] if rows else "Entity"
result = await self._enrich_single(entity_name, entity_type)
if result:
# Return the summary
summary_rows = await self.store.execute_query(
"MATCH (e:Entity {name: $name}) RETURN e.summary as summary",
{"name": entity_name},
)
return summary_rows[0]["summary"] if summary_rows else None
return None
async def get_entity_summary(self, entity_name: str) -> Optional[str]:
"""Get the stored summary for an entity, or None if not enriched."""
rows = await self.store.execute_query(
"MATCH (e:Entity {name: $name}) RETURN e.summary as summary",
{"name": entity_name},
)
if not rows:
return None
return rows[0].get("summary")
# ── Internal ────────────────────────────────────────────────────────────
async def _enrich_single(
self, entity_name: str, entity_type: str
) -> bool:
"""Generate and persist a summary for one entity. Returns True on success."""
try:
# Get the 2-hop neighborhood
neighbors = await self.store.get_neighbors(entity_name, depth=2)
# Also get direct relationship types
rel_query = """
MATCH (e:Entity {name: $name})-[r]-(other:Entity)
RETURN type(r) as rel_type, other.name as other_name,
other.type as other_type
LIMIT 30
"""
rels = await self.store.execute_query(
rel_query, {"name": entity_name}
)
neighborhood_lines: List[str] = []
for rel in rels:
neighborhood_lines.append(
f"- {rel['rel_type']} → {rel['other_name']} ({rel['other_type']})"
)
if not neighborhood_lines and not neighbors:
return False # isolated node — skip
neighborhood_text = "\n".join(neighborhood_lines[:40])
prompt = f"""You are summarizing an entity from a knowledge graph.
Entity: {entity_name}
Type: {entity_type}
Direct relationships:
{neighborhood_text if neighborhood_text else "(no direct relationships found)"}
Write a concise 2-3 sentence factual profile of "{entity_name}" based ONLY
on the graph connections listed above. Be specific, avoid vague language,
and do not add information not implied by the relationships."""
summary = await self.llm.complete(prompt, temperature=0.2)
summary = summary.strip()
if not summary:
return False
# Write summary back to Neo4j
await self.store.execute_query(
"""
MATCH (e:Entity {name: $name})
SET e.summary = $summary,
e.summary_updated_at = datetime()
""",
{"name": entity_name, "summary": summary},
)
return True
except Exception as exc:
logger.info(f"[EntityEnricher] Failed to enrich '{entity_name}': {exc}")
return False
|