949 lines
34 KiB
Python
949 lines
34 KiB
Python
"""
|
|
Temporal + Semantic + Entity Memory System for AI Agents.
|
|
|
|
This implements a sophisticated memory architecture that combines:
|
|
1. Temporal links: Memories connected by time proximity
|
|
2. Semantic links: Memories connected by meaning/similarity
|
|
3. Entity links: Memories connected by shared entities (PERSON, ORG, etc.)
|
|
4. Spreading activation: Search through the graph with activation decay
|
|
5. Dynamic weighting: Recency and frequency-based importance
|
|
"""
|
|
import os
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor, execute_values
|
|
from pgvector.psycopg2 import register_vector
|
|
from sentence_transformers import SentenceTransformer
|
|
from dotenv import load_dotenv
|
|
import asyncio
|
|
import time
|
|
|
|
from .utils import (
|
|
extract_facts,
|
|
calculate_recency_weight,
|
|
calculate_frequency_weight,
|
|
)
|
|
from .entity_resolver import EntityResolver, extract_entities
|
|
from .coref_resolver import resolve_sentences
|
|
|
|
|
|
def utcnow():
|
|
"""Get current UTC time with timezone info."""
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
class TemporalSemanticMemory:
|
|
"""
|
|
Advanced memory system using temporal and semantic linking with PostgreSQL.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
db_url: Optional[str] = None,
|
|
embedding_model: str = "BAAI/bge-small-en-v1.5",
|
|
):
|
|
"""
|
|
Initialize the temporal + semantic memory system.
|
|
|
|
Args:
|
|
db_url: PostgreSQL connection URL (postgresql://user:pass@host:port/dbname)
|
|
embedding_model: Name of the SentenceTransformer model to use
|
|
"""
|
|
load_dotenv()
|
|
|
|
# Initialize PostgreSQL connection
|
|
self.db_url = db_url or os.getenv("DATABASE_URL")
|
|
if not self.db_url:
|
|
raise ValueError(
|
|
"Database URL not found. "
|
|
"Set DATABASE_URL environment variable."
|
|
)
|
|
|
|
self.conn = psycopg2.connect(self.db_url)
|
|
register_vector(self.conn)
|
|
|
|
# Initialize entity resolver
|
|
self.entity_resolver = EntityResolver(self.conn)
|
|
|
|
# Initialize local embedding model (384 dimensions)
|
|
print(f"Loading embedding model: {embedding_model}...")
|
|
self.embedding_model = SentenceTransformer(embedding_model)
|
|
print(f"✓ Model loaded (embedding dim: {self.embedding_model.get_sentence_embedding_dimension()})")
|
|
|
|
def __del__(self):
|
|
"""Clean up database connection."""
|
|
if hasattr(self, 'conn') and self.conn:
|
|
self.conn.close()
|
|
|
|
def _generate_embedding(self, text: str) -> List[float]:
|
|
"""
|
|
Generate embedding for text using local SentenceTransformer model.
|
|
|
|
Args:
|
|
text: Text to embed
|
|
|
|
Returns:
|
|
384-dimensional embedding vector (bge-small-en-v1.5)
|
|
"""
|
|
try:
|
|
embedding = self.embedding_model.encode(text, convert_to_numpy=True, show_progress_bar=False)
|
|
return embedding.tolist()
|
|
except Exception as e:
|
|
raise Exception(f"Failed to generate embedding: {str(e)}")
|
|
|
|
async def _generate_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
|
|
"""
|
|
Generate embeddings for multiple texts using local model (batch processing).
|
|
|
|
Local models are fast and process batches efficiently without needing
|
|
parallel API calls. We run this in asyncio to avoid blocking, but the
|
|
actual embedding generation is synchronous.
|
|
|
|
Args:
|
|
texts: List of texts to embed
|
|
|
|
Returns:
|
|
List of 384-dimensional embeddings in same order as input texts
|
|
"""
|
|
try:
|
|
# Run in thread pool to avoid blocking event loop
|
|
loop = asyncio.get_event_loop()
|
|
embeddings = await loop.run_in_executor(
|
|
None,
|
|
lambda: self.embedding_model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
|
)
|
|
return [emb.tolist() for emb in embeddings]
|
|
except Exception as e:
|
|
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
|
|
|
def _find_duplicate_facts_batch(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
texts: List[str],
|
|
embeddings: List[List[float]],
|
|
event_date: datetime,
|
|
time_window_hours: int = 24,
|
|
similarity_threshold: float = 0.95
|
|
) -> List[bool]:
|
|
"""
|
|
Check which facts are duplicates using semantic similarity + temporal window.
|
|
|
|
For each new fact, checks if a semantically similar fact already exists
|
|
within the time window. Uses pgvector cosine similarity for efficiency.
|
|
|
|
Args:
|
|
cursor: Database cursor
|
|
agent_id: Agent identifier
|
|
texts: List of fact texts to check
|
|
embeddings: Corresponding embeddings
|
|
event_date: Event date for temporal filtering
|
|
time_window_hours: Hours before/after event_date to search (default: 24)
|
|
similarity_threshold: Minimum cosine similarity to consider duplicate (default: 0.95)
|
|
|
|
Returns:
|
|
List of booleans - True if fact is a duplicate (should skip), False if new
|
|
"""
|
|
is_duplicate = []
|
|
|
|
time_lower = event_date - timedelta(hours=time_window_hours)
|
|
time_upper = event_date + timedelta(hours=time_window_hours)
|
|
|
|
for text, embedding in zip(texts, embeddings):
|
|
# Query for similar facts within time window
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, text, 1 - (embedding <=> %s::vector) AS similarity
|
|
FROM memory_units
|
|
WHERE agent_id = %s
|
|
AND event_date BETWEEN %s AND %s
|
|
AND 1 - (embedding <=> %s::vector) > %s
|
|
ORDER BY similarity DESC
|
|
LIMIT 1
|
|
""",
|
|
(embedding, agent_id, time_lower, time_upper, embedding, similarity_threshold)
|
|
)
|
|
|
|
result = cursor.fetchone()
|
|
if result:
|
|
is_duplicate.append(True)
|
|
else:
|
|
is_duplicate.append(False)
|
|
|
|
return is_duplicate
|
|
|
|
def put(
|
|
self,
|
|
agent_id: str,
|
|
content: str,
|
|
context: str = "",
|
|
event_date: Optional[datetime] = None,
|
|
) -> List[str]:
|
|
"""
|
|
Store content as memory units (synchronous wrapper).
|
|
|
|
This is a synchronous wrapper around put_async() for convenience.
|
|
For best performance, use put_async() directly.
|
|
|
|
Args:
|
|
agent_id: Unique identifier for the agent
|
|
content: Text content to store
|
|
context: Context about when/why this memory was formed
|
|
event_date: When the event occurred (defaults to now)
|
|
|
|
Returns:
|
|
List of created unit IDs
|
|
"""
|
|
# Run async version synchronously
|
|
return asyncio.run(self.put_async(agent_id, content, context, event_date))
|
|
|
|
async def put_async(
|
|
self,
|
|
agent_id: str,
|
|
content: str,
|
|
context: str = "",
|
|
event_date: Optional[datetime] = None,
|
|
) -> List[str]:
|
|
"""
|
|
Store content as memory units with temporal and semantic links (ASYNC version).
|
|
|
|
This async version generates ALL embeddings in parallel for maximum speed,
|
|
then uses batch inserts for database operations.
|
|
|
|
Steps:
|
|
1. Split content into sentence units
|
|
2. Resolve coreferences
|
|
3. **Generate ALL embeddings in parallel** (FAST!)
|
|
4. **Batch insert all units and links** (FAST!)
|
|
|
|
Args:
|
|
agent_id: Unique identifier for the agent
|
|
content: Text content to store
|
|
context: Context about when/why this memory was formed
|
|
event_date: When the event occurred (defaults to now)
|
|
|
|
Returns:
|
|
List of created unit IDs
|
|
"""
|
|
start_time = time.time()
|
|
print(f"\n{'='*60}")
|
|
print(f"PUT_ASYNC START: {agent_id}")
|
|
print(f"Content length: {len(content)} chars")
|
|
print(f"{'='*60}")
|
|
|
|
if event_date is None:
|
|
event_date = utcnow()
|
|
|
|
# Step 1: Extract semantic facts using LLM (async)
|
|
step_start = time.time()
|
|
try:
|
|
facts = await extract_facts(content)
|
|
print(f"[1] Extract facts: {len(facts)} facts in {time.time() - step_start:.3f}s")
|
|
except Exception as e:
|
|
print(f"\n{'='*60}")
|
|
print(f"PUT_ASYNC FAILED: Fact extraction error")
|
|
print(f"Error: {e}")
|
|
print(f"{'='*60}\n")
|
|
raise Exception(f"Failed to extract facts from content: {e}")
|
|
|
|
# Step 2: Resolve pronouns to make facts even more self-contained
|
|
step_start = time.time()
|
|
sentences = resolve_sentences(facts)
|
|
print(f"[2] Resolve coreferences: {time.time() - step_start:.3f}s")
|
|
|
|
# Step 3: Generate ALL embeddings in parallel
|
|
step_start = time.time()
|
|
embeddings = await self._generate_embeddings_batch(sentences)
|
|
print(f"[3] Generate embeddings (parallel): {len(embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
|
|
|
# Step 4: Check for duplicates using similarity + temporal window
|
|
cursor = self.conn.cursor()
|
|
step_start = time.time()
|
|
duplicate_flags = self._find_duplicate_facts_batch(
|
|
cursor, agent_id, sentences, embeddings, event_date
|
|
)
|
|
num_duplicates = sum(duplicate_flags)
|
|
|
|
# Filter out duplicates
|
|
filtered_data = [
|
|
(sentence, embedding)
|
|
for sentence, embedding, is_dup in zip(sentences, embeddings, duplicate_flags)
|
|
if not is_dup
|
|
]
|
|
|
|
if filtered_data:
|
|
sentences, embeddings = zip(*filtered_data)
|
|
sentences = list(sentences)
|
|
embeddings = list(embeddings)
|
|
else:
|
|
sentences = []
|
|
embeddings = []
|
|
|
|
print(f"[4] Deduplication check: {num_duplicates} duplicates filtered, {len(sentences)} new facts in {time.time() - step_start:.3f}s")
|
|
|
|
# If all facts were duplicates, return empty list
|
|
if not sentences:
|
|
cursor.close()
|
|
print(f"\n{'='*60}")
|
|
print(f"PUT_ASYNC COMPLETE: All facts were duplicates, nothing stored")
|
|
print(f"{'='*60}\n")
|
|
return []
|
|
|
|
# Step 5: Batch insert everything
|
|
try:
|
|
# Batch INSERT all memory units
|
|
step_start = time.time()
|
|
from psycopg2.extras import execute_values
|
|
unit_data = [
|
|
(agent_id, sentence, embedding, context, event_date, 0)
|
|
for sentence, embedding in zip(sentences, embeddings)
|
|
]
|
|
|
|
unit_ids = execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_units (agent_id, text, embedding, context, event_date, access_count)
|
|
VALUES %s
|
|
RETURNING id
|
|
""",
|
|
unit_data,
|
|
fetch=True
|
|
)
|
|
created_unit_ids = [str(row[0]) for row in unit_ids]
|
|
print(f"[5] Batch insert units: {time.time() - step_start:.3f}s")
|
|
|
|
# Process entities for all units
|
|
step_start = time.time()
|
|
all_entity_links = []
|
|
for unit_id, sentence in zip(created_unit_ids, sentences):
|
|
entity_links = self._extract_entities_for_batch(cursor, agent_id, unit_id, sentence, context, event_date, sentences)
|
|
all_entity_links.extend(entity_links)
|
|
print(f"[6] Extract entities: {time.time() - step_start:.3f}s")
|
|
|
|
# Create ALL temporal links in batch
|
|
step_start = time.time()
|
|
self._create_temporal_links_batch(cursor, agent_id, created_unit_ids, event_date)
|
|
print(f"[7] Batch create temporal links: {time.time() - step_start:.3f}s")
|
|
|
|
# Create ALL semantic links in batch
|
|
step_start = time.time()
|
|
self._create_semantic_links_batch(cursor, agent_id, created_unit_ids, embeddings)
|
|
print(f"[8] Batch create semantic links: {time.time() - step_start:.3f}s")
|
|
|
|
# Insert all entity links in batch
|
|
step_start = time.time()
|
|
if all_entity_links:
|
|
self._insert_entity_links_batch(cursor, all_entity_links)
|
|
print(f"[9] Batch insert entity links: {time.time() - step_start:.3f}s")
|
|
|
|
commit_start = time.time()
|
|
self.conn.commit()
|
|
print(f"[10] Commit: {time.time() - commit_start:.3f}s")
|
|
|
|
total_time = time.time() - start_time
|
|
print(f"\n{'='*60}")
|
|
print(f"PUT_ASYNC COMPLETE: {len(created_unit_ids)} units stored in {total_time:.3f}s")
|
|
print(f"{'='*60}\n")
|
|
|
|
return created_unit_ids
|
|
|
|
except Exception as e:
|
|
self.conn.rollback()
|
|
raise Exception(f"Failed to store memory: {str(e)}")
|
|
finally:
|
|
cursor.close()
|
|
|
|
def _create_temporal_links(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
unit_id: str,
|
|
event_date: datetime,
|
|
time_window_hours: int = 24,
|
|
):
|
|
"""
|
|
Create temporal links to recent memories.
|
|
|
|
Links this unit to other units that occurred within a time window.
|
|
|
|
Args:
|
|
cursor: Database cursor
|
|
agent_id: Agent ID
|
|
unit_id: ID of the current unit
|
|
event_date: When this event occurred
|
|
time_window_hours: Size of the temporal window
|
|
"""
|
|
try:
|
|
# Get recent units within time window
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, event_date
|
|
FROM memory_units
|
|
WHERE agent_id = %s
|
|
AND id != %s
|
|
AND event_date >= %s
|
|
ORDER BY event_date DESC
|
|
LIMIT 10
|
|
""",
|
|
(agent_id, unit_id, event_date - timedelta(hours=time_window_hours))
|
|
)
|
|
|
|
recent_units = cursor.fetchall()
|
|
|
|
# Create links to recent units
|
|
links = []
|
|
for recent_id, recent_event_date in recent_units:
|
|
# Calculate temporal proximity weight
|
|
time_diff_hours = abs((event_date - recent_event_date).total_seconds() / 3600)
|
|
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
|
|
|
|
links.append((unit_id, recent_id, 'temporal', weight, None))
|
|
|
|
if links:
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES %s
|
|
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 e:
|
|
print(f"Warning: Failed to create temporal links: {str(e)}")
|
|
|
|
def _create_semantic_links(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
unit_id: str,
|
|
embedding: List[float],
|
|
top_k: int = 5,
|
|
threshold: float = 0.7,
|
|
):
|
|
"""
|
|
Create semantic links to similar memories.
|
|
|
|
Links this unit to other units with similar meaning.
|
|
|
|
Args:
|
|
cursor: Database cursor
|
|
agent_id: Agent ID
|
|
unit_id: ID of the current unit
|
|
embedding: Embedding of the current unit
|
|
top_k: Number of similar units to link to
|
|
threshold: Minimum similarity threshold
|
|
"""
|
|
try:
|
|
# Find similar units using vector similarity
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, 1 - (embedding <=> %s::vector) AS similarity
|
|
FROM memory_units
|
|
WHERE agent_id = %s
|
|
AND id != %s
|
|
AND embedding IS NOT NULL
|
|
AND (1 - (embedding <=> %s::vector)) >= %s
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT %s
|
|
""",
|
|
(embedding, agent_id, unit_id, embedding, threshold, embedding, top_k)
|
|
)
|
|
|
|
similar_units = cursor.fetchall()
|
|
|
|
# Create links to similar units
|
|
links = []
|
|
for similar_id, similarity in similar_units:
|
|
links.append((unit_id, similar_id, 'semantic', float(similarity), None))
|
|
|
|
if links:
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES %s
|
|
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 e:
|
|
print(f"Warning: Failed to create semantic links: {str(e)}")
|
|
|
|
def _extract_and_link_entities(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
unit_id: str,
|
|
text: str,
|
|
context: str,
|
|
event_date,
|
|
all_sentences: List[str],
|
|
):
|
|
"""
|
|
Extract entities from text, resolve them, and create entity links.
|
|
|
|
Args:
|
|
cursor: Database cursor
|
|
agent_id: Agent ID
|
|
unit_id: Current unit ID
|
|
text: Unit text
|
|
context: Context
|
|
event_date: When created
|
|
all_sentences: All sentences from the same PUT (for context)
|
|
"""
|
|
try:
|
|
# Extract entities from this unit
|
|
entities = extract_entities(text)
|
|
|
|
if not entities:
|
|
return
|
|
|
|
# Resolve each entity and link
|
|
entity_ids = []
|
|
for entity in entities:
|
|
entity_id = self.entity_resolver.resolve_entity(
|
|
agent_id=agent_id,
|
|
entity_text=entity['text'],
|
|
entity_type=entity['type'],
|
|
context=context,
|
|
nearby_entities=entities,
|
|
unit_event_date=event_date
|
|
)
|
|
entity_ids.append(entity_id)
|
|
|
|
# Link unit to entity
|
|
self.entity_resolver.link_unit_to_entity(unit_id, entity_id)
|
|
|
|
# Create entity links to other units that mention the same entities
|
|
for entity_id in set(entity_ids):
|
|
# Get other units that mention this entity
|
|
related_units = self.entity_resolver.get_units_by_entity(entity_id, limit=50)
|
|
|
|
# Create entity links
|
|
links = []
|
|
for related_unit_id in related_units:
|
|
if str(related_unit_id) != str(unit_id):
|
|
links.append((unit_id, related_unit_id, 'entity', 1.0, entity_id))
|
|
|
|
if links:
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES %s
|
|
ON CONFLICT DO NOTHING
|
|
""",
|
|
links
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to extract/link entities: {str(e)}")
|
|
|
|
def search(
|
|
self,
|
|
agent_id: str,
|
|
query: str,
|
|
thinking_budget: int = 50,
|
|
top_k: int = 10,
|
|
live_tracer=None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search memories using spreading activation.
|
|
|
|
This implements the core SEARCH operation:
|
|
1. Find entry points (most relevant units via vector search)
|
|
2. Spread activation through the graph
|
|
3. Weight results by activation + recency + frequency
|
|
4. Return top results
|
|
|
|
Args:
|
|
agent_id: Agent ID to search for
|
|
query: Search query
|
|
thinking_budget: How many units to explore (computational budget)
|
|
top_k: Number of results to return
|
|
live_tracer: Optional LiveSearchTracer for visualization
|
|
|
|
Returns:
|
|
List of memory units with their weights, sorted by relevance
|
|
"""
|
|
cursor = self.conn.cursor(cursor_factory=RealDictCursor)
|
|
|
|
try:
|
|
# Step 1: Generate query embedding
|
|
query_embedding = self._generate_embedding(query)
|
|
|
|
# Step 2: Find entry points
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, text, context, event_date, access_count,
|
|
1 - (embedding <=> %s::vector) AS similarity
|
|
FROM memory_units
|
|
WHERE agent_id = %s
|
|
AND embedding IS NOT NULL
|
|
AND (1 - (embedding <=> %s::vector)) >= 0.5
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT 3
|
|
""",
|
|
(query_embedding, agent_id, query_embedding, query_embedding)
|
|
)
|
|
|
|
entry_points = cursor.fetchall()
|
|
if not entry_points:
|
|
return []
|
|
|
|
# Step 3: Spreading activation with budget
|
|
visited = set()
|
|
results = []
|
|
budget_remaining = thinking_budget
|
|
queue = [(dict(unit), 1.0, True) for unit in entry_points] # (unit, activation, is_entry)
|
|
|
|
while queue and budget_remaining > 0:
|
|
current_unit, activation, is_entry_point = queue.pop(0)
|
|
unit_id = str(current_unit["id"])
|
|
|
|
if unit_id in visited:
|
|
continue
|
|
|
|
visited.add(unit_id)
|
|
budget_remaining -= 1
|
|
|
|
# Increment access count
|
|
cursor.execute(
|
|
"UPDATE memory_units SET access_count = access_count + 1 WHERE id = %s",
|
|
(unit_id,)
|
|
)
|
|
|
|
# Calculate combined weight
|
|
event_date = current_unit["event_date"]
|
|
days_since = (utcnow() - event_date).total_seconds() / 86400
|
|
|
|
recency_weight = calculate_recency_weight(days_since)
|
|
frequency_weight = calculate_frequency_weight(current_unit.get("access_count", 0))
|
|
|
|
# Combined weight: activation * recency * frequency
|
|
final_weight = activation * recency_weight * frequency_weight
|
|
|
|
# Notify tracer
|
|
if live_tracer:
|
|
live_tracer.visit_node(
|
|
node_id=unit_id,
|
|
text=current_unit["text"],
|
|
activation=activation,
|
|
recency=recency_weight,
|
|
frequency=frequency_weight,
|
|
weight=final_weight,
|
|
is_entry_point=is_entry_point,
|
|
)
|
|
import time
|
|
time.sleep(0.15) # Slow down for visualization
|
|
|
|
results.append({
|
|
"id": unit_id,
|
|
"text": current_unit["text"],
|
|
"context": current_unit.get("context", ""),
|
|
"event_date": event_date.isoformat(),
|
|
"weight": final_weight,
|
|
"activation": activation,
|
|
"recency": recency_weight,
|
|
"frequency": frequency_weight,
|
|
})
|
|
|
|
# Spread to neighbors
|
|
cursor.execute(
|
|
"""
|
|
SELECT ml.to_unit_id, ml.weight, mu.text, mu.context, mu.event_date, mu.access_count
|
|
FROM memory_links ml
|
|
JOIN memory_units mu ON ml.to_unit_id = mu.id
|
|
WHERE ml.from_unit_id = %s
|
|
AND ml.weight >= 0.1
|
|
ORDER BY ml.weight DESC
|
|
""",
|
|
(unit_id,)
|
|
)
|
|
|
|
neighbors = cursor.fetchall()
|
|
for neighbor in neighbors:
|
|
neighbor_id = str(neighbor["to_unit_id"])
|
|
if neighbor_id not in visited:
|
|
link_weight = neighbor["weight"]
|
|
new_activation = activation * link_weight * 0.8 # 0.8 = decay factor
|
|
|
|
if new_activation > 0.1:
|
|
queue.append(({
|
|
"id": neighbor["to_unit_id"],
|
|
"text": neighbor["text"],
|
|
"context": neighbor.get("context", ""),
|
|
"event_date": neighbor["event_date"],
|
|
"access_count": neighbor["access_count"],
|
|
}, new_activation, False)) # Not an entry point
|
|
|
|
self.conn.commit()
|
|
|
|
# Step 4: Sort by final weight and return top results
|
|
results.sort(key=lambda x: x["weight"], reverse=True)
|
|
return results[:top_k]
|
|
|
|
except Exception as e:
|
|
self.conn.rollback()
|
|
raise Exception(f"Failed to search memories: {str(e)}")
|
|
finally:
|
|
cursor.close()
|
|
|
|
def get_memory_graph_data(self, agent_id: str = None) -> Tuple[List[Dict], List[Dict]]:
|
|
"""
|
|
Get memory graph data for visualization.
|
|
|
|
Args:
|
|
agent_id: Optional agent ID (if None, returns all data)
|
|
|
|
Returns:
|
|
Tuple of (units, links) for visualization
|
|
"""
|
|
cursor = self.conn.cursor(cursor_factory=RealDictCursor)
|
|
|
|
try:
|
|
# Get all units (optionally filtered by agent)
|
|
if agent_id:
|
|
cursor.execute(
|
|
"SELECT id, text, context, event_date, access_count FROM memory_units WHERE agent_id = %s",
|
|
(agent_id,)
|
|
)
|
|
else:
|
|
cursor.execute(
|
|
"SELECT id, text, context, event_date, access_count FROM memory_units"
|
|
)
|
|
units = [dict(row) for row in cursor.fetchall()]
|
|
|
|
# Get all links (optionally filtered by agent)
|
|
if agent_id:
|
|
cursor.execute(
|
|
"""
|
|
SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight
|
|
FROM memory_links ml
|
|
JOIN memory_units mu1 ON ml.from_unit_id = mu1.id
|
|
JOIN memory_units mu2 ON ml.to_unit_id = mu2.id
|
|
WHERE mu1.agent_id = %s
|
|
""",
|
|
(agent_id,)
|
|
)
|
|
else:
|
|
cursor.execute(
|
|
"SELECT from_unit_id, to_unit_id, link_type, weight FROM memory_links"
|
|
)
|
|
links = [dict(row) for row in cursor.fetchall()]
|
|
|
|
return units, links
|
|
|
|
except Exception as e:
|
|
raise Exception(f"Failed to get memory graph data: {str(e)}")
|
|
finally:
|
|
cursor.close()
|
|
|
|
def _extract_entities_for_batch(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
unit_id: str,
|
|
text: str,
|
|
context: str,
|
|
event_date,
|
|
all_sentences: List[str],
|
|
) -> List[tuple]:
|
|
"""
|
|
Extract entities and return entity links (doesn't insert yet).
|
|
|
|
Returns list of tuples for batch insertion: (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
"""
|
|
from .entity_resolver import extract_entities
|
|
|
|
try:
|
|
# Extract entities from this unit
|
|
entities = extract_entities(text)
|
|
|
|
if not entities:
|
|
return []
|
|
|
|
# Resolve each entity
|
|
entity_ids = []
|
|
for entity in entities:
|
|
entity_id = self.entity_resolver.resolve_entity(
|
|
agent_id=agent_id,
|
|
entity_text=entity['text'],
|
|
entity_type=entity['type'],
|
|
context=context,
|
|
nearby_entities=entities,
|
|
unit_event_date=event_date
|
|
)
|
|
entity_ids.append(entity_id)
|
|
|
|
# Link unit to entity (this inserts into entity_units)
|
|
self.entity_resolver.link_unit_to_entity(unit_id, entity_id)
|
|
|
|
# Now collect entity links for batch insertion
|
|
# After link_unit_to_entity has been called, entity_units should exist
|
|
links = []
|
|
for entity_id in set(entity_ids):
|
|
# Find all other units with this entity (cursor must be fresh)
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
SELECT unit_id
|
|
FROM unit_entities
|
|
WHERE entity_id = %s AND unit_id != %s
|
|
""",
|
|
(entity_id, unit_id)
|
|
)
|
|
|
|
related_units = cursor.fetchall()
|
|
for (related_unit_id,) in related_units:
|
|
# Bidirectional links
|
|
links.append((unit_id, related_unit_id, 'entity', 1.0, entity_id))
|
|
links.append((related_unit_id, unit_id, 'entity', 1.0, entity_id))
|
|
except Exception as query_error:
|
|
# If there's an error querying, just skip this entity
|
|
print(f"Warning: Failed to query entity_units for {entity_id}: {str(query_error)}")
|
|
continue
|
|
|
|
return links
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to extract entities: {str(e)}")
|
|
return []
|
|
|
|
def _create_temporal_links_batch(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
unit_ids: List[str],
|
|
event_date: datetime,
|
|
time_window_hours: int = 24,
|
|
):
|
|
"""
|
|
Create temporal links for multiple units in one batch query.
|
|
|
|
Uses a single query to find all relevant temporal connections.
|
|
"""
|
|
if not unit_ids:
|
|
return
|
|
|
|
try:
|
|
from psycopg2.extras import execute_values
|
|
|
|
# Get ALL recent units within time window (single query)
|
|
# Cast string IDs to UUIDs for comparison
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, event_date
|
|
FROM memory_units
|
|
WHERE agent_id = %s
|
|
AND id::text != ALL(%s)
|
|
AND event_date >= %s
|
|
ORDER BY event_date DESC
|
|
""",
|
|
(agent_id, unit_ids, event_date - timedelta(hours=time_window_hours))
|
|
)
|
|
|
|
recent_units = cursor.fetchall()
|
|
|
|
# Create links from each new unit to all recent units
|
|
links = []
|
|
for unit_id in unit_ids:
|
|
for recent_id, recent_event_date in recent_units:
|
|
# Calculate temporal proximity weight
|
|
time_diff_hours = abs((event_date - recent_event_date).total_seconds() / 3600)
|
|
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
|
|
links.append((unit_id, recent_id, 'temporal', weight, None))
|
|
|
|
if links:
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES %s
|
|
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 e:
|
|
print(f"Warning: Failed to create temporal links: {str(e)}")
|
|
|
|
def _create_semantic_links_batch(
|
|
self,
|
|
cursor,
|
|
agent_id: str,
|
|
unit_ids: List[str],
|
|
embeddings: List[List[float]],
|
|
top_k: int = 5,
|
|
threshold: float = 0.7,
|
|
):
|
|
"""
|
|
Create semantic links for multiple units efficiently.
|
|
|
|
For each unit, finds similar units and creates links.
|
|
"""
|
|
if not unit_ids or not embeddings:
|
|
return
|
|
|
|
try:
|
|
from psycopg2.extras import execute_values
|
|
|
|
all_links = []
|
|
|
|
for unit_id, embedding in zip(unit_ids, embeddings):
|
|
# Find similar units using vector similarity
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, 1 - (embedding <=> %s::vector) AS similarity
|
|
FROM memory_units
|
|
WHERE agent_id = %s
|
|
AND id != %s
|
|
AND embedding IS NOT NULL
|
|
AND (1 - (embedding <=> %s::vector)) >= %s
|
|
ORDER BY embedding <=> %s::vector
|
|
LIMIT %s
|
|
""",
|
|
(embedding, agent_id, unit_id, embedding, threshold, embedding, top_k)
|
|
)
|
|
|
|
similar_units = cursor.fetchall()
|
|
|
|
for similar_id, similarity in similar_units:
|
|
all_links.append((unit_id, similar_id, 'semantic', float(similarity), None))
|
|
|
|
if all_links:
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES %s
|
|
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
|
""",
|
|
all_links
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to create semantic links: {str(e)}")
|
|
|
|
def _insert_entity_links_batch(self, cursor, links: List[tuple]):
|
|
"""Insert all entity links in a single batch."""
|
|
if not links:
|
|
return
|
|
|
|
try:
|
|
from psycopg2.extras import execute_values
|
|
execute_values(
|
|
cursor,
|
|
"""
|
|
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
|
VALUES %s
|
|
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 e:
|
|
print(f"Warning: Failed to insert entity links: {str(e)}")
|