634 lines
22 KiB
Python
634 lines
22 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 spacy
|
|
from typing import List, Dict, Optional, Set
|
|
from difflib import SequenceMatcher
|
|
|
|
|
|
# Load spaCy model (singleton)
|
|
_nlp = None
|
|
|
|
|
|
def get_nlp():
|
|
"""Get or load spaCy model."""
|
|
global _nlp
|
|
if _nlp is None:
|
|
_nlp = spacy.load("en_core_web_sm")
|
|
return _nlp
|
|
|
|
|
|
def extract_entities(text: str) -> List[Dict[str, any]]:
|
|
"""
|
|
Extract entities from text using spaCy.
|
|
|
|
Args:
|
|
text: Input text
|
|
|
|
Returns:
|
|
List of entities with text, type, and span info
|
|
"""
|
|
nlp = get_nlp()
|
|
doc = nlp(text)
|
|
|
|
entities = []
|
|
for ent in doc.ents:
|
|
# Filter to important entity types
|
|
if ent.label_ in ['PERSON', 'ORG', 'GPE', 'LOC', 'PRODUCT', 'EVENT']:
|
|
entities.append({
|
|
'text': ent.text,
|
|
'type': ent.label_,
|
|
'start': ent.start_char,
|
|
'end': ent.end_char,
|
|
})
|
|
|
|
return entities
|
|
|
|
|
|
def extract_entities_batch(texts: List[str]) -> List[List[Dict[str, any]]]:
|
|
"""
|
|
Extract entities from multiple texts in batch (MUCH faster than sequential).
|
|
|
|
Uses spaCy's nlp.pipe() for efficient batch processing.
|
|
|
|
Args:
|
|
texts: List of input texts
|
|
|
|
Returns:
|
|
List of entity lists, one per input text
|
|
"""
|
|
if not texts:
|
|
return []
|
|
|
|
nlp = get_nlp()
|
|
|
|
# Process all texts in batch using nlp.pipe (significantly faster!)
|
|
docs = list(nlp.pipe(texts, batch_size=50))
|
|
|
|
all_entities = []
|
|
for doc in docs:
|
|
entities = []
|
|
for ent in doc.ents:
|
|
# Filter to important entity types
|
|
if ent.label_ in ['PERSON', 'ORG', 'GPE', 'LOC', 'PRODUCT', 'EVENT']:
|
|
entities.append({
|
|
'text': ent.text,
|
|
'type': ent.label_,
|
|
'start': ent.start_char,
|
|
'end': ent.end_char,
|
|
})
|
|
all_entities.append(entities)
|
|
|
|
return all_entities
|
|
|
|
|
|
class EntityResolver:
|
|
"""
|
|
Resolves entities to canonical IDs with disambiguation.
|
|
"""
|
|
|
|
def __init__(self, db_conn):
|
|
"""
|
|
Initialize entity resolver.
|
|
|
|
Args:
|
|
db_conn: psycopg2 database connection
|
|
"""
|
|
self.conn = db_conn
|
|
|
|
def resolve_entities_batch(
|
|
self,
|
|
agent_id: str,
|
|
entities_data: List[Dict],
|
|
context: str,
|
|
unit_event_date,
|
|
) -> 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:
|
|
agent_id: Agent ID
|
|
entities_data: List of dicts with 'text', 'type', 'nearby_entities'
|
|
context: Context where entities appear
|
|
unit_event_date: When this unit was created
|
|
|
|
Returns:
|
|
List of entity IDs in same order as input
|
|
"""
|
|
if not entities_data:
|
|
return []
|
|
|
|
cursor = self.conn.cursor()
|
|
|
|
try:
|
|
import time
|
|
start = time.time()
|
|
|
|
# Group entities by type for efficient querying
|
|
entities_by_type = {}
|
|
for idx, entity_data in enumerate(entities_data):
|
|
entity_type = entity_data['type']
|
|
if entity_type not in entities_by_type:
|
|
entities_by_type[entity_type] = []
|
|
entities_by_type[entity_type].append((idx, entity_data))
|
|
|
|
# Query ALL candidates for each type in batch
|
|
all_candidates = {} # Maps (entity_type, entity_text) -> list of candidates
|
|
for entity_type, entities_list in entities_by_type.items():
|
|
# Extract unique entity texts for this type
|
|
entity_texts = list(set(e[1]['text'] for e in entities_list))
|
|
|
|
# Query candidates for all texts at once
|
|
from psycopg2.extras import execute_values
|
|
cursor.execute(
|
|
"""
|
|
SELECT canonical_name, id, metadata, last_seen, mention_count
|
|
FROM entities
|
|
WHERE agent_id = %s AND entity_type = %s
|
|
""",
|
|
(agent_id, entity_type)
|
|
)
|
|
type_candidates = cursor.fetchall()
|
|
|
|
# Filter candidates in memory (faster than complex SQL for small datasets)
|
|
for entity_text in entity_texts:
|
|
matching = []
|
|
entity_text_lower = entity_text.lower()
|
|
for canonical_name, ent_id, metadata, last_seen, mention_count in type_candidates:
|
|
canonical_lower = canonical_name.lower()
|
|
# Same matching logic as before
|
|
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_type, entity_text)] = matching
|
|
|
|
# Resolve each entity using pre-fetched candidates
|
|
entity_ids = [None] * len(entities_data)
|
|
entities_to_update = [] # (entity_id, unit_event_date)
|
|
entities_to_create = [] # (idx, entity_data)
|
|
|
|
for idx, entity_data in enumerate(entities_data):
|
|
entity_text = entity_data['text']
|
|
entity_type = entity_data['type']
|
|
nearby_entities = entity_data.get('nearby_entities', [])
|
|
|
|
candidates = all_candidates.get((entity_type, entity_text), [])
|
|
|
|
if not candidates:
|
|
# Will create new entity
|
|
entities_to_create.append((idx, entity_data))
|
|
continue
|
|
|
|
# Score candidates (same logic as before but with pre-fetched data)
|
|
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 candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
|
score = 0.0
|
|
|
|
# Name similarity
|
|
name_similarity = SequenceMatcher(
|
|
None,
|
|
entity_text.lower(),
|
|
canonical_name.lower()
|
|
).ratio()
|
|
score += name_similarity * 0.5
|
|
|
|
# Temporal proximity
|
|
if last_seen:
|
|
days_diff = abs((unit_event_date - last_seen).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
|
|
best_name_similarity = name_similarity
|
|
|
|
# Apply threshold
|
|
threshold = 0.4 if entity_type == 'PERSON' and best_name_similarity >= 0.95 else 0.6
|
|
|
|
if best_score > threshold:
|
|
entity_ids[idx] = best_candidate
|
|
entities_to_update.append((best_candidate, unit_event_date))
|
|
else:
|
|
entities_to_create.append((idx, entity_data))
|
|
|
|
# Batch update existing entities
|
|
if entities_to_update:
|
|
from psycopg2.extras import execute_values
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
UPDATE entities SET
|
|
mention_count = mention_count + 1,
|
|
last_seen = data.last_seen
|
|
FROM (VALUES %s) AS data(id, last_seen)
|
|
WHERE entities.id = data.id::uuid
|
|
""",
|
|
entities_to_update
|
|
)
|
|
|
|
# Batch create new entities
|
|
if entities_to_create:
|
|
for idx, entity_data in entities_to_create:
|
|
entity_id = self._create_entity(
|
|
cursor, agent_id, entity_data['text'],
|
|
entity_data['type'], unit_event_date
|
|
)
|
|
entity_ids[idx] = entity_id
|
|
|
|
return entity_ids
|
|
|
|
finally:
|
|
cursor.close()
|
|
|
|
def resolve_entity(
|
|
self,
|
|
agent_id: str,
|
|
entity_text: str,
|
|
entity_type: str,
|
|
context: str,
|
|
nearby_entities: List[Dict],
|
|
unit_event_date,
|
|
) -> str:
|
|
"""
|
|
Resolve an entity to a canonical entity ID.
|
|
|
|
Args:
|
|
agent_id: Agent ID (entities are scoped to agents)
|
|
entity_text: Entity text ("Alice", "Google", etc.)
|
|
entity_type: Entity type (PERSON, ORG, 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)
|
|
"""
|
|
cursor = self.conn.cursor()
|
|
|
|
try:
|
|
# Find candidate entities with same type and similar name
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, canonical_name, metadata, last_seen
|
|
FROM entities
|
|
WHERE agent_id = %s
|
|
AND entity_type = %s
|
|
AND (
|
|
canonical_name ILIKE %s
|
|
OR canonical_name ILIKE %s
|
|
OR %s ILIKE canonical_name || '%%'
|
|
)
|
|
ORDER BY mention_count DESC
|
|
""",
|
|
(agent_id, entity_type, entity_text, f"%{entity_text}%", entity_text)
|
|
)
|
|
|
|
candidates = cursor.fetchall()
|
|
|
|
if not candidates:
|
|
# New entity - create it
|
|
return self._create_entity(
|
|
cursor, agent_id, entity_text, entity_type, 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 candidate_id, canonical_name, metadata, last_seen in candidates:
|
|
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
|
|
cursor.execute(
|
|
"""
|
|
SELECT e.canonical_name, ec.cooccurrence_count
|
|
FROM entity_cooccurrences ec
|
|
JOIN entities e ON (
|
|
CASE
|
|
WHEN ec.entity_id_1 = %s THEN ec.entity_id_2
|
|
WHEN ec.entity_id_2 = %s THEN ec.entity_id_1
|
|
END = e.id
|
|
)
|
|
WHERE ec.entity_id_1 = %s OR ec.entity_id_2 = %s
|
|
""",
|
|
(candidate_id, candidate_id, candidate_id, candidate_id)
|
|
)
|
|
co_entities = {row[0].lower() for row in cursor.fetchall()}
|
|
|
|
# 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
|
|
# For PERSON entities with exact name match, use lower threshold
|
|
threshold = 0.4 if entity_type == 'PERSON' and best_name_similarity >= 0.95 else 0.6
|
|
|
|
if best_score > threshold:
|
|
# Update entity
|
|
cursor.execute(
|
|
"""
|
|
UPDATE entities
|
|
SET mention_count = mention_count + 1,
|
|
last_seen = %s
|
|
WHERE id = %s
|
|
""",
|
|
(unit_event_date, best_candidate)
|
|
)
|
|
return best_candidate
|
|
else:
|
|
# Not confident - create new entity
|
|
return self._create_entity(
|
|
cursor, agent_id, entity_text, entity_type, unit_event_date
|
|
)
|
|
|
|
finally:
|
|
cursor.close()
|
|
|
|
def _create_entity(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
entity_text: str,
|
|
entity_type: str,
|
|
event_date,
|
|
) -> str:
|
|
"""
|
|
Create a new entity.
|
|
|
|
Args:
|
|
cursor: Database cursor
|
|
agent_id: Agent ID
|
|
entity_text: Entity text
|
|
entity_type: Entity type
|
|
event_date: When first seen
|
|
|
|
Returns:
|
|
Entity ID
|
|
"""
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO entities (agent_id, canonical_name, entity_type, first_seen, last_seen, mention_count)
|
|
VALUES (%s, %s, %s, %s, %s, 1)
|
|
RETURNING id
|
|
""",
|
|
(agent_id, entity_text, entity_type, event_date, event_date)
|
|
)
|
|
entity_id = cursor.fetchone()[0]
|
|
return entity_id
|
|
|
|
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
|
|
"""
|
|
cursor = self.conn.cursor()
|
|
try:
|
|
# Insert unit-entity link
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO unit_entities (unit_id, entity_id)
|
|
VALUES (%s, %s)
|
|
ON CONFLICT DO NOTHING
|
|
""",
|
|
(unit_id, entity_id)
|
|
)
|
|
|
|
# Update co-occurrence cache: find other entities in this unit
|
|
cursor.execute(
|
|
"""
|
|
SELECT entity_id
|
|
FROM unit_entities
|
|
WHERE unit_id = %s AND entity_id != %s
|
|
""",
|
|
(unit_id, entity_id)
|
|
)
|
|
|
|
other_entities = [row[0] for row in cursor.fetchall()]
|
|
|
|
# Update co-occurrences for each pair
|
|
for other_entity_id in other_entities:
|
|
self._update_cooccurrence(cursor, entity_id, other_entity_id)
|
|
|
|
finally:
|
|
cursor.close()
|
|
|
|
def _update_cooccurrence(self, cursor, 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:
|
|
cursor: Database cursor
|
|
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
|
|
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
|
VALUES (%s, %s, 1, NOW())
|
|
ON CONFLICT (entity_id_1, entity_id_2)
|
|
DO UPDATE SET
|
|
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
|
|
last_cooccurred = NOW()
|
|
""",
|
|
(entity_id_1, entity_id_2)
|
|
)
|
|
|
|
def link_units_to_entities_batch(self, unit_entity_pairs: List[tuple[str, str]]):
|
|
"""
|
|
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
|
|
"""
|
|
if not unit_entity_pairs:
|
|
return
|
|
|
|
cursor = self.conn.cursor()
|
|
try:
|
|
# Batch insert all unit-entity links
|
|
from psycopg2.extras import execute_values
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO unit_entities (unit_id, entity_id)
|
|
VALUES %s
|
|
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))
|
|
|
|
# Batch update co-occurrences
|
|
if cooccurrence_pairs:
|
|
from datetime import datetime, timezone
|
|
now = datetime.now(timezone.utc)
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
|
VALUES %s
|
|
ON CONFLICT (entity_id_1, entity_id_2)
|
|
DO UPDATE SET
|
|
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
|
|
last_cooccurred = EXCLUDED.last_cooccurred
|
|
""",
|
|
[(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs]
|
|
)
|
|
|
|
finally:
|
|
cursor.close()
|
|
|
|
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
|
|
"""
|
|
cursor = self.conn.cursor()
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
SELECT unit_id
|
|
FROM unit_entities
|
|
WHERE entity_id = %s
|
|
ORDER BY unit_id
|
|
LIMIT %s
|
|
""",
|
|
(entity_id, limit)
|
|
)
|
|
return [row[0] for row in cursor.fetchall()]
|
|
finally:
|
|
cursor.close()
|
|
|
|
def get_entity_by_text(
|
|
self,
|
|
agent_id: str,
|
|
entity_text: str,
|
|
entity_type: Optional[str] = None
|
|
) -> Optional[str]:
|
|
"""
|
|
Find an entity by text (for query resolution).
|
|
|
|
Args:
|
|
agent_id: Agent ID
|
|
entity_text: Entity text to search for
|
|
entity_type: Optional entity type filter
|
|
|
|
Returns:
|
|
Entity ID if found, None otherwise
|
|
"""
|
|
cursor = self.conn.cursor()
|
|
try:
|
|
if entity_type:
|
|
cursor.execute(
|
|
"""
|
|
SELECT id FROM entities
|
|
WHERE agent_id = %s
|
|
AND entity_type = %s
|
|
AND canonical_name ILIKE %s
|
|
ORDER BY mention_count DESC
|
|
LIMIT 1
|
|
""",
|
|
(agent_id, entity_type, entity_text)
|
|
)
|
|
else:
|
|
cursor.execute(
|
|
"""
|
|
SELECT id FROM entities
|
|
WHERE agent_id = %s
|
|
AND canonical_name ILIKE %s
|
|
ORDER BY mention_count DESC
|
|
LIMIT 1
|
|
""",
|
|
(agent_id, entity_text)
|
|
)
|
|
|
|
row = cursor.fetchone()
|
|
return row[0] if row else None
|
|
finally:
|
|
cursor.close()
|