holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

improvements

Browse source
This commit is contained in:
Nicolò Boschi 2025-10-30 18:54:06 +01:00
parent 6d7223b14e
commit bc7d9fe07f
26 changed files with 7675 additions and 3433 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.gitignore vendored
View file

@ -21,3 +21,9 @@ wheels/
# NLTK data (will be downloaded automatically)
nltk_data/
# Large benchmark datasets (will be downloaded automatically)
benchmarks/longmemeval/longmemeval_s_cleaned.json
# Debug logs
logs/

3
CLAUDE.md
View file

@ -2,6 +2,7 @@
Do not write any markdown file, just write the code.
# Workflow
After your changes, make sure everything is working fine by running the main script.
- After your changes, make sure everything is working fine by running the tests.
- keep the readme.md architecture section up to date when you change the implementation
- when changing an implemetation, do not keep the old one as fallback
- to run test, use uv run pytest tests

42
README.md
View file

@ -90,43 +90,44 @@ The combination of these three networks enables powerful memory retrieval that g
The search algorithm explores the memory graph using spreading activation:
1. **Entry Points**: Find top-3 semantically similar memories to the query (vector search)
2. **Activation Spreading**: Start with activation = 1.0 at entry points
1. **Entry Points**: Find top-3 semantically similar memories to the query (vector search, similarity ≥ 0.5)
2. **Activation Spreading**: Start with activation = actual similarity score (0.5 to 1.0) at entry points
3. **Graph Traversal**: Follow links to neighbors, spreading activation with decay (0.8 factor)
4. **Thinking Budget**: Limit exploration to N units (controls computational cost)
5. **Dynamic Weighting**: Combine activation with recency and frequency:
5. **Dynamic Weighting**: Combine activation, semantic similarity, recency, and frequency:
```
final_weight = activation × recency × frequency
final_weight = 0.30 × activation + 0.30 × semantic_similarity + 0.25 × recency + 0.15 × frequency
semantic_similarity = cosine_similarity(query_embedding, memory_embedding)
recency = exp(-0.1 × days_since)
frequency = 1.0 + log(access_count + 1) / log(10)
frequency = normalized to [0, 1] from log(access_count + 1) / log(10)
```
6. **Return Top-K**: Sort by final weight and return top results
This approach ensures:
- Recently accessed memories get boosted (recency bias)
- Frequently accessed memories get boosted (importance signal)
- Graph structure influences results (not just vector similarity)
- Semantic relevance to query is always considered (30% weight)
- Graph structure influences results through activation (30% weight)
- Recently accessed memories get boosted (25% weight - recency bias)
- Frequently accessed memories get boosted (15% weight - importance signal)
### Self-Contained Memory Units
Every memory unit is processed to be self-contained through coreference resolution:
Every memory unit is self-contained through LLM fact extraction:
**Problem**: "She joined Google last year" - unclear who "she" is
**Solution**: Fast batch coreference resolution that:
- Replaces personal pronouns (he, she, it, they) with actual referents
- Processes all sentences in one batch (O(n) instead of O(n²))
- Uses neural coreference model for high accuracy
- Provides fallback to custom spaCy-based resolution if needed
**Solution**: LLM-based fact extraction that:
- Resolves pronouns to actual referents during extraction
- Makes facts readable without original context
- Includes all relevant details (WHO, WHAT, WHERE, WHEN, WHY, HOW)
- Processes facts in parallel for speed
**Result**: "Alice joined Google last year" - fully self-contained
**Technology**:
- **FastCoref** - Fast, accurate neural coreference resolution
- Processes 2.8K documents in 25 seconds on GPU
- Significant speedup over sequential spaCy approach
- Fallback to custom spaCy implementation if needed
- LLM fact extraction with detailed prompts for pronoun resolution
- Structured output using Pydantic models
- Batch processing for efficiency
### LLM-Based Fact Extraction
@ -165,9 +166,8 @@ Raw content is processed through an LLM to extract meaningful facts before stora
- `psycopg2-binary` - PostgreSQL client
- `sentence-transformers` - Local embedding model (bge-small-en-v1.5)
- `torch` - Deep learning framework (for embeddings)
- `fastcoref` - Fast neural coreference resolution
- `spacy` - NLP (NER, dependency parsing, tokenization)
- `nltk` - Sentence tokenization
- `langchain-text-splitters` - Intelligent text chunking
- `networkx` - Graph operations
- `pyvis` - Interactive HTML graph visualization
- `matplotlib` - Static graph visualization
@ -247,7 +247,7 @@ memory-poc/
├── memory/ # Core memory system package
│ ├── temporal_semantic_memory.py # Main memory system class
│ ├── entity_resolver.py # Entity extraction and disambiguation
│ ├── coref_resolver.py # Coreference resolution
│ ├── llm_client.py # LLM-based fact extraction
│ └── utils.py # Utility functions
│
├── demos/ # Demo scripts

2323
benchmarks/locomo/benchmark_results.json
View file

File diff suppressed because it is too large Load diff

370
benchmarks/locomo/run_benchmark.py
View file

