2033 lines
88 KiB
Python
2033 lines
88 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 asyncpg
|
|
from dotenv import load_dotenv
|
|
import asyncio
|
|
from .embeddings import Embeddings, SentenceTransformersEmbeddings
|
|
import time
|
|
import numpy as np
|
|
import uuid
|
|
import logging
|
|
|
|
from .utils import (
|
|
extract_facts,
|
|
calculate_recency_weight,
|
|
calculate_frequency_weight,
|
|
)
|
|
from .entity_resolver import EntityResolver
|
|
from .operations import EmbeddingOperationsMixin, LinkOperationsMixin
|
|
|
|
|
|
def utcnow():
|
|
"""Get current UTC time with timezone info."""
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
# Logger for memory system
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TemporalSemanticMemory(
|
|
EmbeddingOperationsMixin,
|
|
LinkOperationsMixin,
|
|
):
|
|
"""
|
|
Advanced memory system using temporal and semantic linking with PostgreSQL.
|
|
|
|
Uses mixin architecture for code organization:
|
|
- EmbeddingOperationsMixin: Embedding generation
|
|
- LinkOperationsMixin: Entity, temporal, and semantic link creation
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
db_url: Optional[str] = None,
|
|
embeddings: Optional[Embeddings] = None,
|
|
embedding_model: Optional[str] = None,
|
|
):
|
|
"""
|
|
Initialize the temporal + semantic memory system.
|
|
|
|
Args:
|
|
db_url: PostgreSQL connection URL (postgresql://user:pass@host:port/dbname)
|
|
embeddings: Embeddings implementation to use. If not provided, uses SentenceTransformersEmbeddings
|
|
embedding_model: (Deprecated) Name of the SentenceTransformer model to use. Use embeddings parameter instead.
|
|
"""
|
|
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 (will be created in initialize())
|
|
self._pool = None
|
|
self._initialized = False
|
|
|
|
# Initialize entity resolver (will be created in initialize())
|
|
self.entity_resolver = None
|
|
|
|
# Initialize embeddings
|
|
if embeddings is not None:
|
|
self.embeddings = embeddings
|
|
else:
|
|
# Default to SentenceTransformersEmbeddings
|
|
model_name = embedding_model or "BAAI/bge-small-en-v1.5"
|
|
self.embeddings = SentenceTransformersEmbeddings(model_name)
|
|
|
|
# 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:
|
|
# Convert string UUIDs to UUID type for faster matching
|
|
uuid_list = [uuid.UUID(nid) for nid in node_id_list]
|
|
async with pool.acquire() as conn:
|
|
await conn.execute(
|
|
"UPDATE memory_units SET access_count = access_count + 1 WHERE id = ANY($1::uuid[])",
|
|
uuid_list
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Access count worker: Error updating access counts: {e}")
|
|
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Access count worker: Unexpected error: {e}")
|
|
await asyncio.sleep(1) # Backoff on error
|
|
|
|
async def initialize(self):
|
|
"""Initialize the connection pool and background workers."""
|
|
if self._initialized:
|
|
return
|
|
|
|
# Create connection pool
|
|
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
|
|
self.entity_resolver = EntityResolver(self._pool)
|
|
|
|
# Start access count worker
|
|
self._access_count_worker_task = asyncio.create_task(self._access_count_worker())
|
|
|
|
self._initialized = True
|
|
logger.info("Memory system initialized (pool and workers started)")
|
|
|
|
async def _get_pool(self) -> asyncpg.Pool:
|
|
"""Get the connection pool (must call initialize() first)."""
|
|
if not self._initialized:
|
|
await self.initialize()
|
|
return self._pool
|
|
|
|
async def close(self):
|
|
"""Close the connection pool and shutdown background workers."""
|
|
logger.info("close() started")
|
|
|
|
# Signal shutdown to worker
|
|
self._shutdown_event.set()
|
|
logger.info("shutdown event set")
|
|
|
|
# Cancel and wait for worker task
|
|
if self._access_count_worker_task is not None:
|
|
logger.debug("cancelling worker task")
|
|
self._access_count_worker_task.cancel()
|
|
try:
|
|
logger.debug("waiting for worker task to finish")
|
|
await self._access_count_worker_task
|
|
logger.debug("worker task finished")
|
|
except asyncio.CancelledError:
|
|
logger.debug("worker task cancelled successfully")
|
|
else:
|
|
logger.debug("no worker task to cancel")
|
|
|
|
# Close pool
|
|
if self._pool is not None:
|
|
logger.debug("closing connection pool")
|
|
self._pool.terminate()
|
|
logger.debug("connection pool closed")
|
|
self._pool = None
|
|
else:
|
|
logger.debug("no pool to close")
|
|
|
|
logger.debug("close() completed")
|
|
|
|
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
|
|
"""
|
|
if not texts:
|
|
return []
|
|
|
|
time_lower = event_date - timedelta(hours=time_window_hours)
|
|
time_upper = event_date + timedelta(hours=time_window_hours)
|
|
|
|
# Fetch ALL existing facts in time window ONCE (much faster than N queries)
|
|
import time as time_mod
|
|
fetch_start = time_mod.time()
|
|
existing_facts = await conn.fetch(
|
|
"""
|
|
SELECT id, text, embedding
|
|
FROM memory_units
|
|
WHERE agent_id = $1
|
|
AND event_date BETWEEN $2 AND $3
|
|
""",
|
|
agent_id, time_lower, time_upper
|
|
)
|
|
logger.debug(f" [3.X] Fetched {len(existing_facts)} existing facts in {time_mod.time() - fetch_start:.3f}s")
|
|
|
|
# If no existing facts, nothing is duplicate
|
|
if not existing_facts:
|
|
return [False] * len(texts)
|
|
|
|
# Compute similarities in Python (vectorized with numpy)
|
|
import numpy as np
|
|
is_duplicate = []
|
|
|
|
# Convert existing embeddings to numpy for faster computation
|
|
embedding_arrays = []
|
|
for row in existing_facts:
|
|
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
|
|
emb = np.array(raw_emb, dtype=np.float32)
|
|
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)
|
|
|
|
comp_start = time_mod.time()
|
|
for embedding in embeddings:
|
|
# Compute cosine similarity with all existing facts
|
|
emb_array = np.array(embedding)
|
|
# Cosine similarity = 1 - cosine distance
|
|
# For normalized vectors: cosine_sim = dot product
|
|
similarities = np.dot(existing_embeddings, emb_array)
|
|
|
|
# Check if any existing fact is too similar
|
|
max_similarity = np.max(similarities) if len(similarities) > 0 else 0
|
|
is_duplicate.append(max_similarity > similarity_threshold)
|
|
|
|
logger.debug(f" [3.X] Computed {len(texts)} x {len(existing_facts)} similarities in {time_mod.time() - comp_start:.3f}s")
|
|
|
|
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,
|
|
document_id: Optional[str] = None,
|
|
document_metadata: Optional[Dict[str, Any]] = None,
|
|
upsert: bool = False,
|
|
fact_type_override: Optional[str] = None,
|
|
confidence_score: Optional[float] = 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)
|
|
document_id: Optional document ID for tracking and upsert
|
|
document_metadata: Optional metadata about the document
|
|
upsert: If True and document_id exists, delete old units and create new ones
|
|
fact_type_override: Override fact type ('world', 'agent', 'opinion')
|
|
confidence_score: Confidence score for opinions (0.0 to 1.0)
|
|
|
|
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
|
|
}],
|
|
document_id=document_id,
|
|
document_metadata=document_metadata,
|
|
upsert=upsert,
|
|
fact_type_override=fact_type_override,
|
|
confidence_score=confidence_score
|
|
)
|
|
|
|
# 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]],
|
|
document_id: Optional[str] = None,
|
|
document_metadata: Optional[Dict[str, Any]] = None,
|
|
upsert: bool = False,
|
|
fact_type_override: Optional[str] = None,
|
|
confidence_score: Optional[float] = None,
|
|
) -> 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
|
|
document_id: Optional document ID for tracking and upsert
|
|
document_metadata: Optional metadata about the document
|
|
upsert: If True and document_id exists, delete old units and create new ones
|
|
fact_type_override: Override fact type for all facts ('world', 'agent', 'opinion')
|
|
confidence_score: Confidence score for opinions (0.0 to 1.0)
|
|
|
|
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"},
|
|
],
|
|
document_id="meeting-2024-01-15",
|
|
upsert=True
|
|
)
|
|
# Returns: [["unit-id-1"], ["unit-id-2"]]
|
|
"""
|
|
start_time = time.time()
|
|
logger.info(f"\n{'='*60}")
|
|
logger.info(f"PUT_BATCH_ASYNC START: {agent_id}")
|
|
logger.info(f"Batch size: {len(contents)} content items")
|
|
logger.info(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])
|
|
logger.info(f"[1] Extract facts (parallel): {len(fact_extraction_tasks)} contents in {time.time() - step_start:.3f}s")
|
|
|
|
# Flatten and track which facts belong to which content
|
|
all_fact_texts = []
|
|
all_fact_dates = []
|
|
all_contexts = []
|
|
all_fact_entities = [] # NEW: Store LLM-extracted entities per fact
|
|
all_fact_types = [] # Store fact type (world or agent)
|
|
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)
|
|
# Extract entities from fact (default to empty list if not present)
|
|
all_fact_entities.append(fact_dict.get('entities', []))
|
|
# Extract fact type (use override if provided, else use extracted type or default to 'world')
|
|
if fact_type_override:
|
|
all_fact_types.append(fact_type_override)
|
|
else:
|
|
all_fact_types.append(fact_dict.get('fact_type', 'world'))
|
|
|
|
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)
|
|
logger.info(f"[2] Generate embeddings (parallel): {len(all_embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
|
|
|
# Step 3: Process everything in ONE database transaction
|
|
logger.debug("Getting connection pool")
|
|
pool = await self._get_pool()
|
|
logger.debug("Acquiring connection from pool")
|
|
async with pool.acquire() as conn:
|
|
logger.debug("Starting transaction")
|
|
async with conn.transaction():
|
|
logger.debug("Inside transaction")
|
|
try:
|
|
# Handle document tracking and upsert
|
|
if document_id:
|
|
logger.debug(f"Handling document tracking for {document_id}")
|
|
import hashlib
|
|
import json
|
|
|
|
# Calculate content hash from all content items
|
|
combined_content = "\n".join([c.get("content", "") for c in contents])
|
|
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
|
|
|
# If upsert, delete old document first (cascades to units and links)
|
|
if upsert:
|
|
deleted = await conn.fetchval(
|
|
"DELETE FROM documents WHERE id = $1 AND agent_id = $2 RETURNING id",
|
|
document_id, agent_id
|
|
)
|
|
if deleted:
|
|
logger.debug(f"[3.1] Upsert: Deleted existing document '{document_id}' and all its units")
|
|
|
|
# Insert or update document
|
|
# Always use ON CONFLICT for idempotent behavior
|
|
await conn.execute(
|
|
"""
|
|
INSERT INTO documents (id, agent_id, original_text, content_hash, metadata)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
ON CONFLICT (id, agent_id) DO UPDATE
|
|
SET original_text = EXCLUDED.original_text,
|
|
content_hash = EXCLUDED.content_hash,
|
|
metadata = EXCLUDED.metadata,
|
|
updated_at = NOW()
|
|
""",
|
|
document_id,
|
|
agent_id,
|
|
combined_content,
|
|
content_hash,
|
|
json.dumps(document_metadata or {})
|
|
)
|
|
logger.debug(f"[3.2] Document '{document_id}' stored/updated")
|
|
|
|
# Deduplication check for all facts (batched by time window)
|
|
logger.debug("Starting deduplication check")
|
|
step_start = time.time()
|
|
|
|
# Group facts by event_date (rounded to 12-hour buckets) for batching
|
|
from collections import defaultdict
|
|
time_buckets = defaultdict(list)
|
|
for idx, (sentence, embedding, fact_date) in enumerate(zip(all_fact_texts, all_embeddings, all_fact_dates)):
|
|
# Round to 12-hour bucket to group similar times
|
|
bucket_key = fact_date.replace(hour=(fact_date.hour // 12) * 12, minute=0, second=0, microsecond=0)
|
|
time_buckets[bucket_key].append((idx, sentence, embedding, fact_date))
|
|
|
|
# Process each bucket in batch
|
|
all_is_duplicate = [False] * total_facts # Initialize all as not duplicate
|
|
for bucket_date, bucket_items in time_buckets.items():
|
|
indices = [item[0] for item in bucket_items]
|
|
sentences = [item[1] for item in bucket_items]
|
|
embeddings = [item[2] for item in bucket_items]
|
|
# Use bucket_date as representative for time window
|
|
dup_flags = await self._find_duplicate_facts_batch(
|
|
conn, agent_id, sentences, embeddings, bucket_date, time_window_hours=24
|
|
)
|
|
# Map results back to original indices
|
|
for idx, is_dup in zip(indices, dup_flags):
|
|
all_is_duplicate[idx] = is_dup
|
|
|
|
duplicates_filtered = sum(all_is_duplicate)
|
|
new_facts = total_facts - duplicates_filtered
|
|
logger.debug(f"Deduplication complete: {duplicates_filtered} duplicates filtered, {new_facts} new facts ({len(time_buckets)} time buckets)")
|
|
logger.info(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]
|
|
filtered_entities = [ents for ents, is_dup in zip(all_fact_entities, all_is_duplicate) if not is_dup]
|
|
filtered_fact_types = [ft for ft, is_dup in zip(all_fact_types, all_is_duplicate) if not is_dup]
|
|
|
|
if not filtered_sentences:
|
|
logger.debug(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]
|
|
# Prepare confidence scores (only for opinions)
|
|
confidence_scores = [confidence_score if ft == 'opinion' else None for ft in filtered_fact_types]
|
|
results = await conn.fetch(
|
|
"""
|
|
INSERT INTO memory_units (agent_id, document_id, text, context, embedding, event_date, fact_type, confidence_score, access_count)
|
|
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::vector[], $6::timestamptz[], $7::text[], $8::float[], $9::integer[])
|
|
RETURNING id
|
|
""",
|
|
[agent_id] * len(filtered_sentences),
|
|
[document_id] * len(filtered_sentences) if document_id else [None] * len(filtered_sentences),
|
|
filtered_sentences,
|
|
filtered_contexts,
|
|
filtered_embeddings_str,
|
|
filtered_dates,
|
|
filtered_fact_types,
|
|
confidence_scores,
|
|
[0] * len(filtered_sentences)
|
|
)
|
|
|
|
created_unit_ids = [str(row['id']) for row in results]
|
|
logger.debug(f"Batch insert complete: {len(created_unit_ids)} units created")
|
|
logger.info(f"[5] Batch insert units: {len(created_unit_ids)} units in {time.time() - step_start:.3f}s")
|
|
|
|
# Process entities for ALL units
|
|
logger.debug("Processing entities")
|
|
step_start = time.time()
|
|
all_entity_links = await self._extract_entities_batch_optimized(
|
|
conn, agent_id, created_unit_ids, filtered_sentences, "", filtered_dates, filtered_entities
|
|
)
|
|
logger.debug(f"Entity processing complete: {len(all_entity_links)} links")
|
|
logger.info(f"[6] Process entities (batched): {time.time() - step_start:.3f}s")
|
|
|
|
# Create temporal links
|
|
logger.debug("Creating temporal links")
|
|
step_start = time.time()
|
|
await self._create_temporal_links_batch_per_fact(conn, agent_id, created_unit_ids)
|
|
logger.debug("Temporal links complete")
|
|
logger.info(f"[7] Batch create temporal links: {time.time() - step_start:.3f}s")
|
|
|
|
# Create semantic links
|
|
logger.debug("Creating semantic links")
|
|
step_start = time.time()
|
|
await self._create_semantic_links_batch(conn, agent_id, created_unit_ids, filtered_embeddings)
|
|
logger.debug("Semantic links complete")
|
|
logger.info(f"[8] Batch create semantic links: {time.time() - step_start:.3f}s")
|
|
|
|
# Insert entity links
|
|
logger.debug("Inserting entity links")
|
|
step_start = time.time()
|
|
if all_entity_links:
|
|
await self._insert_entity_links_batch(conn, all_entity_links)
|
|
logger.debug("Entity links inserted")
|
|
logger.info(f"[9] Batch insert entity links: {time.time() - step_start:.3f}s")
|
|
|
|
# Transaction auto-commits on success
|
|
commit_start = time.time()
|
|
logger.debug(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
|
|
logger.info(f"\n{'='*60}")
|
|
logger.info(f"PUT_BATCH_ASYNC COMPLETE: {len(created_unit_ids)} units from {len(contents)} contents in {total_time:.3f}s")
|
|
logger.info(f"{'='*60}\n")
|
|
|
|
# Trigger opinion reinforcement in background (non-blocking)
|
|
# Only trigger if there are entities in the new units
|
|
if any(filtered_entities):
|
|
asyncio.create_task(
|
|
self._reinforce_opinions_async(
|
|
agent_id=agent_id,
|
|
created_unit_ids=created_unit_ids,
|
|
unit_texts=filtered_sentences,
|
|
unit_entities=filtered_entities
|
|
)
|
|
)
|
|
logger.debug("[PUT_BATCH_ASYNC] Opinion reinforcement task queued in background")
|
|
|
|
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,
|
|
mmr_lambda: float = 0.5,
|
|
fact_type: Optional[str] = None,
|
|
) -> 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)
|
|
mmr_lambda: Lambda for MMR diversification (0=max diversity, 1=no diversity, default: 0.5)
|
|
fact_type: Optional filter for fact type ('world' or 'agent')
|
|
|
|
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, mmr_lambda, fact_type
|
|
))
|
|
|
|
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,
|
|
mmr_lambda: float = 0.5,
|
|
fact_type: Optional[str] = None,
|
|
) -> 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()
|
|
|
|
# Buffer logs for clean output in concurrent scenarios
|
|
search_id = f"{agent_id[:8]}-{int(time.time() * 1000) % 100000}"
|
|
log_buffer = []
|
|
log_buffer.append(f"[SEARCH {search_id}] Query: '{query[:50]}...' (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
|
|
log_buffer.append(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
|
|
log_buffer.append(f" [2.1] Waited {conn_acquire_time:.3f}s for connection (pool busy)")
|
|
|
|
# Build entry point query with optional fact_type filter
|
|
if fact_type:
|
|
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 fact_type = $3
|
|
AND (1 - (embedding <=> $1::vector)) >= 0.5
|
|
ORDER BY embedding <=> $1::vector
|
|
LIMIT 3
|
|
""",
|
|
query_embedding_str, agent_id, fact_type
|
|
)
|
|
else:
|
|
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
|
|
log_buffer.append(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:
|
|
logger.debug(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
|
|
log_buffer.append(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)
|
|
# Convert string UUIDs to UUID type for faster matching
|
|
substep_start = time.time()
|
|
uuid_array = [uuid.UUID(nid) for nid in node_ids]
|
|
|
|
# Build neighbor query with optional fact_type filter
|
|
if fact_type:
|
|
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 = ANY($1::uuid[])
|
|
AND ml.weight >= 0.1
|
|
AND mu.fact_type = $2
|
|
ORDER BY ml.from_unit_id, ml.weight DESC
|
|
""",
|
|
uuid_array, fact_type
|
|
)
|
|
else:
|
|
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 = ANY($1::uuid[])
|
|
AND ml.weight >= 0.1
|
|
ORDER BY ml.from_unit_id, ml.weight DESC
|
|
""",
|
|
uuid_array
|
|
)
|
|
neighbor_query_time = time.time() - substep_start
|
|
if neighbor_query_time > 1.0: # Log slow neighbor queries
|
|
log_buffer.append(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 = ANY($1::uuid[])",
|
|
uuid_array
|
|
)
|
|
embedding_map = {str(row["id"]): row["embedding"] for row in embeddings}
|
|
fetch_embeddings_time = time.time() - substep_start
|
|
if fetch_embeddings_time > 0.5:
|
|
log_buffer.append(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,
|
|
"embedding": memory_embedding, # Store for MMR
|
|
})
|
|
|
|
# Spread to neighbors (from batch query results)
|
|
neighbors = neighbors_by_node.get(unit_id, [])
|
|
|
|
# Group neighbors by to_unit_id to handle multiple connections
|
|
neighbors_grouped = {}
|
|
for neighbor in neighbors:
|
|
neighbor_id = str(neighbor["to_unit_id"])
|
|
if neighbor_id not in neighbors_grouped:
|
|
neighbors_grouped[neighbor_id] = []
|
|
neighbors_grouped[neighbor_id].append(neighbor)
|
|
|
|
# Process each unique neighbor (aggregating multiple links)
|
|
for neighbor_id, neighbor_links in neighbors_grouped.items():
|
|
if neighbor_id in visited:
|
|
continue
|
|
|
|
# Sort links by weight descending to identify primary link
|
|
neighbor_links_sorted = sorted(neighbor_links, key=lambda x: x["weight"], reverse=True)
|
|
primary_link = neighbor_links_sorted[0]
|
|
|
|
# Aggregate link weights: max + 30% bonus for additional links
|
|
max_weight = primary_link["weight"]
|
|
bonus_weight = sum(link["weight"] for link in neighbor_links_sorted[1:]) * 0.3
|
|
combined_weight = max_weight + bonus_weight
|
|
|
|
# Calculate new activation using combined weight
|
|
new_activation = activation * combined_weight * 0.8 # 0.8 = decay factor
|
|
|
|
# Use primary link metadata for queue and trace
|
|
primary_link_type = primary_link["link_type"]
|
|
primary_entity_id = str(primary_link["entity_id"]) if primary_link["entity_id"] else None
|
|
|
|
if new_activation > 0.1:
|
|
queue.append(({
|
|
"id": primary_link["to_unit_id"],
|
|
"text": primary_link["text"],
|
|
"context": primary_link.get("context", ""),
|
|
"event_date": primary_link["event_date"],
|
|
"access_count": primary_link["access_count"],
|
|
}, new_activation, False, unit_id, primary_link_type, combined_weight)) # parent_id, link_type, combined_weight
|
|
|
|
# Record all links in trace (primary + additional)
|
|
if tracer:
|
|
# Add primary link with combined activation
|
|
tracer.add_neighbor_link(
|
|
from_node_id=unit_id,
|
|
to_node_id=neighbor_id,
|
|
link_type=primary_link_type,
|
|
link_weight=combined_weight,
|
|
entity_id=primary_entity_id,
|
|
new_activation=new_activation,
|
|
followed=True
|
|
)
|
|
|
|
# Add additional links as supplementary (if multiple connections exist)
|
|
for additional_link in neighbor_links_sorted[1:]:
|
|
additional_link_type = additional_link["link_type"]
|
|
additional_entity_id = str(additional_link["entity_id"]) if additional_link["entity_id"] else None
|
|
tracer.add_neighbor_link(
|
|
from_node_id=unit_id,
|
|
to_node_id=neighbor_id,
|
|
link_type=additional_link_type,
|
|
link_weight=additional_link["weight"],
|
|
entity_id=additional_entity_id,
|
|
new_activation=None, # Don't show activation for supplementary links
|
|
followed=True,
|
|
is_supplementary=True # Mark as supplementary link
|
|
)
|
|
elif tracer:
|
|
# Record pruned link
|
|
tracer.add_neighbor_link(
|
|
from_node_id=unit_id,
|
|
to_node_id=neighbor_id,
|
|
link_type=primary_link_type,
|
|
link_weight=combined_weight,
|
|
entity_id=primary_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
|
|
log_buffer.append(f" [3] Spreading activation: {len(visited)} nodes visited in {spreading_activation_time:.3f}s")
|
|
log_buffer.append(f" [3.1] Calculate weights: {calculate_weight_time:.3f}s")
|
|
log_buffer.append(f" [3.2] Query neighbors: {query_neighbors_time:.3f}s ({num_batches} batched queries)")
|
|
log_buffer.append(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)
|
|
log_buffer.append(f" [4] Queued access count updates for {len(visited_node_ids)} nodes")
|
|
|
|
# Step 5: Sort by final weight and apply MMR for diversity
|
|
step_start = time.time()
|
|
results.sort(key=lambda x: x["weight"], reverse=True)
|
|
|
|
# Apply MMR (Maximal Marginal Relevance) for diversity if lambda < 1.0
|
|
if mmr_lambda < 1.0 and len(results) > top_k:
|
|
top_results = self._apply_mmr(results, top_k, mmr_lambda, log_buffer)
|
|
log_buffer.append(f" [5] MMR diversification (λ={mmr_lambda}): {time.time() - step_start:.3f}s")
|
|
else:
|
|
top_results = results[:top_k]
|
|
# Add original rank and remove embeddings from results
|
|
for idx, result in enumerate(top_results):
|
|
result["original_rank"] = idx + 1
|
|
result["mmr_score"] = None
|
|
result["mmr_relevance"] = None
|
|
result["mmr_max_similarity"] = None
|
|
result["mmr_diversified"] = False
|
|
result.pop("embedding", None)
|
|
log_buffer.append(f" [5] Sort and return top {top_k} (no MMR): {time.time() - step_start:.3f}s")
|
|
|
|
total_time = time.time() - search_start
|
|
log_buffer.append(f"[SEARCH {search_id}] Complete: {len(top_results)} results in {total_time:.3f}s")
|
|
|
|
# Log all buffered logs at once
|
|
logger.info("\n" + "\n".join(log_buffer))
|
|
|
|
# Finalize trace if enabled
|
|
if tracer:
|
|
trace = tracer.finalize(top_results)
|
|
return top_results, trace
|
|
return top_results, None
|
|
|
|
except Exception as e:
|
|
log_buffer.append(f"[SEARCH {search_id}] ERROR after {time.time() - search_start:.3f}s: {str(e)}")
|
|
logger.error("\n" + "\n".join(log_buffer))
|
|
raise Exception(f"Failed to search memories: {str(e)}")
|
|
|
|
def _apply_mmr(
|
|
self,
|
|
results: List[Dict[str, Any]],
|
|
top_k: int,
|
|
mmr_lambda: float,
|
|
log_buffer: List[str]
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Apply Maximal Marginal Relevance (MMR) to diversify search results.
|
|
|
|
MMR balances relevance with diversity by selecting results that are:
|
|
1. Relevant to the query (high score)
|
|
2. Different from already selected results (low similarity)
|
|
|
|
Formula: MMR = λ * relevance - (1-λ) * max_similarity_to_selected
|
|
|
|
Args:
|
|
results: Sorted list of all results with embeddings
|
|
top_k: Number of results to select
|
|
mmr_lambda: Balance parameter (0=max diversity, 1=max relevance)
|
|
log_buffer: Logging buffer
|
|
|
|
Returns:
|
|
Diversified list of top_k results
|
|
"""
|
|
if not results or top_k <= 0:
|
|
return []
|
|
|
|
# Normalize weights to [0, 1] for fair comparison with similarity
|
|
max_weight = max(r["weight"] for r in results)
|
|
min_weight = min(r["weight"] for r in results)
|
|
weight_range = max_weight - min_weight if max_weight > min_weight else 1.0
|
|
|
|
# Pre-compute normalized relevance scores for all results
|
|
for idx, result in enumerate(results):
|
|
result["original_rank"] = idx + 1
|
|
result["normalized_relevance"] = (result["weight"] - min_weight) / weight_range
|
|
|
|
# Extract embeddings as a numpy array for vectorized operations
|
|
# Shape: (num_results, embedding_dim)
|
|
embeddings_list = []
|
|
valid_indices = []
|
|
for idx, result in enumerate(results):
|
|
if result.get("embedding") is not None:
|
|
embeddings_list.append(result["embedding"])
|
|
valid_indices.append(idx)
|
|
|
|
if not embeddings_list:
|
|
# No embeddings available, just return top-k by relevance
|
|
return results[:top_k]
|
|
|
|
# Stack embeddings into a matrix (num_results, embedding_dim)
|
|
embeddings_matrix = np.array(embeddings_list, dtype=np.float32)
|
|
|
|
# Normalize embeddings for faster cosine similarity (just dot product after normalization)
|
|
norms = np.linalg.norm(embeddings_matrix, axis=1, keepdims=True)
|
|
norms[norms == 0] = 1.0 # Avoid division by zero
|
|
embeddings_matrix = embeddings_matrix / norms
|
|
|
|
selected_indices = []
|
|
remaining_indices = list(range(len(results)))
|
|
diversified_count = 0
|
|
|
|
for selection_round in range(min(top_k, len(results))):
|
|
if not remaining_indices:
|
|
break
|
|
|
|
best_mmr_score = float('-inf')
|
|
best_remaining_idx = 0
|
|
|
|
# Vectorized computation for all remaining candidates
|
|
for remaining_idx, candidate_idx in enumerate(remaining_indices):
|
|
candidate = results[candidate_idx]
|
|
normalized_relevance = candidate["normalized_relevance"]
|
|
|
|
# Calculate max similarity to selected results
|
|
max_similarity = 0.0
|
|
if selected_indices and candidate_idx in valid_indices:
|
|
# Find position in embeddings_matrix
|
|
embedding_idx = valid_indices.index(candidate_idx)
|
|
candidate_embedding = embeddings_matrix[embedding_idx]
|
|
|
|
# Vectorized similarity calculation with all selected embeddings
|
|
if selected_indices:
|
|
selected_embedding_indices = [valid_indices.index(idx) for idx in selected_indices if idx in valid_indices]
|
|
if selected_embedding_indices:
|
|
selected_embeddings = embeddings_matrix[selected_embedding_indices]
|
|
# Compute cosine similarities in one operation (already normalized, so just dot product)
|
|
similarities = np.dot(selected_embeddings, candidate_embedding)
|
|
max_similarity = float(np.max(similarities))
|
|
|
|
# MMR score: balance relevance and diversity
|
|
mmr_score = mmr_lambda * normalized_relevance - (1 - mmr_lambda) * max_similarity
|
|
|
|
if mmr_score > best_mmr_score:
|
|
best_mmr_score = mmr_score
|
|
best_remaining_idx = remaining_idx
|
|
best_max_similarity = max_similarity
|
|
|
|
# Select the best candidate
|
|
best_candidate_idx = remaining_indices.pop(best_remaining_idx)
|
|
best_candidate = results[best_candidate_idx]
|
|
|
|
# Store MMR metadata
|
|
best_candidate["mmr_score"] = best_mmr_score
|
|
best_candidate["mmr_relevance"] = best_candidate["normalized_relevance"]
|
|
best_candidate["mmr_max_similarity"] = best_max_similarity
|
|
best_candidate["mmr_diversified"] = best_remaining_idx > 0
|
|
|
|
selected_indices.append(best_candidate_idx)
|
|
|
|
if best_remaining_idx > 0:
|
|
diversified_count += 1
|
|
|
|
log_buffer.append(f" MMR: Selected {len(selected_indices)} results, {diversified_count} diversified picks")
|
|
|
|
# Return selected results in order
|
|
selected_results = [results[idx] for idx in selected_indices]
|
|
|
|
# Remove embeddings from final results (not needed in response)
|
|
for result in selected_results:
|
|
result.pop("embedding", None)
|
|
result.pop("normalized_relevance", None) # Clean up temp field
|
|
|
|
return selected_results
|
|
|
|
async def get_document(self, document_id: str, agent_id: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Retrieve document metadata and statistics.
|
|
|
|
Args:
|
|
document_id: Document ID to retrieve
|
|
agent_id: Agent ID that owns the document
|
|
|
|
Returns:
|
|
Dictionary with document info or None if not found
|
|
"""
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
doc = await conn.fetchrow(
|
|
"""
|
|
SELECT d.id, d.agent_id, d.original_text, d.content_hash, d.metadata,
|
|
d.created_at, d.updated_at, COUNT(mu.id) as unit_count
|
|
FROM documents d
|
|
LEFT JOIN memory_units mu ON mu.document_id = d.id
|
|
WHERE d.id = $1 AND d.agent_id = $2
|
|
GROUP BY d.id, d.agent_id, d.original_text, d.content_hash, d.metadata, d.created_at, d.updated_at
|
|
""",
|
|
document_id, agent_id
|
|
)
|
|
|
|
if not doc:
|
|
return None
|
|
|
|
import json
|
|
return {
|
|
"id": doc["id"],
|
|
"agent_id": doc["agent_id"],
|
|
"original_text": doc["original_text"],
|
|
"content_hash": doc["content_hash"],
|
|
"metadata": json.loads(doc["metadata"]) if doc["metadata"] else {},
|
|
"unit_count": doc["unit_count"],
|
|
"created_at": doc["created_at"],
|
|
"updated_at": doc["updated_at"]
|
|
}
|
|
|
|
async def delete_document(self, document_id: str, agent_id: str) -> Dict[str, int]:
|
|
"""
|
|
Delete a document and all its associated memory units and links.
|
|
|
|
Args:
|
|
document_id: Document ID to delete
|
|
agent_id: Agent ID that owns the document
|
|
|
|
Returns:
|
|
Dictionary with counts of deleted items
|
|
"""
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
async with conn.transaction():
|
|
# Count units before deletion
|
|
units_count = await conn.fetchval(
|
|
"SELECT COUNT(*) FROM memory_units WHERE document_id = $1",
|
|
document_id
|
|
)
|
|
|
|
# Delete document (cascades to memory_units and all their links)
|
|
deleted = await conn.fetchval(
|
|
"DELETE FROM documents WHERE id = $1 AND agent_id = $2 RETURNING id",
|
|
document_id, agent_id
|
|
)
|
|
|
|
return {
|
|
"document_deleted": 1 if deleted else 0,
|
|
"memory_units_deleted": units_count if deleted else 0
|
|
}
|
|
|
|
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 list_agents(self) -> List[str]:
|
|
"""
|
|
Get list of all agent IDs in the database.
|
|
|
|
Returns:
|
|
List of agent IDs
|
|
"""
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
# Get distinct agent IDs from memory_units
|
|
agents = await conn.fetch("""
|
|
SELECT DISTINCT agent_id
|
|
FROM memory_units
|
|
WHERE agent_id IS NOT NULL
|
|
ORDER BY agent_id
|
|
""")
|
|
|
|
return [row['agent_id'] for row in agents]
|
|
|
|
async def get_graph_data(self, agent_id: Optional[str] = None, fact_type: Optional[str] = None):
|
|
"""
|
|
Get graph data for visualization.
|
|
|
|
Args:
|
|
agent_id: Filter by agent ID
|
|
fact_type: Filter by fact type (world, agent, opinion)
|
|
|
|
Returns:
|
|
Dict with nodes, edges, and table_rows
|
|
"""
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
# Get memory units, optionally filtered by agent_id and fact_type
|
|
query_conditions = []
|
|
query_params = []
|
|
param_count = 0
|
|
|
|
if agent_id:
|
|
param_count += 1
|
|
query_conditions.append(f"agent_id = ${param_count}")
|
|
query_params.append(agent_id)
|
|
|
|
if fact_type:
|
|
param_count += 1
|
|
query_conditions.append(f"fact_type = ${param_count}")
|
|
query_params.append(fact_type)
|
|
|
|
where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""
|
|
|
|
units = await conn.fetch(f"""
|
|
SELECT id, text, event_date, context
|
|
FROM memory_units
|
|
{where_clause}
|
|
ORDER BY event_date
|
|
""", *query_params)
|
|
|
|
# Get links, filtering to only include links between units of the selected agent
|
|
unit_ids = [row['id'] for row in units]
|
|
if unit_ids:
|
|
links = await conn.fetch("""
|
|
SELECT
|
|
ml.from_unit_id,
|
|
ml.to_unit_id,
|
|
ml.link_type,
|
|
ml.weight,
|
|
e.canonical_name as entity_name
|
|
FROM memory_links ml
|
|
LEFT JOIN entities e ON ml.entity_id = e.id
|
|
WHERE ml.from_unit_id = ANY($1::uuid[]) AND ml.to_unit_id = ANY($1::uuid[])
|
|
ORDER BY ml.link_type, ml.weight DESC
|
|
""", unit_ids)
|
|
else:
|
|
links = []
|
|
|
|
# Get entity information
|
|
unit_entities = await conn.fetch("""
|
|
SELECT ue.unit_id, e.canonical_name, e.entity_type
|
|
FROM unit_entities ue
|
|
JOIN entities e ON ue.entity_id = e.id
|
|
ORDER BY ue.unit_id
|
|
""")
|
|
|
|
# Build entity mapping
|
|
entity_map = {}
|
|
for row in unit_entities:
|
|
unit_id = row['unit_id']
|
|
entity_name = row['canonical_name']
|
|
entity_type = row['entity_type']
|
|
if unit_id not in entity_map:
|
|
entity_map[unit_id] = []
|
|
entity_map[unit_id].append(f"{entity_name} ({entity_type})")
|
|
|
|
# Build nodes
|
|
nodes = []
|
|
for row in units:
|
|
unit_id = row['id']
|
|
text = row['text']
|
|
event_date = row['event_date']
|
|
context = row['context']
|
|
|
|
entities = entity_map.get(unit_id, [])
|
|
entity_count = len(entities)
|
|
|
|
# Color by entity count
|
|
if entity_count == 0:
|
|
color = "#e0e0e0"
|
|
elif entity_count == 1:
|
|
color = "#90caf9"
|
|
else:
|
|
color = "#42a5f5"
|
|
|
|
nodes.append({
|
|
"data": {
|
|
"id": str(unit_id),
|
|
"label": f"{text[:30]}..." if len(text) > 30 else text,
|
|
"text": text,
|
|
"date": event_date.isoformat() if event_date else "",
|
|
"context": context if context else "",
|
|
"entities": ", ".join(entities) if entities else "None",
|
|
"color": color
|
|
}
|
|
})
|
|
|
|
# Build edges
|
|
edges = []
|
|
for row in links:
|
|
from_id = str(row['from_unit_id'])
|
|
to_id = str(row['to_unit_id'])
|
|
link_type = row['link_type']
|
|
weight = row['weight']
|
|
entity_name = row['entity_name']
|
|
|
|
# Color by link type
|
|
if link_type == 'temporal':
|
|
color = "#00bcd4"
|
|
line_style = "dashed"
|
|
elif link_type == 'semantic':
|
|
color = "#ff69b4"
|
|
line_style = "solid"
|
|
elif link_type == 'entity':
|
|
color = "#ffd700"
|
|
line_style = "solid"
|
|
else:
|
|
color = "#999999"
|
|
line_style = "solid"
|
|
|
|
edges.append({
|
|
"data": {
|
|
"id": f"{from_id}-{to_id}-{link_type}",
|
|
"source": from_id,
|
|
"target": to_id,
|
|
"linkType": link_type,
|
|
"weight": weight,
|
|
"entityName": entity_name if entity_name else "",
|
|
"color": color,
|
|
"lineStyle": line_style
|
|
}
|
|
})
|
|
|
|
# Build table rows
|
|
table_rows = []
|
|
for row in units:
|
|
unit_id = row['id']
|
|
entities = entity_map.get(unit_id, [])
|
|
|
|
table_rows.append({
|
|
"id": str(unit_id)[:8] + "...",
|
|
"text": row['text'],
|
|
"context": row['context'] if row['context'] else "N/A",
|
|
"date": row['event_date'].strftime("%Y-%m-%d %H:%M") if row['event_date'] else "N/A",
|
|
"entities": ", ".join(entities) if entities else "None"
|
|
})
|
|
|
|
return {
|
|
"nodes": nodes,
|
|
"edges": edges,
|
|
"table_rows": table_rows,
|
|
"total_units": len(units)
|
|
}
|
|
|
|
async def think_async(
|
|
self,
|
|
agent_id: str,
|
|
query: str,
|
|
thinking_budget: int = 50,
|
|
top_k: int = 10,
|
|
model: str = "openai/gpt-oss-120b",
|
|
temperature: float = 0.7,
|
|
max_tokens: int = 1000,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Think and formulate an answer using agent identity, world facts, and opinions.
|
|
|
|
This method:
|
|
1. Retrieves agent facts (agent's identity and past actions)
|
|
2. Retrieves world facts (general knowledge)
|
|
3. Retrieves existing opinions (agent's formed perspectives)
|
|
4. Uses Groq LLM to formulate an answer
|
|
5. Extracts and stores any new opinions formed during thinking
|
|
6. Returns plain text answer and the facts used
|
|
|
|
Args:
|
|
agent_id: Agent identifier
|
|
query: Question to answer
|
|
thinking_budget: Number of memory units to explore
|
|
top_k: Maximum facts to retrieve
|
|
model: LLM model to use (default: llama-3.3-70b-versatile)
|
|
temperature: Sampling temperature
|
|
max_tokens: Maximum tokens in response
|
|
|
|
Returns:
|
|
Dict with:
|
|
- text: Plain text answer (no markdown)
|
|
- based_on: Dict with 'world', 'agent', and 'opinion' fact lists
|
|
- new_opinions: List of newly formed opinions
|
|
"""
|
|
from openai import AsyncOpenAI
|
|
from datetime import datetime, timezone
|
|
|
|
# Initialize Groq client
|
|
groq_api_key = os.getenv("GROQ_API_KEY")
|
|
if not groq_api_key:
|
|
raise ValueError("GROQ_API_KEY environment variable not set")
|
|
|
|
client = AsyncOpenAI(
|
|
api_key=groq_api_key,
|
|
base_url="https://api.groq.com/openai/v1"
|
|
)
|
|
|
|
# Step 1: Get agent facts (identity)
|
|
agent_results, _ = await self.search_async(
|
|
agent_id=agent_id,
|
|
query=query,
|
|
thinking_budget=thinking_budget,
|
|
top_k=top_k,
|
|
enable_trace=False,
|
|
fact_type='agent'
|
|
)
|
|
|
|
# Step 2: Get world facts
|
|
world_results, _ = await self.search_async(
|
|
agent_id=agent_id,
|
|
query=query,
|
|
thinking_budget=thinking_budget,
|
|
top_k=top_k,
|
|
enable_trace=False,
|
|
fact_type='world'
|
|
)
|
|
|
|
# Step 3: Get existing opinions
|
|
opinion_results, _ = await self.search_async(
|
|
agent_id=agent_id,
|
|
query=query,
|
|
thinking_budget=thinking_budget,
|
|
top_k=top_k,
|
|
enable_trace=False,
|
|
fact_type='opinion'
|
|
)
|
|
|
|
# Step 4: Format facts for LLM
|
|
agent_facts_text = "\n".join([f"- {fact['text']}" for fact in agent_results]) if agent_results else "None"
|
|
world_facts_text = "\n".join([f"- {fact['text']}" for fact in world_results]) if world_results else "None"
|
|
opinion_facts_text = "\n".join([f"- {fact['text']}" for fact in opinion_results]) if opinion_results else "None"
|
|
|
|
# Step 5: Call Groq to formulate answer
|
|
prompt = f"""You are an AI assistant answering a question based on retrieved facts.
|
|
|
|
AGENT IDENTITY (what the agent has done):
|
|
{agent_facts_text}
|
|
|
|
WORLD FACTS (general knowledge):
|
|
{world_facts_text}
|
|
|
|
YOUR EXISTING OPINIONS (perspectives you've formed):
|
|
{opinion_facts_text}
|
|
|
|
QUESTION: {query}
|
|
|
|
Provide a helpful, accurate answer based on the facts above. Be consistent with your existing opinions. If the facts don't contain enough information to answer the question, say so clearly. Do not use markdown formatting - respond in plain text only.
|
|
|
|
If you form any new opinions while thinking about this question, state them clearly in your answer."""
|
|
|
|
response = await client.chat.completions.create(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": "You are a helpful AI assistant. Always respond in plain text without markdown formatting. You can form and express opinions based on facts."},
|
|
{"role": "user", "content": prompt}
|
|
],
|
|
temperature=temperature,
|
|
max_tokens=max_tokens
|
|
)
|
|
|
|
answer_text = response.choices[0].message.content.strip()
|
|
|
|
# Step 6: Extract new opinions from the answer
|
|
new_opinions = await self._extract_opinions_from_text(
|
|
client=client,
|
|
text=answer_text,
|
|
model=model
|
|
)
|
|
|
|
# Step 7: Store new opinions
|
|
if new_opinions:
|
|
current_time = datetime.now(timezone.utc)
|
|
for opinion_dict in new_opinions:
|
|
await self.put_async(
|
|
agent_id=agent_id,
|
|
content=opinion_dict["text"],
|
|
context=f"formed during thinking about: {query}",
|
|
event_date=current_time,
|
|
fact_type_override='opinion',
|
|
confidence_score=opinion_dict["confidence"]
|
|
)
|
|
|
|
# Step 8: Return response with facts split by type
|
|
return {
|
|
"text": answer_text,
|
|
"based_on": {
|
|
"world": world_results,
|
|
"agent": agent_results,
|
|
"opinion": opinion_results
|
|
},
|
|
"new_opinions": new_opinions
|
|
}
|
|
|
|
async def _extract_opinions_from_text(
|
|
self,
|
|
client,
|
|
text: str,
|
|
model: str
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Extract opinions with reasons and confidence from text using LLM.
|
|
|
|
Args:
|
|
client: OpenAI client
|
|
text: Text to extract opinions from
|
|
model: LLM model to use
|
|
|
|
Returns:
|
|
List of dicts with keys: 'text' (opinion with reasons), 'confidence' (score 0-1)
|
|
"""
|
|
from pydantic import BaseModel, Field
|
|
|
|
class Opinion(BaseModel):
|
|
"""An opinion formed by the agent."""
|
|
opinion: str = Field(description="The opinion or perspective formed")
|
|
reasons: str = Field(description="The reasons supporting this opinion")
|
|
confidence: float = Field(description="Confidence score for this opinion (0.0 to 1.0, where 1.0 is very confident)")
|
|
|
|
class OpinionExtractionResponse(BaseModel):
|
|
"""Response containing extracted opinions."""
|
|
opinions: List[Opinion] = Field(
|
|
default_factory=list,
|
|
description="List of opinions formed with their supporting reasons and confidence scores"
|
|
)
|
|
|
|
extraction_prompt = f"""Extract any opinions or perspectives that were formed in the following text.
|
|
An opinion is a judgment, viewpoint, or conclusion that goes beyond just stating facts.
|
|
|
|
TEXT:
|
|
{text}
|
|
|
|
For each opinion found, provide:
|
|
1. The opinion itself
|
|
2. The reasons or facts that support it
|
|
3. A confidence score (0.0 to 1.0) indicating how confident the agent is in this opinion based on the available information
|
|
|
|
If no clear opinions are expressed, return an empty list."""
|
|
|
|
try:
|
|
response = await client.beta.chat.completions.parse(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": "You extract opinions and perspectives from text."},
|
|
{"role": "user", "content": extraction_prompt}
|
|
],
|
|
response_format=OpinionExtractionResponse
|
|
)
|
|
|
|
result = response.choices[0].message.parsed
|
|
|
|
# Format opinions with reasons included in the text and confidence score
|
|
formatted_opinions = []
|
|
for op in result.opinions:
|
|
# Combine opinion and reasons into a single statement
|
|
opinion_with_reasons = f"{op.opinion} (Reasons: {op.reasons})"
|
|
formatted_opinions.append({
|
|
"text": opinion_with_reasons,
|
|
"confidence": op.confidence
|
|
})
|
|
|
|
return formatted_opinions
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to extract opinions: {str(e)}")
|
|
return []
|
|
|
|
async def _evaluate_opinion_update_async(
|
|
self,
|
|
client,
|
|
opinion_text: str,
|
|
opinion_confidence: float,
|
|
new_event_text: str,
|
|
entity_name: str,
|
|
model: str = "llama-3.3-70b-versatile",
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Evaluate if an opinion should be updated based on a new event.
|
|
|
|
Args:
|
|
client: OpenAI client
|
|
opinion_text: Current opinion text (includes reasons)
|
|
opinion_confidence: Current confidence score (0.0-1.0)
|
|
new_event_text: Text of the new event
|
|
entity_name: Name of the entity this opinion is about
|
|
model: LLM model to use
|
|
|
|
Returns:
|
|
Dict with 'action' ('keep'|'update'), 'new_confidence', 'new_text' (if action=='update')
|
|
or None if no changes needed
|
|
"""
|
|
from pydantic import BaseModel, Field
|
|
|
|
class OpinionEvaluation(BaseModel):
|
|
"""Evaluation of whether an opinion should be updated."""
|
|
action: str = Field(description="Action to take: 'keep' (no change) or 'update' (modify opinion)")
|
|
reasoning: str = Field(description="Brief explanation of why this action was chosen")
|
|
new_confidence: float = Field(description="New confidence score (0.0-1.0). Can be higher, lower, or same as before.")
|
|
new_opinion_text: Optional[str] = Field(
|
|
default=None,
|
|
description="If action is 'update', the revised opinion text that acknowledges the previous view. Otherwise None."
|
|
)
|
|
|
|
evaluation_prompt = f"""You are evaluating whether an existing opinion should be updated based on new information.
|
|
|
|
ENTITY: {entity_name}
|
|
|
|
EXISTING OPINION:
|
|
{opinion_text}
|
|
Current confidence: {opinion_confidence:.2f}
|
|
|
|
NEW EVENT:
|
|
{new_event_text}
|
|
|
|
Evaluate whether this new event:
|
|
1. REINFORCES the opinion (increase confidence, keep text)
|
|
2. WEAKENS the opinion (decrease confidence, keep text)
|
|
3. CHANGES the opinion (update both text and confidence, noting "Previously I thought X, but now Y...")
|
|
4. IRRELEVANT (keep everything as is)
|
|
|
|
Guidelines:
|
|
- Only suggest 'update' action if the new event genuinely contradicts or significantly modifies the opinion
|
|
- If updating the text, acknowledge the previous opinion and explain the change
|
|
- Confidence should reflect accumulated evidence (0.0 = no confidence, 1.0 = very confident)
|
|
- Small changes in confidence are normal; large jumps should be rare"""
|
|
|
|
try:
|
|
response = await client.beta.chat.completions.parse(
|
|
model=model,
|
|
messages=[
|
|
{"role": "system", "content": "You evaluate and update opinions based on new information."},
|
|
{"role": "user", "content": evaluation_prompt}
|
|
],
|
|
response_format=OpinionEvaluation,
|
|
temperature=0.3 # Lower temperature for more consistent evaluation
|
|
)
|
|
|
|
result = response.choices[0].message.parsed
|
|
|
|
# Only return updates if something actually changed
|
|
if result.action == 'keep' and abs(result.new_confidence - opinion_confidence) < 0.01:
|
|
return None
|
|
|
|
return {
|
|
'action': result.action,
|
|
'reasoning': result.reasoning,
|
|
'new_confidence': result.new_confidence,
|
|
'new_text': result.new_opinion_text if result.action == 'update' else None
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Failed to evaluate opinion update: {str(e)}")
|
|
return None
|
|
|
|
async def _reinforce_opinions_async(
|
|
self,
|
|
agent_id: str,
|
|
created_unit_ids: List[str],
|
|
unit_texts: List[str],
|
|
unit_entities: List[List[Dict[str, str]]],
|
|
model: str = "llama-3.3-70b-versatile",
|
|
):
|
|
"""
|
|
Background task to reinforce opinions based on newly ingested events.
|
|
|
|
This runs asynchronously and does not block the put operation.
|
|
|
|
Args:
|
|
agent_id: Agent ID
|
|
created_unit_ids: List of newly created memory unit IDs
|
|
unit_texts: Texts of the newly created units
|
|
unit_entities: Entities extracted from each unit
|
|
model: LLM model to use for evaluation
|
|
"""
|
|
try:
|
|
# Extract all unique entity names from the new units
|
|
entity_names = set()
|
|
for entities_list in unit_entities:
|
|
for entity in entities_list:
|
|
entity_names.add(entity['text'])
|
|
|
|
if not entity_names:
|
|
logger.debug("[REINFORCE] No entities found in new units, skipping opinion reinforcement")
|
|
return
|
|
|
|
logger.debug(f"[REINFORCE] Starting opinion reinforcement for {len(entity_names)} entities")
|
|
|
|
pool = await self._get_pool()
|
|
async with pool.acquire() as conn:
|
|
# Find all opinions related to these entities
|
|
opinions = await conn.fetch(
|
|
"""
|
|
SELECT DISTINCT mu.id, mu.text, mu.confidence_score, e.canonical_name
|
|
FROM memory_units mu
|
|
JOIN unit_entities ue ON mu.id = ue.unit_id
|
|
JOIN entities e ON ue.entity_id = e.id
|
|
WHERE mu.agent_id = $1
|
|
AND mu.fact_type = 'opinion'
|
|
AND e.canonical_name = ANY($2::text[])
|
|
""",
|
|
agent_id,
|
|
list(entity_names)
|
|
)
|
|
|
|
if not opinions:
|
|
logger.debug("[REINFORCE] No existing opinions found for these entities")
|
|
return
|
|
|
|
logger.debug(f"[REINFORCE] Found {len(opinions)} opinions to potentially reinforce")
|
|
|
|
# Get OpenAI client
|
|
from openai import AsyncOpenAI
|
|
groq_api_key = os.getenv("GROQ_API_KEY")
|
|
client = AsyncOpenAI(
|
|
api_key=groq_api_key,
|
|
base_url="https://api.groq.com/openai/v1"
|
|
)
|
|
|
|
# Evaluate each opinion against the new events
|
|
updates_to_apply = []
|
|
for opinion in opinions:
|
|
opinion_id = str(opinion['id'])
|
|
opinion_text = opinion['text']
|
|
opinion_confidence = opinion['confidence_score']
|
|
entity_name = opinion['canonical_name']
|
|
|
|
# Find all new events mentioning this entity
|
|
relevant_events = []
|
|
for unit_text, entities_list in zip(unit_texts, unit_entities):
|
|
if any(e['text'] == entity_name for e in entities_list):
|
|
relevant_events.append(unit_text)
|
|
|
|
if not relevant_events:
|
|
continue
|
|
|
|
# Combine all relevant events
|
|
combined_events = "\n".join(relevant_events)
|
|
|
|
# Evaluate if opinion should be updated
|
|
evaluation = await self._evaluate_opinion_update_async(
|
|
client,
|
|
opinion_text,
|
|
opinion_confidence,
|
|
combined_events,
|
|
entity_name,
|
|
model
|
|
)
|
|
|
|
if evaluation:
|
|
updates_to_apply.append({
|
|
'opinion_id': opinion_id,
|
|
'evaluation': evaluation
|
|
})
|
|
|
|
# Apply all updates in a single transaction
|
|
if updates_to_apply:
|
|
async with conn.transaction():
|
|
for update in updates_to_apply:
|
|
opinion_id = update['opinion_id']
|
|
evaluation = update['evaluation']
|
|
|
|
if evaluation['action'] == 'update' and evaluation['new_text']:
|
|
# Update both text and confidence
|
|
await conn.execute(
|
|
"""
|
|
UPDATE memory_units
|
|
SET text = $1, confidence_score = $2, updated_at = NOW()
|
|
WHERE id = $3
|
|
""",
|
|
evaluation['new_text'],
|
|
evaluation['new_confidence'],
|
|
uuid.UUID(opinion_id)
|
|
)
|
|
logger.debug(f"[REINFORCE] Updated opinion {opinion_id[:8]}... (action: {evaluation['action']}, confidence: {evaluation['new_confidence']:.2f})")
|
|
else:
|
|
# Only update confidence
|
|
await conn.execute(
|
|
"""
|
|
UPDATE memory_units
|
|
SET confidence_score = $1, updated_at = NOW()
|
|
WHERE id = $2
|
|
""",
|
|
evaluation['new_confidence'],
|
|
uuid.UUID(opinion_id)
|
|
)
|
|
logger.debug(f"[REINFORCE] Updated confidence for opinion {opinion_id[:8]}... (confidence: {evaluation['new_confidence']:.2f})")
|
|
|
|
logger.debug(f"[REINFORCE] Applied {len(updates_to_apply)} opinion updates")
|
|
else:
|
|
logger.debug("[REINFORCE] No opinion updates needed")
|
|
|
|
except Exception as e:
|
|
logger.error(f"[REINFORCE] Error during opinion reinforcement: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|