"""LOCOMO benchmark experiment runner for Mem0 agent.""" import asyncio import json import random import time from pathlib import Path from typing import List, Dict, Any, Tuple from datetime import datetime import argparse import numpy as np import pandas as pd from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn from rich.table import Table from rich.panel import Panel from rich.layout import Layout import matplotlib.pyplot as plt import seaborn as sns from agent import Mem0Agent, AgentContext from config import Config console = Console() class LOCOMOBenchmark: """LOCOMO benchmark implementation for evaluating long-context multi-agent communication.""" def __init__(self, agent: Mem0Agent, config: Config): self.agent = agent self.config = config self.results = [] self.scenario_data = [] def generate_scenario(self, scenario_id: int, num_agents: int = 3) -> Dict[str, Any]: """Generate a LOCOMO benchmark scenario.""" scenarios = [ { "type": "collaborative_planning", "description": "Multiple agents collaborate to plan a complex project", "topics": ["project timeline", "resource allocation", "task dependencies", "risk assessment"], "context_requirements": ["maintain consistency across planning decisions", "remember previous constraints", "coordinate between agents"] }, { "type": "information_sharing", "description": "Agents share and synthesize information across sessions", "topics": ["research findings", "data analysis", "hypothesis formation", "conclusion drawing"], "context_requirements": ["retain factual information", "build upon previous insights", "cross-reference between sources"] }, { "type": "problem_solving", "description": "Agents work together to solve multi-step problems", "topics": ["problem decomposition", "solution strategies", "intermediate results", "final synthesis"], "context_requirements": ["remember partial solutions", "maintain logical consistency", "track progress across sessions"] }, { "type": "negotiation", "description": "Agents engage in multi-round negotiations", "topics": ["initial positions", "concessions", "agreements", "conflict resolution"], "context_requirements": ["remember previous offers", "maintain negotiation stance", "track agreement points"] }, { "type": "teaching_learning", "description": "Agents engage in educational dialogue", "topics": ["concept explanation", "question answering", "knowledge verification", "skill progression"], "context_requirements": ["track learning progress", "adapt to understanding level", "remember misconceptions"] } ] scenario = scenarios[scenario_id % len(scenarios)].copy() scenario["scenario_id"] = f"scenario_{scenario_id:03d}" scenario["num_agents"] = num_agents scenario["agents"] = [f"agent_{i:02d}" for i in range(num_agents)] scenario["num_sessions"] = random.randint(3, 8) scenario["turns_per_session"] = random.randint(5, 15) return scenario def generate_conversation_prompts(self, scenario: Dict[str, Any], session_num: int) -> List[str]: """Generate conversation prompts for a scenario session.""" prompts = [] topic = random.choice(scenario["topics"]) base_prompts = { "collaborative_planning": [ f"Let's discuss the {topic} for our project. What are your thoughts?", f"Based on our previous discussion, how should we adjust the {topic}?", f"Can you summarize what we've decided about {topic} so far?", f"What challenges do you foresee with the current {topic}?", f"How does the {topic} align with our overall objectives?" ], "information_sharing": [ f"What new information do you have about {topic}?", f"How does this relate to what we discussed about {topic} before?", f"Can you integrate the findings about {topic} with our previous data?", f"What patterns are emerging from our {topic} analysis?", f"What conclusions can we draw about {topic} at this point?" ], "problem_solving": [ f"What's our current approach to {topic}?", f"Have we made progress on {topic} since last time?", f"What obstacles are we facing with {topic}?", f"Can you propose an alternative solution for {topic}?", f"How can we validate our solution for {topic}?" ], "negotiation": [ f"What's your position on {topic}?", f"Can we find middle ground on {topic}?", f"What concessions are you willing to make regarding {topic}?", f"How does this affect our previous agreement on {topic}?", f"Let's finalize our agreement on {topic}." ], "teaching_learning": [ f"Can you explain {topic} in simple terms?", f"What questions do you have about {topic}?", f"How would you apply {topic} in practice?", f"What did we learn about {topic} last time?", f"Can you give an example of {topic}?" ] } scenario_prompts = base_prompts.get(scenario["type"], base_prompts["information_sharing"]) # Add session-specific context for i in range(scenario["turns_per_session"]): if session_num == 0: prompt = f"[Session {session_num + 1}, Turn {i + 1}] {random.choice(scenario_prompts)}" else: prompt = f"[Session {session_num + 1}, Turn {i + 1}] Continuing from our previous session, {random.choice(scenario_prompts)}" prompts.append(prompt) return prompts async def run_scenario_session(self, scenario: Dict[str, Any], session_num: int) -> Dict[str, Any]: """Run a single session of a scenario.""" session_id = f"{scenario['scenario_id']}_session_{session_num:02d}" user_id = f"user_{scenario['scenario_id']}" # Create contexts for all agents agent_contexts = {} for agent_id in scenario["agents"]: context = self.agent.create_context( agent_id=agent_id, user_id=user_id, session_id=f"{session_id}_{agent_id}" ) agent_contexts[agent_id] = context # Generate conversation prompts prompts = self.generate_conversation_prompts(scenario, session_num) session_results = { "session_id": session_id, "session_num": session_num, "turns": [], "metrics": {} } # Run conversation turns for turn_idx, prompt in enumerate(prompts): # Randomly select which agent responds responding_agent = random.choice(scenario["agents"]) agent_session_id = f"{session_id}_{responding_agent}" # Process turn response, metrics = await self.agent.process_turn_async(agent_session_id, prompt) session_results["turns"].append({ "turn": turn_idx, "agent": responding_agent, "prompt": prompt, "response": response, "metrics": metrics }) # Small delay to simulate realistic conversation await asyncio.sleep(0.1) # Calculate session-level metrics session_results["metrics"] = { "total_turns": len(prompts), "avg_response_time": np.mean([t["metrics"]["generation_time"] for t in session_results["turns"]]), "avg_response_length": np.mean([t["metrics"]["response_length"] for t in session_results["turns"]]), "consistency_scores": { agent_id: self.agent.evaluate_consistency(f"{session_id}_{agent_id}") for agent_id in scenario["agents"] }, "coherence_scores": { agent_id: self.agent.evaluate_coherence(f"{session_id}_{agent_id}") for agent_id in scenario["agents"] } } return session_results async def run_scenario(self, scenario_id: int) -> Dict[str, Any]: """Run a complete LOCOMO scenario.""" scenario = self.generate_scenario(scenario_id) console.print(f"\n[cyan]Running Scenario {scenario['scenario_id']}[/cyan]") console.print(f"Type: {scenario['type']}") console.print(f"Agents: {', '.join(scenario['agents'])}") console.print(f"Sessions: {scenario['num_sessions']}") scenario_results = { "scenario": scenario, "sessions": [], "start_time": datetime.now().isoformat() } # Run all sessions with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), TimeElapsedColumn(), console=console ) as progress: task = progress.add_task( f"Running {scenario['num_sessions']} sessions...", total=scenario["num_sessions"] ) for session_num in range(scenario["num_sessions"]): session_results = await self.run_scenario_session(scenario, session_num) scenario_results["sessions"].append(session_results) progress.update(task, advance=1) # Delay between sessions await asyncio.sleep(0.5) scenario_results["end_time"] = datetime.now().isoformat() # Calculate scenario-level metrics scenario_results["overall_metrics"] = self.calculate_scenario_metrics(scenario_results) return scenario_results def calculate_scenario_metrics(self, scenario_results: Dict[str, Any]) -> Dict[str, Any]: """Calculate overall metrics for a scenario.""" all_turns = [] for session in scenario_results["sessions"]: all_turns.extend(session["turns"]) consistency_scores = [] coherence_scores = [] for session in scenario_results["sessions"]: consistency_scores.extend(list(session["metrics"]["consistency_scores"].values())) coherence_scores.extend(list(session["metrics"]["coherence_scores"].values())) metrics = { "total_turns": len(all_turns), "total_sessions": len(scenario_results["sessions"]), "avg_response_time": np.mean([t["metrics"]["generation_time"] for t in all_turns]), "std_response_time": np.std([t["metrics"]["generation_time"] for t in all_turns]), "avg_response_length": np.mean([t["metrics"]["response_length"] for t in all_turns]), "avg_consistency": np.mean(consistency_scores) if consistency_scores else 0, "avg_coherence": np.mean(coherence_scores) if coherence_scores else 0, "memory_utilization": len(self.agent.get_all_memories( f"user_{scenario_results['scenario']['scenario_id']}", top_k=100 )) } return metrics async def run_benchmark(self, num_scenarios: int = 5) -> Dict[str, Any]: """Run the complete LOCOMO benchmark.""" console.print(Panel.fit( f"[bold cyan]LOCOMO Benchmark[/bold cyan]\n" f"Scenarios: {num_scenarios}\n" f"Model: {self.config.kimi.model_name}\n" f"Memory Backend: {self.config.mem0.backend}", title="Benchmark Configuration" )) benchmark_results = { "config": self.config.to_dict(), "start_time": datetime.now().isoformat(), "scenarios": [] } for i in range(num_scenarios): scenario_results = await self.run_scenario(i) benchmark_results["scenarios"].append(scenario_results) self.results.append(scenario_results) # Display interim results self.display_scenario_results(scenario_results) benchmark_results["end_time"] = datetime.now().isoformat() benchmark_results["overall_metrics"] = self.calculate_overall_metrics(benchmark_results) return benchmark_results def calculate_overall_metrics(self, benchmark_results: Dict[str, Any]) -> Dict[str, Any]: """Calculate overall benchmark metrics.""" all_metrics = [s["overall_metrics"] for s in benchmark_results["scenarios"]] return { "total_scenarios": len(benchmark_results["scenarios"]), "avg_response_time": np.mean([m["avg_response_time"] for m in all_metrics]), "std_response_time": np.std([m["avg_response_time"] for m in all_metrics]), "avg_consistency": np.mean([m["avg_consistency"] for m in all_metrics]), "std_consistency": np.std([m["avg_consistency"] for m in all_metrics]), "avg_coherence": np.mean([m["avg_coherence"] for m in all_metrics]), "std_coherence": np.std([m["avg_coherence"] for m in all_metrics]), "avg_memory_utilization": np.mean([m["memory_utilization"] for m in all_metrics]), "total_turns": sum([m["total_turns"] for m in all_metrics]) } def display_scenario_results(self, scenario_results: Dict[str, Any]) -> None: """Display results for a single scenario.""" metrics = scenario_results["overall_metrics"] table = Table(title=f"Scenario {scenario_results['scenario']['scenario_id']} Results") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") table.add_row("Type", scenario_results["scenario"]["type"]) table.add_row("Sessions", str(metrics["total_sessions"])) table.add_row("Total Turns", str(metrics["total_turns"])) table.add_row("Avg Response Time", f"{metrics['avg_response_time']:.3f}s") table.add_row("Consistency Score", f"{metrics['avg_consistency']:.3f}") table.add_row("Coherence Score", f"{metrics['avg_coherence']:.3f}") table.add_row("Memory Utilization", str(metrics["memory_utilization"])) console.print(table) def display_overall_results(self, benchmark_results: Dict[str, Any]) -> None: """Display overall benchmark results.""" metrics = benchmark_results["overall_metrics"] console.print("\n") console.print(Panel.fit( f"[bold green]Benchmark Complete![/bold green]\n" f"Total Scenarios: {metrics['total_scenarios']}\n" f"Total Turns: {metrics['total_turns']}\n" f"Avg Response Time: {metrics['avg_response_time']:.3f}s ± {metrics['std_response_time']:.3f}s\n" f"Avg Consistency: {metrics['avg_consistency']:.3f} ± {metrics['std_consistency']:.3f}\n" f"Avg Coherence: {metrics['avg_coherence']:.3f} ± {metrics['std_coherence']:.3f}\n" f"Avg Memory Utilization: {metrics['avg_memory_utilization']:.1f}", title="Overall Results" )) def save_results(self, benchmark_results: Dict[str, Any], filepath: Path) -> None: """Save benchmark results to file.""" filepath.parent.mkdir(parents=True, exist_ok=True) with open(filepath, "w") as f: json.dump(benchmark_results, f, indent=2) console.print(f"[green]Results saved to {filepath}[/green]") def generate_report(self, benchmark_results: Dict[str, Any], output_dir: Path) -> None: """Generate a detailed report with visualizations.""" output_dir.mkdir(parents=True, exist_ok=True) # Extract data for visualization scenarios_df = pd.DataFrame([ { "scenario_id": s["scenario"]["scenario_id"], "type": s["scenario"]["type"], "consistency": s["overall_metrics"]["avg_consistency"], "coherence": s["overall_metrics"]["avg_coherence"], "response_time": s["overall_metrics"]["avg_response_time"], "memory_utilization": s["overall_metrics"]["memory_utilization"] } for s in benchmark_results["scenarios"] ]) # Create visualizations fig, axes = plt.subplots(2, 2, figsize=(12, 10)) # Consistency scores by scenario type sns.boxplot(data=scenarios_df, x="type", y="consistency", ax=axes[0, 0]) axes[0, 0].set_title("Consistency Scores by Scenario Type") axes[0, 0].set_xticklabels(axes[0, 0].get_xticklabels(), rotation=45) # Coherence scores by scenario type sns.boxplot(data=scenarios_df, x="type", y="coherence", ax=axes[0, 1]) axes[0, 1].set_title("Coherence Scores by Scenario Type") axes[0, 1].set_xticklabels(axes[0, 1].get_xticklabels(), rotation=45) # Response time distribution axes[1, 0].hist(scenarios_df["response_time"], bins=20, edgecolor='black') axes[1, 0].set_title("Response Time Distribution") axes[1, 0].set_xlabel("Response Time (s)") axes[1, 0].set_ylabel("Frequency") # Memory utilization vs performance axes[1, 1].scatter(scenarios_df["memory_utilization"], scenarios_df["consistency"], alpha=0.6, label="Consistency") axes[1, 1].scatter(scenarios_df["memory_utilization"], scenarios_df["coherence"], alpha=0.6, label="Coherence") axes[1, 1].set_title("Memory Utilization vs Performance") axes[1, 1].set_xlabel("Memory Utilization") axes[1, 1].set_ylabel("Score") axes[1, 1].legend() plt.tight_layout() plt.savefig(output_dir / "benchmark_results.png", dpi=300) console.print(f"[green]Report generated in {output_dir}[/green]") async def main(): """Main function to run the LOCOMO benchmark.""" parser = argparse.ArgumentParser(description="Run LOCOMO benchmark for Mem0 agent") parser.add_argument("--scenarios", type=int, default=5, help="Number of scenarios to run") parser.add_argument("--output", type=str, default="results", help="Output directory for results") parser.add_argument("--config", type=str, help="Path to configuration file") args = parser.parse_args() # Initialize configuration config = Config.from_env() # Initialize agent console.print("[yellow]Initializing Mem0 agent...[/yellow]") agent = Mem0Agent(config) # Initialize benchmark benchmark = LOCOMOBenchmark(agent, config) # Run benchmark console.print(f"[yellow]Starting benchmark with {args.scenarios} scenarios...[/yellow]") benchmark_results = await benchmark.run_benchmark(num_scenarios=args.scenarios) # Save results output_dir = Path(args.output) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") results_file = output_dir / f"locomo_results_{timestamp}.json" benchmark.save_results(benchmark_results, results_file) # Generate report report_dir = output_dir / f"report_{timestamp}" benchmark.generate_report(benchmark_results, report_dir) # Display overall results benchmark.display_overall_results(benchmark_results) if __name__ == "__main__": asyncio.run(main())