@ -13,10 +13,12 @@ import json
from datetime import datetime, timezone, timedelta
from memory import TemporalSemanticMemory
from typing import List, Dict
from openai import AsyncOpenAI
import openai
from dotenv import load_dotenv
import os
import asyncio
import pydantic
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
from rich.table import Table
@ -27,6 +29,24 @@ load_dotenv()
console = Console()
def get_groq_client() -> AsyncOpenAI:
"""
Get configured async Groq client for LLM judge.
Returns:
Configured AsyncOpenAI client pointing to Groq
"""
groq_api_key = os.getenv('GROQ_API_KEY')
if not groq_api_key:
raise ValueError("GROQ_API_KEY environment variable not set")
base_url = os.getenv('GROQ_BASE_URL', 'https://api.groq.com/openai/v1')
return AsyncOpenAI(
api_key=groq_api_key,
base_url=base_url
)
def parse_date(date_string: str) -> datetime:
"""Parse LoComo date format to datetime."""
# Format: "1:56 pm on 8 May, 2023"
@ -41,7 +61,7 @@ async def ingest_conversation(memory: TemporalSemanticMemory, conversation_data:
"""
Ingest a LoComo conversation into the memory system (ASYNC version).
Ingests entire conversation as a single large document for maximum efficiency.
Ingests ALL sessions in ONE batch for maximum efficiency.
Args:
memory: Memory system instance
@ -55,45 +75,55 @@ async def ingest_conversation(memory: TemporalSemanticMemory, conversation_data:
# Get all session keys sorted
session_keys = sorted([k for k in conv.keys() if k.startswith('session_') and not k.endswith('_date_time')])
# Collect all sessions as batch items
batch_contents = []
total_turns = 0
# Build entire conversation as one large text
conversation_parts = []
for session_key in session_keys:
if session_key not in conv or not isinstance(conv[session_key], list):
continue
session_data = conv[session_key]
# Add all turns from this session
# Build session content from all turns
session_parts = []
for turn in session_data:
speaker = turn['speaker']
text = turn['text']
conversation_parts.append(f"{speaker} said: {text}")
session_parts.append(f"{speaker}: {text}")
total_turns += 1
# Ingest entire conversation in ONE put_async call
# Use the first session date as the event date
first_session_key = session_keys[0] if session_keys else "session_1"
date_key = f"{first_session_key}_date_time"
conversation_date = parse_date(conv.get(date_key, "1:00 pm on 1 January, 2023"))
if not session_parts:
continue
full_conversation = " ".join(conversation_parts)
# Get session date
date_key = f"{session_key}_date_time"
session_date = parse_date(conv.get(date_key, "1:00 pm on 1 January, 2023"))
await memory.put_async(
agent_id=agent_id,
content=full_conversation,
context=f"Full conversation between {speaker_a} and {speaker_b}",
event_date=conversation_date
)
# Add to batch
session_content = "\n".join(session_parts)
batch_contents.append({
"content": session_content,
"context": f"Conversation session between {speaker_a} and {speaker_b}",
"event_date": session_date
})
# Ingest ALL sessions in ONE batch call (MUCH faster!)
if batch_contents:
await memory.put_batch_async(
agent_id=agent_id,
contents=batch_contents
)
return total_turns
class QuestionAnswer(pydantic.BaseModel):
answer: str
reasoning: str
def answer_question(memory: TemporalSemanticMemory, agent_id: str, question: str, thinking_budget: int = 100) -> str:
async def answer_question(memory: TemporalSemanticMemory, agent_id: str, question: str, thinking_budget: int = 500) -> tuple[str, str, List[Dict]]:
"""
Answer a question using the memory system.
Answer a question using the memory system (ASYNC version).
Args:
memory: Memory system instance
@ -102,36 +132,33 @@ def answer_question(memory: TemporalSemanticMemory, agent_id: str, question: str
thinking_budget: How many memory units to explore
Returns:
Answer string
Tuple of (answer string, reasoning string, retrieved memories list)
"""
# Search memory
results = memory.search(
results = await memory.search_async(
agent_id=agent_id,
query=question,
thinking_budget=thinking_budget,
top_k=20 # Get more results for better context
)
print("question:", question)
print("Got results:", results)
if not results:
return "I don't have enough information to answer that question."
return "I don't have enough information to answer that question.", "No relevant memories found.", []
# Build context from top results
context_parts = []
for i, result in enumerate(results[:10], 1):
for i, result in enumerate(results):
context_parts.append(f"{i}. {result['text']}")
context = "\n".join(context_parts)
# Use OpenAI to generate answer from context
# Use AsyncOpenAI to generate answer from context
try:
response = openai.chat.completions.create(
client = AsyncOpenAI()
response = await client.beta.chat.completions.parse(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": "You are a helpful assistant. Answer the question based ONLY on the provided context. If the context doesn't contain the answer, say 'I don't know'."
"content": "You are a helpful assistant. Answer the question based ONLY on the provided context. If the context doesn't contain the answer, say 'I don't know'. In the reasoning, explain why you choose or not choose the context items for the answer."
},
{
"role": "user",
@ -139,14 +166,16 @@ def answer_question(memory: TemporalSemanticMemory, agent_id: str, question: str
}
],
temperature=0,
max_tokens=150
max_tokens=8000,
response_format=QuestionAnswer
)
return response.choices[0].message.content.strip()
answer = response.choices[0].message.parsed
return answer.answer, answer.reasoning, results
except Exception as e:
return f"Error generating answer: {str(e)}"
return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", results
def evaluate_qa_task(
async def evaluate_qa_task(
memory: TemporalSemanticMemory,
agent_id: str,
qa_pairs: List[Dict],
@ -154,13 +183,11 @@ def evaluate_qa_task(
max_questions: int = None
) -> Dict:
"""
Evaluate the QA task.
Evaluate the QA task (ASYNC version - processes questions in parallel).
Returns:
Dict with evaluation metrics
"""
results = []
questions_to_eval = qa_pairs[:max_questions] if max_questions else qa_pairs
with Progress(
@ -170,38 +197,97 @@ def evaluate_qa_task(
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
console=console
) as progress:
task = progress.add_task(f"[cyan]Evaluating QA for sample {sample_id}...", total=len(questions_to_eval))
task = progress.add_task(f"[cyan]Evaluating QA for sample {sample_id} (parallel)...", total=len(questions_to_eval))
for qa in questions_to_eval:
# Create tasks for all questions
async def process_question(qa):
question = qa['question']
correct_answer = qa['answer']
category = qa.get('category', 0)
# Get predicted answer
predicted_answer = answer_question(memory, agent_id, question)
# Get predicted answer, reasoning, and retrieved memories
predicted_answer, reasoning, retrieved_memories = await answer_question(memory, agent_id, question)
results.append({
return {
'question': question,
'correct_answer': correct_answer,
'predicted_answer': predicted_answer,
'category': category
})
'reasoning': reasoning,
'category': category,
'retrieved_memories': retrieved_memories
}
# Process all questions in parallel
question_tasks = [process_question(qa) for qa in questions_to_eval]
# Use as_completed to update progress as results come in
results = []
for coro in asyncio.as_completed(question_tasks):
result = await coro
results.append(result)
progress.update(task, advance=1)
return results
class JudgeResponse(pydantic.BaseModel):
correct: bool
reasoning: str
def calculate_metrics(results: List[Dict]) -> Dict:
async def judge_single_answer(client: AsyncOpenAI, result: Dict, semaphore: asyncio.Semaphore) -> Dict:
"""
Calculate evaluation metrics.
Judge a single answer using LLM (with concurrency control).
Uses LLM-as-judge to evaluate answer quality.
Args:
client: Async OpenAI client (Groq)
result: Result dict with question, correct_answer, predicted_answer, category
semaphore: Semaphore to limit concurrent requests
Returns:
Updated result dict with is_correct field
"""
async with semaphore:
try:
response = await client.beta.chat.completions.parse(
model="openai/gpt-oss-120b",
messages=[
{
"role": "system",
"content":
"You are an objective judge. Determine if the predicted answer contains the correct answer or they are the same content (with different form is fine)."
},
{
"role": "user",
"content": f"Question: {result['question']}\nCorrect answer: {result['correct_answer']}\nPredicted answer: {result['predicted_answer']}\n\nAre they equivalent?"
}
],
temperature=0,
max_tokens=512,
response_format=JudgeResponse
)
judgement = response.choices[0].message.parsed
result['is_correct'] = judgement.correct
result['correctness_reasoning'] = judgement.reasoning
except Exception as e:
console.print(f"[red]Error judging answer: {e}[/red]")
result['is_correct'] = False
return result
async def calculate_metrics(results: List[Dict]) -> Dict:
"""
Calculate evaluation metrics using parallel LLM-as-judge.
Processes up to 8 judgments concurrently for speed.
"""
correct = 0
total = len(results)
client = get_groq_client()
category_stats = {}
# Semaphore to limit to 8 concurrent requests
semaphore = asyncio.Semaphore(8)
with Progress(
SpinnerColumn(),
@ -210,49 +296,33 @@ def calculate_metrics(results: List[Dict]) -> Dict:
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
console=console
) as progress:
task = progress.add_task("[yellow]Judging answers with LLM...", total=total)
task = progress.add_task("[yellow]Judging answers with LLM (parallel, max 8)...", total=total)
# Create all judgment tasks
judgment_tasks = []
for result in results:
# Use LLM as judge
try:
response = openai.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": "You are an objective judge. Determine if the predicted answer is semantically equivalent to the correct answer. Answer with ONLY 'yes' or 'no'."
},
{
"role": "user",
"content": f"Question: {result['question']}\nCorrect answer: {result['correct_answer']}\nPredicted answer: {result['predicted_answer']}\n\nAre they equivalent?"
}
],
temperature=0,
max_tokens=5
)
judgment = response.choices[0].message.content.strip().lower()
is_correct = 'yes' in judgment
if is_correct:
correct += 1
result['is_correct'] = is_correct
# Track by category
category = result['category']
if category not in category_stats:
category_stats[category] = {'correct': 0, 'total': 0}
category_stats[category]['total'] += 1
if is_correct:
category_stats[category]['correct'] += 1
except Exception as e:
console.print(f"[red]Error judging answer: {e}[/red]")
result['is_correct'] = False
judgment_task = judge_single_answer(client, result, semaphore)
judgment_tasks.append(judgment_task)
# Process in parallel with progress updates
judged_results = []
for coro in asyncio.as_completed(judgment_tasks):
judged_result = await coro
judged_results.append(judged_result)
progress.update(task, advance=1)
# Calculate stats
correct = sum(1 for r in judged_results if r.get('is_correct', False))
category_stats = {}
for result in judged_results:
category = result['category']
if category not in category_stats:
category_stats[category] = {'correct': 0, 'total': 0}
category_stats[category]['total'] += 1
if result.get('is_correct', False):
category_stats[category]['correct'] += 1
accuracy = (correct / total * 100) if total > 0 else 0
return {
@ -260,17 +330,82 @@ def calculate_metrics(results: List[Dict]) -> Dict:
'correct': correct,
'total': total,
'category_stats': category_stats,
'detailed_results': results
'detailed_results': judged_results
}
def run_benchmark(max_conversations: int = None, max_questions_per_conv: int = None):
async def process_single_conversation(
memory: TemporalSemanticMemory,
conv_data: Dict,
i: int,
total_convs: int,
max_questions_per_conv: int,
skip_ingestion: bool
) -> Dict:
"""
Process a single conversation (ingest + evaluate).
Args:
memory: Memory system instance
conv_data: Conversation data
i: Conversation index (1-based)
total_convs: Total number of conversations
max_questions_per_conv: Max questions to evaluate per conversation
skip_ingestion: Whether to skip ingestion
Returns:
Result dict with sample_id, metrics, total_turns
"""
sample_id = conv_data['sample_id']
agent_id = "locomo" # Single agent for all Locomo benchmark data
console.print(f"\n[bold blue]Conversation {i}/{total_convs}[/bold blue] (Sample ID: {sample_id})")
if not skip_ingestion:
# Clear previous locomo agent data only (multi-tenant safe)
if i == 1: # Only cleanup on first conversation
console.print(" [2] Clearing previous 'locomo' agent data...")
memory.delete_agent(agent_id)
console.print(f" [green]✓[/green] Cleared 'locomo' agent data")
# Ingest conversation (sessions processed in parallel)
console.print(" [3] Ingesting conversation (sessions in parallel)...")
total_turns = await ingest_conversation(memory, conv_data, agent_id)
console.print(f" [green]✓[/green] Ingested {total_turns} turns across multiple sessions")
else:
total_turns = -1
# Evaluate QA (async - questions processed in parallel)
console.print(f" [4] Evaluating {len(conv_data['qa'])} QA pairs (parallel)...")
qa_results = await evaluate_qa_task(
memory,
agent_id,
conv_data['qa'],
sample_id,
max_questions=max_questions_per_conv
)
# Calculate metrics (async with parallel LLM judging)
console.print(" [5] Calculating metrics...")
metrics = await calculate_metrics(qa_results)
console.print(f" [green]✓[/green] Accuracy: {metrics['accuracy']:.2f}% ({metrics['correct']}/{metrics['total']})")
return {
'sample_id': sample_id,
'metrics': metrics,
'total_turns': total_turns
}
def run_benchmark(max_conversations: int = None, max_questions_per_conv: int = None, skip_ingestion: bool = False):
"""
Run the LoComo benchmark.
Args:
max_conversations: Maximum number of conversations to evaluate (None for all)
max_questions_per_conv: Maximum questions per conversation (None for all)
skip_ingestion: Whether to skip ingestion and use existing data
"""
console.print("\n[bold cyan]LoComo Benchmark - Entity-Aware Memory System[/bold cyan]")
console.print("=" * 80)
@ -288,54 +423,17 @@ def run_benchmark(max_conversations: int = None, max_questions_per_conv: int = N
memory = TemporalSemanticMemory()
console.print(" [green]✓[/green] Memory system initialized")
# Run evaluation for each conversation
# Run evaluation (conversations sequential, sessions within each conversation parallel)
all_results = []
for i, conv_data in enumerate(conversations_to_eval, 1):
sample_id = conv_data['sample_id']
agent_id = f"locomo_{sample_id}"
console.print(f"\n[bold blue]Conversation {i}/{len(conversations_to_eval)}[/bold blue] (Sample ID: {sample_id})")
# Clear previous data
import psycopg2
conn = psycopg2.connect(os.getenv('DATABASE_URL'))
cursor = conn.cursor()
cursor.execute("DELETE FROM memory_units WHERE agent_id = %s", (agent_id,))
cursor.execute("DELETE FROM memory_links WHERE agent_id = %s", (agent_id,))
cursor.execute("DELETE FROM entity_cooccurrences WHERE agent_id = %s", (agent_id,))
cursor.execute("DELETE FROM unit_entities WHERE agent_id = %s", (agent_id,))
cursor.execute("DELETE FROM entities WHERE agent_id = %s", (agent_id,))
conn.commit()
cursor.close()
conn.close()
# Ingest conversation (using async for parallel embedding generation)
console.print(" [3] Ingesting conversation (async with parallel embeddings)...")
total_turns = asyncio.run(ingest_conversation(memory, conv_data, agent_id))
console.print(f" [green]✓[/green] Ingested {total_turns} conversation turns")
# Evaluate QA
console.print(f" [4] Evaluating {len(conv_data['qa'])} QA pairs...")
qa_results = evaluate_qa_task(
memory,
agent_id,
conv_data['qa'],
sample_id,
max_questions=max_questions_per_conv
result = asyncio.run(
process_single_conversation(
memory, conv_data, i, len(conversations_to_eval),
max_questions_per_conv, skip_ingestion
)
)
# Calculate metrics
console.print(" [5] Calculating metrics...")
metrics = calculate_metrics(qa_results)
console.print(f" [green]✓[/green] Accuracy: {metrics['accuracy']:.2f}% ({metrics['correct']}/{metrics['total']})")
all_results.append({
'sample_id': sample_id,
'metrics': metrics,
'total_turns': total_turns
})
all_results.append(result)
# Overall results
console.print("\n[bold green]✓ Benchmark Complete![/bold green]\n")
@ -387,12 +485,14 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Run LoComo benchmark')
parser.add_argument('--max-conversations', type=int, default=None, help='Maximum conversations to evaluate')
parser.add_argument('--max-questions', type=int, default=None, help='Maximum questions per conversation')
parser.add_argument('--skip-ingestion', action='store_true', help='Skip ingestion and use existing data')
args = parser.parse_args()
results = run_benchmark(
max_conversations=args.max_conversations,
max_questions_per_conv=args.max_questions
max_questions_per_conv=args.max_questions,
skip_ingestion=args.skip_ingestion
)
# Save results

62
benchmarks/longmemeval/run_benchmark.py
View file

@ -22,6 +22,7 @@ from typing import Dict, List, Any
from pathlib import Path
import time
import asyncio
import subprocess
from dotenv import load_dotenv
# Load environment variables from .env
@ -74,6 +75,43 @@ def parse_args():
return parser.parse_args()
def download_dataset(dataset_path: Path) -> bool:
"""
Download the LongMemEval dataset if it doesn't exist.
Returns:
True if successful, False otherwise
"""
url = "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json"
console.print(f"[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
console.print(f"[dim]URL: {url}[/dim]")
console.print(f"[dim]Destination: {dataset_path}[/dim]")
try:
# Use curl to download with progress
result = subprocess.run(
["curl", "-L", "-o", str(dataset_path), url],
capture_output=True,
text=True,
timeout=300 # 5 minute timeout
)
if result.returncode == 0 and dataset_path.exists():
console.print(f"[green]✓ Dataset downloaded successfully[/green]")
return True
else:
console.print(f"[red]✗ Download failed: {result.stderr}[/red]")
return False
except subprocess.TimeoutExpired:
console.print(f"[red]✗ Download timed out after 5 minutes[/red]")
return False
except Exception as e:
console.print(f"[red]✗ Download error: {e}[/red]")
return False
def load_dataset(dataset_path: str) -> List[Dict[str, Any]]:
"""Load LongMemEval dataset from JSON file."""
with open(dataset_path, 'r') as f:
@ -158,7 +196,7 @@ async def ingest_conversation(memory: TemporalSemanticMemory, agent_id: str, ins
console.print(f"[yellow]Warning: Failed to ingest session {session_id}: {e}[/yellow]")
def retrieve_memories(
async def retrieve_memories(
memory: TemporalSemanticMemory,
agent_id: str,
query: str,
@ -179,7 +217,7 @@ def retrieve_memories(
List of retrieved memory units
"""
try:
results = memory.search(
results = await memory.search_async(
agent_id=agent_id,
query=query,
thinking_budget=thinking_budget,
@ -323,12 +361,13 @@ def run_benchmark(args):
"""Run the LongMemEval benchmark evaluation."""
console.print("\n[bold cyan]LongMemEval Benchmark Evaluation[/bold cyan]\n")
# Load dataset
# Load dataset - download if needed
dataset_path = Path(__file__).parent / "longmemeval_s_cleaned.json"
if not dataset_path.exists():
console.print(f"[red]Error: Dataset not found at {dataset_path}[/red]")
console.print("[yellow]Run: curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o longmemeval_s_cleaned.json[/yellow]")
return
if not download_dataset(dataset_path):
console.print(f"[red]Failed to download dataset. Please download manually:[/red]")
console.print("[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/longmemeval_s_cleaned.json[/yellow]")
return
console.print(f"[green]Loading dataset from {dataset_path}[/green]")
dataset = load_dataset(dataset_path)
@ -374,8 +413,11 @@ def run_benchmark(args):
progress.update(instance_task, description=f"[cyan]Instance {idx+1}/{len(dataset)}: {question_id}")
# Create unique agent ID for this instance
agent_id = f"longmemeval_{question_id}"
# Use single agent for all LongMemEval data (cleared per question for isolation)
agent_id = "longmemeval"
# Clear agent data for this question (each question needs fresh isolated context)
memory.delete_agent(agent_id)
# Ingest conversation history
try:
@ -385,13 +427,13 @@ def run_benchmark(args):
continue
# Retrieve memories
memories = retrieve_memories(
memories = asyncio.run(retrieve_memories(
memory,
agent_id,
question,
args.thinking_budget,
args.top_k
)
))
# Generate answer
predicted_answer = generate_answer(client, question, memories)

336
memory/coref_resolver.py
View file

@ -1,336 +0,0 @@
"""
Coreference resolution for memory units.
Ensures every memory unit is self-contained by replacing pronouns
with their actual referents.
"""
import spacy
from typing import List, Dict, Optional
from fastcoref import FCoref
import threading
def get_nlp():
"""Get or load spaCy model."""
try:
return spacy.load("en_core_web_sm")
except OSError:
raise Exception("spaCy model not found. Run: uv pip install https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl")
# Global fastcoref model instance (singleton pattern)
_fastcoref_model = None
_fastcoref_lock = threading.Lock()
def get_fastcoref_model():
"""Get or load FastCoref model (singleton pattern)."""
global _fastcoref_model
if _fastcoref_model is None:
with _fastcoref_lock:
if _fastcoref_model is None:
# Use CPU by default, can be configured with device='cuda:0' for GPU
_fastcoref_model = FCoref(device='cpu')
return _fastcoref_model
def resolve_pronouns_in_text(text: str, context_sentences: List[str] = None) -> str:
"""
Resolve pronouns to their referents to make text self-contained.
Strategy:
1. Identify pronouns in the text
2. Look for named entities in the same sentence or previous sentences
3. Replace pronouns with the most likely referent based on:
- Gender agreement
- Number agreement (singular/plural)
- Proximity (closer entities more likely)
Args:
text: The sentence to resolve
context_sentences: Previous sentences for context (optional)
Returns:
Text with pronouns resolved
"""
nlp = get_nlp()
# Parse the target sentence
doc = nlp(text)
# Collect all sentences for context
all_text = text
if context_sentences:
# Add previous sentences for context
all_text = " ".join(context_sentences) + " " + text
full_doc = nlp(all_text)
# Extract entities with their positions
entities = []
for ent in full_doc.ents:
if ent.label_ in ['PERSON', 'ORG', 'GPE', 'PRODUCT']:
entities.append({
'text': ent.text,
'label': ent.label_,
'start': ent.start_char,
'end': ent.end_char,
})
# Check if sentence already has a named entity subject
has_named_subject = False
for token in doc:
if token.dep_ in ['nsubj', 'nsubjpass'] and token.pos_ == 'PROPN':
has_named_subject = True
break
# Find pronouns and anaphoric references that need resolution
pronouns_to_replace = []
for token in doc:
# Handle pronouns (he, she, it, they)
if token.pos_ == 'PRON' and token.dep_ in ['nsubj', 'nsubjpass']:
# Subject pronouns that need resolution
pron_lower = token.text.lower()
# Skip if sentence already has a named subject earlier
if has_named_subject and any(
t.dep_ in ['nsubj', 'nsubjpass'] and t.pos_ == 'PROPN' and t.i < token.i
for t in doc
):
continue
# Skip if it's already a proper name or demonstrative
if pron_lower in ['i', 'you', 'we', 'this', 'that', 'these', 'those']:
continue
# Find the best entity to replace it with
referent = find_best_referent(
pronoun=token,
entities=entities,
doc=full_doc
)
if referent:
pronouns_to_replace.append({
'pronoun': token,
'referent': referent,
'start': token.idx,
'end': token.idx + len(token.text)
})
# Handle definite noun phrases (e.g., "The project")
elif token.text.lower() == 'the' and token.head.pos_ == 'NOUN':
# Check if this "the X" phrase is a subject
if token.head.dep_ in ['nsubj', 'nsubjpass']:
# Try to find what "the X" refers to
noun = token.head.text
# Look for indefinite mentions earlier ("a project", "an organization")
for ent_token in reversed(list(full_doc)):
if ent_token.text.lower() == noun.lower():
# Found a matching noun - check if it has indefinite article
if any(child.text.lower() in ['a', 'an'] for child in ent_token.children):
# Replace "the project" with "the Python project" or similar
# Get the full noun phrase
descriptors = []
for child in ent_token.children:
if child.pos_ in ['ADJ', 'PROPN', 'NOUN'] and child.i < ent_token.i:
descriptors.append(child.text)
if descriptors:
full_phrase = ' '.join(descriptors) + ' ' + noun
# Calculate span to replace
span_start = token.idx
span_end = token.head.idx + len(token.head.text)
pronouns_to_replace.append({
'pronoun': token,
'referent': 'the ' + full_phrase,
'start': span_start,
'end': span_end
})
break
# Replace pronouns with referents (in reverse order to maintain indices)
result = text
for item in reversed(pronouns_to_replace):
start = item['start']
end = item['end']
result = result[:start] + item['referent'] + result[end:]
return result
def find_best_referent(
pronoun,
entities: List[Dict],
doc
) -> Optional[str]:
"""
Find the best entity referent for a pronoun.
Uses:
- Gender agreement (he/she -> PERSON)
- Number agreement (singular/plural)
- Entity type (he/she -> PERSON, it -> ORG/PRODUCT)
- Proximity (closer entities preferred)
"""
pron_text = pronoun.text.lower()
# Determine pronoun properties
is_singular = pron_text in ['he', 'she', 'it', 'him', 'her']
is_plural = pron_text in ['they', 'them']
is_person = pron_text in ['he', 'she', 'him', 'her']
is_thing = pron_text in ['it']
# Score each entity
candidates = []
for entity in entities:
score = 0.0
# Proximity score (entities closer to pronoun are better)
# Since entities come from context, those appearing later (higher start position) are closer
proximity_score = entity['start'] / 1000.0 # Normalize by position
score += proximity_score
# Type matching
if is_person and entity['label'] == 'PERSON':
score += 2.0 # Strong preference for person entities
elif is_thing and entity['label'] in ['ORG', 'PRODUCT', 'GPE']:
score += 2.0 # Organizations/products for "it"
# Recency: prefer entities that appear just before the pronoun
if entity['end'] < pronoun.idx:
distance = pronoun.idx - entity['end']
recency = 1.0 / (1.0 + distance / 100.0)
score += recency
candidates.append((entity['text'], score))
# Return the highest scoring candidate
if candidates:
candidates.sort(key=lambda x: x[1], reverse=True)
return candidates[0][0]
return None
def resolve_sentences_fast(sentences: List[str]) -> List[str]:
"""
Fast batch coreference resolution using FastCoref.
This is significantly faster than the sequential spaCy-based approach:
- Processes entire document in one pass (O(n) instead of O(n²))
- Uses efficient batching and neural model
- Can process 2.8K documents in 25 seconds on GPU
Args:
sentences: List of sentences to resolve
Returns:
List of resolved sentences (self-contained)
"""
if not sentences:
return []
# Join sentences into a single document for batch processing
# Add markers to track sentence boundaries
full_text = " ".join(sentences)
# Get the fastcoref model
model = get_fastcoref_model()
# Predict coreferences in batch
preds = model.predict(texts=[full_text])
if not preds or len(preds) == 0:
# No coreferences found, return original sentences
return sentences
# Get the first (and only) result
result = preds[0]
# Get clusters as text strings
clusters = result.get_clusters(as_strings=True)
if not clusters:
return sentences
# Build a replacement map: pronoun -> main referent
replacements = {}
for cluster in clusters:
if len(cluster) < 2:
continue
# The first mention is typically the most complete referent
main_referent = cluster[0]
# Map all other mentions (pronouns/short references) to the main referent
for mention in cluster[1:]:
mention_lower = mention.lower()
# Only replace if it's likely a pronoun or short reference
if len(mention.split()) <= 2 and any(
pron in mention_lower
for pron in ['he', 'she', 'it', 'they', 'him', 'her', 'them', 'his', 'her', 'their', 'the']
):
replacements[mention] = main_referent
# Apply replacements to each sentence
resolved = []
for sentence in sentences:
resolved_sentence = sentence
for mention, referent in replacements.items():
# Case-insensitive replacement but preserve capitalization context
if mention in resolved_sentence:
resolved_sentence = resolved_sentence.replace(mention, referent)
resolved.append(resolved_sentence)
return resolved
def resolve_sentences(sentences: List[str]) -> List[str]:
"""
Resolve pronouns across a list of sentences.
Uses FastCoref for efficient batch processing.
Falls back to legacy spaCy method if FastCoref fails.
Args:
sentences: List of sentences to resolve
Returns:
List of resolved sentences (self-contained)
"""
try:
return resolve_sentences_fast(sentences)
except Exception as e:
# Fallback to legacy method
print(f"FastCoref failed ({e}), falling back to spaCy method")
return resolve_sentences_legacy(sentences)
def resolve_sentences_legacy(sentences: List[str]) -> List[str]:
"""
Legacy sequential pronoun resolution (slower, O(n²) complexity).
Kept as fallback in case FastCoref is unavailable or fails.
Args:
sentences: List of sentences to resolve
Returns:
List of resolved sentences (self-contained)
"""
resolved = []
for i, sentence in enumerate(sentences):
# Use all previous sentences as context
context = resolved[:i] if i > 0 else []
# Resolve pronouns in this sentence
resolved_sentence = resolve_pronouns_in_text(sentence, context)
resolved.append(resolved_sentence)
return resolved

261
memory/entity_resolver.py
View file

@ -48,6 +48,43 @@ def extract_entities(text: str) -> List[Dict[str, any]]:
return entities
def extract_entities_batch(texts: List[str]) -> List[List[Dict[str, any]]]:
"""
Extract entities from multiple texts in batch (MUCH faster than sequential).
Uses spaCy's nlp.pipe() for efficient batch processing.
Args:
texts: List of input texts
Returns:
List of entity lists, one per input text
"""
if not texts:
return []
nlp = get_nlp()
# Process all texts in batch using nlp.pipe (significantly faster!)
docs = list(nlp.pipe(texts, batch_size=50))
all_entities = []
for doc in docs:
entities = []
for ent in doc.ents:
# Filter to important entity types
if ent.label_ in ['PERSON', 'ORG', 'GPE', 'LOC', 'PRODUCT', 'EVENT']:
entities.append({
'text': ent.text,
'type': ent.label_,
'start': ent.start_char,
'end': ent.end_char,
})
all_entities.append(entities)
return all_entities
class EntityResolver:
"""
Resolves entities to canonical IDs with disambiguation.
@ -62,6 +99,161 @@ class EntityResolver:
"""
self.conn = db_conn
def resolve_entities_batch(
self,
agent_id: str,
entities_data: List[Dict],
context: str,
unit_event_date,
) -> List[str]:
"""
Resolve multiple entities in batch (MUCH faster than sequential).
Groups entities by type, queries candidates in bulk, and resolves
all entities with minimal DB queries.
Args:
agent_id: Agent ID
entities_data: List of dicts with 'text', 'type', 'nearby_entities'
context: Context where entities appear
unit_event_date: When this unit was created
Returns:
List of entity IDs in same order as input
"""
if not entities_data:
return []
cursor = self.conn.cursor()
try:
import time
start = time.time()
# Group entities by type for efficient querying
entities_by_type = {}
for idx, entity_data in enumerate(entities_data):
entity_type = entity_data['type']
if entity_type not in entities_by_type:
entities_by_type[entity_type] = []
entities_by_type[entity_type].append((idx, entity_data))
# Query ALL candidates for each type in batch
all_candidates = {} # Maps (entity_type, entity_text) -> list of candidates
for entity_type, entities_list in entities_by_type.items():
# Extract unique entity texts for this type
entity_texts = list(set(e[1]['text'] for e in entities_list))
# Query candidates for all texts at once
from psycopg2.extras import execute_values
cursor.execute(
"""
SELECT canonical_name, id, metadata, last_seen, mention_count
FROM entities
WHERE agent_id = %s AND entity_type = %s
""",
(agent_id, entity_type)
)
type_candidates = cursor.fetchall()
# Filter candidates in memory (faster than complex SQL for small datasets)
for entity_text in entity_texts:
matching = []
entity_text_lower = entity_text.lower()
for canonical_name, ent_id, metadata, last_seen, mention_count in type_candidates:
canonical_lower = canonical_name.lower()
# Same matching logic as before
if (entity_text_lower == canonical_lower or
entity_text_lower in canonical_lower or
canonical_lower in entity_text_lower):
matching.append((ent_id, canonical_name, metadata, last_seen, mention_count))
all_candidates[(entity_type, entity_text)] = matching
# Resolve each entity using pre-fetched candidates
entity_ids = [None] * len(entities_data)
entities_to_update = [] # (entity_id, unit_event_date)
entities_to_create = [] # (idx, entity_data)
for idx, entity_data in enumerate(entities_data):
entity_text = entity_data['text']
entity_type = entity_data['type']
nearby_entities = entity_data.get('nearby_entities', [])
candidates = all_candidates.get((entity_type, entity_text), [])
if not candidates:
# Will create new entity
entities_to_create.append((idx, entity_data))
continue
# Score candidates (same logic as before but with pre-fetched data)
best_candidate = None
best_score = 0.0
best_name_similarity = 0.0
nearby_entity_set = {e['text'].lower() for e in nearby_entities if e['text'] != entity_text}
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
score = 0.0
# Name similarity
name_similarity = SequenceMatcher(
None,
entity_text.lower(),
canonical_name.lower()
).ratio()
score += name_similarity * 0.5
# Temporal proximity
if last_seen:
days_diff = abs((unit_event_date - last_seen).total_seconds() / 86400)
if days_diff < 7:
temporal_score = max(0, 1.0 - (days_diff / 7))
score += temporal_score * 0.2
if score > best_score:
best_score = score
best_candidate = candidate_id
best_name_similarity = name_similarity
# Apply threshold
threshold = 0.4 if entity_type == 'PERSON' and best_name_similarity >= 0.95 else 0.6
if best_score > threshold:
entity_ids[idx] = best_candidate
entities_to_update.append((best_candidate, unit_event_date))
else:
entities_to_create.append((idx, entity_data))
# Batch update existing entities
if entities_to_update:
from psycopg2.extras import execute_values
execute_values(
cursor,
"""
UPDATE entities SET
mention_count = mention_count + 1,
last_seen = data.last_seen
FROM (VALUES %s) AS data(id, last_seen)
WHERE entities.id = data.id::uuid
""",
entities_to_update
)
# Batch create new entities
if entities_to_create:
for idx, entity_data in entities_to_create:
entity_id = self._create_entity(
cursor, agent_id, entity_data['text'],
entity_data['type'], unit_event_date
)
entity_ids[idx] = entity_id
return entity_ids
finally:
cursor.close()
def resolve_entity(
self,
agent_id: str,
@ -297,6 +489,75 @@ class EntityResolver:
(entity_id_1, entity_id_2)
)
def link_units_to_entities_batch(self, unit_entity_pairs: List[tuple[str, str]]):
"""
Link multiple memory units to entities in batch (MUCH faster than sequential).
Also updates co-occurrence cache for entities that appear in the same unit.
Args:
unit_entity_pairs: List of (unit_id, entity_id) tuples
"""
if not unit_entity_pairs:
return
cursor = self.conn.cursor()
try:
# Batch insert all unit-entity links
from psycopg2.extras import execute_values
execute_values(
cursor,
"""
INSERT INTO unit_entities (unit_id, entity_id)
VALUES %s
ON CONFLICT DO NOTHING
""",
unit_entity_pairs
)
# Build map of unit -> entities for co-occurrence calculation
# Use sets to avoid duplicate entities in the same unit
unit_to_entities = {}
for unit_id, entity_id in unit_entity_pairs:
if unit_id not in unit_to_entities:
unit_to_entities[unit_id] = set()
unit_to_entities[unit_id].add(entity_id)
# Update co-occurrences for all pairs in each unit
cooccurrence_pairs = set() # Use set to avoid duplicates
for unit_id, entity_ids in unit_to_entities.items():
entity_list = list(entity_ids) # Convert set to list for iteration
# For each pair of entities in this unit, create co-occurrence
for i, entity_id_1 in enumerate(entity_list):
for entity_id_2 in entity_list[i+1:]:
# Skip if same entity (shouldn't happen with set, but be safe)
if entity_id_1 == entity_id_2:
continue
# Ensure consistent ordering (entity_id_1 < entity_id_2)
if entity_id_1 > entity_id_2:
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
cooccurrence_pairs.add((entity_id_1, entity_id_2))
# Batch update co-occurrences
if cooccurrence_pairs:
from datetime import datetime, timezone
now = datetime.now(timezone.utc)
execute_values(
cursor,
"""
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
VALUES %s
ON CONFLICT (entity_id_1, entity_id_2)
DO UPDATE SET
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
last_cooccurred = EXCLUDED.last_cooccurred
""",
[(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs]
)
finally:
cursor.close()
def get_units_by_entity(self, entity_id: str, limit: int = 100) -> List[str]:
"""
Get all units that mention an entity.

303
memory/llm_client.py
View file

@ -6,6 +6,8 @@ Uses OpenAI-compatible API (works with Groq, OpenAI, etc.)
import os
import json
import re
import asyncio
from datetime import datetime
from typing import List, Dict, Optional, Literal
from openai import AsyncOpenAI
from pydantic import BaseModel, Field
@ -16,16 +18,8 @@ class ExtractedFact(BaseModel):
fact: str = Field(
description="Self-contained factual statement with subject + action + context"
)
speaker: str = Field(
default="narrator",
description="Who said this (name or 'narrator' if not a conversation)"
)
type: Literal["biographical", "event", "opinion", "recommendation", "description", "relationship"] = Field(
description="Category of the fact"
)
confidence: Literal["high", "medium", "low"] = Field(
default="medium",
description="Confidence level of the extraction"
date: str = Field(
description="Absolute date/time when this fact occurred in ISO format (YYYY-MM-DDTHH:MM:SSZ). If text mentions relative time (yesterday, last week, this morning), calculate absolute date from the provided context date."
)
@ -36,75 +30,48 @@ class FactExtractionResponse(BaseModel):
)
def split_into_sentences(text: str) -> List[str]:
"""
Fast sentence splitter using regex.
Splits on periods, exclamation marks, and question marks followed by whitespace or end of string.
Args:
text: Input text to split
Returns:
List of sentences
"""
# Split on sentence boundaries: .!? followed by space/newline/end
sentences = re.split(r'(?<=[.!?])\s+', text)
return [s.strip() for s in sentences if s.strip()]
def chunk_text(text: str, max_chars: int = 120000) -> List[str]:
"""
Split text into chunks at sentence boundaries.
Split text into chunks at sentence boundaries using LangChain's text splitter.
Keeps chunks under max_chars (~30k tokens assuming 1 token ≈ 4 chars).
This prevents hitting output token limits on large documents.
Uses RecursiveCharacterTextSplitter which intelligently splits at sentence boundaries
and allows chunks to slightly exceed max_chars to finish sentences naturally.
Args:
text: Input text to chunk
max_chars: Maximum characters per chunk (default 120k ≈ 30k tokens)
Note: chunks may slightly exceed this to complete sentences
Returns:
List of text chunks, each under max_chars
List of text chunks, roughly under max_chars
"""
from langchain_text_splitters import RecursiveCharacterTextSplitter
# If text is small enough, return as-is
if len(text) <= max_chars:
return [text]
sentences = split_into_sentences(text)
chunks = []
current_chunk = []
current_length = 0
# Configure splitter to split at sentence boundaries first
# Separators in order of preference: paragraphs, newlines, sentences, words
splitter = RecursiveCharacterTextSplitter(
chunk_size=max_chars,
chunk_overlap=0,
length_function=len,
is_separator_regex=False,
separators=[
"\n\n", # Paragraph breaks
"\n", # Line breaks
". ", # Sentence endings
"! ", # Exclamations
"? ", # Questions
"; ", # Semicolons
", ", # Commas
" ", # Words
"", # Characters (last resort)
],
)
for sentence in sentences:
sentence_length = len(sentence)
# If single sentence exceeds max_chars, split it forcefully
if sentence_length > max_chars:
# Save current chunk if any
if current_chunk:
chunks.append(' '.join(current_chunk))
current_chunk = []
current_length = 0
# Split long sentence into smaller pieces
for i in range(0, len(sentence), max_chars):
chunks.append(sentence[i:i + max_chars])
continue
# If adding this sentence would exceed limit, start new chunk
if current_length + sentence_length + 1 > max_chars:
chunks.append(' '.join(current_chunk))
current_chunk = [sentence]
current_length = sentence_length
else:
current_chunk.append(sentence)
current_length += sentence_length + 1 # +1 for space
# Add remaining chunk
if current_chunk:
chunks.append(' '.join(current_chunk))
return chunks
return splitter.split_text(text)
def get_llm_client() -> AsyncOpenAI:
@ -137,22 +104,28 @@ def get_llm_client() -> AsyncOpenAI:
)
async def extract_facts_from_text(
text: str,
model: str = "openai/gpt-oss-20b",
temperature: float = 0.1,
max_tokens: int = 65000,
chunk_size: int = 60000
async def _extract_facts_from_chunk(
chunk: str,
chunk_index: int,
total_chunks: int,
event_date: datetime,
context: str,
model: str,
temperature: float,
max_tokens: int,
client: AsyncOpenAI
) -> List[Dict[str, str]]:
client = get_llm_client()
"""
Extract facts from a single chunk (internal helper for parallel processing).
"""
# Format event_date for the prompt
event_date_str = event_date.strftime("%Y-%m-%dT%H:%M:%SZ")
# Chunk text if necessary
chunks = chunk_text(text, max_chars=chunk_size)
prompt = f"""You are extracting facts from text for an AI memory system. Each fact will be stored and retrieved later.
all_facts = []
for i, chunk in enumerate(chunks):
prompt = f"""You are extracting facts from text for an AI memory system. Each fact will be stored and retrieved later.
## CONTEXT INFORMATION
- Current reference date/time: {event_date_str}
- Context: {context if context else 'no context provided'}
## CRITICAL: Facts must be DETAILED and COMPREHENSIVE
@ -164,66 +137,158 @@ Each fact should:
5. Include surrounding context that makes the fact meaningful
6. Capture nuances, reasons, causes, and implications
## What to EXTRACT:
- Biographical information (jobs, roles, backgrounds, experiences)
- Events (what happened, when, where, who was involved, why)
- Opinions and beliefs (who believes what and why)
- Recommendations and advice (specific suggestions with reasoning)
- Descriptions (detailed explanations of how things work)
- Relationships (connections between people, organizations, concepts)
## TEMPORAL INFORMATION (VERY IMPORTANT)
For each fact, extract the ABSOLUTE date/time when it occurred:
- If text mentions ABSOLUTE dates ("on March 15, 2024", "last Tuesday"), use that date
- If text mentions RELATIVE times ("yesterday", "last week", "this morning", "3 days ago"), calculate the absolute date using the reference date above.
- if text mentions a vague relative time without a specific day ("last week", "this morning"), transform the date in relative with absolute context ("last week" + " 2 june 2024" -> "week before June 2 2024") in the text and use the absolute date for the 'date' field
- If NO specific time is mentioned, use the reference date
- Always output dates in ISO format: YYYY-MM-DDTHH:MM:SSZ
## What to SKIP:
- Greetings, thank yous, acknowledgments
Examples of date extraction:
- Reference: 2024-03-20T10:00:00Z
- "Yesterday I went hiking" → date: 2024-03-19T10:00:00Z
- "Last week I joined Google" → date: 2024-03-13T10:00:00Z (approximately)
- "This morning I had coffee" → date: 2024-03-20T08:00:00Z
- "I work at Google" (no time mentioned) → date: 2024-03-20T10:00:00Z (use reference)
## What to EXTRACT (BE EXHAUSTIVE - DO NOT SKIP ANYTHING):
- **Biographical information**: jobs, roles, backgrounds, experiences, skills
- **Events (NEVER MISS THESE)**:
- ANY action that happened (went, did, attended, joined, started, finished, etc.)
- Photos, images, videos shared or taken ("here's a photo", "took a picture", "captured")
- Social activities (meetups, gatherings, meals, conversations)
- Achievements, milestones, accomplishments
- Travels, visits, locations visited
- Purchases, acquisitions, creations
- **Opinions and beliefs**: who believes what and why
- **Recommendations and advice**: specific suggestions with reasoning
- **Descriptions**: detailed explanations of how things work
- **Relationships**: connections between people, organizations, concepts
- **States and conditions**: current status, ongoing situations
## CRITICAL: Extract EVERY event mentioned, even casual ones
- "here's a photo of X" = someone took/shared a photo of X
- "I was with friends last week" = meetup/gathering with friends last week
- "sent you that link" = action of sending a link
- DO NOT skip events just because they seem minor or casual
## What to SKIP (ONLY these):
- Greetings, thank yous, acknowledgments (unless they reveal information)
- Filler words ("um", "uh", "like")
- Pure reactions without content ("wow", "cool")
- Incomplete thoughts
- Pure reactions without content ("wow", "cool", "nice")
- Incomplete thoughts or sentence fragments with no meaning
## EXAMPLES of GOOD facts (detailed, comprehensive):
Input: "Alice mentioned she works at Google in Mountain View. She joined the AI team last year and loves working on large language models."
GOOD: "Alice works at Google in Mountain View on the AI team, which she joined last year, and she loves working on large language models"
BAD: "Alice works at Google" (too short, missing context)
Input: "Alice mentioned she works at Google in Mountain View. She joined the AI team last year."
GOOD fact: "Alice works at Google in Mountain View on the AI team, which she joined last year"
GOOD date: Calculate based on reference date (if reference is 2024-03-20, "last year" = 2023-03-20)
Input: "Bob said he's been hiking every weekend in Yosemite because it helps him clear his mind after stressful work weeks."
GOOD: "Bob has been hiking every weekend in Yosemite because it helps him clear his mind after stressful work weeks"
BAD: "Bob hikes in Yosemite" (missing frequency, reason, and context)
Input: "Yesterday Bob went hiking in Yosemite because it helps him clear his mind."
GOOD fact: "Bob went hiking in Yosemite because it helps him clear his mind"
GOOD date: Reference date minus 1 day
Input: "The new algorithm reduced latency by 40% compared to the baseline by using a novel caching strategy."
GOOD: "The new algorithm reduced latency by 40% compared to the baseline by using a novel caching strategy"
BAD: "The algorithm is faster" (missing numbers, comparison, and method)
Input: "Here's a photo of me with my friends taken last week at the beach."
GOOD fact: "Someone shared/took a photo with their friends at the beach"
GOOD date: Reference date minus 7 days (last week)
NOTE: Extract the event (photo taken/shared with friends at beach), NOT just that a photo exists
Input: "I sent you that article about AI last Tuesday."
GOOD fact: "Someone sent an article about AI"
GOOD date: Calculate last Tuesday from reference date
## TEXT TO EXTRACT FROM:
{chunk}
Remember: Include ALL details, names, numbers, reasons, and context. Facts should be rich and informative, not summaries."""
Remember:
1. BE EXHAUSTIVE - Extract EVERY event, action, and fact mentioned
2. DO NOT skip casual mentions like "here's a photo", "I was with X", "sent you Y"
3. Include ALL details, names, numbers, reasons, and context in the fact text
4. Extract the absolute date for EACH fact by calculating relative times from the reference date
5. When in doubt, EXTRACT IT - better to have too many facts than miss important events"""
# Use parse() for structured outputs with Pydantic models
response = await client.beta.chat.completions.parse(
response = await client.beta.chat.completions.parse(
model=model,
messages=[
{
"role": "system",
"content": "You are an EXHAUSTIVE fact extractor. Extract EVERY event, action, and fact mentioned - never skip anything. This includes casual mentions like photos shared, things sent, meetups, gatherings, or any action. Preserve all context, details, and nuances. Calculate absolute dates from relative time expressions. When in doubt, extract it - missing facts is worse than extracting too many."
},
{
"role": "user",
"content": prompt
}
],
temperature=temperature,
max_tokens=max_tokens,
response_format=FactExtractionResponse,
extra_body={"service_tier": "auto"},
)
# Extract the parsed response
extraction_response = response.choices[0].message.parsed
# Convert to dict format
chunk_facts = [fact.model_dump() for fact in extraction_response.facts]
return chunk_facts
async def extract_facts_from_text(
text: str,
event_date: datetime,
context: str = "",
model: str = "openai/gpt-oss-120b",
temperature: float = 0.1,
max_tokens: int = 65000,
chunk_size: int = 5000
) -> List[Dict[str, str]]:
"""
Extract semantic facts from conversational or narrative text using LLM.
For large texts (>chunk_size chars), automatically chunks at sentence boundaries
to avoid hitting output token limits. Processes ALL chunks in PARALLEL for speed.
Args:
text: Input text (conversation, article, etc.)
event_date: Reference date for resolving relative times
context: Context about the conversation/document
model: LLM model to use
temperature: Sampling temperature (lower = more focused)
max_tokens: Maximum tokens in response
chunk_size: Maximum characters per chunk
Returns:
List of fact dictionaries with 'fact' and 'date' keys
"""
client = get_llm_client()
# Chunk text if necessary
chunks = chunk_text(text, max_chars=chunk_size)
# Process all chunks in parallel using asyncio.gather
tasks = [
_extract_facts_from_chunk(
chunk=chunk,
chunk_index=i,
total_chunks=len(chunks),
event_date=event_date,
context=context,
model=model,
messages=[
{
"role": "system",
"content": "You extract detailed, comprehensive facts from text. Preserve all context, details, and nuances. Never summarize or shorten - include everything relevant."
},
{
"role": "user",
"content": prompt
}
],
temperature=temperature,
max_tokens=max_tokens,
response_format=FactExtractionResponse
client=client
)
for i, chunk in enumerate(chunks)
]
# Extract the parsed response
extraction_response = response.choices[0].message.parsed
# Wait for all chunks to complete in parallel
chunk_results = await asyncio.gather(*tasks)
# Convert to dict format and add to aggregate
chunk_facts = [fact.model_dump() for fact in extraction_response.facts]
# Flatten results from all chunks
all_facts = []
for chunk_facts in chunk_results:
all_facts.extend(chunk_facts)
# Log progress for large documents
if len(chunks) > 1:
print(f"Processed chunk {i + 1}/{len(chunks)}: extracted {len(chunk_facts)} facts")
return all_facts

952
memory/temporal_semantic_memory.py
View file

File diff suppressed because it is too large Load diff

20
memory/utils.py
View file

@ -1,24 +1,28 @@
"""
Utility functions for memory system.
"""
from typing import List
from datetime import datetime
from typing import List, Dict
from .llm_client import extract_facts_from_text
async def extract_facts(text: str) -> List[str]:
async def extract_facts(text: str, event_date: datetime, context: str = "") -> List[Dict[str, str]]:
"""
Extract semantic facts from text using LLM.
Uses LLM for intelligent fact extraction that:
- Filters out social pleasantries and filler words
- Creates self-contained statements
- Creates self-contained statements with absolute dates
- Handles conversational text well
- Resolves relative time expressions to absolute dates
Args:
text: Input text (conversation, article, etc.)
event_date: Reference date for resolving relative times
context: Context about the conversation/document
Returns:
List of factual statements
List of fact dictionaries with keys: 'fact' (text) and 'date' (ISO string)
Raises:
Exception: If LLM fact extraction fails
@ -26,14 +30,12 @@ async def extract_facts(text: str) -> List[str]:
if not text or not text.strip():
return []
fact_dicts = await extract_facts_from_text(text)
# Extract just the fact text
facts = [f['fact'] for f in fact_dicts if f.get('fact')]
fact_dicts = await extract_facts_from_text(text, event_date, context)
if not facts:
if not fact_dicts:
raise Exception(f"LLM extracted 0 facts from text of length {len(text)}. This may indicate the text contains no meaningful information, or the LLM failed to extract facts.")
return facts
return fact_dicts
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:

514
memory_graph_interactive.html
View file

File diff suppressed because one or more lines are too long

3
pyproject.toml
View file

@ -15,10 +15,9 @@ dependencies = [
"matplotlib>=3.7.0",
"rich>=13.0.0",
"spacy>=3.7.0",
"pyvis>=0.3.0",
"sentence-transformers>=2.2.0",
"torch>=2.0.0",
"pytest>=7.0.0",
"pytest-asyncio>=0.21.0",
"fastcoref>=2.1.0",
"langchain-text-splitters>=0.3.0",
]

21
tests/conftest.py
View file

@ -2,10 +2,10 @@
Pytest configuration and shared fixtures.
"""
import pytest
import psycopg2
import os
from dotenv import load_dotenv
from memory import TemporalSemanticMemory
import psycopg2
load_dotenv()
@ -24,28 +24,17 @@ def memory():
def clean_agent(memory):
"""
Provide a clean agent ID and clean up data after test.
Uses agent_id='test' for all tests (multi-tenant isolation).
"""
agent_id = "test_agent"
agent_id = "test"
# Clean up before test
conn = psycopg2.connect(os.getenv('DATABASE_URL'))
cursor = conn.cursor()
cursor.execute("DELETE FROM memory_units WHERE agent_id = %s", (agent_id,))
cursor.execute("DELETE FROM entities WHERE agent_id = %s", (agent_id,))
conn.commit()
cursor.close()
conn.close()
memory.delete_agent(agent_id)
yield agent_id
# Clean up after test
conn = psycopg2.connect(os.getenv('DATABASE_URL'))
cursor = conn.cursor()
cursor.execute("DELETE FROM memory_units WHERE agent_id = %s", (agent_id,))
cursor.execute("DELETE FROM entities WHERE agent_id = %s", (agent_id,))
conn.commit()
cursor.close()
conn.close()
memory.delete_agent(agent_id)
@pytest.fixture

14
tests/test_chunking.py
View file

@ -2,19 +2,7 @@
Test chunking functionality for large documents.
"""
import pytest
from memory.llm_client import chunk_text, split_into_sentences
def test_split_into_sentences():
"""Test sentence splitting."""
text = "This is sentence one. This is sentence two! Is this sentence three? Yes it is."
sentences = split_into_sentences(text)
assert len(sentences) == 4, f"Expected 4 sentences, got {len(sentences)}"
assert "This is sentence one" in sentences[0]
assert "This is sentence two" in sentences[1]
assert "Is this sentence three" in sentences[2]
assert "Yes it is" in sentences[3]
from memory.llm_client import chunk_text
def test_chunk_text_small():

57
tests/test_coref_performance.py
View file

@ -1,57 +0,0 @@
"""
Performance test for coreference resolution.
"""
import time
from memory.coref_resolver import resolve_sentences, resolve_sentences_fast, resolve_sentences_legacy
def test_coref_performance():
"""Compare performance of fast vs legacy coreference resolution."""
# Sample sentences with coreferences
test_sentences = [
"John is a software engineer.",
"He works at a tech company.",
"The company is based in San Francisco.",
"He enjoys working on AI projects.",
"The projects involve machine learning.",
"John believes AI will transform the industry.",
"He has been working on this for 5 years.",
"The experience has been valuable.",
"John plans to continue his research.",
"He is passionate about the field.",
] * 10 # Repeat 10 times to make it 100 sentences
print(f"\nTesting with {len(test_sentences)} sentences...")
# Test fast method
start = time.time()
resolved_fast = resolve_sentences_fast(test_sentences)
fast_time = time.time() - start
print(f"FastCoref: {fast_time:.3f} seconds")
# Test legacy method (with smaller dataset to avoid timeout)
small_test = test_sentences[:20]
start = time.time()
resolved_legacy = resolve_sentences_legacy(small_test)
legacy_time = time.time() - start
print(f"Legacy (20 sentences): {legacy_time:.3f} seconds")
# Extrapolate legacy time
extrapolated_legacy = legacy_time * (len(test_sentences) / len(small_test)) ** 2
print(f"Legacy (extrapolated for {len(test_sentences)}): {extrapolated_legacy:.3f} seconds")
speedup = extrapolated_legacy / fast_time if fast_time > 0 else float('inf')
print(f"Speedup: {speedup:.1f}x faster")
# Verify resolution worked
print("\nSample resolved sentences (FastCoref):")
for i, sent in enumerate(resolved_fast[:3]):
print(f" {i+1}. {sent}")
assert len(resolved_fast) == len(test_sentences)
assert fast_time < extrapolated_legacy
if __name__ == "__main__":
test_coref_performance()

111
tests/test_deduplication.py
View file

@ -1,111 +0,0 @@
"""
Test deduplication of identical puts.
"""
import pytest
from datetime import datetime, timezone
from memory.temporal_semantic_memory import TemporalSemanticMemory
@pytest.fixture
def memory():
"""Create a memory instance for testing."""
mem = TemporalSemanticMemory()
yield mem
# Cleanup after test
cursor = mem.conn.cursor()
cursor.execute("DELETE FROM memory_units WHERE agent_id LIKE 'test_%'")
mem.conn.commit()
cursor.close()
@pytest.mark.asyncio
async def test_duplicate_put_filters_identical_content(memory):
"""Test that putting the same content twice doesn't create duplicates."""
agent_id = "test_dedup_agent"
content = "Alice works at Google as a software engineer. She joined last year and loves Python."
event_date = datetime(2024, 1, 15, 10, 0, 0, tzinfo=timezone.utc)
# First put - should create units
print("\n--- FIRST PUT ---")
units_1 = await memory.put_async(agent_id, content, "Test context", event_date)
assert len(units_1) > 0, "First put should create units"
print(f"First put created {len(units_1)} units")
# Second put with identical content and same date - should be filtered as duplicates
print("\n--- SECOND PUT (identical) ---")
units_2 = await memory.put_async(agent_id, content, "Test context", event_date)
assert len(units_2) == 0, "Second identical put should create no new units (all duplicates)"
print(f"Second put created {len(units_2)} units (expected 0)")
# Verify database has only the first set of units
cursor = memory.conn.cursor()
cursor.execute(
"SELECT COUNT(*) FROM memory_units WHERE agent_id = %s",
(agent_id,)
)
total_units = cursor.fetchone()[0]
cursor.close()
assert total_units == len(units_1), f"Database should have {len(units_1)} units, found {total_units}"
print(f"✅ Deduplication working: {total_units} total units in database")
@pytest.mark.asyncio
async def test_duplicate_put_with_paraphrased_content(memory):
"""Test that similar but paraphrased content is also deduplicated."""
agent_id = "test_paraphrase_agent"
event_date = datetime(2024, 1, 15, 10, 0, 0, tzinfo=timezone.utc)
# First put
content_1 = "Bob is a chef in New York. He owns a restaurant."
print("\n--- FIRST PUT ---")
units_1 = await memory.put_async(agent_id, content_1, "Test", event_date)
assert len(units_1) > 0, "First put should create units"
print(f"First put created {len(units_1)} units")
# Second put with paraphrased content - should be mostly deduplicated
# The LLM will extract similar facts that should match via embeddings
content_2 = "Bob works as a chef in New York City. He is the owner of a restaurant."
print("\n--- SECOND PUT (paraphrased) ---")
units_2 = await memory.put_async(agent_id, content_2, "Test", event_date)
# May create 0 or very few new units (depending on how LLM extracts facts)
print(f"Second put created {len(units_2)} units")
print(f"Deduplication ratio: {len(units_2)}/{len(units_1)} new units from paraphrase")
# Just verify it doesn't create the same number of units (some deduplication should happen)
assert len(units_2) < len(units_1), "Paraphrased content should have fewer new units due to deduplication"
@pytest.mark.asyncio
async def test_different_dates_not_deduplicated(memory):
"""Test that same content with different dates is NOT deduplicated."""
agent_id = "test_dates_agent"
content = "Charlie went hiking in Yosemite."
# First put at date 1
date_1 = datetime(2024, 1, 1, 10, 0, 0, tzinfo=timezone.utc)
print("\n--- FIRST PUT (Jan 1) ---")
units_1 = await memory.put_async(agent_id, content, "Test", date_1)
assert len(units_1) > 0
print(f"First put created {len(units_1)} units")
# Second put at date 2 (outside 24-hour window)
date_2 = datetime(2024, 2, 1, 10, 0, 0, tzinfo=timezone.utc)
print("\n--- SECOND PUT (Feb 1, outside time window) ---")
units_2 = await memory.put_async(agent_id, content, "Test", date_2)
# Should create new units because dates are far apart
assert len(units_2) > 0, "Same content with different dates (outside window) should create new units"
print(f"Second put created {len(units_2)} units (not deduplicated due to date difference)")
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])

107
tests/test_detailed_extraction.py
View file

@ -1,107 +0,0 @@
"""
Test that the improved prompt extracts detailed, comprehensive facts.
"""
import pytest
from memory.llm_client import extract_facts_from_text
@pytest.mark.asyncio
async def test_detailed_extraction_preserves_context():
"""Test that facts preserve all context and details."""
text = """
Alice mentioned she works at Google in Mountain View on the AI research team.
She joined last year after finishing her PhD at Stanford, and she's currently
focused on improving large language model safety through red teaming and
adversarial testing. She said the work is challenging but very rewarding because
it directly impacts millions of users.
"""
facts = await extract_facts_from_text(text)
print(f"\nExtracted {len(facts)} facts:")
for i, fact in enumerate(facts, 1):
print(f"{i}. {fact['fact']}")
print(f" Type: {fact['type']}, Speaker: {fact['speaker']}, Confidence: {fact['confidence']}\n")
# Verify we got facts
assert len(facts) > 0, "Should extract at least one fact"
# Check that facts contain detailed information
fact_texts = [f['fact'].lower() for f in facts]
combined_facts = ' '.join(fact_texts)
# Should preserve location details
assert 'mountain view' in combined_facts, "Should preserve specific location 'Mountain View'"
# Should preserve team/department
assert 'ai' in combined_facts or 'research' in combined_facts, "Should preserve team information"
# Should preserve educational background
assert 'stanford' in combined_facts or 'phd' in combined_facts, "Should preserve educational background"
# Should preserve work details
assert 'safety' in combined_facts or 'red teaming' in combined_facts or 'adversarial' in combined_facts, \
"Should preserve specific work focus details"
# Check that at least one fact is reasonably detailed (not just "Alice works at Google")
detailed_fact_found = any(len(f['fact'].split()) >= 10 for f in facts)
assert detailed_fact_found, "At least one fact should be detailed (10+ words)"
@pytest.mark.asyncio
async def test_numbers_and_metrics_preserved():
"""Test that numbers, percentages, and metrics are preserved."""
text = """
Bob explained that the new caching algorithm reduced API latency by 40%
compared to the baseline, processing 10,000 requests per second instead
of the previous 7,000. This improvement was achieved by implementing a
two-tier LRU cache with 1GB memory allocation.
"""
facts = await extract_facts_from_text(text)
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
print(f"- {fact['fact']}")
combined = ' '.join([f['fact'] for f in facts])
# Should preserve specific numbers
assert '40' in combined or 'forty' in combined.lower(), "Should preserve percentage"
assert '10,000' in combined or '10000' in combined or 'ten thousand' in combined.lower(), \
"Should preserve request rate"
assert 'cache' in combined.lower(), "Should preserve technical details"
@pytest.mark.asyncio
async def test_reasons_and_causality_preserved():
"""Test that reasons, causes, and explanations are preserved."""
text = """
Sarah has been meditating every morning for the past 6 months because
she found it significantly reduced her anxiety levels and improved her
focus during work hours. She started this practice after reading a research
paper on mindfulness benefits.
"""
facts = await extract_facts_from_text(text)
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
print(f"- {fact['fact']}")
combined = ' '.join([f['fact'] for f in facts])
# Should preserve the causal relationship (because/reason)
assert any(keyword in combined.lower() for keyword in ['because', 'reduced', 'anxiety', 'improved']), \
"Should preserve the reason/causality"
# Should preserve frequency
assert 'morning' in combined.lower() or 'every' in combined.lower(), \
"Should preserve frequency information"
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])

212
tests/test_entity_linking.py
View file

@ -1,212 +0,0 @@
"""
Test entity-aware memory linking functionality.
Tests that entity resolution connects memories about the same person/place/thing.
"""
import pytest
from datetime import datetime, timedelta, timezone
def utcnow():
"""Get current UTC time with timezone info."""
return datetime.now(timezone.utc)
def test_entity_extraction_and_linking(memory, clean_agent, db_connection):
"""Test that entities are extracted and linked correctly."""
agent_id = clean_agent
# Store memories about Alice's hiking hobby
memory.put(
agent_id=agent_id,
content="Alice told me she loves hiking in the mountains. "
"She goes hiking every weekend in Yosemite.",
context="Casual conversation about hobbies",
event_date=utcnow() - timedelta(days=7),
)
# Store memories about Alice's work (different context!)
memory.put(
agent_id=agent_id,
content="Alice works at Google as a software engineer. "
"She joined Google last year and loves the culture.",
context="Discussion about careers",
event_date=utcnow() - timedelta(days=3),
)
# Store more about hiking (no Alice mention)
memory.put(
agent_id=agent_id,
content="Bob mentioned he enjoys rock climbing. "
"He climbs in Yosemite too, on weekends.",
context="Outdoor activities discussion",
event_date=utcnow() - timedelta(days=1),
)
# Store another Alice memory
memory.put(
agent_id=agent_id,
content="Alice is working on a Python project at Google. "
"The project uses machine learning.",
context="Technical discussion",
event_date=utcnow(),
)
# Verify entities were extracted
cursor = db_connection.cursor()
cursor.execute("""
SELECT canonical_name, entity_type, mention_count
FROM entities
WHERE agent_id = %s
ORDER BY mention_count DESC
""", (agent_id,))
entities = cursor.fetchall()
entity_names = [e[0] for e in entities]
# Should have Alice, Google, Yosemite, Bob
assert "Alice" in entity_names, "Alice entity should be extracted"
assert "Google" in entity_names, "Google entity should be extracted"
assert "Yosemite" in entity_names, "Yosemite entity should be extracted"
assert "Bob" in entity_names, "Bob entity should be extracted"
# Alice should have multiple mentions
alice_entity = next((e for e in entities if e[0] == "Alice"), None)
assert alice_entity is not None
assert alice_entity[2] >= 3, "Alice should have at least 3 mentions"
# Verify entity links exist
cursor.execute("""
SELECT COUNT(*)
FROM memory_links
WHERE link_type = 'entity'
AND from_unit_id IN (
SELECT id FROM memory_units WHERE agent_id = %s
)
""", (agent_id,))
entity_link_count = cursor.fetchone()[0]
assert entity_link_count > 0, "Entity links should be created"
cursor.close()
def test_entity_search_retrieves_all_related_memories(memory, clean_agent):
"""Test that searching for an entity retrieves ALL memories about that entity."""
agent_id = clean_agent
# Store diverse memories about Alice
memory.put(
agent_id=agent_id,
content="Alice loves hiking in the mountains.",
context="Hobbies",
event_date=utcnow() - timedelta(days=7),
)
memory.put(
agent_id=agent_id,
content="Alice works at Google as a software engineer.",
context="Career",
event_date=utcnow() - timedelta(days=3),
)
memory.put(
agent_id=agent_id,
content="Alice is working on a Python machine learning project.",
context="Technical",
event_date=utcnow(),
)
# Query about Alice - should get ALL Alice memories via entity links
results = memory.search(
agent_id=agent_id,
query="What does Alice do?",
thinking_budget=30,
top_k=10,
)
# Should retrieve multiple memories about Alice
assert len(results) >= 2, "Should find multiple memories about Alice"
# Check that results contain Alice-related content
alice_mentions = sum(1 for r in results if "Alice" in r['text'])
assert alice_mentions >= 2, "Multiple results should mention Alice"
def test_entity_disambiguation(memory, clean_agent, db_connection):
"""Test that entity disambiguation correctly identifies same vs different entities."""
agent_id = clean_agent
# Store two memories about "Alice" in different contexts
memory.put(
agent_id=agent_id,
content="Alice from engineering loves Python.",
context="Tech team",
event_date=utcnow() - timedelta(days=2),
)
memory.put(
agent_id=agent_id,
content="Alice from engineering is working on a new project.",
context="Tech team",
event_date=utcnow(),
)
# Check that only ONE Alice entity was created (not two)
cursor = db_connection.cursor()
cursor.execute("""
SELECT COUNT(*)
FROM entities
WHERE agent_id = %s AND canonical_name = 'Alice'
""", (agent_id,))
alice_count = cursor.fetchone()[0]
assert alice_count == 1, "Should create only one Alice entity (disambiguation)"
cursor.close()
def test_link_type_distribution(memory, clean_agent, db_connection):
"""Test that all three link types (temporal, semantic, entity) are created."""
agent_id = clean_agent
# Store related memories
memory.put(
agent_id=agent_id,
content="Alice works at Google. She loves her job.",
context="Career",
event_date=utcnow() - timedelta(hours=2),
)
memory.put(
agent_id=agent_id,
content="Bob also works at Google. He is in sales.",
context="Career",
event_date=utcnow() - timedelta(hours=1),
)
memory.put(
agent_id=agent_id,
content="Google is a great company to work for.",
context="Career",
event_date=utcnow(),
)
# Check link types
cursor = db_connection.cursor()
cursor.execute("""
SELECT link_type, COUNT(*) as count
FROM memory_links ml
JOIN memory_units mu ON ml.from_unit_id = mu.id
WHERE mu.agent_id = %s
GROUP BY link_type
ORDER BY count DESC
""", (agent_id,))
link_types = {row[0]: row[1] for row in cursor.fetchall()}
# Should have at least temporal and entity links (semantic depends on similarity threshold)
assert 'temporal' in link_types, "Should create temporal links"
assert 'entity' in link_types or 'semantic' in link_types, "Should create entity or semantic links"
cursor.close()

130
tests/test_fact_extraction.py
View file

@ -1,130 +0,0 @@
"""
Test LLM-based fact extraction.
"""
import pytest
from memory.llm_client import extract_facts_from_text
from memory.utils import extract_facts
async def test_fact_extraction_filters_pleasantries():
"""Test that fact extraction filters out social pleasantries."""
conversation = """
Host: Welcome to the show, Marta! Thanks for joining us.
Marta: Oh, thank you so much for having me!
Host: So tell us, what do you do?
Marta: I work at Google as a software engineer. I've been there for 3 years now.
Host: That's amazing!
Marta: Yeah, I really enjoy it. I mostly work on AI infrastructure.
Host: Uh-huh, interesting.
Marta: And I'm also passionate about hiking. I go to Yosemite almost every weekend.
Host: Wow, that sounds great!
"""
facts = await extract_facts_from_text(conversation)
# Extract just the fact texts
fact_texts = [f['fact'].lower() for f in facts]
print("\nExtracted facts:")
for fact in facts:
print(f" - {fact['fact']} (speaker: {fact['speaker']}, type: {fact['type']})")
# Should extract meaningful facts
assert any('google' in fact and 'software engineer' in fact for fact in fact_texts), \
"Should extract Marta's job at Google"
assert any('yosemite' in fact and 'hiking' in fact for fact in fact_texts), \
"Should extract Marta's hiking hobby"
# Should NOT extract pleasantries
assert not any('thank you' in fact for fact in fact_texts), \
"Should not extract 'thank you'"
assert not any('amazing' in fact and len(fact.split()) < 5 for fact in fact_texts), \
"Should not extract simple reactions like 'that's amazing'"
assert not any('uh-huh' in fact for fact in fact_texts), \
"Should not extract acknowledgments"
async def test_fact_extraction_makes_self_contained():
"""Test that facts are self-contained (pronouns resolved)."""
conversation = """
Alice told me she works at Microsoft.
She mentioned that she's been there for 5 years.
She really enjoys her team.
"""
facts = await extract_facts_from_text(conversation)
print("\nExtracted facts:")
for fact in facts:
print(f" - {fact['fact']}")
# All facts should mention "Alice" explicitly, not "she"
for fact in facts:
fact_text = fact['fact'].lower()
# If it's about Alice, it should say "alice" not "she"
if 'microsoft' in fact_text or 'team' in fact_text:
assert 'alice' in fact_text, \
f"Fact should be self-contained with 'Alice', not pronouns: {fact['fact']}"
async def test_extract_facts_util_function():
"""Test the utils.extract_facts() wrapper function."""
text = """
Bob is a chef in New York. He owns a restaurant called "The Kitchen".
Thank you! Yeah, uh-huh.
"""
facts = await extract_facts(text)
print("\nExtracted facts:")
for fact in facts:
print(f" - {fact}")
assert len(facts) > 0, "Should extract at least one fact"
assert any('bob' in fact.lower() and 'chef' in fact.lower() for fact in facts), \
"Should extract Bob's profession"
assert not any('thank you' in fact.lower() for fact in facts), \
"Should filter out pleasantries"
async def test_extract_facts_basic():
"""Test basic fact extraction."""
text = "Alice works at Google. She loves Python programming."
facts = await extract_facts(text)
assert len(facts) > 0, "Should extract at least one fact"
assert any('alice' in fact.lower() for fact in facts), "Should extract facts about Alice"
if __name__ == "__main__":
import asyncio
# Run a manual test
async def main():
conversation = """
Host: Welcome to the AI podcast! Today we have Dr. Sarah Chen with us.
Sarah: Hi! Thanks for having me.
Host: So Sarah, tell us about your work.
Sarah: I'm a researcher at Stanford focusing on large language models.
Host: Oh wow!
Sarah: Yeah, I've been studying how LLMs handle reasoning tasks. It's fascinating.
Sarah: We published a paper last month showing that chain-of-thought prompting improves accuracy by 40%.
Host: That's incredible!
Sarah: And I'm also advising a startup called MemoryAI that's building long-term memory systems.
Host: Cool, cool.
"""
print("Testing fact extraction with podcast conversation:")
print("=" * 60)
facts = await extract_facts_from_text(conversation)
print(f"\nExtracted {len(facts)} facts:\n")
for i, fact in enumerate(facts, 1):
print(f"{i}. {fact['fact']}")
print(f" Speaker: {fact['speaker']}, Type: {fact['type']}, Confidence: {fact['confidence']}\n")
asyncio.run(main())

223
tests/test_memory_operations.py
View file

@ -1,223 +0,0 @@
"""
Test basic memory operations: PUT, SEARCH, GET_RECENT.
Tests the core functionality of the temporal + semantic memory system.
"""
import pytest
import asyncio
from datetime import datetime, timedelta, timezone
def utcnow():
"""Get current UTC time with timezone info."""
return datetime.now(timezone.utc)
@pytest.mark.asyncio
async def test_put_creates_memory_units(memory, clean_agent, db_connection):
"""Test that PUT operation creates memory units."""
agent_id = clean_agent
# Store a conversation
await memory.put_async(
agent_id=agent_id,
content="Alice told me she loves hiking in the mountains. "
"She mentioned that she goes hiking every weekend. "
"Her favorite trail is in Yosemite National Park.",
context="Casual conversation about hobbies",
event_date=utcnow() - timedelta(hours=2),
)
# Verify memory units were created
cursor = db_connection.cursor()
cursor.execute("SELECT COUNT(*) FROM memory_units WHERE agent_id = %s", (agent_id,))
count = cursor.fetchone()[0]
assert count > 0, "Memory units should be created"
assert count <= 3, "Should create approximately 3 units (one per sentence)"
cursor.close()
@pytest.mark.asyncio
async def test_put_creates_temporal_links(memory, clean_agent, db_connection):
"""Test that temporal links are created between recent memories."""
agent_id = clean_agent
# Store two memories close in time
await memory.put_async(
agent_id=agent_id,
content="Alice loves hiking.",
context="Hobbies",
event_date=utcnow() - timedelta(hours=2),
)
await memory.put_async(
agent_id=agent_id,
content="Bob enjoys climbing.",
context="Sports",
event_date=utcnow() - timedelta(hours=1),
)
# Verify temporal links were created
cursor = db_connection.cursor()
cursor.execute("""
SELECT COUNT(*)
FROM memory_links
WHERE link_type = 'temporal'
AND from_unit_id IN (
SELECT id FROM memory_units WHERE agent_id = %s
)
""", (agent_id,))
temporal_link_count = cursor.fetchone()[0]
assert temporal_link_count > 0, "Temporal links should be created"
cursor.close()
@pytest.mark.asyncio
async def test_put_creates_semantic_links(memory, clean_agent, db_connection):
"""Test that semantic links are created between similar memories."""
agent_id = clean_agent
# Store semantically similar memories
await memory.put_async(
agent_id=agent_id,
content="Alice loves hiking in the mountains.",
context="Hobbies",
event_date=utcnow() - timedelta(days=2),
)
await memory.put_async(
agent_id=agent_id,
content="Bob enjoys climbing mountains.",
context="Sports",
event_date=utcnow(),
)
# Verify semantic links were created
cursor = db_connection.cursor()
cursor.execute("""
SELECT COUNT(*)
FROM memory_links
WHERE link_type = 'semantic'
AND from_unit_id IN (
SELECT id FROM memory_units WHERE agent_id = %s
)
""", (agent_id,))
semantic_link_count = cursor.fetchone()[0]
# Semantic links may or may not be created depending on similarity threshold
# So we just check that the query works
assert semantic_link_count >= 0, "Query should execute successfully"
cursor.close()
@pytest.mark.asyncio
async def test_search_with_spreading_activation(memory, clean_agent):
"""Test search using spreading activation algorithm."""
agent_id = clean_agent
# Store memories about outdoor activities
await memory.put_async(
agent_id=agent_id,
content="Alice told me she loves hiking in the mountains. "
"She goes hiking every weekend.",
context="Casual conversation about hobbies",
event_date=utcnow() - timedelta(hours=2),
)
await memory.put_async(
agent_id=agent_id,
content="Bob mentioned he enjoys rock climbing. "
"He climbs mountains on weekends too.",
context="Discussion about outdoor sports",
event_date=utcnow() - timedelta(hours=1),
)
# Search for outdoor activities
results = memory.search(
agent_id=agent_id,
query="outdoor mountain activities",
thinking_budget=50,
top_k=5,
)
assert len(results) > 0, "Search should return results"
# Verify result structure
for result in results:
assert 'id' in result, "Result should have id"
assert 'text' in result, "Result should have text"
assert 'weight' in result, "Result should have weight"
assert 'activation' in result, "Result should have activation"
assert 'recency' in result, "Result should have recency"
assert 'frequency' in result, "Result should have frequency"
# Results should be sorted by weight (descending)
weights = [r['weight'] for r in results]
assert weights == sorted(weights, reverse=True), "Results should be sorted by weight"
@pytest.mark.asyncio
async def test_search_returns_relevant_memories(memory, clean_agent):
"""Test that search returns semantically relevant memories."""
agent_id = clean_agent
# Store memories about different topics
await memory.put_async(
agent_id=agent_id,
content="Alice loves hiking in the mountains.",
context="Hobbies",
event_date=utcnow() - timedelta(hours=2),
)
await memory.put_async(
agent_id=agent_id,
content="Bob is working on a Python web application.",
context="Tech",
event_date=utcnow() - timedelta(hours=1),
)
# Search for programming-related memories
results = memory.search(
agent_id=agent_id,
query="software development",
thinking_budget=50,
top_k=3,
)
# Should find the programming-related memory
assert len(results) > 0, "Search should return results"
# Top result should be about programming (more relevant)
top_result_text = results[0]['text'].lower()
assert 'python' in top_result_text or 'application' in top_result_text or 'working' in top_result_text, \
"Top result should be about programming"
@pytest.mark.asyncio
async def test_search_with_no_results(memory, clean_agent):
"""Test search behavior when no relevant memories exist."""
agent_id = clean_agent
# Store unrelated memories
await memory.put_async(
agent_id=agent_id,
content="Alice loves cooking pasta.",
context="Food",
event_date=utcnow(),
)
# Search for something completely unrelated
results = memory.search(
agent_id=agent_id,
query="quantum physics theories",
thinking_budget=20,
top_k=5,
)
# May return low-scoring results or empty list
assert isinstance(results, list), "Search should return a list"

94
tests/test_temporal_extraction.py Normal file
View file

@ -0,0 +1,94 @@
"""
Test temporal extraction and per-fact dating.
"""
import pytest
from datetime import datetime, timezone, timedelta
from memory.llm_client import extract_facts_from_text
@pytest.mark.asyncio
async def test_extract_facts_with_relative_dates():
"""Test that relative dates are converted to absolute dates."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=timezone.utc)
text = """
Yesterday I went hiking in Yosemite.
Last week I started my new job at Google.
This morning I had coffee with Alice.
"""
facts = await extract_facts_from_text(text, reference_date, "Personal diary")
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
print(f"- {fact['fact']}")
print(f" Date: {fact['date']}")
# Verify we got facts
assert len(facts) > 0, "Should extract at least one fact"
# Check that all facts have dates
for fact in facts:
assert 'fact' in fact, "Each fact should have 'fact' field"
assert 'date' in fact, "Each fact should have 'date' field"
assert fact['date'], f"Date should not be empty for fact: {fact['fact']}"
# Verify dates are different (not all using reference date)
dates = [f['date'] for f in facts]
unique_dates = set(dates)
if len(facts) >= 3:
assert len(unique_dates) >= 2, "Should have different dates for different temporal facts"
print(f"\n✅ All facts have absolute dates")
@pytest.mark.asyncio
async def test_extract_facts_with_no_temporal_info():
"""Test that facts without temporal info use the reference date."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=timezone.utc)
text = "Alice works at Google. She loves Python programming."
facts = await extract_facts_from_text(text, reference_date, "General info")
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
print(f"- {fact['fact']}")
print(f" Date: {fact['date']}")
assert len(facts) > 0, "Should extract at least one fact"
# All facts should use the reference date since no temporal info is mentioned
for fact in facts:
assert fact['date'], f"Fact should have a date: {fact['fact']}"
@pytest.mark.asyncio
async def test_extract_facts_with_absolute_dates():
"""Test that absolute dates in text are preserved."""
reference_date = datetime(2024, 3, 20, 14, 0, 0, tzinfo=timezone.utc)
text = """
On March 15, 2024, Alice joined Google.
Bob will start his vacation on April 1st.
"""
facts = await extract_facts_from_text(text, reference_date, "Calendar events")
print(f"\nExtracted {len(facts)} facts:")
for fact in facts:
print(f"- {fact['fact']}")
print(f" Date: {fact['date']}")
assert len(facts) > 0, "Should extract at least one fact"
# Check that dates are present
for fact in facts:
assert fact['date'], f"Fact should have a date: {fact['fact']}"
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])

