852 lines
33 KiB
Python
852 lines
33 KiB
Python
"""
|
|
Entity extraction and resolution for memory system.
|
|
|
|
Uses spaCy for entity extraction and implements resolution logic
|
|
to disambiguate entities across memory units.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass, field
|
|
from datetime import UTC, datetime
|
|
from difflib import SequenceMatcher
|
|
|
|
import asyncpg
|
|
|
|
from .db_utils import acquire_with_retry
|
|
from .memory_engine import fq_table
|
|
from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class _EntityToCreate:
|
|
"""An entity that needs to be inserted (no matching candidate found)."""
|
|
|
|
idx: int
|
|
name: str
|
|
event_date: datetime | None
|
|
|
|
|
|
@dataclass
|
|
class _EntityStat:
|
|
"""Stat accumulation entry for a resolved entity (post-transaction update)."""
|
|
|
|
entity_id: str
|
|
event_date: datetime | None
|
|
|
|
|
|
@dataclass
|
|
class _EntityStatAgg:
|
|
"""Aggregated stats used when flushing pending updates."""
|
|
|
|
count: int = 0
|
|
max_date: datetime | None = None
|
|
|
|
|
|
@dataclass
|
|
class _CooccurrencePair:
|
|
"""A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
|
|
|
|
entity_id_1: str
|
|
entity_id_2: str
|
|
|
|
|
|
# Load spaCy model (singleton)
|
|
_nlp = None
|
|
|
|
|
|
class EntityResolver:
|
|
"""
|
|
Resolves entities to canonical IDs with disambiguation.
|
|
"""
|
|
|
|
def __init__(self, pool: asyncpg.Pool, entity_lookup: str = "full"):
|
|
"""
|
|
Initialize entity resolver.
|
|
|
|
Args:
|
|
pool: asyncpg connection pool
|
|
entity_lookup: Lookup strategy — "full" loads all bank entities then
|
|
matches in Python; "trigram" uses pg_trgm GIN index to fetch only
|
|
similar candidates per entity name (much faster for large banks).
|
|
"""
|
|
self.pool = pool
|
|
self.entity_lookup = entity_lookup
|
|
# Keyed by asyncio task id so concurrent retain batches never mix their
|
|
# pending updates. flush_pending_stats() pops only the calling task's items.
|
|
self._pending_stats: dict[int, list[_EntityStat]] = {}
|
|
self._pending_cooccurrences: dict[int, list[_CooccurrencePair]] = {}
|
|
|
|
def _task_key(self) -> int:
|
|
"""Return a unique key for the current asyncio task (or 0 for non-task context)."""
|
|
task = asyncio.current_task()
|
|
return id(task) if task is not None else 0
|
|
|
|
async def flush_pending_stats(self) -> None:
|
|
"""
|
|
Flush accumulated entity stats and co-occurrence counts for the current task.
|
|
|
|
Must be called AFTER the retain transaction commits. Pops only the items
|
|
accumulated by the calling asyncio task so concurrent retain batches never
|
|
flush each other's uncommitted entity IDs.
|
|
"""
|
|
if self.pool is None:
|
|
return
|
|
|
|
key = self._task_key()
|
|
stats = self._pending_stats.pop(key, [])
|
|
cooccurrences = self._pending_cooccurrences.pop(key, [])
|
|
|
|
if not stats and not cooccurrences:
|
|
return
|
|
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
if stats:
|
|
# Aggregate: sum counts and find max date per entity_id.
|
|
agg: dict[str, _EntityStatAgg] = defaultdict(_EntityStatAgg)
|
|
for s in stats:
|
|
entry = agg[s.entity_id]
|
|
entry.count += 1
|
|
if s.event_date is not None:
|
|
entry.max_date = s.event_date if entry.max_date is None else max(entry.max_date, s.event_date)
|
|
|
|
# Sort by entity_id so all concurrent workers acquire row locks in
|
|
# the same order — prevents circular lock dependencies (deadlocks).
|
|
rows = sorted((eid, a.count, a.max_date) for eid, a in agg.items())
|
|
await conn.executemany(
|
|
f"""
|
|
UPDATE {fq_table("entities")} SET
|
|
mention_count = mention_count + $2,
|
|
last_seen = GREATEST(last_seen, $3)
|
|
WHERE id = $1::uuid
|
|
""",
|
|
rows,
|
|
)
|
|
|
|
if cooccurrences:
|
|
# Aggregate: count occurrences per (entity_id_1, entity_id_2) pair.
|
|
coo_agg: dict[tuple[str, str], int] = {}
|
|
for c in cooccurrences:
|
|
pair = (c.entity_id_1, c.entity_id_2)
|
|
coo_agg[pair] = coo_agg.get(pair, 0) + 1
|
|
|
|
now = datetime.now(UTC)
|
|
# Sort by (entity_id_1, entity_id_2) for consistent lock ordering.
|
|
await conn.executemany(
|
|
f"""
|
|
INSERT INTO {fq_table("entity_cooccurrences")}
|
|
(entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
|
VALUES ($1, $2, $3, $4)
|
|
ON CONFLICT (entity_id_1, entity_id_2)
|
|
DO UPDATE SET
|
|
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + EXCLUDED.cooccurrence_count,
|
|
last_cooccurred = GREATEST({fq_table("entity_cooccurrences")}.last_cooccurred, EXCLUDED.last_cooccurred)
|
|
""",
|
|
sorted((e1, e2, count, now) for (e1, e2), count in coo_agg.items()),
|
|
)
|
|
|
|
@staticmethod
|
|
def _build_labels_lookup(entity_labels: list | None) -> set[str]:
|
|
"""Build a set of valid 'key:value' entity label strings for fast lookup."""
|
|
return _build_labels_lookup_from_config(entity_labels)
|
|
|
|
async def resolve_entities_batch(
|
|
self,
|
|
bank_id: str,
|
|
entities_data: list[dict],
|
|
context: str,
|
|
unit_event_date,
|
|
conn=None,
|
|
entity_labels: list | None = None,
|
|
) -> list[str]:
|
|
"""
|
|
Resolve multiple entities in batch (MUCH faster than sequential).
|
|
|
|
Groups entities by type, queries candidates in bulk, and resolves
|
|
all entities with minimal DB queries.
|
|
|
|
Args:
|
|
bank_id: bank ID
|
|
entities_data: List of dicts with 'text', 'type', 'nearby_entities'
|
|
context: Context where entities appear
|
|
unit_event_date: When this unit was created
|
|
conn: Optional connection to use (if None, acquires from pool)
|
|
|
|
Returns:
|
|
List of entity IDs in same order as input
|
|
"""
|
|
if not entities_data:
|
|
return []
|
|
|
|
taxonomy_lookup = self._build_labels_lookup(entity_labels)
|
|
if conn is None:
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
return await self._resolve_entities_batch_impl(
|
|
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
|
)
|
|
else:
|
|
return await self._resolve_entities_batch_impl(
|
|
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
|
)
|
|
|
|
async def _resolve_entities_batch_impl(
|
|
self,
|
|
conn,
|
|
bank_id: str,
|
|
entities_data: list[dict],
|
|
context: str,
|
|
unit_event_date,
|
|
taxonomy_lookup: set[str] | None = None,
|
|
) -> list[str]:
|
|
if self.entity_lookup == "trigram":
|
|
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
|
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
|
|
|
async def _resolve_entities_batch_full(
|
|
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
|
) -> list[str]:
|
|
"""Original strategy: load all bank entities then match in Python."""
|
|
# Query ALL candidates for this bank
|
|
all_entities = await conn.fetch(
|
|
f"""
|
|
SELECT canonical_name, id, metadata, last_seen, mention_count
|
|
FROM {fq_table("entities")}
|
|
WHERE bank_id = $1
|
|
""",
|
|
bank_id,
|
|
)
|
|
|
|
# Build entity ID to name mapping for co-occurrence lookups
|
|
entity_id_to_name = {row["id"]: row["canonical_name"].lower() for row in all_entities}
|
|
|
|
# Query ALL co-occurrences for this bank's entities in one query
|
|
# This builds a map of entity_id -> set of co-occurring entity names
|
|
all_cooccurrences = await conn.fetch(
|
|
f"""
|
|
SELECT ec.entity_id_1, ec.entity_id_2, ec.cooccurrence_count
|
|
FROM {fq_table("entity_cooccurrences")} ec
|
|
WHERE ec.entity_id_1 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
|
|
OR ec.entity_id_2 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
|
|
""",
|
|
bank_id,
|
|
)
|
|
|
|
# Build co-occurrence map: entity_id -> set of co-occurring entity names (lowercase)
|
|
cooccurrence_map: dict[str, set[str]] = {}
|
|
for row in all_cooccurrences:
|
|
eid1, eid2 = row["entity_id_1"], row["entity_id_2"]
|
|
# Add both directions
|
|
if eid1 not in cooccurrence_map:
|
|
cooccurrence_map[eid1] = set()
|
|
if eid2 not in cooccurrence_map:
|
|
cooccurrence_map[eid2] = set()
|
|
# Map to canonical names for comparison with nearby_entities
|
|
if eid2 in entity_id_to_name:
|
|
cooccurrence_map[eid1].add(entity_id_to_name[eid2])
|
|
if eid1 in entity_id_to_name:
|
|
cooccurrence_map[eid2].add(entity_id_to_name[eid1])
|
|
|
|
# Build candidate map for each entity text
|
|
all_candidates = {} # Maps entity_text -> list of candidates
|
|
entity_texts = list(set(e["text"] for e in entities_data))
|
|
|
|
for entity_text in entity_texts:
|
|
matching = []
|
|
entity_text_lower = entity_text.lower()
|
|
for row in all_entities:
|
|
canonical_name = row["canonical_name"]
|
|
ent_id = row["id"]
|
|
metadata = row["metadata"]
|
|
last_seen = row["last_seen"]
|
|
mention_count = row["mention_count"]
|
|
canonical_lower = canonical_name.lower()
|
|
# Match if exact or substring match
|
|
if (
|
|
entity_text_lower == canonical_lower
|
|
or entity_text_lower in canonical_lower
|
|
or canonical_lower in entity_text_lower
|
|
):
|
|
matching.append((ent_id, canonical_name, metadata, last_seen, mention_count))
|
|
all_candidates[entity_text] = matching
|
|
|
|
return await self._resolve_from_candidates(
|
|
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
|
)
|
|
|
|
async def _resolve_entities_batch_trigram(
|
|
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
|
) -> list[str]:
|
|
"""
|
|
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
|
|
|
Instead of loading all bank entities (O(N)), uses a GIN trigram index to fetch
|
|
only the small set of candidates that are textually similar to each input name.
|
|
Reduces DB data transfer from 165K rows to ~5-20 rows per entity.
|
|
"""
|
|
entity_texts = list(set(e["text"] for e in entities_data))
|
|
|
|
# Fetch candidates for all unique entity texts in a single batched query.
|
|
# The trigram % operator uses the GIN index; the substring conditions cover
|
|
# exact prefix/suffix matches that trigrams might miss at low similarity.
|
|
rows = await conn.fetch(
|
|
f"""
|
|
SELECT DISTINCT ON (e.id)
|
|
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
|
q.query_text
|
|
FROM unnest($2::text[]) AS q(query_text)
|
|
JOIN {fq_table("entities")} e ON (
|
|
e.bank_id = $1
|
|
AND (
|
|
e.canonical_name % q.query_text
|
|
OR LOWER(e.canonical_name) LIKE '%' || LOWER(q.query_text) || '%'
|
|
OR LOWER(q.query_text) LIKE '%' || LOWER(e.canonical_name) || '%'
|
|
)
|
|
)
|
|
""",
|
|
bank_id,
|
|
entity_texts,
|
|
)
|
|
|
|
# Group candidates by query_text
|
|
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
|
candidate_ids: set = set()
|
|
for row in rows:
|
|
query_text = row["query_text"]
|
|
all_candidates[query_text].append(
|
|
(row["id"], row["canonical_name"], row["metadata"], row["last_seen"], row["mention_count"])
|
|
)
|
|
candidate_ids.add(row["id"])
|
|
|
|
# Fetch co-occurrences only for the candidate entities (not all bank entities)
|
|
cooccurrence_map: dict[str, set[str]] = {}
|
|
if candidate_ids:
|
|
candidate_id_list = list(candidate_ids)
|
|
cooc_rows = await conn.fetch(
|
|
f"""
|
|
SELECT ec.entity_id_1, ec.entity_id_2
|
|
FROM {fq_table("entity_cooccurrences")} ec
|
|
WHERE ec.entity_id_1 = ANY($1::uuid[])
|
|
OR ec.entity_id_2 = ANY($1::uuid[])
|
|
""",
|
|
candidate_id_list,
|
|
)
|
|
# Build name lookup for co-occurrence mapping
|
|
id_to_name = {
|
|
row["id"]: row["canonical_name"].lower()
|
|
for cands in all_candidates.values()
|
|
for row in [{"id": c[0], "canonical_name": c[1]} for c in cands]
|
|
}
|
|
for row in cooc_rows:
|
|
eid1, eid2 = row["entity_id_1"], row["entity_id_2"]
|
|
if eid1 not in cooccurrence_map:
|
|
cooccurrence_map[eid1] = set()
|
|
if eid2 not in cooccurrence_map:
|
|
cooccurrence_map[eid2] = set()
|
|
if eid2 in id_to_name:
|
|
cooccurrence_map[eid1].add(id_to_name[eid2])
|
|
if eid1 in id_to_name:
|
|
cooccurrence_map[eid2].add(id_to_name[eid1])
|
|
|
|
return await self._resolve_from_candidates(
|
|
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
|
)
|
|
|
|
async def _resolve_from_candidates(
|
|
self,
|
|
conn,
|
|
bank_id: str,
|
|
entities_data: list[dict],
|
|
unit_event_date,
|
|
all_candidates: dict[str, list],
|
|
cooccurrence_map: dict[str, set[str]],
|
|
) -> list[str]:
|
|
"""Shared scoring + upsert logic used by both lookup strategies."""
|
|
|
|
# Resolve each entity using pre-fetched candidates
|
|
entity_ids = [None] * len(entities_data)
|
|
entities_to_update: list[_EntityStat] = []
|
|
entities_to_create: list[_EntityToCreate] = []
|
|
|
|
for idx, entity_data in enumerate(entities_data):
|
|
entity_text = entity_data["text"]
|
|
nearby_entities = entity_data.get("nearby_entities", [])
|
|
# Use per-entity date if available, otherwise fall back to batch-level date
|
|
entity_event_date = entity_data.get("event_date", unit_event_date)
|
|
|
|
candidates = all_candidates.get(entity_text, [])
|
|
|
|
if not candidates:
|
|
# Will create new entity
|
|
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
|
continue
|
|
|
|
# Score candidates
|
|
best_candidate = None
|
|
best_score = 0.0
|
|
|
|
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
|
|
|
|
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
|
score = 0.0
|
|
|
|
# 1. Name similarity (0-0.5)
|
|
name_similarity = SequenceMatcher(None, entity_text.lower(), canonical_name.lower()).ratio()
|
|
score += name_similarity * 0.5
|
|
|
|
# 2. Co-occurring entities (0-0.3)
|
|
if nearby_entity_set:
|
|
co_entities = cooccurrence_map.get(candidate_id, set())
|
|
overlap = len(nearby_entity_set & co_entities)
|
|
co_entity_score = overlap / len(nearby_entity_set)
|
|
score += co_entity_score * 0.3
|
|
|
|
# 3. Temporal proximity (0-0.2)
|
|
if last_seen and entity_event_date:
|
|
# Normalize timezone awareness for comparison
|
|
event_date_utc = (
|
|
entity_event_date if entity_event_date.tzinfo else entity_event_date.replace(tzinfo=UTC)
|
|
)
|
|
last_seen_utc = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=UTC)
|
|
days_diff = abs((event_date_utc - last_seen_utc).total_seconds() / 86400)
|
|
if days_diff < 7:
|
|
temporal_score = max(0, 1.0 - (days_diff / 7))
|
|
score += temporal_score * 0.2
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
best_candidate = candidate_id
|
|
|
|
# Apply unified threshold
|
|
threshold = 0.6
|
|
|
|
if best_score > threshold:
|
|
entity_ids[idx] = best_candidate
|
|
entities_to_update.append(_EntityStat(entity_id=best_candidate, event_date=entity_event_date))
|
|
else:
|
|
entities_to_create.append(
|
|
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
|
|
)
|
|
|
|
# Existing entities: IDs already known from the candidate SELECT above.
|
|
# No in-transaction UPDATE — mention_count/last_seen are stats deferred to
|
|
# flush_pending_stats() which the orchestrator calls after the transaction.
|
|
pending: list[_EntityStat] = list(entities_to_update)
|
|
|
|
# New entities: INSERT with DO NOTHING to avoid row locks on concurrent races.
|
|
# ON CONFLICT DO NOTHING returns nothing for rows that conflicted; we handle
|
|
# that rare case with a fallback SELECT.
|
|
if entities_to_create:
|
|
# Group by lowercase name — deduplicate within the batch.
|
|
@dataclass
|
|
class _NameGroup:
|
|
name: str
|
|
event_date: datetime | None
|
|
indices: list[int] = field(default_factory=list)
|
|
|
|
groups: dict[str, _NameGroup] = {}
|
|
for e in entities_to_create:
|
|
name_lower = e.name.lower()
|
|
if name_lower not in groups:
|
|
groups[name_lower] = _NameGroup(name=e.name, event_date=e.event_date)
|
|
groups[name_lower].indices.append(e.idx)
|
|
|
|
# Sort by lowercase name for deterministic ordering.
|
|
sorted_groups = sorted(groups.items())
|
|
entity_names = [g.name for _, g in sorted_groups]
|
|
entity_dates = [g.event_date for _, g in sorted_groups]
|
|
|
|
# INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
|
|
inserted_rows = await conn.fetch(
|
|
f"""
|
|
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
|
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 1
|
|
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
|
ON CONFLICT (bank_id, LOWER(canonical_name))
|
|
DO NOTHING
|
|
RETURNING id, LOWER(canonical_name) AS name_lower
|
|
""",
|
|
bank_id,
|
|
entity_names,
|
|
entity_dates,
|
|
)
|
|
id_by_name: dict[str, str] = {row["name_lower"]: row["id"] for row in inserted_rows}
|
|
|
|
# Fallback SELECT for names that conflicted (another worker won the race).
|
|
missing = [n for n, _ in sorted_groups if n not in id_by_name]
|
|
if missing:
|
|
existing_rows = await conn.fetch(
|
|
f"""
|
|
SELECT id, LOWER(canonical_name) AS name_lower
|
|
FROM {fq_table("entities")}
|
|
WHERE bank_id = $1 AND LOWER(canonical_name) = ANY($2::text[])
|
|
""",
|
|
bank_id,
|
|
missing,
|
|
)
|
|
for row in existing_rows:
|
|
id_by_name[row["name_lower"]] = row["id"]
|
|
|
|
# Assign entity IDs back and queue for post-txn stats flush.
|
|
for name_lower, g in sorted_groups:
|
|
entity_id = id_by_name.get(name_lower)
|
|
if entity_id:
|
|
for original_idx in g.indices:
|
|
entity_ids[original_idx] = entity_id
|
|
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
|
|
|
|
# Accumulate into the resolver's pending list; the orchestrator flushes
|
|
# these with await entity_resolver.flush_pending_stats() after the txn.
|
|
key = self._task_key()
|
|
self._pending_stats.setdefault(key, []).extend(pending)
|
|
|
|
return entity_ids
|
|
|
|
async def resolve_entity(
|
|
self,
|
|
bank_id: str,
|
|
entity_text: str,
|
|
context: str,
|
|
nearby_entities: list[dict],
|
|
unit_event_date,
|
|
) -> str:
|
|
"""
|
|
Resolve an entity to a canonical entity ID.
|
|
|
|
Args:
|
|
bank_id: bank ID (entities are scoped to agents)
|
|
entity_text: Entity text ("Alice", "Google", etc.)
|
|
context: Context where entity appears
|
|
nearby_entities: Other entities in the same unit
|
|
unit_event_date: When this unit was created
|
|
|
|
Returns:
|
|
Entity ID (creates new entity if needed)
|
|
"""
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
# Find candidate entities with similar name
|
|
candidates = await conn.fetch(
|
|
f"""
|
|
SELECT id, canonical_name, metadata, last_seen
|
|
FROM {fq_table("entities")}
|
|
WHERE bank_id = $1
|
|
AND (
|
|
canonical_name ILIKE $2
|
|
OR canonical_name ILIKE $3
|
|
OR $2 ILIKE canonical_name || '%%'
|
|
)
|
|
ORDER BY mention_count DESC
|
|
""",
|
|
bank_id,
|
|
entity_text,
|
|
f"%{entity_text}%",
|
|
)
|
|
|
|
if not candidates:
|
|
# New entity - create it
|
|
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
|
|
|
# Score candidates based on:
|
|
# 1. Name similarity
|
|
# 2. Context overlap (TODO: could use embeddings)
|
|
# 3. Co-occurring entities
|
|
# 4. Temporal proximity
|
|
|
|
best_candidate = None
|
|
best_score = 0.0
|
|
best_name_similarity = 0.0
|
|
|
|
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
|
|
|
|
for row in candidates:
|
|
candidate_id = row["id"]
|
|
canonical_name = row["canonical_name"]
|
|
metadata = row["metadata"]
|
|
last_seen = row["last_seen"]
|
|
score = 0.0
|
|
|
|
# 1. Name similarity (0-1)
|
|
name_similarity = SequenceMatcher(None, entity_text.lower(), canonical_name.lower()).ratio()
|
|
score += name_similarity * 0.5
|
|
|
|
# 2. Co-occurring entities (0-0.5)
|
|
# Get entities that co-occurred with this candidate before
|
|
# Use the materialized co-occurrence cache for fast lookup
|
|
co_entity_rows = await conn.fetch(
|
|
f"""
|
|
SELECT e.canonical_name, ec.cooccurrence_count
|
|
FROM {fq_table("entity_cooccurrences")} ec
|
|
JOIN {fq_table("entities")} e ON (
|
|
CASE
|
|
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
|
|
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
|
|
END = e.id
|
|
)
|
|
WHERE ec.entity_id_1 = $1 OR ec.entity_id_2 = $1
|
|
""",
|
|
candidate_id,
|
|
)
|
|
co_entities = {r["canonical_name"].lower() for r in co_entity_rows}
|
|
|
|
# Check overlap with nearby entities
|
|
overlap = len(nearby_entity_set & co_entities)
|
|
if nearby_entity_set:
|
|
co_entity_score = overlap / len(nearby_entity_set)
|
|
score += co_entity_score * 0.3
|
|
|
|
# 3. Temporal proximity (0-0.2)
|
|
if last_seen:
|
|
days_diff = abs((unit_event_date - last_seen).total_seconds() / 86400)
|
|
if days_diff < 7: # Within a week
|
|
temporal_score = max(0, 1.0 - (days_diff / 7))
|
|
score += temporal_score * 0.2
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
best_candidate = candidate_id
|
|
best_name_similarity = name_similarity
|
|
|
|
# Threshold for considering it the same entity
|
|
threshold = 0.6
|
|
|
|
if best_score > threshold:
|
|
# Update entity
|
|
await conn.execute(
|
|
f"""
|
|
UPDATE {fq_table("entities")}
|
|
SET mention_count = mention_count + 1,
|
|
last_seen = $1
|
|
WHERE id = $2
|
|
""",
|
|
unit_event_date,
|
|
best_candidate,
|
|
)
|
|
return best_candidate
|
|
else:
|
|
# Not confident - create new entity
|
|
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
|
|
|
async def _create_entity(
|
|
self,
|
|
conn,
|
|
bank_id: str,
|
|
entity_text: str,
|
|
event_date,
|
|
) -> str:
|
|
"""
|
|
Create a new entity or get existing one if it already exists.
|
|
|
|
Uses INSERT ... ON CONFLICT to handle race conditions where
|
|
two concurrent transactions try to create the same entity.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
bank_id: bank ID
|
|
entity_text: Entity text
|
|
event_date: When first seen
|
|
|
|
Returns:
|
|
Entity ID
|
|
"""
|
|
entity_id = await conn.fetchval(
|
|
f"""
|
|
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
|
VALUES ($1, $2, COALESCE($3, now()), COALESCE($4, now()), 1)
|
|
ON CONFLICT (bank_id, LOWER(canonical_name))
|
|
DO UPDATE SET
|
|
mention_count = {fq_table("entities")}.mention_count + 1,
|
|
last_seen = EXCLUDED.last_seen
|
|
RETURNING id
|
|
""",
|
|
bank_id,
|
|
entity_text,
|
|
event_date,
|
|
event_date,
|
|
)
|
|
return entity_id
|
|
|
|
async def link_unit_to_entity(self, unit_id: str, entity_id: str):
|
|
"""
|
|
Link a memory unit to an entity.
|
|
Also updates co-occurrence cache with other entities in the same unit.
|
|
|
|
Args:
|
|
unit_id: Memory unit ID
|
|
entity_id: Entity ID
|
|
"""
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
# Insert unit-entity link
|
|
await conn.execute(
|
|
f"""
|
|
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
|
VALUES ($1, $2)
|
|
ON CONFLICT DO NOTHING
|
|
""",
|
|
unit_id,
|
|
entity_id,
|
|
)
|
|
|
|
# Update co-occurrence cache: find other entities in this unit
|
|
rows = await conn.fetch(
|
|
f"""
|
|
SELECT entity_id
|
|
FROM {fq_table("unit_entities")}
|
|
WHERE unit_id = $1 AND entity_id != $2
|
|
""",
|
|
unit_id,
|
|
entity_id,
|
|
)
|
|
|
|
other_entities = [row["entity_id"] for row in rows]
|
|
|
|
# Update co-occurrences for each pair
|
|
for other_entity_id in other_entities:
|
|
await self._update_cooccurrence(conn, entity_id, other_entity_id)
|
|
|
|
async def _update_cooccurrence(self, conn, entity_id_1: str, entity_id_2: str):
|
|
"""
|
|
Update the co-occurrence cache for two entities.
|
|
|
|
Uses CHECK constraint ordering (entity_id_1 < entity_id_2) to avoid duplicates.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
entity_id_1: First entity ID
|
|
entity_id_2: Second entity ID
|
|
"""
|
|
# Ensure consistent ordering (smaller UUID first)
|
|
if entity_id_1 > entity_id_2:
|
|
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
|
|
|
await conn.execute(
|
|
f"""
|
|
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
|
VALUES ($1, $2, 1, NOW())
|
|
ON CONFLICT (entity_id_1, entity_id_2)
|
|
DO UPDATE SET
|
|
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
|
last_cooccurred = NOW()
|
|
""",
|
|
entity_id_1,
|
|
entity_id_2,
|
|
)
|
|
|
|
async def link_units_to_entities_batch(self, unit_entity_pairs: list[tuple[str, str]], conn=None):
|
|
"""
|
|
Link multiple memory units to entities in batch (MUCH faster than sequential).
|
|
|
|
Also updates co-occurrence cache for entities that appear in the same unit.
|
|
|
|
Args:
|
|
unit_entity_pairs: List of (unit_id, entity_id) tuples
|
|
conn: Optional connection to use (if None, acquires from pool)
|
|
"""
|
|
if not unit_entity_pairs:
|
|
return
|
|
|
|
if conn is None:
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
return await self._link_units_to_entities_batch_impl(conn, unit_entity_pairs)
|
|
else:
|
|
return await self._link_units_to_entities_batch_impl(conn, unit_entity_pairs)
|
|
|
|
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str]]):
|
|
# Batch insert all unit-entity links
|
|
await conn.executemany(
|
|
f"""
|
|
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
|
VALUES ($1, $2)
|
|
ON CONFLICT DO NOTHING
|
|
""",
|
|
unit_entity_pairs,
|
|
)
|
|
|
|
# Build map of unit -> entities for co-occurrence calculation
|
|
# Use sets to avoid duplicate entities in the same unit
|
|
unit_to_entities = {}
|
|
for unit_id, entity_id in unit_entity_pairs:
|
|
if unit_id not in unit_to_entities:
|
|
unit_to_entities[unit_id] = set()
|
|
unit_to_entities[unit_id].add(entity_id)
|
|
|
|
# Update co-occurrences for all pairs in each unit
|
|
cooccurrence_pairs = set() # Use set to avoid duplicates
|
|
for unit_id, entity_ids in unit_to_entities.items():
|
|
entity_list = list(entity_ids) # Convert set to list for iteration
|
|
# For each pair of entities in this unit, create co-occurrence
|
|
for i, entity_id_1 in enumerate(entity_list):
|
|
for entity_id_2 in entity_list[i + 1 :]:
|
|
# Skip if same entity (shouldn't happen with set, but be safe)
|
|
if entity_id_1 == entity_id_2:
|
|
continue
|
|
# Ensure consistent ordering (entity_id_1 < entity_id_2)
|
|
if entity_id_1 > entity_id_2:
|
|
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
|
cooccurrence_pairs.add((entity_id_1, entity_id_2))
|
|
|
|
# Accumulate co-occurrence pairs for post-transaction flush.
|
|
# The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
|
|
# row-level lock contention (ON CONFLICT DO UPDATE inside a long transaction
|
|
# serialises concurrent writers on popular entity pairs).
|
|
if cooccurrence_pairs:
|
|
key = self._task_key()
|
|
self._pending_cooccurrences.setdefault(key, []).extend(
|
|
_CooccurrencePair(entity_id_1=e1, entity_id_2=e2) for e1, e2 in cooccurrence_pairs
|
|
)
|
|
|
|
async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]:
|
|
"""
|
|
Get all units that mention an entity.
|
|
|
|
Args:
|
|
entity_id: Entity ID
|
|
limit: Max results
|
|
|
|
Returns:
|
|
List of unit IDs
|
|
"""
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
rows = await conn.fetch(
|
|
f"""
|
|
SELECT unit_id
|
|
FROM {fq_table("unit_entities")}
|
|
WHERE entity_id = $1
|
|
ORDER BY unit_id
|
|
LIMIT $2
|
|
""",
|
|
entity_id,
|
|
limit,
|
|
)
|
|
return [row["unit_id"] for row in rows]
|
|
|
|
async def get_entity_by_text(
|
|
self,
|
|
bank_id: str,
|
|
entity_text: str,
|
|
) -> str | None:
|
|
"""
|
|
Find an entity by text (for query resolution).
|
|
|
|
Args:
|
|
bank_id: bank ID
|
|
entity_text: Entity text to search for
|
|
|
|
Returns:
|
|
Entity ID if found, None otherwise
|
|
"""
|
|
async with acquire_with_retry(self.pool) as conn:
|
|
row = await conn.fetchrow(
|
|
f"""
|
|
SELECT id FROM {fq_table("entities")}
|
|
WHERE bank_id = $1
|
|
AND canonical_name ILIKE $2
|
|
ORDER BY mention_count DESC
|
|
LIMIT 1
|
|
""",
|
|
bank_id,
|
|
entity_text,
|
|
)
|
|
|
|
return row["id"] if row else None
|