679 lines
27 KiB
Python
679 lines
27 KiB
Python
"""
|
|
Link creation utilities for temporal, semantic, and entity links.
|
|
"""
|
|
|
|
import time
|
|
import logging
|
|
from typing import List
|
|
from datetime import timedelta, datetime, timezone
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _normalize_datetime(dt):
|
|
"""Normalize datetime to be timezone-aware (UTC) for consistent comparison."""
|
|
if dt is None:
|
|
return None
|
|
if dt.tzinfo is None:
|
|
# Naive datetime - assume UTC
|
|
return dt.replace(tzinfo=timezone.utc)
|
|
return dt
|
|
|
|
|
|
def compute_temporal_links(
|
|
new_units: dict,
|
|
candidates: list,
|
|
time_window_hours: int = 24,
|
|
) -> list:
|
|
"""
|
|
Compute temporal links between new units and candidate neighbors.
|
|
|
|
This is a pure function that takes query results and returns link tuples,
|
|
making it easy to test without database access.
|
|
|
|
Args:
|
|
new_units: Dict mapping unit_id (str) to event_date (datetime)
|
|
candidates: List of dicts with 'id' and 'event_date' keys (candidate neighbors)
|
|
time_window_hours: Time window in hours for temporal links
|
|
|
|
Returns:
|
|
List of tuples: (from_unit_id, to_unit_id, 'temporal', weight, None)
|
|
"""
|
|
if not new_units:
|
|
return []
|
|
|
|
links = []
|
|
for unit_id, unit_event_date in new_units.items():
|
|
# Normalize unit_event_date for consistent comparison
|
|
unit_event_date_norm = _normalize_datetime(unit_event_date)
|
|
|
|
# Calculate time window bounds with overflow protection
|
|
try:
|
|
time_lower = unit_event_date_norm - timedelta(hours=time_window_hours)
|
|
except OverflowError:
|
|
time_lower = datetime.min.replace(tzinfo=timezone.utc)
|
|
try:
|
|
time_upper = unit_event_date_norm + timedelta(hours=time_window_hours)
|
|
except OverflowError:
|
|
time_upper = datetime.max.replace(tzinfo=timezone.utc)
|
|
|
|
# Filter candidates within this unit's time window
|
|
matching_neighbors = [
|
|
(row['id'], row['event_date'])
|
|
for row in candidates
|
|
if time_lower <= _normalize_datetime(row['event_date']) <= time_upper
|
|
][:10] # Limit to top 10
|
|
|
|
for recent_id, recent_event_date in matching_neighbors:
|
|
# Calculate temporal proximity weight
|
|
time_diff_hours = abs((unit_event_date_norm - _normalize_datetime(recent_event_date)).total_seconds() / 3600)
|
|
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
|
|
links.append((unit_id, str(recent_id), 'temporal', weight, None))
|
|
|
|
return links
|
|
|
|
|
|
def compute_temporal_query_bounds(
|
|
new_units: dict,
|
|
time_window_hours: int = 24,
|
|
) -> tuple:
|
|
"""
|
|
Compute the min/max date bounds for querying temporal neighbors.
|
|
|
|
Args:
|
|
new_units: Dict mapping unit_id (str) to event_date (datetime)
|
|
time_window_hours: Time window in hours
|
|
|
|
Returns:
|
|
Tuple of (min_date, max_date) with overflow protection
|
|
"""
|
|
if not new_units:
|
|
return None, None
|
|
|
|
# Normalize all dates to be timezone-aware to avoid comparison issues
|
|
all_dates = [_normalize_datetime(d) for d in new_units.values()]
|
|
|
|
try:
|
|
min_date = min(all_dates) - timedelta(hours=time_window_hours)
|
|
except OverflowError:
|
|
min_date = datetime.min.replace(tzinfo=timezone.utc)
|
|
|
|
try:
|
|
max_date = max(all_dates) + timedelta(hours=time_window_hours)
|
|
except OverflowError:
|
|
max_date = datetime.max.replace(tzinfo=timezone.utc)
|
|
|
|
return min_date, max_date
|
|
|
|
|
|
def _log(log_buffer, message, level='info'):
|
|
"""Helper to log to buffer if available, otherwise use logger."""
|
|
if log_buffer is not None:
|
|
log_buffer.append(message)
|
|
else:
|
|
if level == 'info':
|
|
logger.info(message)
|
|
else:
|
|
logger.log(logging.WARNING if level == 'warning' else logging.ERROR, message)
|
|
|
|
|
|
async def extract_entities_batch_optimized(
|
|
entity_resolver,
|
|
conn,
|
|
bank_id: str,
|
|
unit_ids: List[str],
|
|
sentences: List[str],
|
|
context: str,
|
|
fact_dates: List,
|
|
llm_entities: List[List[dict]],
|
|
log_buffer: List[str] = None,
|
|
) -> List[tuple]:
|
|
"""
|
|
Process LLM-extracted entities for ALL facts in batch.
|
|
|
|
Uses entities provided by the LLM (no spaCy needed), then resolves
|
|
and links them in bulk.
|
|
|
|
Args:
|
|
entity_resolver: EntityResolver instance for entity resolution
|
|
conn: Database connection
|
|
agent_id: bank IDentifier
|
|
unit_ids: List of unit IDs
|
|
sentences: List of fact sentences
|
|
context: Context string
|
|
fact_dates: List of fact dates
|
|
llm_entities: List of entity lists from LLM extraction
|
|
log_buffer: Optional buffer for logging
|
|
|
|
Returns:
|
|
List of tuples for batch insertion: (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
"""
|
|
try:
|
|
# Step 1: Convert LLM entities to the format expected by entity resolver
|
|
substep_start = time.time()
|
|
all_entities = []
|
|
for entity_list in llm_entities:
|
|
# Convert List[Entity] or List[dict] to List[Dict] format
|
|
formatted_entities = []
|
|
for ent in entity_list:
|
|
# Handle both Entity objects and dicts
|
|
if hasattr(ent, 'text'):
|
|
# Entity objects only have 'text', default type to 'CONCEPT'
|
|
formatted_entities.append({'text': ent.text, 'type': 'CONCEPT'})
|
|
elif isinstance(ent, dict):
|
|
formatted_entities.append({'text': ent.get('text', ''), 'type': ent.get('type', 'CONCEPT')})
|
|
all_entities.append(formatted_entities)
|
|
|
|
total_entities = sum(len(ents) for ents in all_entities)
|
|
_log(log_buffer, f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s")
|
|
|
|
# Step 2: Resolve entities in BATCH (much faster!)
|
|
substep_start = time.time()
|
|
step_6_2_start = time.time()
|
|
|
|
# [6.2.1] Prepare all entities for batch resolution
|
|
substep_6_2_1_start = time.time()
|
|
all_entities_flat = []
|
|
entity_to_unit = [] # Maps flat index to (unit_id, local_index)
|
|
|
|
for unit_id, entities, fact_date in zip(unit_ids, all_entities, fact_dates):
|
|
if not entities:
|
|
continue
|
|
|
|
for local_idx, entity in enumerate(entities):
|
|
all_entities_flat.append({
|
|
'text': entity['text'],
|
|
'type': entity['type'],
|
|
'nearby_entities': entities,
|
|
})
|
|
entity_to_unit.append((unit_id, local_idx, fact_date))
|
|
_log(log_buffer, f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s")
|
|
|
|
# Resolve ALL entities in one batch call
|
|
if all_entities_flat:
|
|
# [6.2.2] Batch resolve entities
|
|
substep_6_2_2_start = time.time()
|
|
# Group by date for batch resolution (round to hour to reduce buckets)
|
|
entities_by_date = {}
|
|
for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit):
|
|
# Round to hour to group facts from same time period
|
|
date_key = fact_date.replace(minute=0, second=0, microsecond=0)
|
|
if date_key not in entities_by_date:
|
|
entities_by_date[date_key] = []
|
|
entities_by_date[date_key].append((idx, all_entities_flat[idx]))
|
|
|
|
_log(log_buffer, f" [6.2.2] Grouped into {len(entities_by_date)} date buckets, resolving in parallel...")
|
|
|
|
# Resolve all date groups in PARALLEL using asyncio.gather
|
|
resolved_entity_ids = [None] * len(all_entities_flat)
|
|
|
|
# Prepare all resolution tasks
|
|
async def resolve_date_bucket(date_idx, date_key, entities_group):
|
|
date_bucket_start = time.time()
|
|
indices = [idx for idx, _ in entities_group]
|
|
entities_data = [entity_data for _, entity_data in entities_group]
|
|
# Use the first fact's date for this bucket (all should be in same hour)
|
|
fact_date = entity_to_unit[indices[0]][2]
|
|
|
|
# Pass conn=None to let each parallel task acquire its own connection
|
|
batch_resolved = await entity_resolver.resolve_entities_batch(
|
|
bank_id=bank_id,
|
|
entities_data=entities_data,
|
|
context=context,
|
|
unit_event_date=fact_date,
|
|
conn=None # Each task gets its own connection from pool
|
|
)
|
|
|
|
if len(entities_by_date) <= 10: # Only log individual buckets if there aren't too many
|
|
_log(log_buffer, f" [6.2.2.{date_idx}] Resolved {len(entities_data)} entities in {time.time() - date_bucket_start:.3f}s")
|
|
|
|
return indices, batch_resolved
|
|
|
|
# Execute all resolution tasks in parallel
|
|
import asyncio
|
|
tasks = [
|
|
resolve_date_bucket(date_idx, date_key, entities_group)
|
|
for date_idx, (date_key, entities_group) in enumerate(entities_by_date.items(), 1)
|
|
]
|
|
results = await asyncio.gather(*tasks)
|
|
|
|
# Map results back to resolved_entity_ids
|
|
for indices, batch_resolved in results:
|
|
for idx, entity_id in zip(indices, batch_resolved):
|
|
resolved_entity_ids[idx] = entity_id
|
|
|
|
_log(log_buffer, f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities across {len(entities_by_date)} buckets in {time.time() - substep_6_2_2_start:.3f}s")
|
|
|
|
# [6.2.3] Create unit-entity links in BATCH
|
|
substep_6_2_3_start = time.time()
|
|
# Map resolved entities back to units and collect all (unit, entity) pairs
|
|
unit_to_entity_ids = {}
|
|
unit_entity_pairs = []
|
|
for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit):
|
|
if unit_id not in unit_to_entity_ids:
|
|
unit_to_entity_ids[unit_id] = []
|
|
|
|
entity_id = resolved_entity_ids[idx]
|
|
unit_to_entity_ids[unit_id].append(entity_id)
|
|
unit_entity_pairs.append((unit_id, entity_id))
|
|
|
|
# Batch insert all unit-entity links (MUCH faster!)
|
|
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
|
_log(log_buffer, f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s")
|
|
|
|
_log(log_buffer, f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s")
|
|
else:
|
|
unit_to_entity_ids = {}
|
|
_log(log_buffer, f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s")
|
|
|
|
# Step 3: Create entity links between units that share entities
|
|
substep_start = time.time()
|
|
# Collect all unique entity IDs
|
|
all_entity_ids = set()
|
|
for entity_ids in unit_to_entity_ids.values():
|
|
all_entity_ids.update(entity_ids)
|
|
|
|
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...")
|
|
|
|
# Find all units that reference these entities (ONE batched query)
|
|
entity_to_units = {}
|
|
if all_entity_ids:
|
|
query_start = time.time()
|
|
import uuid
|
|
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
|
rows = await conn.fetch(
|
|
"""
|
|
SELECT entity_id, unit_id
|
|
FROM unit_entities
|
|
WHERE entity_id = ANY($1::uuid[])
|
|
""",
|
|
entity_id_list
|
|
)
|
|
_log(log_buffer, f" [6.3.1] Query unit_entities: {len(rows)} rows in {time.time() - query_start:.3f}s")
|
|
|
|
# Group by entity_id
|
|
group_start = time.time()
|
|
for row in rows:
|
|
entity_id = row['entity_id']
|
|
if entity_id not in entity_to_units:
|
|
entity_to_units[entity_id] = []
|
|
entity_to_units[entity_id].append(row['unit_id'])
|
|
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s")
|
|
|
|
# Create bidirectional links between units that share entities
|
|
link_gen_start = time.time()
|
|
links = []
|
|
for entity_id, units_with_entity in entity_to_units.items():
|
|
# For each pair of units with this entity, create bidirectional links
|
|
for i, unit_id_1 in enumerate(units_with_entity):
|
|
for unit_id_2 in units_with_entity[i+1:]:
|
|
# Bidirectional links
|
|
links.append((unit_id_1, unit_id_2, 'entity', 1.0, entity_id))
|
|
links.append((unit_id_2, unit_id_1, 'entity', 1.0, entity_id))
|
|
|
|
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s")
|
|
_log(log_buffer, f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s")
|
|
|
|
return links
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to extract entities in batch: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
|
|
async def create_temporal_links_batch_per_fact(
|
|
conn,
|
|
bank_id: str,
|
|
unit_ids: List[str],
|
|
time_window_hours: int = 24,
|
|
log_buffer: List[str] = None,
|
|
):
|
|
"""
|
|
Create temporal links for multiple units, each with their own event_date.
|
|
|
|
Queries the event_date for each unit from the database and creates temporal
|
|
links based on individual dates (supports per-fact dating).
|
|
|
|
Args:
|
|
conn: Database connection
|
|
agent_id: bank IDentifier
|
|
unit_ids: List of unit IDs
|
|
time_window_hours: Time window in hours for temporal links
|
|
log_buffer: Optional buffer for logging
|
|
"""
|
|
if not unit_ids:
|
|
return
|
|
|
|
try:
|
|
import time as time_mod
|
|
|
|
# Get the event_date for each new unit
|
|
fetch_dates_start = time_mod.time()
|
|
rows = await conn.fetch(
|
|
"""
|
|
SELECT id, event_date
|
|
FROM memory_units
|
|
WHERE id::text = ANY($1)
|
|
""",
|
|
unit_ids
|
|
)
|
|
new_units = {str(row['id']): row['event_date'] for row in rows}
|
|
_log(log_buffer, f" [7.1] Fetch event_dates for {len(unit_ids)} units: {time_mod.time() - fetch_dates_start:.3f}s")
|
|
|
|
# Fetch ALL potential temporal neighbors in ONE query (much faster!)
|
|
# Get time range across all units with overflow protection
|
|
min_date, max_date = compute_temporal_query_bounds(new_units, time_window_hours)
|
|
|
|
fetch_neighbors_start = time_mod.time()
|
|
all_candidates = await conn.fetch(
|
|
"""
|
|
SELECT id, event_date
|
|
FROM memory_units
|
|
WHERE bank_id = $1
|
|
AND event_date BETWEEN $2 AND $3
|
|
AND id::text != ALL($4)
|
|
ORDER BY event_date DESC
|
|
""",
|
|
bank_id,
|
|
min_date,
|
|
max_date,
|
|
unit_ids
|
|
)
|
|
_log(log_buffer, f" [7.2] Fetch {len(all_candidates)} candidate neighbors (1 query): {time_mod.time() - fetch_neighbors_start:.3f}s")
|
|
|
|
# Filter and create links in memory (much faster than N queries)
|
|
link_gen_start = time_mod.time()
|
|
links = compute_temporal_links(new_units, all_candidates, time_window_hours)
|
|
_log(log_buffer, f" [7.3] Generate {len(links)} temporal links: {time_mod.time() - link_gen_start:.3f}s")
|
|
|
|
if links:
|
|
insert_start = time_mod.time()
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
|
""",
|
|
links
|
|
)
|
|
_log(log_buffer, f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to create temporal links: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
|
|
async def create_semantic_links_batch(
|
|
conn,
|
|
bank_id: str,
|
|
unit_ids: List[str],
|
|
embeddings: List[List[float]],
|
|
top_k: int = 5,
|
|
threshold: float = 0.7,
|
|
log_buffer: List[str] = None,
|
|
):
|
|
"""
|
|
Create semantic links for multiple units efficiently.
|
|
|
|
For each unit, finds similar units and creates links.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
agent_id: bank IDentifier
|
|
unit_ids: List of unit IDs
|
|
embeddings: List of embedding vectors
|
|
top_k: Number of top similar units to link
|
|
threshold: Minimum similarity threshold
|
|
log_buffer: Optional buffer for logging
|
|
"""
|
|
if not unit_ids or not embeddings:
|
|
return
|
|
|
|
try:
|
|
import time as time_mod
|
|
import numpy as np
|
|
|
|
# Fetch ALL existing units with embeddings in ONE query
|
|
fetch_start = time_mod.time()
|
|
all_existing = await conn.fetch(
|
|
"""
|
|
SELECT id, embedding
|
|
FROM memory_units
|
|
WHERE bank_id = $1
|
|
AND embedding IS NOT NULL
|
|
AND id::text != ALL($2)
|
|
""",
|
|
bank_id,
|
|
unit_ids
|
|
)
|
|
_log(log_buffer, f" [8.1] Fetch {len(all_existing)} existing embeddings (1 query): {time_mod.time() - fetch_start:.3f}s")
|
|
|
|
# Convert to numpy for vectorized similarity computation
|
|
compute_start = time_mod.time()
|
|
all_links = []
|
|
|
|
if all_existing:
|
|
# Convert existing embeddings to numpy array
|
|
existing_ids = [str(row['id']) for row in all_existing]
|
|
# Stack embeddings as 2D array: (num_embeddings, embedding_dim)
|
|
embedding_arrays = []
|
|
for row in all_existing:
|
|
raw_emb = row['embedding']
|
|
# Handle different pgvector formats
|
|
if isinstance(raw_emb, str):
|
|
# Parse string format: "[1.0, 2.0, ...]"
|
|
import json
|
|
emb = np.array(json.loads(raw_emb), dtype=np.float32)
|
|
elif isinstance(raw_emb, (list, tuple)):
|
|
emb = np.array(raw_emb, dtype=np.float32)
|
|
else:
|
|
# Try direct conversion (works for numpy arrays, pgvector objects, etc.)
|
|
emb = np.array(raw_emb, dtype=np.float32)
|
|
|
|
# Ensure it's 1D
|
|
if emb.ndim != 1:
|
|
raise ValueError(f"Expected 1D embedding, got shape {emb.shape}")
|
|
embedding_arrays.append(emb)
|
|
|
|
if not embedding_arrays:
|
|
existing_embeddings = np.array([])
|
|
elif len(embedding_arrays) == 1:
|
|
# Single embedding: reshape to (1, dim)
|
|
existing_embeddings = embedding_arrays[0].reshape(1, -1)
|
|
else:
|
|
# Multiple embeddings: vstack
|
|
existing_embeddings = np.vstack(embedding_arrays)
|
|
|
|
# For each new unit, compute similarities with ALL existing units
|
|
for unit_id, new_embedding in zip(unit_ids, embeddings):
|
|
new_emb_array = np.array(new_embedding)
|
|
|
|
# Compute cosine similarities (dot product for normalized vectors)
|
|
similarities = np.dot(existing_embeddings, new_emb_array)
|
|
|
|
# Find top-k above threshold
|
|
# Get indices of similarities above threshold
|
|
above_threshold = np.where(similarities >= threshold)[0]
|
|
|
|
if len(above_threshold) > 0:
|
|
# Sort by similarity (descending) and take top-k
|
|
sorted_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
|
|
|
|
for idx in sorted_indices:
|
|
similar_id = existing_ids[idx]
|
|
similarity = float(similarities[idx])
|
|
all_links.append((unit_id, similar_id, 'semantic', similarity, None))
|
|
|
|
_log(log_buffer, f" [8.2] Compute similarities & generate {len(all_links)} semantic links: {time_mod.time() - compute_start:.3f}s")
|
|
|
|
if all_links:
|
|
insert_start = time_mod.time()
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
|
""",
|
|
all_links
|
|
)
|
|
_log(log_buffer, f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to create semantic links: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
|
|
async def insert_entity_links_batch(conn, links: List[tuple], chunk_size: int = 5000):
|
|
"""
|
|
Insert all entity links in bulk using unnest for efficiency.
|
|
|
|
Uses PostgreSQL unnest() to insert many rows in a single query,
|
|
which is much faster than executemany over high-latency connections.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
links: List of tuples (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
chunk_size: Number of rows per batch (default 5000)
|
|
"""
|
|
if not links:
|
|
return
|
|
|
|
import uuid as uuid_mod
|
|
|
|
# Process in chunks to avoid query size limits
|
|
for i in range(0, len(links), chunk_size):
|
|
chunk = links[i:i + chunk_size]
|
|
|
|
# Separate into arrays for unnest
|
|
from_ids = []
|
|
to_ids = []
|
|
link_types = []
|
|
weights = []
|
|
entity_ids = []
|
|
|
|
for from_id, to_id, link_type, weight, entity_id in chunk:
|
|
from_ids.append(uuid_mod.UUID(from_id) if isinstance(from_id, str) else from_id)
|
|
to_ids.append(uuid_mod.UUID(to_id) if isinstance(to_id, str) else to_id)
|
|
link_types.append(link_type)
|
|
weights.append(weight)
|
|
entity_ids.append(
|
|
uuid_mod.UUID(str(entity_id)) if entity_id and not isinstance(entity_id, uuid_mod.UUID)
|
|
else entity_id
|
|
)
|
|
|
|
# Use unnest to insert all rows in one query
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
SELECT * FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float[], $5::uuid[])
|
|
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
|
""",
|
|
from_ids, to_ids, link_types, weights, entity_ids
|
|
)
|
|
|
|
|
|
async def create_causal_links_batch(
|
|
conn,
|
|
unit_ids: List[str],
|
|
causal_relations_per_fact: List[List[dict]],
|
|
) -> int:
|
|
"""
|
|
Create causal links between facts based on LLM-extracted causal relationships.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
unit_ids: List of unit IDs (in same order as causal_relations_per_fact)
|
|
causal_relations_per_fact: List of causal relations for each fact.
|
|
Each element is a list of dicts with:
|
|
- target_fact_index: Index into unit_ids for the target fact
|
|
- relation_type: "causes", "caused_by", "enables", or "prevents"
|
|
- strength: Float in [0.0, 1.0] representing relationship strength
|
|
|
|
Returns:
|
|
Number of causal links created
|
|
|
|
Causal link types:
|
|
- "causes": This fact directly causes the target fact (forward causation)
|
|
- "caused_by": This fact was caused by the target fact (backward causation)
|
|
- "enables": This fact enables/allows the target fact (enablement)
|
|
- "prevents": This fact prevents/blocks the target fact (prevention)
|
|
"""
|
|
if not unit_ids or not causal_relations_per_fact:
|
|
return 0
|
|
|
|
try:
|
|
import time as time_mod
|
|
create_start = time_mod.time()
|
|
|
|
# Build links list
|
|
links = []
|
|
for fact_idx, causal_relations in enumerate(causal_relations_per_fact):
|
|
if not causal_relations:
|
|
continue
|
|
|
|
from_unit_id = unit_ids[fact_idx]
|
|
|
|
for relation in causal_relations:
|
|
target_idx = relation['target_fact_index']
|
|
relation_type = relation['relation_type']
|
|
strength = relation.get('strength', 1.0)
|
|
|
|
# Validate relation_type - must match database constraint
|
|
valid_types = {'causes', 'caused_by', 'enables', 'prevents'}
|
|
if relation_type not in valid_types:
|
|
logger.error(
|
|
f"Invalid relation_type '{relation_type}' (type: {type(relation_type).__name__}) "
|
|
f"from fact {fact_idx}. Must be one of: {valid_types}. "
|
|
f"Relation data: {relation}"
|
|
)
|
|
continue
|
|
|
|
# Validate target index
|
|
if target_idx < 0 or target_idx >= len(unit_ids):
|
|
logger.warning(f"Invalid target_fact_index {target_idx} in causal relation from fact {fact_idx}")
|
|
continue
|
|
|
|
to_unit_id = unit_ids[target_idx]
|
|
|
|
# Don't create self-links
|
|
if from_unit_id == to_unit_id:
|
|
continue
|
|
|
|
# Add the causal link
|
|
# link_type is the relation_type (e.g., "causes", "caused_by")
|
|
# weight is the strength of the relationship
|
|
links.append((from_unit_id, to_unit_id, relation_type, strength, None))
|
|
|
|
|
|
if links:
|
|
insert_start = time_mod.time()
|
|
try:
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
|
""",
|
|
links
|
|
)
|
|
except Exception as db_error:
|
|
# Log the actual data being inserted for debugging
|
|
logger.error(f"Database insert failed for causal links. Error: {db_error}")
|
|
logger.error(f"Attempted to insert {len(links)} links. First few:")
|
|
for i, link in enumerate(links[:3]):
|
|
logger.error(f" Link {i}: from={link[0]}, to={link[1]}, type='{link[2]}' (repr={repr(link[2])}), weight={link[3]}, entity={link[4]}")
|
|
raise
|
|
|
|
return len(links)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to create causal links: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
raise
|