96
tests/test_visualization.py
View file

@ -1,96 +0,0 @@
"""
Test visualization functionality.
Tests memory graph data retrieval (not actual rendering).
"""
import pytest
from datetime import datetime, timedelta, timezone
def utcnow():
"""Get current UTC time with timezone info."""
return datetime.now(timezone.utc)
def test_get_memory_graph_data(memory, clean_agent):
"""Test retrieval of memory graph data for visualization."""
agent_id = clean_agent
# Store some memories
memory.put(
agent_id=agent_id,
content="Alice loves hiking in the mountains.",
context="Hobbies",
event_date=utcnow() - timedelta(hours=2),
)
memory.put(
agent_id=agent_id,
content="Bob enjoys rock climbing.",
context="Sports",
event_date=utcnow() - timedelta(hours=1),
)
memory.put(
agent_id=agent_id,
content="Alice is working on a Python project.",
context="Tech",
event_date=utcnow(),
)
# Get graph data
units, links = memory.get_memory_graph_data(agent_id)
assert isinstance(units, list), "Units should be a list"
assert isinstance(links, list), "Links should be a list"
assert len(units) > 0, "Should have memory units"
# Verify unit structure
for unit in units:
assert 'id' in unit, "Unit should have id"
assert 'text' in unit, "Unit should have text"
assert 'context' in unit, "Unit should have context"
assert 'event_date' in unit, "Unit should have event_date"
assert 'access_count' in unit, "Unit should have access_count"
# Links may or may not exist depending on similarity/proximity
if len(links) > 0:
# Verify link structure
for link in links:
assert 'from_unit_id' in link, "Link should have from_unit_id"
assert 'to_unit_id' in link, "Link should have to_unit_id"
assert 'link_type' in link, "Link should have link_type"
assert 'weight' in link, "Link should have weight"
assert link['link_type'] in ['temporal', 'semantic', 'entity'], \
"Link type should be temporal, semantic, or entity"
def test_memory_graph_has_correct_agent_data(memory, clean_agent):
"""Test that graph data only includes data for the specified agent."""
agent_id = clean_agent
other_agent_id = "other_agent"
# Store memories for test agent
memory.put(
agent_id=agent_id,
content="Alice loves hiking.",
context="Hobbies",
event_date=utcnow(),
)
# Store memories for another agent
memory.put(
agent_id=other_agent_id,
content="Charlie enjoys swimming.",
context="Sports",
event_date=utcnow(),
)
# Get graph data for test agent
units, links = memory.get_memory_graph_data(agent_id)
# Should only include test agent's data
for unit in units:
# Verify by checking text content (Alice should be present, Charlie should not)
unit_text = unit['text']
assert 'Charlie' not in unit_text, "Should not include other agent's memories"

