""" 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 asyncpg from sentence_transformers import SentenceTransformer from dotenv import load_dotenv import asyncio import time from concurrent.futures import ProcessPoolExecutor import numpy as np from .utils import ( extract_facts, calculate_recency_weight, calculate_frequency_weight, ) from .entity_resolver import EntityResolver def utcnow(): """Get current UTC time with timezone info.""" return datetime.now(timezone.utc) # Global process pool for parallel embedding generation # Each process loads its own copy of the embedding model # This provides TRUE parallelism for CPU-bound embedding operations _PROCESS_POOL = None _EMBEDDING_MODEL_NAME = "BAAI/bge-small-en-v1.5" # Process-local model cache (one per worker process) _worker_model = None def _get_worker_model(): """Get or load the embedding model in worker process.""" global _worker_model if _worker_model is None: _worker_model = SentenceTransformer(_EMBEDDING_MODEL_NAME) return _worker_model def _encode_batch_worker(texts: List[str]) -> List[List[float]]: """ Worker function for process pool - encodes texts to embeddings. This function runs in a separate process and loads its own model. """ model = _get_worker_model() embeddings = model.encode(texts, convert_to_numpy=True, show_progress_bar=False) return [emb.tolist() for emb in embeddings] def _get_process_pool(): """Get or create the global process pool.""" global _PROCESS_POOL if _PROCESS_POOL is None: # Use 4 worker processes for true parallelism # Adjust based on your CPU cores (each process loads ~500MB model) _PROCESS_POOL = ProcessPoolExecutor(max_workers=4) return _PROCESS_POOL 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 URL 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." ) # Connection pool (created lazily on first use) self._pool = None self._pool_lock = asyncio.Lock() # Initialize entity resolver (will be created with pool) self.entity_resolver = None # 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()})") # Background queue for access count updates (to avoid blocking searches) self._access_count_queue = asyncio.Queue() self._access_count_worker_task = None self._shutdown_event = asyncio.Event() async def _access_count_worker(self): """Background worker that processes access count updates in batches.""" pool = self._pool # Pool is guaranteed to exist when worker starts while not self._shutdown_event.is_set(): try: # Collect updates for up to 1 second or 1000 items updates = {} deadline = asyncio.get_event_loop().time() + 1.0 while len(updates) < 1000 and asyncio.get_event_loop().time() < deadline: try: # Wait for items with short timeout remaining_time = max(0.1, deadline - asyncio.get_event_loop().time()) node_ids = await asyncio.wait_for( self._access_count_queue.get(), timeout=remaining_time ) # Deduplicate by adding to set for node_id in node_ids: updates[node_id] = True except asyncio.TimeoutError: break # Process batch if we have updates if updates: node_id_list = list(updates.keys()) try: async with pool.acquire() as conn: await conn.execute( "UPDATE memory_units SET access_count = access_count + 1 WHERE id::text = ANY($1)", node_id_list ) except Exception as e: print(f"[ACCESS_COUNT_WORKER] Error updating access counts: {e}") except asyncio.CancelledError: break except Exception as e: print(f"[ACCESS_COUNT_WORKER] Unexpected error: {e}") await asyncio.sleep(1) # Backoff on error async def _get_pool(self) -> asyncpg.Pool: """Get or create the connection pool (lazy initialization).""" if self._pool is None: async with self._pool_lock: if self._pool is None: self._pool = await asyncpg.create_pool( self.db_url, min_size=2, max_size=10, command_timeout=60, statement_cache_size=0 # Disable prepared statement cache ) # Initialize entity resolver with pool if self.entity_resolver is None: self.entity_resolver = EntityResolver(self._pool) # Start access count worker (outside lock, after pool is created) if self._access_count_worker_task is None and self._pool is not None: self._access_count_worker_task = asyncio.create_task(self._access_count_worker()) return self._pool async def close(self): """Close the connection pool and shutdown background workers.""" # Signal shutdown to worker self._shutdown_event.set() # Cancel and wait for worker task if self._access_count_worker_task is not None: self._access_count_worker_task.cancel() try: await self._access_count_worker_task except asyncio.CancelledError: pass # Close pool if self._pool is not None: await self._pool.close() self._pool = None 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 in parallel. Uses a ProcessPoolExecutor to achieve TRUE parallelism for CPU-bound embedding generation. Each worker process loads its own model copy. When multiple put_async calls run in parallel, each can generate embeddings concurrently in separate processes (no GIL contention). Args: texts: List of texts to embed Returns: List of 384-dimensional embeddings in same order as input texts """ try: # Run in process pool for true parallelism loop = asyncio.get_event_loop() pool = _get_process_pool() embeddings = await loop.run_in_executor( pool, _encode_batch_worker, texts ) return embeddings except Exception as e: raise Exception(f"Failed to generate batch embeddings: {str(e)}") async def _find_duplicate_facts_batch( self, conn, 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: conn: Database connection 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 # Convert embedding list to string for asyncpg vector type embedding_str = str(embedding) result = await conn.fetchrow( """ SELECT id, text, 1 - (embedding <=> $1::vector) AS similarity FROM memory_units WHERE agent_id = $2 AND event_date BETWEEN $3 AND $4 AND 1 - (embedding <=> $1::vector) > $5 ORDER BY similarity DESC LIMIT 1 """, embedding_str, agent_id, time_lower, time_upper, similarity_threshold ) 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 is a convenience wrapper around put_batch_async for a single content item. 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 """ # Use put_batch_async with a single item (avoids code duplication) result = await self.put_batch_async( agent_id=agent_id, contents=[{ "content": content, "context": context, "event_date": event_date }] ) # Return the first (and only) list of unit IDs return result[0] if result else [] async def put_batch_async( self, agent_id: str, contents: List[Dict[str, Any]], ) -> List[List[str]]: """ Store multiple content items as memory units in ONE batch operation. This is MUCH more efficient than calling put_async multiple times: - Extracts facts from all contents in parallel - Generates ALL embeddings in ONE batch - Does ALL database operations in ONE transaction Args: agent_id: Unique identifier for the agent contents: List of dicts with keys: - "content" (required): Text content to store - "context" (optional): Context about the memory - "event_date" (optional): When the event occurred Returns: List of lists of unit IDs (one list per content item) Example: unit_ids = await memory.put_batch_async( agent_id="user123", contents=[ {"content": "Alice works at Google", "context": "conversation"}, {"content": "Bob loves Python", "context": "conversation"}, ] ) # Returns: [["unit-id-1"], ["unit-id-2"]] """ start_time = time.time() print(f"\n{'='*60}") print(f"PUT_BATCH_ASYNC START: {agent_id}") print(f"Batch size: {len(contents)} content items") print(f"{'='*60}") if not contents: return [] # Step 1: Extract facts from ALL contents in parallel step_start = time.time() # Create tasks for parallel fact extraction fact_extraction_tasks = [] for item in contents: content = item["content"] context = item.get("context", "") event_date = item.get("event_date") or utcnow() task = extract_facts(content, event_date, context) fact_extraction_tasks.append((task, event_date, context)) # Wait for all fact extractions to complete all_fact_results = await asyncio.gather(*[task for task, _, _ in fact_extraction_tasks]) # Flatten and track which facts belong to which content all_fact_texts = [] all_fact_dates = [] all_contexts = [] content_boundaries = [] # [(start_idx, end_idx), ...] current_idx = 0 for i, ((_, event_date, context), fact_dicts) in enumerate(zip(fact_extraction_tasks, all_fact_results)): start_idx = current_idx for fact_dict in fact_dicts: all_fact_texts.append(fact_dict['fact']) try: from dateutil import parser as date_parser fact_date = date_parser.isoparse(fact_dict['date']) all_fact_dates.append(fact_date) except Exception: all_fact_dates.append(event_date) all_contexts.append(context) end_idx = current_idx + len(fact_dicts) content_boundaries.append((start_idx, end_idx)) current_idx = end_idx total_facts = len(all_fact_texts) if total_facts == 0: return [[] for _ in contents] # Step 2: Generate ALL embeddings in ONE batch (HUGE speedup!) step_start = time.time() all_embeddings = await self._generate_embeddings_batch(all_fact_texts) print(f"[2] Generate embeddings (parallel): {len(all_embeddings)} embeddings in {time.time() - step_start:.3f}s") # Step 3: Process everything in ONE database transaction pool = await self._get_pool() async with pool.acquire() as conn: async with conn.transaction(): try: # Deduplication check for all facts step_start = time.time() all_is_duplicate = [] for sentence, embedding, fact_date in zip(all_fact_texts, all_embeddings, all_fact_dates): dup_flags = await self._find_duplicate_facts_batch( conn, agent_id, [sentence], [embedding], fact_date ) all_is_duplicate.extend(dup_flags) duplicates_filtered = sum(all_is_duplicate) new_facts = total_facts - duplicates_filtered print(f"[3] Deduplication check: {duplicates_filtered} duplicates filtered, {new_facts} new facts in {time.time() - step_start:.3f}s") # Filter out duplicates filtered_sentences = [s for s, is_dup in zip(all_fact_texts, all_is_duplicate) if not is_dup] filtered_embeddings = [e for e, is_dup in zip(all_embeddings, all_is_duplicate) if not is_dup] filtered_dates = [d for d, is_dup in zip(all_fact_dates, all_is_duplicate) if not is_dup] filtered_contexts = [c for c, is_dup in zip(all_contexts, all_is_duplicate) if not is_dup] if not filtered_sentences: print(f"[PUT_BATCH_ASYNC] All facts were duplicates, returning empty") return [[] for _ in contents] # Batch insert ALL units step_start = time.time() # Convert embeddings to strings for asyncpg vector type filtered_embeddings_str = [str(emb) for emb in filtered_embeddings] results = await conn.fetch( """ INSERT INTO memory_units (agent_id, text, context, embedding, event_date, access_count) SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::vector[], $5::timestamptz[], $6::integer[]) RETURNING id """, [agent_id] * len(filtered_sentences), filtered_sentences, filtered_contexts, filtered_embeddings_str, filtered_dates, [0] * len(filtered_sentences) ) created_unit_ids = [str(row['id']) for row in results] print(f"[5] Batch insert units: {len(created_unit_ids)} units in {time.time() - step_start:.3f}s") # Process entities for ALL units step_start = time.time() all_entity_links = await self._extract_entities_batch_optimized( conn, agent_id, created_unit_ids, filtered_sentences, "", filtered_dates ) print(f"[6] Extract entities (batched): {time.time() - step_start:.3f}s") # Create temporal links step_start = time.time() await self._create_temporal_links_batch_per_fact(conn, agent_id, created_unit_ids) print(f"[7] Batch create temporal links: {time.time() - step_start:.3f}s") # Create semantic links step_start = time.time() await self._create_semantic_links_batch(conn, agent_id, created_unit_ids, filtered_embeddings) print(f"[8] Batch create semantic links: {time.time() - step_start:.3f}s") # Insert entity links step_start = time.time() if all_entity_links: await self._insert_entity_links_batch(conn, all_entity_links) print(f"[9] Batch insert entity links: {time.time() - step_start:.3f}s") # Transaction auto-commits on success commit_start = time.time() print(f"[10] Commit: {time.time() - commit_start:.3f}s") # Map created unit IDs back to original content items # Account for duplicates when mapping back result_unit_ids = [] filtered_idx = 0 for start_idx, end_idx in content_boundaries: content_unit_ids = [] for i in range(start_idx, end_idx): if not all_is_duplicate[i]: content_unit_ids.append(created_unit_ids[filtered_idx]) filtered_idx += 1 result_unit_ids.append(content_unit_ids) total_time = time.time() - start_time print(f"\n{'='*60}") print(f"PUT_BATCH_ASYNC COMPLETE: {len(created_unit_ids)} units from {len(contents)} contents in {total_time:.3f}s") print(f"{'='*60}\n") return result_unit_ids except Exception as e: # Transaction auto-rolls back on exception import traceback traceback.print_exc() raise Exception(f"Failed to store batch memory: {str(e)}") def search( self, agent_id: str, query: str, thinking_budget: int = 50, top_k: int = 10, enable_trace: bool = False, weight_activation: float = 0.30, weight_semantic: float = 0.30, weight_recency: float = 0.25, weight_frequency: float = 0.15, ) -> tuple[List[Dict[str, Any]], Optional[Any]]: """ Search memories using spreading activation (synchronous wrapper). This is a synchronous wrapper around search_async() for convenience. For best performance, use search_async() directly. 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 enable_trace: If True, returns detailed SearchTrace object weight_activation: Weight for activation component (default: 0.30) weight_semantic: Weight for semantic similarity component (default: 0.30) weight_recency: Weight for recency component (default: 0.25) weight_frequency: Weight for frequency component (default: 0.15) Returns: Tuple of (results, trace) """ # Run async version synchronously return asyncio.run(self.search_async( agent_id, query, thinking_budget, top_k, enable_trace, weight_activation, weight_semantic, weight_recency, weight_frequency )) async def search_async( self, agent_id: str, query: str, thinking_budget: int = 50, top_k: int = 10, enable_trace: bool = False, weight_activation: float = 0.30, weight_semantic: float = 0.30, weight_recency: float = 0.25, weight_frequency: float = 0.15, ) -> tuple[List[Dict[str, Any]], Optional[Any]]: """ Search memories using spreading activation (ASYNC version). 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 """ # Initialize tracer if requested from .search_tracer import SearchTracer tracer = SearchTracer(query, thinking_budget, top_k) if enable_trace else None if tracer: tracer.start() pool = await self._get_pool() search_start = time.time() print(f"\n[SEARCH] Starting search for query: '{query[:50]}...' (thinking_budget={thinking_budget}, top_k={top_k})") try: # Step 1: Generate query embedding (CPU-bound, no DB needed) step_start = time.time() query_embedding = self._generate_embedding(query) step_duration = time.time() - step_start print(f" [1] Generate query embedding: {step_duration:.3f}s") if tracer: tracer.record_query_embedding(query_embedding) tracer.add_phase_metric("generate_query_embedding", step_duration) # Step 2: Find entry points (acquire connection only for this query) step_start = time.time() query_embedding_str = str(query_embedding) # Log connection acquisition conn_acquire_start = time.time() async with pool.acquire() as conn: conn_acquire_time = time.time() - conn_acquire_start if conn_acquire_time > 0.1: # Log if waiting > 100ms print(f" [2.1] Waited {conn_acquire_time:.3f}s for connection (pool busy)") entry_points = await conn.fetch( """ SELECT id, text, context, event_date, access_count, embedding, 1 - (embedding <=> $1::vector) AS similarity FROM memory_units WHERE agent_id = $2 AND embedding IS NOT NULL AND (1 - (embedding <=> $1::vector)) >= 0.5 ORDER BY embedding <=> $1::vector LIMIT 3 """, query_embedding_str, agent_id ) step_duration = time.time() - step_start print(f" [2] Find entry points: {len(entry_points)} found in {step_duration:.3f}s") if tracer: tracer.add_phase_metric("find_entry_points", step_duration, {"count": len(entry_points)}) for rank, ep in enumerate(entry_points, 1): tracer.add_entry_point( node_id=str(ep["id"]), text=ep["text"], similarity=ep["similarity"], rank=rank ) if not entry_points: print(f"[SEARCH] Complete: 0 results in {time.time() - search_start:.3f}s") if tracer: trace = tracer.finalize([]) return [], trace return [], None # Step 3: Spreading activation with budget (in-memory processing) step_start = time.time() visited = set() results = [] budget_remaining = thinking_budget # Initialize entry points with their actual similarity scores instead of 1.0 # Format: (unit, activation, is_entry, parent_node_id, link_type, link_weight) queue = [(dict(unit), unit["similarity"], True, None, None, None) for unit in entry_points] # Track substep timings calculate_weight_time = 0 query_neighbors_time = 0 process_neighbors_time = 0 # Track which nodes were visited for deferred access count update visited_node_ids = [] # Process nodes in batches for efficient neighbor querying BATCH_SIZE = 50 nodes_to_process = [] # (unit, activation, is_entry_point, parent_node_id, link_type, link_weight) while queue and budget_remaining > 0: # Collect a batch of nodes to process (in-memory, no DB) while queue and len(nodes_to_process) < BATCH_SIZE and budget_remaining > 0: current_unit, activation, is_entry_point, parent_node_id, link_type, link_weight = queue.pop(0) unit_id = str(current_unit["id"]) if unit_id not in visited: visited.add(unit_id) budget_remaining -= 1 nodes_to_process.append((current_unit, activation, is_entry_point, parent_node_id, link_type, link_weight)) visited_node_ids.append(unit_id) # Track for deferred update elif tracer: # Node already visited - prune tracer.prune_node(unit_id, "already_visited", activation) if not nodes_to_process: break # Acquire connection ONLY for neighbor queries (defer access count updates) node_ids = [str(node[0]["id"]) for node in nodes_to_process] # Log connection acquisition for batch queries batch_conn_start = time.time() async with pool.acquire() as conn: batch_conn_acquire = time.time() - batch_conn_start if batch_conn_acquire > 0.1: # Log if waiting > 100ms print(f" [3.3.1] Waited {batch_conn_acquire:.3f}s for connection (pool busy) - batch size: {len(node_ids)}") # Query neighbors for ALL nodes in batch at once (without embeddings for speed) substep_start = time.time() all_neighbors = await conn.fetch( """ SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id, mu.text, mu.context, mu.event_date, mu.access_count, mu.id as neighbor_id FROM memory_links ml JOIN memory_units mu ON ml.to_unit_id = mu.id WHERE ml.from_unit_id::text = ANY($1) AND ml.weight >= 0.1 ORDER BY ml.from_unit_id, ml.weight DESC """, node_ids ) neighbor_query_time = time.time() - substep_start if neighbor_query_time > 1.0: # Log slow neighbor queries print(f" [3.3.3] Slow NEIGHBOR query: {neighbor_query_time:.3f}s for {len(node_ids)} nodes → {len(all_neighbors)} neighbors") query_neighbors_time += neighbor_query_time # Fetch embeddings for current batch nodes (needed for weight calculation) substep_start = time.time() embeddings = await conn.fetch( "SELECT id, embedding FROM memory_units WHERE id::text = ANY($1)", node_ids ) embedding_map = {str(row["id"]): row["embedding"] for row in embeddings} fetch_embeddings_time = time.time() - substep_start if fetch_embeddings_time > 0.5: print(f" [3.3.4] Slow EMBEDDING fetch: {fetch_embeddings_time:.3f}s for {len(node_ids)} nodes") query_neighbors_time += fetch_embeddings_time # Group neighbors by from_unit_id (in-memory, no DB) substep_start = time.time() neighbors_by_node = {} for neighbor in all_neighbors: from_id = str(neighbor["from_unit_id"]) if from_id not in neighbors_by_node: neighbors_by_node[from_id] = [] neighbors_by_node[from_id].append(neighbor) # Process each node in the batch (CPU-bound, no DB) for current_unit, activation, is_entry_point, parent_node_id, parent_link_type, parent_link_weight in nodes_to_process: unit_id = str(current_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)) # Normalize frequency to [0, 1] range frequency_normalized = (frequency_weight - 1.0) / 1.0 # Calculate semantic similarity between query and this memory # Get embedding from the map we fetched memory_embedding = embedding_map.get(unit_id) if memory_embedding is not None: # Convert embedding to list of floats if it's a string or other type if isinstance(memory_embedding, str): import json memory_embedding = json.loads(memory_embedding) elif not isinstance(memory_embedding, (list, np.ndarray)): # If it's some other type, try to convert it memory_embedding = list(memory_embedding) # Cosine similarity = 1 - cosine distance query_vec = np.array(query_embedding, dtype=np.float64) memory_vec = np.array(memory_embedding, dtype=np.float64) # Cosine similarity dot_product = np.dot(query_vec, memory_vec) norm_query = np.linalg.norm(query_vec) norm_memory = np.linalg.norm(memory_vec) semantic_similarity = dot_product / (norm_query * norm_memory) if norm_query > 0 and norm_memory > 0 else 0.0 else: semantic_similarity = 0.0 # Combined weight using configurable parameters final_weight = ( weight_activation * activation + weight_semantic * semantic_similarity + weight_recency * recency_weight + weight_frequency * frequency_normalized ) # Notify tracer if tracer: tracer.visit_node( node_id=unit_id, text=current_unit["text"], context=current_unit.get("context", ""), event_date=event_date, access_count=current_unit.get("access_count", 0), is_entry_point=is_entry_point, parent_node_id=parent_node_id, link_type=parent_link_type, link_weight=parent_link_weight, activation=activation, semantic_similarity=semantic_similarity, recency=recency_weight, frequency=frequency_normalized, final_weight=final_weight, ) results.append({ "id": unit_id, "text": current_unit["text"], "context": current_unit.get("context", ""), "event_date": event_date.isoformat(), "weight": final_weight, "activation": activation, "semantic_similarity": semantic_similarity, "recency": recency_weight, "frequency": frequency_weight, }) # Spread to neighbors (from batch query results) neighbors = neighbors_by_node.get(unit_id, []) for neighbor in neighbors: neighbor_id = str(neighbor["to_unit_id"]) link_weight = neighbor["weight"] link_type = neighbor["link_type"] entity_id = str(neighbor["entity_id"]) if neighbor["entity_id"] else None new_activation = activation * link_weight * 0.8 # 0.8 = decay factor if neighbor_id not in visited: 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, unit_id, link_type, link_weight)) # parent_id, link_type, link_weight if tracer: tracer.add_neighbor_link( from_node_id=unit_id, to_node_id=neighbor_id, link_type=link_type, link_weight=link_weight, entity_id=entity_id, new_activation=new_activation, followed=True ) elif tracer: tracer.add_neighbor_link( from_node_id=unit_id, to_node_id=neighbor_id, link_type=link_type, link_weight=link_weight, entity_id=entity_id, new_activation=new_activation, followed=False, prune_reason="activation_too_low" ) calculate_weight_time += time.time() - substep_start process_neighbors_time += time.time() - substep_start # Clear batch for next iteration nodes_to_process = [] spreading_activation_time = time.time() - step_start num_batches = (len(visited) + BATCH_SIZE - 1) // BATCH_SIZE # Ceiling division print(f" [3] Spreading activation: {len(visited)} nodes visited in {spreading_activation_time:.3f}s") print(f" [3.1] Calculate weights: {calculate_weight_time:.3f}s") print(f" [3.2] Query neighbors: {query_neighbors_time:.3f}s ({num_batches} batched queries)") print(f" [3.3] Process neighbors: {process_neighbors_time:.3f}s") if tracer: tracer.add_phase_metric("spreading_activation", spreading_activation_time, { "nodes_visited": len(visited), "num_batches": num_batches }) # Step 4: Queue access count updates (background worker will process them) if visited_node_ids: await self._access_count_queue.put(visited_node_ids) print(f" [4] Queued access count updates for {len(visited_node_ids)} nodes") # Step 5: Sort by final weight and return top results step_start = time.time() results.sort(key=lambda x: x["weight"], reverse=True) top_results = results[:top_k] print(f" [5] Sort and return top {top_k}: {time.time() - step_start:.3f}s") print(f"[SEARCH] Complete: {len(top_results)} results in {time.time() - search_start:.3f}s\n") # Finalize trace if enabled if tracer: trace = tracer.finalize(top_results) return top_results, trace return top_results, None except Exception as e: print(f"[SEARCH] ERROR after {time.time() - search_start:.3f}s: {str(e)}") raise Exception(f"Failed to search memories: {str(e)}") async def delete_agent(self, agent_id: str) -> Dict[str, int]: """ Delete all data for a specific agent (multi-tenant cleanup). This is much more efficient than dropping all tables and allows multiple agents to coexist in the same database. Deletes (with CASCADE): - All memory units for this agent - All entities for this agent - All associated links, unit-entity associations, and co-occurrences Args: agent_id: Agent ID to delete Returns: Dictionary with counts of deleted items """ pool = await self._get_pool() async with pool.acquire() as conn: async with conn.transaction(): try: # Count before deletion for reporting units_count = await conn.fetchval("SELECT COUNT(*) FROM memory_units WHERE agent_id = $1", agent_id) entities_count = await conn.fetchval("SELECT COUNT(*) FROM entities WHERE agent_id = $1", agent_id) # Delete memory units (cascades to unit_entities, memory_links) await conn.execute("DELETE FROM memory_units WHERE agent_id = $1", agent_id) # Delete entities (cascades to unit_entities, entity_cooccurrences, memory_links with entity_id) await conn.execute("DELETE FROM entities WHERE agent_id = $1", agent_id) return { "memory_units_deleted": units_count, "entities_deleted": entities_count } except Exception as e: raise Exception(f"Failed to delete agent data: {str(e)}") async def _extract_entities_batch_optimized( self, conn, agent_id: str, unit_ids: List[str], sentences: List[str], context: str, fact_dates: List, ) -> List[tuple]: """ Extract entities from ALL sentences in one batch (MUCH faster than sequential). Uses spaCy's batch processing to extract entities from all texts at once, then resolves and links them in bulk. Returns list of tuples for batch insertion: (from_unit_id, to_unit_id, link_type, weight, entity_id) """ from .entity_resolver import extract_entities_batch try: # Step 1: Extract entities from ALL sentences in one batch (fast!) substep_start = time.time() all_entities = extract_entities_batch(sentences) total_entities = sum(len(ents) for ents in all_entities) print(f" [6.1] spaCy NER (batch): {total_entities} entities from {len(sentences)} sentences 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)) print(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 (most will have same date) entities_by_date = {} for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit): date_key = fact_date if date_key not in entities_by_date: entities_by_date[date_key] = [] entities_by_date[date_key].append((idx, all_entities_flat[idx])) # Resolve each date group in batch resolved_entity_ids = [None] * len(all_entities_flat) for fact_date, entities_group in entities_by_date.items(): indices = [idx for idx, _ in entities_group] entities_data = [entity_data for _, entity_data in entities_group] batch_resolved = await self.entity_resolver.resolve_entities_batch( agent_id=agent_id, entities_data=entities_data, context=context, unit_event_date=fact_date, conn=conn ) for idx, entity_id in zip(indices, batch_resolved): resolved_entity_ids[idx] = entity_id print(f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities 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 self.entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn) print(f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s") print(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 = {} print(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) # For each entity, find all units that reference it (one query per entity) entity_to_units = {} for entity_id in all_entity_ids: rows = await conn.fetch( """ SELECT unit_id FROM unit_entities WHERE entity_id = $1 """, entity_id ) entity_to_units[entity_id] = [row['unit_id'] for row in rows] # Create bidirectional links between units that share entities 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)) print(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: print(f"ERROR: Failed to extract entities in batch: {str(e)}") import traceback traceback.print_exc() # Re-raise to trigger rollback at put_async level raise async def _create_temporal_links_batch_per_fact( self, conn, agent_id: str, unit_ids: List[str], time_window_hours: int = 24, ): """ 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). """ if not unit_ids: return try: # Get the event_date for each new unit 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} # Create links based on each unit's individual event_date links = [] for unit_id, unit_event_date in new_units.items(): # Find units within the time window of THIS specific unit recent_units = await conn.fetch( """ SELECT id, event_date FROM memory_units WHERE agent_id = $1 AND id != $2 AND event_date BETWEEN $3 AND $4 ORDER BY event_date DESC LIMIT 10 """, agent_id, unit_id, unit_event_date - timedelta(hours=time_window_hours), unit_event_date + timedelta(hours=time_window_hours) ) for recent_row in recent_units: recent_id = recent_row['id'] recent_event_date = recent_row['event_date'] # Calculate temporal proximity weight time_diff_hours = abs((unit_event_date - 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)) if links: 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 e: print(f"ERROR: Failed to create temporal links: {str(e)}") import traceback traceback.print_exc() # Re-raise to trigger rollback at put_async level raise async def _create_semantic_links_batch( self, conn, 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: all_links = [] for unit_id, embedding in zip(unit_ids, embeddings): # Find similar units using vector similarity # Convert embedding to string for asyncpg embedding_str = str(embedding) similar_units = await conn.fetch( """ SELECT id, 1 - (embedding <=> $1::vector) AS similarity FROM memory_units WHERE agent_id = $2 AND id != $3 AND embedding IS NOT NULL AND (1 - (embedding <=> $1::vector)) >= $4 ORDER BY embedding <=> $1::vector LIMIT $5 """, embedding_str, agent_id, unit_id, threshold, top_k ) for row in similar_units: similar_id = row['id'] similarity = row['similarity'] all_links.append((unit_id, str(similar_id), 'semantic', float(similarity), None)) if all_links: 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 ) except Exception as e: print(f"ERROR: Failed to create semantic links: {str(e)}") import traceback traceback.print_exc() # Re-raise to trigger rollback at put_async level raise async def _insert_entity_links_batch(self, conn, links: List[tuple]): """Insert all entity links in a single batch.""" if not links: return 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 e: print(f"Warning: Failed to insert entity links: {str(e)}")