251 lines
7.8 KiB
Python
251 lines
7.8 KiB
Python
"""
|
|
Observation regeneration for retain pipeline.
|
|
|
|
Regenerates entity observations as part of the retain transaction.
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
|
|
from ..search import observation_utils
|
|
from . import embedding_utils
|
|
from .types import EntityLink
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def utcnow():
|
|
"""Get current UTC time."""
|
|
return datetime.now(UTC)
|
|
|
|
|
|
# Simple dataclass-like container for facts (avoid importing from memory_engine)
|
|
class MemoryFactForObservation:
|
|
def __init__(self, id: str, text: str, fact_type: str, context: str, occurred_start: str | None):
|
|
self.id = id
|
|
self.text = text
|
|
self.fact_type = fact_type
|
|
self.context = context
|
|
self.occurred_start = occurred_start
|
|
|
|
|
|
async def regenerate_observations_batch(
|
|
conn, embeddings_model, llm_config, bank_id: str, entity_links: list[EntityLink], log_buffer: list[str] = None
|
|
) -> None:
|
|
"""
|
|
Regenerate observations for top entities in this batch.
|
|
|
|
Called INSIDE the retain transaction for atomicity - if observations
|
|
fail, the entire retain batch is rolled back.
|
|
|
|
Args:
|
|
conn: Database connection (from the retain transaction)
|
|
embeddings_model: Embeddings model for generating observation embeddings
|
|
llm_config: LLM configuration for observation extraction
|
|
bank_id: Bank identifier
|
|
entity_links: Entity links from this batch
|
|
log_buffer: Optional log buffer for timing
|
|
"""
|
|
TOP_N_ENTITIES = 5
|
|
MIN_FACTS_THRESHOLD = 5
|
|
|
|
if not entity_links:
|
|
return
|
|
|
|
# Count mentions per entity in this batch
|
|
entity_mention_counts: dict[str, int] = {}
|
|
for link in entity_links:
|
|
if link.entity_id:
|
|
entity_id = str(link.entity_id)
|
|
entity_mention_counts[entity_id] = entity_mention_counts.get(entity_id, 0) + 1
|
|
|
|
if not entity_mention_counts:
|
|
return
|
|
|
|
# Sort by mention count descending and take top N
|
|
sorted_entities = sorted(entity_mention_counts.items(), key=lambda x: x[1], reverse=True)
|
|
entities_to_process = [e[0] for e in sorted_entities[:TOP_N_ENTITIES]]
|
|
|
|
obs_start = time.time()
|
|
|
|
# Convert to UUIDs
|
|
entity_uuids = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in entities_to_process]
|
|
|
|
# Batch query for entity names
|
|
entity_rows = await conn.fetch(
|
|
"""
|
|
SELECT id, canonical_name FROM entities
|
|
WHERE id = ANY($1) AND bank_id = $2
|
|
""",
|
|
entity_uuids,
|
|
bank_id,
|
|
)
|
|
entity_names = {row["id"]: row["canonical_name"] for row in entity_rows}
|
|
|
|
# Batch query for fact counts
|
|
fact_counts = await conn.fetch(
|
|
"""
|
|
SELECT ue.entity_id, COUNT(*) as cnt
|
|
FROM unit_entities ue
|
|
JOIN memory_units mu ON ue.unit_id = mu.id
|
|
WHERE ue.entity_id = ANY($1) AND mu.bank_id = $2
|
|
GROUP BY ue.entity_id
|
|
""",
|
|
entity_uuids,
|
|
bank_id,
|
|
)
|
|
entity_fact_counts = {row["entity_id"]: row["cnt"] for row in fact_counts}
|
|
|
|
# Filter entities that meet the threshold
|
|
entities_with_names = []
|
|
for entity_id in entities_to_process:
|
|
entity_uuid = uuid.UUID(entity_id) if isinstance(entity_id, str) else entity_id
|
|
if entity_uuid not in entity_names:
|
|
continue
|
|
fact_count = entity_fact_counts.get(entity_uuid, 0)
|
|
if fact_count >= MIN_FACTS_THRESHOLD:
|
|
entities_with_names.append((entity_id, entity_names[entity_uuid]))
|
|
|
|
if not entities_with_names:
|
|
return
|
|
|
|
# Process entities SEQUENTIALLY (asyncpg doesn't allow concurrent queries on same connection)
|
|
# We must use the same connection to stay in the retain transaction
|
|
total_observations = 0
|
|
|
|
for entity_id, entity_name in entities_with_names:
|
|
try:
|
|
obs_ids = await _regenerate_entity_observations(
|
|
conn, embeddings_model, llm_config, bank_id, entity_id, entity_name
|
|
)
|
|
total_observations += len(obs_ids)
|
|
except Exception as e:
|
|
logger.error(f"[OBSERVATIONS] Error processing entity {entity_id}: {e}")
|
|
|
|
obs_time = time.time() - obs_start
|
|
if log_buffer is not None:
|
|
log_buffer.append(
|
|
f"[11] Observations: {total_observations} observations for {len(entities_with_names)} entities in {obs_time:.3f}s"
|
|
)
|
|
|
|
|
|
async def _regenerate_entity_observations(
|
|
conn, embeddings_model, llm_config, bank_id: str, entity_id: str, entity_name: str
|
|
) -> list[str]:
|
|
"""
|
|
Regenerate observations for a single entity.
|
|
|
|
Uses the provided connection (part of retain transaction).
|
|
|
|
Args:
|
|
conn: Database connection (from the retain transaction)
|
|
embeddings_model: Embeddings model
|
|
llm_config: LLM configuration
|
|
bank_id: Bank identifier
|
|
entity_id: Entity UUID
|
|
entity_name: Canonical name of the entity
|
|
|
|
Returns:
|
|
List of created observation IDs
|
|
"""
|
|
entity_uuid = uuid.UUID(entity_id) if isinstance(entity_id, str) else entity_id
|
|
|
|
# Get all facts mentioning this entity (exclude observations themselves)
|
|
rows = await conn.fetch(
|
|
"""
|
|
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.fact_type
|
|
FROM memory_units mu
|
|
JOIN unit_entities ue ON mu.id = ue.unit_id
|
|
WHERE mu.bank_id = $1
|
|
AND ue.entity_id = $2
|
|
AND mu.fact_type IN ('world', 'experience')
|
|
ORDER BY mu.occurred_start DESC
|
|
LIMIT 50
|
|
""",
|
|
bank_id,
|
|
entity_uuid,
|
|
)
|
|
|
|
if not rows:
|
|
return []
|
|
|
|
# Convert to fact objects for observation extraction
|
|
facts = []
|
|
for row in rows:
|
|
occurred_start = row["occurred_start"].isoformat() if row["occurred_start"] else None
|
|
facts.append(
|
|
MemoryFactForObservation(
|
|
id=str(row["id"]),
|
|
text=row["text"],
|
|
fact_type=row["fact_type"],
|
|
context=row["context"],
|
|
occurred_start=occurred_start,
|
|
)
|
|
)
|
|
|
|
# Extract observations using LLM
|
|
observations = await observation_utils.extract_observations_from_facts(llm_config, entity_name, facts)
|
|
|
|
if not observations:
|
|
return []
|
|
|
|
# Delete old observations for this entity
|
|
await conn.execute(
|
|
"""
|
|
DELETE FROM memory_units
|
|
WHERE id IN (
|
|
SELECT mu.id
|
|
FROM memory_units mu
|
|
JOIN unit_entities ue ON mu.id = ue.unit_id
|
|
WHERE mu.bank_id = $1
|
|
AND mu.fact_type = 'observation'
|
|
AND ue.entity_id = $2
|
|
)
|
|
""",
|
|
bank_id,
|
|
entity_uuid,
|
|
)
|
|
|
|
# Generate embeddings for new observations
|
|
embeddings = await embedding_utils.generate_embeddings_batch(embeddings_model, observations)
|
|
|
|
# Insert new observations
|
|
current_time = utcnow()
|
|
created_ids = []
|
|
|
|
for obs_text, embedding in zip(observations, embeddings):
|
|
result = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO memory_units (
|
|
bank_id, text, embedding, context, event_date,
|
|
occurred_start, occurred_end, mentioned_at,
|
|
fact_type, access_count
|
|
)
|
|
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'observation', 0)
|
|
RETURNING id
|
|
""",
|
|
bank_id,
|
|
obs_text,
|
|
str(embedding),
|
|
f"observation about {entity_name}",
|
|
current_time,
|
|
current_time,
|
|
current_time,
|
|
current_time,
|
|
)
|
|
obs_id = str(result["id"])
|
|
created_ids.append(obs_id)
|
|
|
|
# Link observation to entity
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO unit_entities (unit_id, entity_id)
|
|
VALUES ($1, $2)
|
|
""",
|
|
uuid.UUID(obs_id),
|
|
entity_uuid,
|
|
)
|
|
|
|
return created_ids
|