331 lines
12 KiB
Python
331 lines
12 KiB
Python
"""
|
|
Consolidation performance benchmark.
|
|
|
|
Measures consolidation throughput (op/sec) and identifies bottlenecks by:
|
|
1. Ingesting a batch of diverse memories
|
|
2. Running consolidation manually with detailed timing
|
|
3. Analyzing timing breakdown to identify bottlenecks
|
|
4. Reporting op/sec and time spent in each component
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from hindsight_api.config import get_config
|
|
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
|
|
from hindsight_api.engine.memory_engine import MemoryEngine
|
|
from hindsight_api.models import RequestContext
|
|
from rich.console import Console
|
|
from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn
|
|
from rich.table import Table
|
|
|
|
console = Console()
|
|
|
|
|
|
# Sample diverse memories to trigger different consolidation patterns
|
|
SAMPLE_MEMORIES = [
|
|
# Similar memories (should merge)
|
|
"Alice loves coffee and drinks it every morning.",
|
|
"Alice prefers coffee over tea for her morning beverage.",
|
|
"Alice switched to decaf coffee recently.",
|
|
# Different person (should NOT merge with Alice)
|
|
"Bob works at Google as a software engineer.",
|
|
"Bob has been at Google for 5 years.",
|
|
# Technical facts
|
|
"Python is a programming language used for data science.",
|
|
"Python supports object-oriented and functional programming.",
|
|
# Product info
|
|
"The new iPhone 15 was released in September 2023.",
|
|
"The iPhone 15 features USB-C charging instead of Lightning.",
|
|
# Contradictions (should merge with conflict resolution)
|
|
"The meeting is scheduled for Tuesday at 2pm.",
|
|
"The meeting was moved to Wednesday at 3pm.",
|
|
# Entity-rich content
|
|
"Sarah Smith works at Microsoft in Seattle.",
|
|
"Sarah graduated from Stanford University in 2015.",
|
|
# Temporal information
|
|
"The project started on January 15, 2024.",
|
|
"The project deadline is March 30, 2024.",
|
|
# Preferences
|
|
"User prefers dark mode in applications.",
|
|
"User uses keyboard shortcuts extensively.",
|
|
# World knowledge
|
|
"Paris is the capital of France.",
|
|
"The Eiffel Tower is located in Paris.",
|
|
# Multiple entities
|
|
"John and Mary went to the Italian restaurant on Main Street.",
|
|
"The Italian restaurant on Main Street has excellent pizza.",
|
|
]
|
|
|
|
|
|
async def create_test_memories(memory_engine: MemoryEngine, bank_id: str, num_memories: int = 100) -> None:
|
|
"""
|
|
Create test memories by repeating and varying the sample memories.
|
|
|
|
Args:
|
|
memory_engine: MemoryEngine instance
|
|
bank_id: Bank ID to ingest into
|
|
num_memories: Number of memories to create
|
|
"""
|
|
console.print(f"\n[cyan]Creating {num_memories} test memories...[/cyan]")
|
|
|
|
# Generate memories by cycling through samples
|
|
memories = []
|
|
for i in range(num_memories):
|
|
base_memory = SAMPLE_MEMORIES[i % len(SAMPLE_MEMORIES)]
|
|
# Add variation to avoid exact duplicates
|
|
memory = f"{base_memory} (context: test {i + 1})"
|
|
memories.append(
|
|
{
|
|
"content": memory,
|
|
"context": f"Test memory {i + 1}",
|
|
}
|
|
)
|
|
|
|
# Batch ingest
|
|
console.print("[yellow]Ingesting memories in batch...[/yellow]")
|
|
start_time = time.time()
|
|
await memory_engine.retain_batch_async(
|
|
bank_id=bank_id,
|
|
contents=memories,
|
|
request_context=RequestContext(),
|
|
)
|
|
ingest_time = time.time() - start_time
|
|
console.print(
|
|
f"[green]✓[/green] Ingested {num_memories} memories in {ingest_time:.2f}s ({num_memories / ingest_time:.2f} mem/sec)"
|
|
)
|
|
|
|
|
|
async def run_consolidation_benchmark(
|
|
memory_engine: MemoryEngine,
|
|
bank_id: str,
|
|
enable_detailed_logs: bool = True,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Run consolidation and measure performance.
|
|
|
|
Args:
|
|
memory_engine: MemoryEngine instance
|
|
bank_id: Bank ID to consolidate
|
|
enable_detailed_logs: Enable detailed consolidation logs
|
|
|
|
Returns:
|
|
Performance metrics dict
|
|
"""
|
|
console.print("\n[cyan]Running consolidation benchmark...[/cyan]")
|
|
|
|
# Set log level to INFO to see consolidation logs
|
|
if enable_detailed_logs:
|
|
# Configure logging for consolidation
|
|
consolidation_logger = logging.getLogger("hindsight_api.engine.consolidation.consolidator")
|
|
consolidation_logger.setLevel(logging.INFO)
|
|
|
|
# Add console handler if not present
|
|
if not consolidation_logger.handlers:
|
|
handler = logging.StreamHandler()
|
|
handler.setLevel(logging.INFO)
|
|
formatter = logging.Formatter("%(message)s")
|
|
handler.setFormatter(formatter)
|
|
consolidation_logger.addHandler(handler)
|
|
|
|
console.print("[yellow]Detailed logging enabled for consolidation[/yellow]")
|
|
|
|
# Run consolidation and measure time
|
|
start_time = time.time()
|
|
result = await run_consolidation_job(
|
|
memory_engine=memory_engine,
|
|
bank_id=bank_id,
|
|
request_context=RequestContext(),
|
|
)
|
|
total_time = time.time() - start_time
|
|
|
|
# Calculate op/sec
|
|
memories_processed = result.get("memories_processed", 0)
|
|
ops_per_sec = memories_processed / total_time if total_time > 0 else 0
|
|
|
|
console.print("\n[green]✓[/green] Consolidation complete!")
|
|
console.print(f" Total time: {total_time:.2f}s")
|
|
console.print(f" Memories processed: {memories_processed}")
|
|
console.print(f" Throughput: {ops_per_sec:.2f} op/sec")
|
|
console.print(f" Avg time per memory: {total_time / memories_processed:.3f}s" if memories_processed > 0 else "")
|
|
|
|
return {
|
|
"total_time": total_time,
|
|
"memories_processed": memories_processed,
|
|
"ops_per_sec": ops_per_sec,
|
|
"consolidation_result": result,
|
|
}
|
|
|
|
|
|
async def analyze_timing_breakdown(bank_id: str) -> None:
|
|
"""
|
|
Analyze the timing breakdown from consolidation logs.
|
|
|
|
NOTE: This relies on the performance logging in ConsolidationPerfLog.
|
|
The logs will show timing breakdowns for: recall, llm, embedding, db_write
|
|
"""
|
|
console.print("\n[cyan]Timing Breakdown Analysis:[/cyan]")
|
|
console.print("Check the logs above for detailed timing breakdown:")
|
|
console.print(" - recall: Time spent finding related observations")
|
|
console.print(" - llm: Time spent in LLM calls for consolidation decisions")
|
|
console.print(" - embedding: Time spent generating embeddings")
|
|
console.print(" - db_write: Time spent writing to database")
|
|
|
|
|
|
async def get_bank_stats(memory_engine: MemoryEngine, bank_id: str) -> dict[str, Any]:
|
|
"""Get memory statistics for the bank."""
|
|
pool = await memory_engine._get_pool()
|
|
from hindsight_api.engine.memory_engine import fq_table
|
|
|
|
async with pool.acquire() as conn:
|
|
# Count memories by fact type
|
|
stats = await conn.fetch(
|
|
f"""
|
|
SELECT fact_type, COUNT(*) as count
|
|
FROM {fq_table("memory_units")}
|
|
WHERE bank_id = $1
|
|
GROUP BY fact_type
|
|
""",
|
|
bank_id,
|
|
)
|
|
|
|
return {row["fact_type"]: row["count"] for row in stats}
|
|
|
|
|
|
def display_results_table(metrics: dict[str, Any], stats_before: dict, stats_after: dict) -> None:
|
|
"""Display benchmark results in a formatted table."""
|
|
table = Table(title="Consolidation Benchmark Results")
|
|
table.add_column("Metric", style="cyan")
|
|
table.add_column("Value", style="green")
|
|
|
|
table.add_row("Total Time", f"{metrics['total_time']:.2f}s")
|
|
table.add_row("Memories Processed", str(metrics["memories_processed"]))
|
|
table.add_row("Throughput", f"{metrics['ops_per_sec']:.2f} op/sec")
|
|
table.add_row(
|
|
"Avg Time/Memory",
|
|
f"{metrics['total_time'] / metrics['memories_processed']:.3f}s" if metrics["memories_processed"] > 0 else "N/A",
|
|
)
|
|
|
|
result = metrics["consolidation_result"]
|
|
table.add_row("", "") # Separator
|
|
table.add_row("Observations Created", str(result.get("observations_created", 0)))
|
|
table.add_row("Observations Updated", str(result.get("observations_updated", 0)))
|
|
table.add_row("Observations Merged", str(result.get("observations_merged", 0)))
|
|
table.add_row("Skipped (No Durable Knowledge)", str(result.get("skipped", 0)))
|
|
|
|
table.add_row("", "") # Separator
|
|
table.add_row("Memories Before", str(stats_before.get("experience", 0) + stats_before.get("world", 0)))
|
|
table.add_row("Observations After", str(stats_after.get("observation", 0)))
|
|
|
|
console.print("\n")
|
|
console.print(table)
|
|
|
|
|
|
async def main():
|
|
"""Run the consolidation benchmark."""
|
|
console.print("\n[bold cyan]Consolidation Performance Benchmark[/bold cyan]")
|
|
console.print("=" * 80)
|
|
|
|
# Configuration
|
|
num_memories = int(os.getenv("NUM_MEMORIES", "100"))
|
|
bank_id = f"consolidation-bench-{uuid.uuid4().hex[:8]}"
|
|
|
|
console.print("\n[cyan]Configuration:[/cyan]")
|
|
console.print(f" Bank ID: {bank_id}")
|
|
console.print(f" Number of memories: {num_memories}")
|
|
console.print(f" LLM Provider: {os.getenv('HINDSIGHT_API_LLM_PROVIDER', 'not set')}")
|
|
console.print(f" LLM Model: {os.getenv('HINDSIGHT_API_LLM_MODEL', 'not set')}")
|
|
|
|
# Check if consolidation is enabled
|
|
config = get_config()
|
|
if not config.enable_observations:
|
|
console.print("\n[red]ERROR: Consolidation is disabled (enable_observations=False)[/red]")
|
|
console.print("Set HINDSIGHT_API_ENABLE_OBSERVATIONS=true to enable consolidation")
|
|
return
|
|
|
|
# Initialize memory engine
|
|
console.print("\n[1] Initializing memory engine...")
|
|
memory = MemoryEngine(
|
|
db_url=os.getenv("HINDSIGHT_API_DATABASE_URL", "pg0"),
|
|
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
|
|
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
|
|
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
|
|
memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None,
|
|
)
|
|
await memory.initialize()
|
|
console.print("[green]✓[/green] Memory engine initialized")
|
|
|
|
try:
|
|
# Create bank
|
|
console.print("\n[2] Creating test bank...")
|
|
await memory.get_bank_profile(bank_id=bank_id, request_context=RequestContext())
|
|
console.print(f"[green]✓[/green] Created bank: {bank_id}")
|
|
|
|
# Get initial stats
|
|
stats_before = await get_bank_stats(memory, bank_id)
|
|
|
|
# Create test memories
|
|
console.print("\n[3] Creating test memories...")
|
|
await create_test_memories(memory, bank_id, num_memories)
|
|
|
|
# Run consolidation benchmark
|
|
console.print("\n[4] Running consolidation benchmark...")
|
|
metrics = await run_consolidation_benchmark(memory, bank_id, enable_detailed_logs=True)
|
|
|
|
# Get final stats
|
|
stats_after = await get_bank_stats(memory, bank_id)
|
|
|
|
# Analyze timing
|
|
console.print("\n[5] Analyzing performance...")
|
|
await analyze_timing_breakdown(bank_id)
|
|
|
|
# Display results
|
|
console.print("\n[6] Results:")
|
|
display_results_table(metrics, stats_before, stats_after)
|
|
|
|
# Save results to file
|
|
output_dir = Path("benchmarks/results")
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_file = output_dir / f"consolidation_benchmark_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
|
|
|
results = {
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
"config": {
|
|
"num_memories": num_memories,
|
|
"bank_id": bank_id,
|
|
"llm_provider": os.getenv("HINDSIGHT_API_LLM_PROVIDER"),
|
|
"llm_model": os.getenv("HINDSIGHT_API_LLM_MODEL"),
|
|
},
|
|
"metrics": metrics,
|
|
"stats_before": stats_before,
|
|
"stats_after": stats_after,
|
|
}
|
|
|
|
with open(output_file, "w") as f:
|
|
json.dump(results, f, indent=2, default=str)
|
|
|
|
console.print(f"\n[green]✓[/green] Results saved to: {output_file}")
|
|
|
|
finally:
|
|
# Cleanup
|
|
console.print("\n[7] Cleaning up...")
|
|
await memory.delete_bank(bank_id, request_context=RequestContext())
|
|
console.print(f"[green]✓[/green] Deleted bank: {bank_id}")
|
|
|
|
# Close memory engine connections
|
|
pool = await memory._get_pool()
|
|
await pool.close()
|
|
console.print("[green]✓[/green] Memory engine connections closed")
|
|
|
|
console.print("\n[bold green]✓ Benchmark Complete![/bold green]\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|