1268
uv.lock
View file

File diff suppressed because it is too large Load diff

680
visualizations/interactive_graph.py
View file

@ -1,20 +1,19 @@
"""
Interactive HTML graph visualization of memory system.
Uses pyvis to create a smooth, interactive network graph that can be
Uses Cytoscape.js to create a performant, interactive network graph that can be
explored in the browser. Shows all memory units and their links with weights.
"""
import psycopg2
from dotenv import load_dotenv
import os
from pyvis.network import Network
import networkx as nx
import json
load_dotenv()
def create_interactive_graph():
"""Create an interactive HTML graph visualization."""
"""Create an interactive HTML graph visualization using Cytoscape.js."""
# Connect to database
conn = psycopg2.connect(os.getenv('DATABASE_URL'))
@ -61,235 +60,532 @@ def create_interactive_graph():
entity_map[unit_id] = []
entity_map[unit_id].append(f"{entity_name} ({entity_type})")
# Create pyvis network
net = Network(
height="900px",
width="100%",
bgcolor="#ffffff",
font_color="#000000",
heading="Entity-Aware Memory Graph - Interactive Visualization"
)
# Configure physics for smooth layout with performance optimizations
net.set_options("""
{
"nodes": {
"font": {
"size": 14,
"face": "Tahoma"
},
"borderWidth": 2,
"borderWidthSelected": 3
},
"edges": {
"smooth": {
"enabled": false
},
"font": {
"size": 10,
"align": "middle"
}
},
"physics": {
"enabled": true,
"stabilization": {
"enabled": true,
"iterations": 100,
"updateInterval": 10
},
"barnesHut": {
"gravitationalConstant": -12000,
"centralGravity": 0.2,
"springLength": 350,
"springConstant": 0.02,
"damping": 0.09,
"avoidOverlap": 0.8
},
"solver": "barnesHut",
"timestep": 0.5,
"adaptiveTimestep": true
},
"interaction": {
"hover": true,
"tooltipDelay": 100,
"navigationButtons": true,
"keyboard": true
}
}
""")
# Build Cytoscape.js graph data
cy_nodes = []
cy_edges = []
# Add nodes
for unit_id, text, event_date, context in units:
# Truncate text for display
display_text = text[:50] + "..." if len(text) > 50 else text
# Get entities
entities = entity_map.get(unit_id, [])
entity_str = "\\n".join(entities) if entities else "No entities"
# Build node label and title (hover)
label = display_text
title = f"""
<b>Text:</b> {text}<br>
<b>Date:</b> {event_date.date()}<br>
<b>Context:</b> {context}<br>
<b>Entities:</b> {entity_str}
"""
entity_count = len(entities)
# Color by entity count
if len(entities) == 0:
color = "#e0e0e0" # Gray
size = 20
elif len(entities) == 1:
color = "#90caf9" # Light blue
size = 25
if entity_count == 0:
color = "#e0e0e0"
elif entity_count == 1:
color = "#90caf9"
else:
color = "#42a5f5" # Dark blue
size = 30
color = "#42a5f5"
net.add_node(
str(unit_id),
label=label,
title=title,
color=color,
size=size,
shape="box",
font={"color": "#000000"}
)
cy_nodes.append({
"data": {
"id": str(unit_id),
"label": text[:50] + "..." if len(text) > 50 else text,
"text": text,
"context": context,
"date": str(event_date.date()),
"entities": ", ".join(entities) if entities else "None",
"color": color
}
})
# Add edges with colors and weights
# Add edges
for from_id, to_id, link_type, weight, entity_name in links:
# Set color and style based on link type
# Set color based on link type
if link_type == 'temporal':
color = "#00bcd4" # Cyan
dashes = [5, 5]
width = 0.5
label = f"T: {weight:.2f}"
color = "#00bcd4"
line_style = "dashed"
elif link_type == 'semantic':
color = "#ff69b4" # Pink
dashes = False
width = 0.5
label = f"S: {weight:.2f}"
color = "#ff69b4"
line_style = "solid"
elif link_type == 'entity':
color = "#ffd700" # Gold
dashes = False
width = 0.8
label = f"{entity_name}: {weight:.2f}"
color = "#ffd700"
line_style = "solid"
else:
color = "#999999"
dashes = False
width = 0.5
label = f"{weight:.2f}"
line_style = "solid"
net.add_edge(
str(from_id),
str(to_id),
value=weight * 1, # Scale for visual thickness
color=color,
dashes=dashes,
width=width,
label=label,
title=f"{link_type.upper()}: {weight:.3f}" + (f" (Entity: {entity_name})" if entity_name else "")
)
cy_edges.append({
"data": {
"id": f"{from_id}-{to_id}-{link_type}",
"source": str(from_id),
"target": str(to_id),
"weight": weight,
"linkType": link_type,
"entityName": entity_name or "",
"color": color,
"lineStyle": line_style
}
})
# Add legend as HTML
legend_html = """
<div style="position: absolute; top: 80px; left: 10px; background: white;
padding: 15px; border: 2px solid #333; border-radius: 8px;
font-family: Tahoma; box-shadow: 2px 2px 8px rgba(0,0,0,0.3); z-index: 1000;">
<h3 style="margin-top: 0; border-bottom: 2px solid #333; padding-bottom: 5px;">Legend</h3>
graph_data = {"nodes": cy_nodes, "edges": cy_edges}
<h4 style="margin-bottom: 5px;">Link Types:</h4>
<div style="margin-left: 10px;">
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 40px; height: 1px;
background: #00bcd4; border-top: 1px dashed #00bcd4;
vertical-align: middle;"></span>
<span style="margin-left: 10px;"><b>Temporal</b> - Time-based (cyan, dashed)</span>
# Build table rows for table view
table_rows = []
for unit_id, text, event_date, context in units:
entities = entity_map.get(unit_id, [])
entity_str = ", ".join(entities) if entities else "None"
table_rows.append(f"""
<tr>
<td style="padding: 8px; border: 1px solid #ddd;">{str(unit_id)[:8]}...</td>
<td style="padding: 8px; border: 1px solid #ddd;">{text}</td>
<td style="padding: 8px; border: 1px solid #ddd;">{context}</td>
<td style="padding: 8px; border: 1px solid #ddd;">{event_date.date()}</td>
<td style="padding: 8px; border: 1px solid #ddd;">{entity_str}</td>
</tr>
""")
# Generate HTML with Cytoscape.js
html_content = f"""
<!DOCTYPE html>
<html>
<head>
<title>Memory Graph - Interactive Visualization</title>
<meta charset="utf-8">
<script src="https://cdnjs.cloudflare.com/ajax/libs/cytoscape/3.28.1/cytoscape.min.js"></script>
<style>
body {{
font-family: Tahoma, sans-serif;
margin: 0;
padding: 0;
background: #f5f5f5;
}}
.tab-container {{
background: white;
}}
.tab-buttons {{
background: #f0f0f0;
border-bottom: 2px solid #333;
padding: 0;
margin: 0;
}}
.tab-button {{
background: #e0e0e0;
border: none;
padding: 12px 24px;
cursor: pointer;
font-size: 16px;
font-weight: bold;
border-top: 2px solid transparent;
border-left: 2px solid transparent;
border-right: 2px solid transparent;
transition: background 0.2s;
}}
.tab-button:hover {{
background: #d0d0d0;
}}
.tab-button.active {{
background: white;
border-top: 2px solid #333;
border-left: 2px solid #333;
border-right: 2px solid #333;
border-bottom: 2px solid white;
margin-bottom: -2px;
}}
.tab-content {{
display: none;
background: white;
}}
.tab-content.active {{
display: block;
}}
#cy {{
width: 100%;
height: 800px;
background: #ffffff;
}}
#graph-tab {{
position: relative;
}}
#table-tab {{
padding: 20px;
}}
.legend {{
position: absolute;
top: 20px;
left: 20px;
background: white;
padding: 15px;
border: 2px solid #333;
border-radius: 8px;
box-shadow: 2px 2px 8px rgba(0,0,0,0.3);
z-index: 1000;
max-width: 250px;
}}
.legend h3 {{
margin-top: 0;
border-bottom: 2px solid #333;
padding-bottom: 5px;
}}
.legend-item {{
margin: 8px 0;
display: flex;
align-items: center;
}}
.legend-line {{
width: 30px;
height: 2px;
margin-right: 10px;
}}
.legend-node {{
width: 20px;
height: 20px;
margin-right: 10px;
border: 1px solid #999;
border-radius: 3px;
}}
#table-filter {{
width: 100%;
max-width: 600px;
padding: 10px;
margin-bottom: 15px;
border: 2px solid #ccc;
border-radius: 4px;
font-size: 14px;
box-sizing: border-box;
}}
#memory-table {{
width: 100%;
border-collapse: collapse;
font-size: 13px;
max-width: 1400px;
}}
#memory-table th {{
padding: 10px;
text-align: left;
border: 1px solid #ddd;
background: #f0f0f0;
}}
#memory-table td {{
padding: 8px;
border: 1px solid #ddd;
}}
.tooltip {{
position: absolute;
background: white;
border: 2px solid #333;
border-radius: 4px;
padding: 10px;
box-shadow: 2px 2px 8px rgba(0,0,0,0.3);
max-width: 300px;
font-size: 12px;
pointer-events: none;
z-index: 9999;
}}
</style>
</head>
<body>
<div class="tab-container">
<div class="tab-buttons">
<button class="tab-button active" onclick="switchTab('graph')">Graph View</button>
<button class="tab-button" onclick="switchTab('table')">Table View</button>
</div>
<div id="graph-tab" class="tab-content active">
<div style="padding: 15px; background: #f9f9f9; border-bottom: 2px solid #333;">
<div style="display: flex; gap: 15px; align-items: center; flex-wrap: wrap;">
<div>
<label style="font-weight: bold; margin-right: 5px;">Limit nodes:</label>
<input type="number" id="node-limit" value="50" min="10" max="1000" step="10"
style="width: 80px; padding: 5px; border: 1px solid #ccc; border-radius: 4px;">
</div>
<div>
<label style="font-weight: bold; margin-right: 5px;">Layout:</label>
<select id="layout-select" style="padding: 5px; border: 1px solid #ccc; border-radius: 4px;">
<option value="circle">Circle (fast)</option>
<option value="grid">Grid (fast)</option>
<option value="cose">Force-directed (slow)</option>
</select>
</div>
<button onclick="reloadGraph()" style="padding: 6px 15px; background: #42a5f5; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold;">
Apply
</button>
<span id="node-count" style="color: #666; font-size: 14px;"></span>
</div>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 40px; height: 1px;
background: #ff69b4; vertical-align: middle;"></span>
<span style="margin-left: 10px;"><b>Semantic</b> - Meaning-based (pink, solid)</span>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 40px; height: 1.5px;
background: #ffd700; vertical-align: middle;"></span>
<span style="margin-left: 10px;"><b>Entity</b> - Same entity (gold)</span>
<div id="cy"></div>
<div class="legend">
<h3>Legend</h3>
<h4 style="margin: 10px 0 5px 0;">Link Types:</h4>
<div class="legend-item">
<div class="legend-line" style="background: #00bcd4; border-top: 1px dashed #00bcd4;"></div>
<span><b>Temporal</b></span>
</div>
<div class="legend-item">
<div class="legend-line" style="background: #ff69b4;"></div>
<span><b>Semantic</b></span>
</div>
<div class="legend-item">
<div class="legend-line" style="background: #ffd700;"></div>
<span><b>Entity</b></span>
</div>
<h4 style="margin: 15px 0 5px 0;">Nodes:</h4>
<div class="legend-item">
<div class="legend-node" style="background: #e0e0e0;"></div>
<span>No entities</span>
</div>
<div class="legend-item">
<div class="legend-node" style="background: #90caf9;"></div>
<span>1 entity</span>
</div>
<div class="legend-item">
<div class="legend-node" style="background: #42a5f5;"></div>
<span>2+ entities</span>
</div>
</div>
</div>
<h4 style="margin-bottom: 5px; margin-top: 15px;">Node Colors:</h4>
<div style="margin-left: 10px;">
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 20px; height: 20px;
background: #e0e0e0; border: 1px solid #999;
vertical-align: middle;"></span>
<span style="margin-left: 10px;">Gray - No entities</span>
<div id="table-tab" class="tab-content">
<h2>Memory Units ({len(units)})</h2>
<input type="text" id="table-filter" placeholder="Filter by text, context, or entities...">
<div style="overflow-x: auto;">
<table id="memory-table">
<thead>
<tr>
<th>ID</th>
<th>Text</th>
<th>Context</th>
<th>Date</th>
<th>Entities</th>
</tr>
</thead>
<tbody>
{''.join(table_rows)}
</tbody>
</table>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 20px; height: 20px;
background: #90caf9; border: 1px solid #999;
vertical-align: middle;"></span>
<span style="margin-left: 10px;">Light Blue - 1 entity</span>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 20px; height: 20px;
background: #42a5f5; border: 1px solid #999;
vertical-align: middle;"></span>
<span style="margin-left: 10px;">Dark Blue - 2+ entities</span>
</div>
</div>
<div style="margin-top: 15px; padding-top: 10px; border-top: 1px solid #ccc;
font-size: 11px; color: #666;">
<b>Tip:</b> Hover over nodes/edges for details<br>
<b>Controls:</b> Drag to move, scroll to zoom
</div>
</div>
"""
# Generate the HTML
output_file = "memory_graph_interactive.html"
net.save_graph(output_file)
<script>
// Graph data
const allGraphData = {json.dumps(graph_data)};
let cy = null;
# Read the generated HTML and inject our legend
with open(output_file, 'r') as f:
html_content = f.read()
// Initialize graph with filtering
function initGraph(nodeLimit, layoutName) {{
// Filter nodes to limit
const limitedNodes = allGraphData.nodes.slice(0, nodeLimit);
const nodeIds = new Set(limitedNodes.map(n => n.data.id));
# Inject legend after the opening body tag
html_content = html_content.replace('<body>', '<body>' + legend_html)
// Filter edges to only include those between visible nodes
const limitedEdges = allGraphData.edges.filter(e =>
nodeIds.has(e.data.source) && nodeIds.has(e.data.target)
);
# Add script to disable physics after stabilization for better performance
physics_script = """
<script type="text/javascript">
// Disable physics after initial stabilization for better performance
network.on("stabilizationIterationsDone", function () {
network.setOptions({ physics: false });
console.log("Physics disabled - graph should be much more responsive now!");
});
// Update count display
document.getElementById('node-count').textContent =
`Showing ${{limitedNodes.length}} of ${{allGraphData.nodes.length}} nodes`;
// Destroy existing graph if any
if (cy) {{
cy.destroy();
}}
// Layout configurations
const layouts = {{
'circle': {{
name: 'circle',
animate: false,
radius: 300,
spacingFactor: 1.5
}},
'grid': {{
name: 'grid',
animate: false,
rows: Math.ceil(Math.sqrt(limitedNodes.length)),
cols: Math.ceil(Math.sqrt(limitedNodes.length)),
spacingFactor: 2
}},
'cose': {{
name: 'cose',
animate: false,
nodeRepulsion: 15000,
idealEdgeLength: 150,
edgeElasticity: 100,
nestingFactor: 1.2,
gravity: 1,
numIter: 1000,
initialTemp: 200,
coolingFactor: 0.95,
minTemp: 1.0
}}
}};
// Initialize Cytoscape
cy = cytoscape({{
container: document.getElementById('cy'),
elements: [
...limitedNodes.map(n => ({{ data: n.data }})),
...limitedEdges.map(e => ({{ data: e.data }}))
],
style: [
{{
selector: 'node',
style: {{
'background-color': 'data(color)',
'label': 'data(label)',
'text-valign': 'center',
'text-halign': 'center',
'font-size': '10px',
'font-weight': 'bold',
'text-wrap': 'wrap',
'text-max-width': '100px',
'width': 40,
'height': 40,
'border-width': 2,
'border-color': '#333'
}}
}},
{{
selector: 'edge',
style: {{
'width': 1,
'line-color': 'data(color)',
'line-style': 'data(lineStyle)',
'target-arrow-shape': 'triangle',
'target-arrow-color': 'data(color)',
'curve-style': 'bezier',
'opacity': 0.7
}}
}},
{{
selector: 'node:selected',
style: {{
'border-width': 4,
'border-color': '#000'
}}
}}
],
layout: layouts[layoutName] || layouts['circle']
}});
// Simple tooltip on hover
let tooltip = null;
cy.on('mouseover', 'node', function(evt) {{
const node = evt.target;
const data = node.data();
const renderedPosition = node.renderedPosition();
// Create tooltip
tooltip = document.createElement('div');
tooltip.className = 'tooltip';
tooltip.innerHTML = `
<b>Text:</b> ${{data.text}}<br>
<b>Context:</b> ${{data.context}}<br>
<b>Date:</b> ${{data.date}}<br>
<b>Entities:</b> ${{data.entities}}
`;
tooltip.style.left = renderedPosition.x + 20 + 'px';
tooltip.style.top = renderedPosition.y + 'px';
document.body.appendChild(tooltip);
}});
cy.on('mouseout', 'node', function(evt) {{
if (tooltip) {{
tooltip.remove();
tooltip = null;
}}
}});
}}
// Reload graph with current settings
function reloadGraph() {{
const nodeLimit = parseInt(document.getElementById('node-limit').value) || 50;
const layoutName = document.getElementById('layout-select').value;
initGraph(nodeLimit, layoutName);
}}
// Initialize with default settings (50 nodes, circle layout)
initGraph(50, 'circle');
// Tab switching
function switchTab(tabName) {{
document.querySelectorAll('.tab-content').forEach(tab => {{
tab.classList.remove('active');
}});
document.querySelectorAll('.tab-button').forEach(btn => {{
btn.classList.remove('active');
}});
if (tabName === 'graph') {{
document.getElementById('graph-tab').classList.add('active');
document.querySelectorAll('.tab-button')[0].classList.add('active');
cy.resize(); // Resize graph when switching to it
}} else if (tabName === 'table') {{
document.getElementById('table-tab').classList.add('active');
document.querySelectorAll('.tab-button')[1].classList.add('active');
}}
}}
// Table filtering
document.getElementById('table-filter').addEventListener('input', function() {{
const filterValue = this.value.toLowerCase();
const rows = document.querySelectorAll('#memory-table tbody tr');
rows.forEach(row => {{
const text = row.textContent.toLowerCase();
if (text.includes(filterValue)) {{
row.style.display = '';
}} else {{
row.style.display = 'none';
}}
}});
}});
</script>
"""
html_content = html_content.replace('</body>', physics_script + '</body>')
</body>
</html>
"""
# Write back
with open(output_file, 'w') as f:
# Write HTML file
output_file = "memory_graph_interactive.html"
with open(output_file, 'w', encoding='utf-8') as f:
f.write(html_content)
# Print summary
print(f"\n{'='*80}")
print("INTERACTIVE GRAPH GENERATED")
print("INTERACTIVE GRAPH GENERATED (Cytoscape.js)")
print(f"{'='*80}")
print(f"\nFile: {output_file}")
print(f"Units: {len(units)}")
print(f"Links: {len(links)}")
print("\nFeatures:")
print(" • Smooth, physics-based layout")
print(" • Interactive - drag nodes, zoom, pan")
print(" • Hover for details on nodes and edges")
print(" • Color-coded by link type and entity count")
print(" • Built-in navigation controls")
print(" • Tab 1: Graph View - Fast interactive network (Cytoscape.js)")
print(" - Limit nodes (default: 50) for better performance")
print(" - Choose layout: Circle (fast), Grid (fast), or Force-directed")
print(" - Drag nodes, zoom, pan")
print(" - Hover for details")
print(" • Tab 2: Table View - Searchable memory units")
print(" - Filter by text, context, or entities")
print(" - Case-insensitive search")
print(" - Shows ALL nodes")
print(f"\n{'='*80}")
print(f"✓ Open {output_file} in your browser to explore!")
print(f" TIP: Start with 50 nodes and Circle layout for best performance")
print(f"{'='*80}\n")

2870
visualizations/memory_graph_interactive.html
View file

File diff suppressed because one or more lines are too long