mem0¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/mem0
项目说明¶
Mem0 Agent with Kimi K3 for LOCOMO Benchmark¶
A sophisticated AI agent implementation that combines the Mem0 memory framework with the Kimi K3 language model, specifically designed for the LOCOMO (Long-Context Multi-Agent) benchmark.
Overview¶
This project implements an advanced conversational AI agent that: - Persistent Memory: Uses Mem0 framework for long-term memory management across sessions - Advanced Language Model: Integrates Kimi K3 model (1M-token context window; experiment caps usage at a smaller budget) - LOCOMO Benchmark: Evaluates agent performance on long-context multi-agent communication tasks - Multi-Session Support: Maintains context and consistency across multiple conversation sessions - Multi-Agent Collaboration: Supports multiple agents working together with shared memory
Features¶
Core Capabilities¶
- Dynamic Memory Management: Automatically extracts, consolidates, and retrieves relevant information
- Context Preservation: Maintains conversation context across sessions and agents
- Performance Metrics: Tracks consistency, coherence, response time, and memory utilization
- Flexible Backend: Supports both local and cloud-based memory storage
Benchmark Scenarios¶
The LOCOMO benchmark includes five scenario types: 1. Collaborative Planning: Multiple agents plan complex projects together 2. Information Sharing: Agents share and synthesize information across sessions 3. Problem Solving: Multi-step problem solving with memory retention 4. Negotiation: Multi-round negotiations with position tracking 5. Teaching & Learning: Educational dialogues with progress tracking
Installation¶
Prerequisites¶
- Python 3.8 or higher
- Kimi API key (from Moonshot AI)
- Optional: Mem0 cloud API key for cloud storage
Setup¶
-
Clone the repository:
-
Install dependencies:
-
Configure environment variables:
Required environment variables:
- KIMI_API_KEY: Your Kimi K3 API key
- MODEL_NAME: Model name (default: kimi-k3). This is the raw Moonshot model
id (e.g. kimi-k3, kimi-k2.5, kimi-k2.6) passed straight to the API — do
NOT use a provider/model slash form here; Mem0 is configured with the
openai provider pointed at Moonshot's base_url, so it forwards the string
verbatim and kimi/k3 returns "Not found the model".
- MEMORY_BACKEND: Storage backend (local/cloud)
- MAX_TOKENS: Maximum token limit (default: 128000)
Quick Start¶
Basic Usage¶
Run the quickstart examples to see the agent in action:
This will demonstrate: - Basic conversation with memory - Multi-session memory persistence - Multi-agent collaboration
Memory Pipeline Demo (提取—对比—决策)¶
The single clearest demonstration of Mem0's value — its ADD / UPDATE / DELETE /
NOOP pipeline and cross-session recall — is the demo mode:
This reproduces the book's centerpiece example (chapter3.md): the user first
says they live in Beijing, a later turn says they moved to Shanghai, and Mem0
resolves the conflict with an UPDATE (revising the existing memory) instead
of storing two contradictory records. In between, a stored memory is recalled
via semantic search, showing memory being used later. (The same routine is
memory_pipeline_example() in quickstart.py, run first by python quickstart.py.)
Direct Memory Operations CLI¶
Mem0's memory API is exposed directly so you can observe each pipeline decision
without the chat loop. All flags have Chinese --help (python main.py --help):
# ADD — write a conversation/utterance; prints the ADD/UPDATE/DELETE events
python main.py --mode memory --op add --text "我住在北京，是一名后端工程师" --user-id u1
# SEARCH — semantic recall
python main.py --mode memory --op search --query "这个用户住在哪里？" --user-id u1
# GET-ALL — list every stored memory (optionally dump to JSON)
python main.py --mode memory --op get-all --user-id u1 --output mem.json
# HISTORY — the change/audit trail of one memory id (shows UPDATE/DELETE over time)
python main.py --mode memory --op history --memory-id <id>
# DELETE — remove one memory by id
python main.py --mode memory --op delete --memory-id <id>
Key flags: --op {add,search,get-all,history,delete}, --text, --query,
--memory-id, --user-id, --agent-id, --model (override MODEL_NAME),
--output (write result JSON). --text accepts either a raw string or a path
to a JSON message list.
These operations,
demomode, and the chat modes all require a working LLM API key (KIMI_API_KEY) and a vector store — Mem0's fact extraction and semantic retrieval are online model calls. With no key the CLI parses arguments and then reports the missing key; no memory output is fabricated.
Interactive Mode¶
Start an interactive conversation session:
Available commands in interactive mode:
- help - Show available commands
- memories - Display stored memories
- metrics - Show performance metrics
- save - Save conversation state
- load - Load previous state
- new - Start a new session
- exit - Exit the program
Batch Processing¶
Process multiple conversations from a JSON file:
Input format:
[
{
"session_id": "session_001",
"user_id": "user_001",
"agent_id": "agent_001",
"turns": [
"First user message",
"Second user message"
]
}
]
Running LOCOMO Benchmark¶
Full Benchmark¶
Run the complete LOCOMO benchmark evaluation:
This will: 1. Generate and run 10 benchmark scenarios 2. Evaluate agent performance on each scenario 3. Calculate metrics (consistency, coherence, memory utilization) 4. Generate detailed reports with visualizations
Benchmark Metrics¶
The benchmark evaluates: - Consistency Score: How well the agent maintains consistent information - Coherence Score: Relevance and logical flow of responses - Memory Retention: Effectiveness of memory storage and retrieval - Response Time: Average generation time per turn - Context Utilization: How well the agent uses available context
Understanding Results¶
Results are saved in JSON format with the following structure:
{
"config": {...},
"scenarios": [
{
"scenario": {...},
"sessions": [...],
"overall_metrics": {
"avg_consistency": 0.92,
"avg_coherence": 0.88,
"memory_utilization": 42
}
}
],
"overall_metrics": {...}
}
Architecture¶
Components¶
- Agent Module (
agent.py) Mem0Agent: Main agent class with memory integrationKimiK3Client: Client for Kimi K3 model API-
AgentContext: Context management for sessions -
Configuration (
config.py) KimiConfig: Kimi model settingsMem0Config: Memory system configuration-
LOCOMOConfig: Benchmark parameters -
Experiment Framework (
experiment.py) LOCOMOBenchmark: Benchmark implementation- Scenario generation and evaluation
- Metrics calculation and reporting
Memory System¶
The Mem0 framework provides: - Vector Storage: Efficient semantic search using embeddings - Memory Consolidation: Automatic extraction of key information - Context Retrieval: Intelligent retrieval of relevant memories - Multi-level Organization: User, agent, and session-level memories
Advanced Usage¶
Custom Scenarios¶
Create custom benchmark scenarios by modifying experiment.py:
custom_scenario = {
"type": "custom_type",
"description": "Your scenario description",
"topics": ["topic1", "topic2"],
"context_requirements": ["requirement1", "requirement2"]
}
Memory Backends¶
Local Storage (Chroma)¶
config.mem0.backend = "local"
config.mem0.vector_store_config = {
"provider": "chroma",
"config": {
"collection_name": "my_collection",
"path": "./data/chroma_db"
}
}
Cloud Storage (Mem0 Cloud)¶
Performance Tuning¶
Optimize performance by adjusting:
- MAX_TOKENS: Reduce for faster responses
- TEMPERATURE: Lower for more consistent outputs
- context_window_size: Balance between context and speed
- Memory retrieval limit: Adjust in _prepare_messages()
Troubleshooting¶
Common Issues¶
- API Key Errors
- Ensure
KIMI_API_KEYis set in.env -
Check API key validity and permissions
-
Memory Backend Issues
- For local: Ensure write permissions in
./data/ -
For cloud: Verify
MEM0_API_KEYis correct -
Performance Issues
- Reduce
MAX_TOKENSfor faster responses - Use local memory backend for lower latency
- Adjust batch sizes in benchmark runs
Debug Mode¶
Enable detailed logging:
Development¶
Project Structure¶
mem0/
├── agent.py # Core agent implementation
├── config.py # Configuration management
├── experiment.py # LOCOMO benchmark
├── main.py # Main entry point
├── quickstart.py # Example demonstrations
├── requirements.txt # Dependencies
├── env.example # Environment template
└── README.md # Documentation
Testing¶
Run tests (when available):
Contributing¶
- Fork the repository
- Create a feature branch
- Make your changes
- Add tests if applicable
- Submit a pull request
Performance Benchmarks¶
Typical performance metrics on standard hardware: - Average response time: 1-3 seconds - Memory retrieval: <100ms - Consistency score: 0.85-0.95 - Coherence score: 0.80-0.90 - Memory utilization: 20-100 items per session
Limitations¶
- Requires active internet connection for API calls
- Memory storage grows with usage (periodic cleanup recommended)
- Context window limited to 128K tokens
- Response quality depends on model availability
License¶
This project is part of the AI Agent Book training materials.
Acknowledgments¶
- Mem0 framework by Mem0 AI
- Kimi K3 model by Moonshot AI
- LOCOMO benchmark concept for long-context evaluation
Support¶
For issues and questions:
- Check the troubleshooting section
- Review example code in quickstart.py
- Refer to the main AI Agent Book documentation
OpenRouter 通用回退 / Universal OpenRouter fallback¶
This experiment now supports a universal OpenRouter fallback for its chat LLM.
- If the primary provider key (e.g.
MOONSHOT_API_KEY/KIMI_API_KEY/OPENAI_API_KEY/DOUBAO_API_KEY…) is present, behavior is unchanged. - Else if
OPENROUTER_API_KEYis set, the chat LLM is automatically routed through OpenRouter (https://openrouter.ai/api/v1). Model names are mapped automatically:gpt-*/o1-*→openai/…,claude-*→anthropic/claude-opus-4.8,kimi-*→moonshotai/kimi-k2.6, ids already containing/are kept as-is, and other provider-native ids (e.g.doubao-*) fall back toopenai/gpt-5.6-luna. SetOPENROUTER_MODELto force a specific OpenRouter model id. - Else a clear error lists the accepted keys.
Add OPENROUTER_API_KEY=... to your .env (see env.example) to enable it.
Note: mem0's embedder still uses OpenAI embeddings (OpenRouter has no embeddings endpoint), so
OPENAI_API_KEYis still required for storing/retrieving memories. The OpenRouter fallback only covers the chat LLM (fact extraction, ADD/UPDATE/DELETE decisions, and answering).
源代码¶
agent.py¶
"""Mem0-powered agent with Kimi K3 integration for LOCOMO benchmark."""
import json
import logging
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
import asyncio
from collections import defaultdict
from mem0 import Memory
from openai import OpenAI
from tenacity import retry, stop_after_attempt, wait_exponential
import numpy as np
from rich.console import Console
from rich.table import Table
from rich.progress import track
from config import Config, config as default_config
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
def _as_memory_list(result: Any) -> List[Dict[str, Any]]:
"""Normalize a mem0 return value to a plain list of memory dicts.
mem0 >=1.0 returns ``{"results": [...], "relations": [...]}`` from
``search``/``get_all``, whereas older versions returned a bare list.
Accepting either form keeps the agent working across mem0 versions
(iterating the dict directly would otherwise yield the string keys).
"""
if isinstance(result, dict):
return result.get("results", []) or []
if isinstance(result, list):
return result
return []
def _extract_memory_events(add_result: Any) -> List[Dict[str, str]]:
"""Extract ADD/UPDATE/DELETE decisions from a mem0 ``add`` return value.
Mem0's "extract-compare-decide" pipeline judges every candidate fact
against existing memories and emits one of ADD (brand-new), UPDATE
(revise an existing memory), DELETE (retract a contradicted memory)
or NOOP (duplicate, nothing changes). ``add`` surfaces the non-NOOP
decisions; this returns them as ``[{"event", "memory", "id"}]`` so the
book's four-way decision is visible to callers.
"""
events = []
for item in _as_memory_list(add_result):
events.append({
"event": item.get("event", "ADD"),
"memory": item.get("memory", item.get("text", "")),
"id": item.get("id", ""),
})
return events
# Set up logging
logging.basicConfig(
level=getattr(logging, default_config.logging.level),
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
console = Console()
@dataclass
class AgentContext:
"""Context information for an agent in the LOCOMO benchmark."""
agent_id: str
user_id: str
session_id: str
turn_count: int = 0
conversation_history: List[Dict[str, str]] = field(default_factory=list)
metadata: Dict[str, Any] = field(default_factory=dict)
def add_turn(self, role: str, content: str) -> None:
"""Add a turn to the conversation history."""
self.conversation_history.append({
"role": role,
"content": content,
"timestamp": datetime.now().isoformat(),
"turn": self.turn_count
})
self.turn_count += 1
class KimiK3Client:
"""Client for interacting with Kimi K3 model."""
def __init__(self, config: Config):
self.config = config
self.client = OpenAI(
api_key=config.kimi.api_key,
base_url=config.kimi.api_base
)
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def generate(self, messages: List[Dict[str, str]], **kwargs) -> str:
"""Generate response using Kimi K3 model."""
try:
response = self.client.chat.completions.create(
model=self.config.kimi.model_name,
messages=messages,
max_tokens=kwargs.get("max_tokens", self.config.kimi.max_tokens),
temperature=_reasoning_safe_temperature(self.config.kimi.model_name, kwargs.get("temperature", self.config.kimi.temperature)),
top_p=kwargs.get("top_p", 0.95),
frequency_penalty=kwargs.get("frequency_penalty", 0),
presence_penalty=kwargs.get("presence_penalty", 0)
)
return response.choices[0].message.content
except Exception as e:
logger.error(f"Error generating response with Kimi K3: {e}")
raise
async def agenerate(self, messages: List[Dict[str, str]], **kwargs) -> str:
"""Async generate response using Kimi K3 model."""
return await asyncio.to_thread(self.generate, messages, **kwargs)
class Mem0Agent:
"""Agent powered by Mem0 memory system and Kimi K3 model."""
def __init__(self, config: Optional[Config] = None):
self.config = config or default_config
self.config.validate()
# Initialize Kimi K3 client
self.llm_client = KimiK3Client(self.config)
# Initialize Mem0 memory system
self._init_memory()
# Agent state management
self.active_contexts: Dict[str, AgentContext] = {}
self.performance_metrics: Dict[str, List[float]] = defaultdict(list)
def _init_memory(self) -> None:
"""Initialize Mem0 memory system."""
# Mem0 runs its own LLM calls for fact extraction and the
# ADD/UPDATE/DELETE decision. Left unset, mem0 defaults to
# max_tokens=2000 / temperature=0.1, which is unsafe for reasoning
# models (Kimi K3 wants temperature=1 and enough room for its thinking
# tokens). Pin both explicitly so the pipeline is reasoning-safe.
mem0_config = {
"llm": {
"provider": "openai",
"config": {
"api_key": self.config.kimi.api_key,
# mem0 >=1.0 names this field openai_base_url (not base_url);
# it points the OpenAI-compatible client at Moonshot.
"openai_base_url": self.config.kimi.api_base,
"model": self.config.kimi.model_name,
"temperature": _reasoning_safe_temperature(
self.config.kimi.model_name, self.config.kimi.temperature
),
"max_tokens": max(self.config.kimi.max_tokens, 2048),
}
},
"vector_store": self.config.mem0.vector_store_config,
"embedder": {
"provider": "openai",
"config": {
"model": self.config.mem0.embedding_model
}
}
}
if self.config.mem0.backend == "local":
self.memory = Memory.from_config(mem0_config)
else:
# For cloud backend
self.memory = Memory(api_key=self.config.mem0.api_key)
logger.info(f"Initialized Mem0 memory system with {self.config.mem0.backend} backend")
def create_context(self, agent_id: str, user_id: str, session_id: str) -> AgentContext:
"""Create a new agent context for a session."""
context = AgentContext(
agent_id=agent_id,
user_id=user_id,
session_id=session_id,
metadata={
"created_at": datetime.now().isoformat(),
"model": self.config.kimi.model_name
}
)
self.active_contexts[session_id] = context
logger.info(f"Created context for agent {agent_id} in session {session_id}")
return context
def get_context(self, session_id: str) -> Optional[AgentContext]:
"""Get agent context for a session."""
return self.active_contexts.get(session_id)
def _prepare_messages(self, context: AgentContext, user_input: str) -> List[Dict[str, str]]:
"""Prepare messages for LLM including memory context."""
messages = []
# System prompt
system_prompt = f"""You are an intelligent agent participating in the LOCOMO benchmark.
Your task is to maintain consistent and coherent conversations across multiple sessions.
You have access to a memory system that helps you remember important information.
Agent ID: {context.agent_id}
User ID: {context.user_id}
Session ID: {context.session_id}
Current Turn: {context.turn_count}
Guidelines:
1. Maintain consistency with previous conversations
2. Reference relevant past information when appropriate
3. Build upon established context naturally
4. Be concise but informative in your responses
"""
messages.append({"role": "system", "content": system_prompt})
# Retrieve relevant memories
memories = self.memory.search(
query=user_input,
user_id=context.user_id,
agent_id=context.agent_id,
limit=5
)
if memories and len(memories) > 0:
memory_context = "\n\nRelevant memories from past interactions:\n"
for mem in memories:
memory_context += f"- {mem.get('memory', mem.get('text', ''))}\n"
messages.append({"role": "system", "content": memory_context})
# Add recent conversation history (last 10 turns)
recent_history = context.conversation_history[-10:] if len(context.conversation_history) > 10 else context.conversation_history
for turn in recent_history:
messages.append({"role": turn["role"], "content": turn["content"]})
# Add current user input
messages.append({"role": "user", "content": user_input})
return messages
def process_turn(self, session_id: str, user_input: str) -> Tuple[str, Dict[str, Any]]:
"""Process a single turn in the conversation."""
context = self.get_context(session_id)
if not context:
raise ValueError(f"No context found for session {session_id}")
# Record user input
context.add_turn("user", user_input)
# Prepare messages with memory context
messages = self._prepare_messages(context, user_input)
# Generate response using Kimi K3
start_time = datetime.now()
response = self.llm_client.generate(messages)
generation_time = (datetime.now() - start_time).total_seconds()
# Record assistant response
context.add_turn("assistant", response)
# Store interaction in memory. mem0 runs its extract-compare-decide
# pipeline here and returns the ADD/UPDATE/DELETE decisions it made.
add_result = self.memory.add(
messages=[
{"role": "user", "content": user_input},
{"role": "assistant", "content": response}
],
user_id=context.user_id,
agent_id=context.agent_id,
metadata={
"session_id": session_id,
"turn": context.turn_count - 1,
"timestamp": datetime.now().isoformat()
}
)
memory_events = _extract_memory_events(add_result)
# Calculate metrics
metrics = {
"generation_time": generation_time,
"response_length": len(response),
"turn_count": context.turn_count,
"memory_count": len(_as_memory_list(self.memory.get_all(user_id=context.user_id))),
"memory_events": memory_events
}
# Store performance metrics
self.performance_metrics[session_id].append(generation_time)
logger.info(f"Processed turn {context.turn_count} for session {session_id} in {generation_time:.2f}s")
return response, metrics
async def process_turn_async(self, session_id: str, user_input: str) -> Tuple[str, Dict[str, Any]]:
"""Async version of process_turn."""
return await asyncio.to_thread(self.process_turn, session_id, user_input)
# ------------------------------------------------------------------
# Direct memory operations (used by the CLI and the pipeline demo)
# ------------------------------------------------------------------
def add_memory(self, messages, user_id: str, agent_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None) -> List[Dict[str, str]]:
"""Add a message/conversation to memory.
Returns the ADD/UPDATE/DELETE decisions from mem0's pipeline. An
empty list means every candidate fact was judged NOOP (duplicate).
``messages`` may be a plain string or an OpenAI-style message list.
"""
add_result = self.memory.add(
messages=messages,
user_id=user_id,
agent_id=agent_id,
metadata=metadata or {}
)
return _extract_memory_events(add_result)
def search_memory(self, query: str, user_id: str, agent_id: Optional[str] = None,
limit: int = 5) -> List[Dict[str, Any]]:
"""Semantically retrieve memories relevant to ``query``."""
return _as_memory_list(self.memory.search(
query=query, user_id=user_id, agent_id=agent_id, limit=limit
))
def get_all_memories(self, user_id: str, agent_id: Optional[str] = None) -> List[Dict[str, Any]]:
"""List every stored memory for a user."""
return _as_memory_list(self.memory.get_all(user_id=user_id, agent_id=agent_id))
def memory_history(self, memory_id: str) -> List[Dict[str, Any]]:
"""Return the change history (ADD/UPDATE/DELETE audit trail) of one memory."""
return self.memory.history(memory_id)
def delete_memory(self, memory_id: str) -> str:
"""Delete a single memory by id."""
self.memory.delete(memory_id)
return memory_id
def evaluate_consistency(self, session_id: str) -> float:
"""Evaluate consistency of responses in a session."""
context = self.get_context(session_id)
if not context or len(context.conversation_history) < 2:
return 1.0
# Simple consistency check based on response patterns
responses = [turn["content"] for turn in context.conversation_history if turn["role"] == "assistant"]
if len(responses) < 2:
return 1.0
# Calculate consistency score based on semantic similarity (simplified)
# In a real implementation, you would use embeddings and cosine similarity
consistency_scores = []
for i in range(1, len(responses)):
# Simplified: check for contradiction keywords
prev_response = responses[i-1].lower()
curr_response = responses[i].lower()
contradiction_words = ["however", "but actually", "correction", "i was wrong", "let me correct"]
has_contradiction = any(word in curr_response for word in contradiction_words)
consistency_scores.append(0.5 if has_contradiction else 1.0)
return np.mean(consistency_scores) if consistency_scores else 1.0
def evaluate_coherence(self, session_id: str) -> float:
"""Evaluate coherence of the conversation."""
context = self.get_context(session_id)
if not context or len(context.conversation_history) < 2:
return 1.0
# Simple coherence check based on response relevance
coherence_scores = []
for i in range(0, len(context.conversation_history) - 1, 2):
if i + 1 < len(context.conversation_history):
user_turn = context.conversation_history[i]["content"]
assistant_turn = context.conversation_history[i + 1]["content"]
# Check if response addresses the user input (simplified)
user_keywords = set(user_turn.lower().split())
assistant_keywords = set(assistant_turn.lower().split())
overlap = len(user_keywords.intersection(assistant_keywords))
score = min(1.0, overlap / max(len(user_keywords), 1) * 2)
coherence_scores.append(score)
return np.mean(coherence_scores) if coherence_scores else 1.0
def evaluate_memory_retention(self, user_id: str) -> float:
"""Evaluate memory retention for a user."""
memories = _as_memory_list(self.memory.get_all(user_id=user_id))
if not memories or len(memories) == 0:
return 0.0
# Calculate retention score based on memory count and recency
now = datetime.now()
retention_scores = []
for memory in memories:
created_at = memory.get("created_at", now.isoformat())
if isinstance(created_at, str):
created_at = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
age_hours = (now - created_at).total_seconds() / 3600
# Decay function: memories lose value over time
retention_score = np.exp(-age_hours / 24) # Half-life of 24 hours
retention_scores.append(retention_score)
return np.mean(retention_scores)
def get_performance_summary(self, session_id: Optional[str] = None) -> Dict[str, Any]:
"""Get performance summary for a session or all sessions."""
if session_id:
context = self.get_context(session_id)
if not context:
return {}
metrics = self.performance_metrics.get(session_id, [])
return {
"session_id": session_id,
"turn_count": context.turn_count,
"avg_response_time": np.mean(metrics) if metrics else 0,
"consistency_score": self.evaluate_consistency(session_id),
"coherence_score": self.evaluate_coherence(session_id),
"memory_retention": self.evaluate_memory_retention(context.user_id)
}
else:
# Aggregate metrics for all sessions
all_metrics = []
for sid in self.active_contexts:
all_metrics.append(self.get_performance_summary(sid))
if not all_metrics:
return {}
return {
"total_sessions": len(all_metrics),
"avg_turn_count": np.mean([m["turn_count"] for m in all_metrics]),
"avg_response_time": np.mean([m["avg_response_time"] for m in all_metrics]),
"avg_consistency": np.mean([m["consistency_score"] for m in all_metrics]),
"avg_coherence": np.mean([m["coherence_score"] for m in all_metrics]),
"avg_memory_retention": np.mean([m["memory_retention"] for m in all_metrics])
}
def display_metrics(self, session_id: Optional[str] = None) -> None:
"""Display performance metrics in a formatted table."""
summary = self.get_performance_summary(session_id)
if not summary:
console.print("[yellow]No metrics available[/yellow]")
return
table = Table(title="Performance Metrics")
table.add_column("Metric", style="cyan")
table.add_column("Value", style="green")
for key, value in summary.items():
if isinstance(value, float):
table.add_row(key.replace("_", " ").title(), f"{value:.4f}")
else:
table.add_row(key.replace("_", " ").title(), str(value))
console.print(table)
def reset(self) -> None:
"""Reset the agent state."""
self.active_contexts.clear()
self.performance_metrics.clear()
logger.info("Agent state reset")
def save_state(self, filepath: str) -> None:
"""Save agent state to file."""
state = {
"contexts": {
sid: {
"agent_id": ctx.agent_id,
"user_id": ctx.user_id,
"session_id": ctx.session_id,
"turn_count": ctx.turn_count,
"conversation_history": ctx.conversation_history,
"metadata": ctx.metadata
}
for sid, ctx in self.active_contexts.items()
},
"metrics": dict(self.performance_metrics),
"timestamp": datetime.now().isoformat()
}
with open(filepath, "w") as f:
json.dump(state, f, indent=2)
logger.info(f"Agent state saved to {filepath}")
def load_state(self, filepath: str) -> None:
"""Load agent state from file."""
with open(filepath, "r") as f:
state = json.load(f)
self.active_contexts.clear()
for sid, ctx_data in state["contexts"].items():
context = AgentContext(
agent_id=ctx_data["agent_id"],
user_id=ctx_data["user_id"],
session_id=ctx_data["session_id"],
turn_count=ctx_data["turn_count"],
conversation_history=ctx_data["conversation_history"],
metadata=ctx_data["metadata"]
)
self.active_contexts[sid] = context
self.performance_metrics = defaultdict(list, state["metrics"])
logger.info(f"Agent state loaded from {filepath}")
config.py¶
"""Configuration module for Mem0 agent with Kimi K3 integration."""
import os
from pathlib import Path
from typing import Optional, Dict, Any
from dataclasses import dataclass, field
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
def _openrouter_model_id(model) -> str:
"""Map a provider-native model name to an OpenRouter model id, used by the
universal OpenRouter fallback. An explicit OPENROUTER_MODEL env var wins."""
override = os.getenv("OPENROUTER_MODEL")
if override:
return override
m = (model or "").strip()
if not m:
return "openai/gpt-5.6-luna"
if "/" in m:
return m # already an OpenRouter-style id (e.g. openai/gpt-5.6-luna)
ml = m.lower()
if ml.startswith(("gpt-", "o1", "o3", "o4", "chatgpt")):
return "openai/" + m
if ml.startswith("claude-"):
return "anthropic/claude-opus-4.8"
if ml.startswith("kimi"):
# kimi-k3 is not on OpenRouter; moonshotai/kimi-k2.6 is the closest hosted id.
return "moonshotai/kimi-k2.6"
# Provider-native ids (kimi-*/doubao-*/qwen/deepseek-*) not hosted on
# OpenRouter under the same name -> a widely-available OpenAI chat model.
return "openai/gpt-5.6-luna"
@dataclass
class KimiConfig:
"""Configuration for Kimi K3 model."""
api_key: str = field(default_factory=lambda: os.getenv("KIMI_API_KEY", ""))
model_name: str = field(default_factory=lambda: os.getenv("MODEL_NAME", "kimi-k3"))
max_tokens: int = field(default_factory=lambda: int(os.getenv("MAX_TOKENS", "128000")))
temperature: float = field(default_factory=lambda: float(os.getenv("TEMPERATURE", "0.7")))
api_base: str = field(default_factory=lambda: os.getenv("KIMI_API_BASE", "https://api.moonshot.cn/v1"))
def __post_init__(self):
"""Universal OpenRouter fallback for the chat LLM: when KIMI_API_KEY is
absent but OPENROUTER_API_KEY is present, route the chat model (used by
KimiK3Client and threaded into mem0's own LLM config) through OpenRouter.
NB: mem0's embedder still uses OpenAI embeddings (OpenRouter has no
embeddings endpoint), so OPENAI_API_KEY remains needed for memory add."""
if not self.api_key and os.getenv("OPENROUTER_API_KEY"):
self.api_key = os.getenv("OPENROUTER_API_KEY")
self.api_base = "https://openrouter.ai/api/v1"
self.model_name = _openrouter_model_id(self.model_name)
def validate(self) -> bool:
"""Validate Kimi configuration."""
if not self.api_key:
raise ValueError("KIMI_API_KEY is required (or set OPENROUTER_API_KEY for the fallback)")
if self.max_tokens <= 0 or self.max_tokens > 128000:
raise ValueError("MAX_TOKENS must be between 1 and 128000")
if self.temperature < 0 or self.temperature > 2:
raise ValueError("TEMPERATURE must be between 0 and 2")
return True
@dataclass
class Mem0Config:
"""Configuration for Mem0 memory system."""
api_key: Optional[str] = field(default_factory=lambda: os.getenv("MEM0_API_KEY"))
backend: str = field(default_factory=lambda: os.getenv("MEMORY_BACKEND", "local"))
collection_name: str = field(default_factory=lambda: os.getenv("MEMORY_COLLECTION", "locomo_benchmark"))
embedding_model: str = field(default_factory=lambda: os.getenv("MEMORY_EMBEDDING_MODEL", "text-embedding-3-small"))
vector_store_config: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
"""Initialize vector store configuration based on backend."""
if self.backend == "local":
# NB: mem0 >=1.0 validates the chroma config against a fixed field
# set (collection_name/path/host/port/api_key/tenant/client). The
# embedding model belongs to the top-level "embedder" block (set in
# agent.py), NOT here — passing embedding_function raises a
# MemoryConfig validation error.
self.vector_store_config = {
"provider": "chroma",
"config": {
"collection_name": self.collection_name,
"path": "./data/chroma_db",
}
}
elif self.backend == "cloud":
if not self.api_key:
raise ValueError("MEM0_API_KEY is required for cloud backend")
self.vector_store_config = {
"provider": "mem0_cloud",
"config": {
"api_key": self.api_key,
"collection_name": self.collection_name
}
}
else:
raise ValueError(f"Invalid backend: {self.backend}. Must be 'local' or 'cloud'")
def validate(self) -> bool:
"""Validate Mem0 configuration."""
if self.backend not in ["local", "cloud"]:
raise ValueError("MEMORY_BACKEND must be 'local' or 'cloud'")
if self.backend == "cloud" and not self.api_key:
raise ValueError("MEM0_API_KEY is required for cloud backend")
return True
@dataclass
class LOCOMOConfig:
"""Configuration for LOCOMO benchmark."""
data_path: Path = field(default_factory=lambda: Path(os.getenv("BENCHMARK_DATA_PATH", "./data/locomo")))
max_sessions: int = field(default_factory=lambda: int(os.getenv("MAX_SESSIONS", "100")))
max_agents: int = field(default_factory=lambda: int(os.getenv("MAX_AGENTS", "10")))
context_window_size: int = field(default_factory=lambda: int(os.getenv("CONTEXT_WINDOW_SIZE", "128000")))
evaluation_metrics: list = field(default_factory=lambda: [
"consistency_score",
"coherence_score",
"memory_retention",
"context_utilization",
"response_relevance"
])
def __post_init__(self):
"""Ensure data path exists."""
self.data_path.mkdir(parents=True, exist_ok=True)
def validate(self) -> bool:
"""Validate LOCOMO configuration."""
if self.max_sessions <= 0:
raise ValueError("MAX_SESSIONS must be positive")
if self.max_agents <= 0:
raise ValueError("MAX_AGENTS must be positive")
if self.context_window_size <= 0:
raise ValueError("CONTEXT_WINDOW_SIZE must be positive")
return True
@dataclass
class LoggingConfig:
"""Configuration for logging."""
level: str = field(default_factory=lambda: os.getenv("LOG_LEVEL", "INFO"))
file_path: Optional[Path] = field(default_factory=lambda: Path(os.getenv("LOG_FILE", "./logs/mem0_agent.log")) if os.getenv("LOG_FILE") else None)
def __post_init__(self):
"""Ensure log directory exists."""
if self.file_path:
self.file_path.parent.mkdir(parents=True, exist_ok=True)
@dataclass
class Config:
"""Main configuration class."""
kimi: KimiConfig = field(default_factory=KimiConfig)
mem0: Mem0Config = field(default_factory=Mem0Config)
locomo: LOCOMOConfig = field(default_factory=LOCOMOConfig)
logging: LoggingConfig = field(default_factory=LoggingConfig)
def validate(self) -> bool:
"""Validate all configurations."""
self.kimi.validate()
self.mem0.validate()
self.locomo.validate()
return True
@classmethod
def from_env(cls) -> "Config":
"""Create configuration from environment variables."""
return cls()
def to_dict(self) -> Dict[str, Any]:
"""Convert configuration to dictionary."""
return {
"kimi": {
"model_name": self.kimi.model_name,
"max_tokens": self.kimi.max_tokens,
"temperature": _reasoning_safe_temperature(self.kimi.model_name, self.kimi.temperature),
"api_base": self.kimi.api_base
},
"mem0": {
"backend": self.mem0.backend,
"collection_name": self.mem0.collection_name,
"embedding_model": self.mem0.embedding_model
},
"locomo": {
"data_path": str(self.locomo.data_path),
"max_sessions": self.locomo.max_sessions,
"max_agents": self.locomo.max_agents,
"context_window_size": self.locomo.context_window_size,
"evaluation_metrics": self.locomo.evaluation_metrics
},
"logging": {
"level": self.logging.level,
"file_path": str(self.logging.file_path) if self.logging.file_path else None
}
}
# Global configuration instance
config = Config.from_env()
experiment.py¶
"""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.memory.get_all(user_id=f"user_{scenario_results['scenario']['scenario_id']}"))
}
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())
main.py¶
"""Main entry point for the Mem0 agent with Kimi K3."""
import asyncio
import argparse
import json
import os
from pathlib import Path
from typing import Optional
import sys
from rich.console import Console
from rich.prompt import Prompt, Confirm
from rich.panel import Panel
from rich.markdown import Markdown
from agent import Mem0Agent
from config import Config
console = Console()
class InteractiveSession:
"""Interactive session manager for Mem0 agent."""
def __init__(self, agent: Mem0Agent):
self.agent = agent
self.current_session = None
self.current_user = None
self.current_agent_id = None
def start_session(self) -> None:
"""Start a new interactive session."""
console.print(Panel.fit(
"[bold cyan]Mem0 Agent Interactive Session[/bold cyan]\n"
"Type 'help' for commands, 'exit' to quit",
title="Welcome"
))
# Get session details
self.current_user = Prompt.ask("Enter user ID", default="user_001")
self.current_agent_id = Prompt.ask("Enter agent ID", default="agent_001")
session_id = Prompt.ask("Enter session ID", default="session_001")
# Create context
context = self.agent.create_context(
agent_id=self.current_agent_id,
user_id=self.current_user,
session_id=session_id
)
self.current_session = session_id
console.print(f"[green]Session started:[/green] {session_id}")
console.print(f"[green]User:[/green] {self.current_user}")
console.print(f"[green]Agent:[/green] {self.current_agent_id}")
def show_help(self) -> None:
"""Display help information."""
help_text = """
# Available Commands
- **help** - Show this help message
- **exit/quit** - Exit the session
- **clear** - Clear the screen
- **metrics** - Show performance metrics
- **memories** - Show stored memories
- **save** - Save conversation state
- **load** - Load conversation state
- **reset** - Reset the agent state
- **new** - Start a new session
"""
console.print(Markdown(help_text))
def show_memories(self) -> None:
"""Display stored memories."""
if not self.current_user:
console.print("[yellow]No active session[/yellow]")
return
memories = self.agent.get_all_memories(user_id=self.current_user)
if not memories:
console.print("[yellow]No memories found[/yellow]")
return
console.print(f"\n[cyan]Memories for {self.current_user}:[/cyan]")
for i, memory in enumerate(memories, 1):
console.print(f"{i}. {memory.get('memory', memory.get('text', 'N/A'))}")
def save_state(self) -> None:
"""Save the current state."""
filepath = Prompt.ask("Enter filepath to save", default="state.json")
self.agent.save_state(filepath)
console.print(f"[green]State saved to {filepath}[/green]")
def load_state(self) -> None:
"""Load a saved state."""
filepath = Prompt.ask("Enter filepath to load", default="state.json")
if Path(filepath).exists():
self.agent.load_state(filepath)
console.print(f"[green]State loaded from {filepath}[/green]")
else:
console.print(f"[red]File not found: {filepath}[/red]")
async def run(self) -> None:
"""Run the interactive session."""
self.start_session()
while True:
try:
# Get user input
user_input = Prompt.ask("\n[bold]You[/bold]")
# Check for commands
if user_input.lower() in ["exit", "quit"]:
if Confirm.ask("Are you sure you want to exit?"):
break
elif user_input.lower() == "help":
self.show_help()
continue
elif user_input.lower() == "clear":
console.clear()
continue
elif user_input.lower() == "metrics":
self.agent.display_metrics(self.current_session)
continue
elif user_input.lower() == "memories":
self.show_memories()
continue
elif user_input.lower() == "save":
self.save_state()
continue
elif user_input.lower() == "load":
self.load_state()
continue
elif user_input.lower() == "reset":
if Confirm.ask("Reset agent state?"):
self.agent.reset()
console.print("[green]Agent state reset[/green]")
continue
elif user_input.lower() == "new":
self.start_session()
continue
# Process the input through the agent
console.print("[dim]Processing...[/dim]")
response, metrics = await self.agent.process_turn_async(
self.current_session,
user_input
)
# Display response
console.print(f"\n[bold cyan]Agent[/bold cyan]: {response}")
# Display metrics (optional)
if metrics.get("generation_time"):
console.print(
f"[dim]Generated in {metrics['generation_time']:.2f}s | "
f"Turn {metrics['turn_count']} | "
f"Memories: {metrics['memory_count']}[/dim]"
)
except KeyboardInterrupt:
console.print("\n[yellow]Interrupted[/yellow]")
if Confirm.ask("Exit session?"):
break
except Exception as e:
console.print(f"[red]Error: {e}[/red]")
console.print("\n[cyan]Session ended. Goodbye![/cyan]")
async def run_batch_mode(agent: Mem0Agent, input_file: Path, output_file: Path) -> None:
"""Run the agent in batch mode."""
console.print(f"[yellow]Processing batch file: {input_file}[/yellow]")
# Read input file
with open(input_file, "r") as f:
batch_data = f.read()
# Parse batch data (assuming JSON format)
import json
try:
sessions = json.loads(batch_data)
except json.JSONDecodeError:
console.print("[red]Invalid JSON in input file[/red]")
return
results = []
# Process each session
for session_data in sessions:
session_id = session_data.get("session_id", "batch_session")
user_id = session_data.get("user_id", "batch_user")
agent_id = session_data.get("agent_id", "batch_agent")
turns = session_data.get("turns", [])
# Create context
context = agent.create_context(
agent_id=agent_id,
user_id=user_id,
session_id=session_id
)
session_results = {
"session_id": session_id,
"user_id": user_id,
"agent_id": agent_id,
"turns": []
}
# Process turns
for turn in turns:
response, metrics = await agent.process_turn_async(session_id, turn)
session_results["turns"].append({
"input": turn,
"response": response,
"metrics": metrics
})
results.append(session_results)
# Save results
with open(output_file, "w") as f:
json.dump(results, f, indent=2)
console.print(f"[green]Results saved to {output_file}[/green]")
def _load_add_messages(text: str):
"""Resolve the --text argument for a memory add operation.
If it points to an existing JSON file, load it (expects a message list
or a string); otherwise treat the argument itself as a user utterance.
"""
if os.path.exists(text):
with open(text, "r", encoding="utf-8") as f:
return json.load(f)
return text
async def run_memory_op(agent: Mem0Agent, args) -> None:
"""Run a single direct memory operation (add/search/get-all/history/delete).
This exposes mem0's memory API on the command line so the extract-
compare-decide pipeline (ADD/UPDATE/DELETE/NOOP) is observable turn by
turn, independent of the chat loop.
"""
op = args.op
if not op:
console.print("[red]memory 模式需要 --op 参数（add/search/get-all/history/delete）[/red]")
sys.exit(1)
result = None
if op == "add":
if not args.text:
console.print("[red]add 操作需要 --text 参数（一段对话文本，或 JSON 消息文件路径）[/red]")
sys.exit(1)
messages = _load_add_messages(args.text)
events = await asyncio.to_thread(agent.add_memory, messages, args.user_id, args.agent_id)
console.print(f"[green]写入完成，记忆流水线（提取—对比—决策）产生的事件：[/green]")
if events:
for ev in events:
console.print(f" [{ev['event']}] {ev['memory']} [dim](id={ev['id']})[/dim]")
else:
console.print(" [dim]（没有 ADD/UPDATE/DELETE —— 候选事实被判定为 NOOP 重复信息）[/dim]")
result = {"op": "add", "user_id": args.user_id, "events": events}
elif op == "search":
if not args.query:
console.print("[red]search 操作需要 --query 参数[/red]")
sys.exit(1)
hits = await asyncio.to_thread(agent.search_memory, args.query, args.user_id, args.agent_id)
console.print(f"[green]检索到 {len(hits)} 条相关记忆：[/green]")
for mem in hits:
console.print(f" - {mem.get('memory', mem.get('text', 'N/A'))} [dim](id={mem.get('id','')})[/dim]")
result = {"op": "search", "query": args.query, "user_id": args.user_id, "memories": hits}
elif op == "get-all":
memories = await asyncio.to_thread(agent.get_all_memories, args.user_id, args.agent_id)
console.print(f"[green]用户 {args.user_id} 共有 {len(memories)} 条记忆：[/green]")
for i, mem in enumerate(memories, 1):
console.print(f" {i}. {mem.get('memory', mem.get('text', 'N/A'))} [dim](id={mem.get('id','')})[/dim]")
result = {"op": "get-all", "user_id": args.user_id, "memories": memories}
elif op == "history":
if not args.memory_id:
console.print("[red]history 操作需要 --memory-id 参数[/red]")
sys.exit(1)
history = await asyncio.to_thread(agent.memory_history, args.memory_id)
console.print(f"[green]记忆 {args.memory_id} 的修改历史：[/green]")
for entry in history:
console.print(f" - {entry}")
result = {"op": "history", "memory_id": args.memory_id, "history": history}
elif op == "delete":
if not args.memory_id:
console.print("[red]delete 操作需要 --memory-id 参数[/red]")
sys.exit(1)
await asyncio.to_thread(agent.delete_memory, args.memory_id)
console.print(f"[green]已删除记忆 {args.memory_id}[/green]")
result = {"op": "delete", "memory_id": args.memory_id}
if args.output and result is not None:
with open(args.output, "w", encoding="utf-8") as f:
json.dump(result, f, ensure_ascii=False, indent=2, default=str)
console.print(f"[green]结果已写入 {args.output}[/green]")
CLI_EPILOG = """\
示例：
python main.py # 默认进入交互式对话（记忆随对话自动写入/检索）
python main.py --mode demo --user-id u1 # 运行“北京→上海”记忆流水线演示（ADD/UPDATE/DELETE/NOOP）
python main.py --mode memory --op add --text "我住在北京，是一名后端工程师" --user-id u1
python main.py --mode memory --op search --query "这个用户住在哪里？" --user-id u1
python main.py --mode memory --op get-all --user-id u1 --output mem.json
python main.py --mode batch --input conversations.json --output results.json
python main.py --mode benchmark --model kimi-k3
说明：memory / demo / interactive / batch / benchmark 均需要可用的 LLM API（KIMI_API_KEY）
及向量存储；Mem0 的记忆提取与检索依赖在线模型调用。
"""
async def main():
"""Main entry point."""
parser = argparse.ArgumentParser(
description="Mem0 记忆智能体（Kimi K3）— 演示 Mem0 的“提取—对比—决策”记忆流水线",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=CLI_EPILOG,
)
parser.add_argument(
"--mode",
choices=["interactive", "batch", "benchmark", "memory", "demo"],
default="interactive",
help="运行模式：interactive 交互对话（默认）/ batch 批量对话 / benchmark 跑 LOCOMO 基准 / "
"memory 直接调用记忆操作 / demo 记忆流水线演示",
)
parser.add_argument(
"--op",
choices=["add", "search", "get-all", "history", "delete"],
help="memory 模式下的记忆操作：add 写入 / search 检索 / get-all 列出全部 / "
"history 查看某条记忆的修改历史 / delete 删除",
)
parser.add_argument("--text", type=str,
help="add 操作的对话输入：一段文本，或指向 JSON 消息列表文件的路径")
parser.add_argument("--query", type=str, help="search 操作的查询语句")
parser.add_argument("--memory-id", type=str, help="history / delete 操作针对的记忆 ID")
parser.add_argument("--user-id", type=str, default="user_001",
help="记忆归属的用户 ID（默认 user_001）")
parser.add_argument("--agent-id", type=str, default="agent_001",
help="智能体 ID（默认 agent_001）")
parser.add_argument("--model", type=str,
help="覆盖 MODEL_NAME，指定对话模型（如 kimi-k3）")
parser.add_argument("--input", type=str, help="batch 模式的输入 JSON 文件")
parser.add_argument("--output", type=str,
help="将结果写入的 JSON 文件（memory / batch 模式）")
parser.add_argument("--config", type=str, help="配置文件路径（预留）")
args = parser.parse_args()
# Initialize configuration
config = Config.from_env()
if args.model:
config.kimi.model_name = args.model
# Initialize agent
console.print("[yellow]Initializing Mem0 agent...[/yellow]")
try:
agent = Mem0Agent(config)
console.print("[green]Agent initialized successfully[/green]")
except Exception as e:
console.print(f"[red]Failed to initialize agent: {e}[/red]")
sys.exit(1)
# Run based on mode
if args.mode == "interactive":
session = InteractiveSession(agent)
await session.run()
elif args.mode == "batch":
if not args.input or not args.output:
console.print("[red]Batch mode requires --input and --output arguments[/red]")
sys.exit(1)
await run_batch_mode(agent, Path(args.input), Path(args.output))
elif args.mode == "memory":
await run_memory_op(agent, args)
elif args.mode == "demo":
from quickstart import memory_pipeline_example
await memory_pipeline_example(agent=agent, user_id=args.user_id)
elif args.mode == "benchmark":
# Import and run benchmark
from experiment import LOCOMOBenchmark
benchmark = LOCOMOBenchmark(agent, config)
results = await benchmark.run_benchmark(num_scenarios=3)
benchmark.display_overall_results(results)
if __name__ == "__main__":
asyncio.run(main())
quickstart.py¶
"""Quick start example for Mem0 agent with Kimi K3."""
import asyncio
import os
from dotenv import load_dotenv
from rich.console import Console
from agent import Mem0Agent
from config import Config
# Load environment variables
load_dotenv()
console = Console()
async def basic_example():
"""Basic example of using Mem0 agent."""
console.print("[bold cyan]Basic Mem0 Agent Example[/bold cyan]\n")
# Initialize configuration
config = Config.from_env()
# Initialize agent
console.print("[yellow]Initializing agent...[/yellow]")
agent = Mem0Agent(config)
# Create a session context
session_id = "quickstart_session"
user_id = "quickstart_user"
agent_id = "quickstart_agent"
context = agent.create_context(
agent_id=agent_id,
user_id=user_id,
session_id=session_id
)
console.print(f"[green]Session created: {session_id}[/green]\n")
# Example conversation
conversations = [
"Hello! I'm interested in learning about machine learning.",
"I prefer Python for programming and have experience with scikit-learn.",
"What would you recommend as the next step in my ML journey?",
"Can you remind me what programming language I mentioned earlier?",
"What libraries have I mentioned using?"
]
for i, user_input in enumerate(conversations, 1):
console.print(f"[bold]Turn {i} - User:[/bold] {user_input}")
# Process the turn
response, metrics = await agent.process_turn_async(session_id, user_input)
console.print(f"[cyan]Agent:[/cyan] {response}")
console.print(f"[dim]Response time: {metrics['generation_time']:.2f}s[/dim]\n")
# Small delay for readability
await asyncio.sleep(0.5)
# Display final metrics
console.print("\n[bold]Session Metrics:[/bold]")
agent.display_metrics(session_id)
# Show stored memories
console.print("\n[bold]Stored Memories:[/bold]")
memories = agent.memory.get_all(user_id=user_id)
for memory in memories:
console.print(f"- {memory.get('memory', memory.get('text', 'N/A'))}")
async def memory_pipeline_example(agent=None, user_id: str = "pipeline_user"):
"""Demonstrate Mem0's extract-compare-decide pipeline (ADD/UPDATE/DELETE/NOOP).
This reproduces the book's centerpiece example: the user first says they
live in Beijing, later says they moved to Shanghai, and Mem0 resolves the
conflict by UPDATE-ing the existing memory instead of keeping two
contradictory records. It also shows a stored memory being recalled in a
*later* session, which is the whole point of a memory framework.
Requires a working LLM API (KIMI_API_KEY) and vector store — Mem0's fact
extraction and semantic retrieval are online model calls.
"""
console.print("\n[bold cyan]Memory Pipeline Example (提取—对比—决策)[/bold cyan]\n")
if agent is None:
agent = Mem0Agent(Config.from_env())
def show_events(label, events):
console.print(f"[bold]{label}[/bold]")
if events:
for ev in events:
console.print(f" [magenta][{ev['event']}][/magenta] {ev['memory']} "
f"[dim](id={ev['id']})[/dim]")
else:
console.print(" [dim](NOOP — 没有产生新的记忆变更)[/dim]")
console.print()
# --- Session 1: establish facts about the user ---------------------------
console.print("[yellow]Session 1 —— 首次对话，建立用户画像[/yellow]")
events = await asyncio.to_thread(
agent.add_memory,
"我住在北京，在一家 AI 创业公司做后端工程师。",
user_id,
)
show_events("写入「我住在北京 / 后端工程师」的记忆决策：", events)
events = await asyncio.to_thread(
agent.add_memory,
"我平时喜欢周末去爬山，也在学弹吉他。",
user_id,
)
show_events("写入「爱好」的记忆决策：", events)
# --- Recall the stored memory (used later, across the session) -----------
console.print("[yellow]检索 —— 从记忆中回忆用户信息（跨轮次复用）[/yellow]")
hits = await asyncio.to_thread(
agent.search_memory, "这个用户住在哪座城市？做什么工作？", user_id
)
console.print(f"[bold]检索到 {len(hits)} 条相关记忆：[/bold]")
for mem in hits:
console.print(f" - {mem.get('memory', mem.get('text', 'N/A'))}")
console.print()
# --- Session 2 (later): conflicting fact triggers UPDATE -----------------
console.print("[yellow]Session 2（一段时间后）—— 用户搬家，出现冲突信息[/yellow]")
events = await asyncio.to_thread(
agent.add_memory,
"更新一下，我上个月从北京搬到上海了。",
user_id,
)
show_events("写入「搬到上海」后的记忆决策（预期出现 UPDATE，而非新增矛盾条目）：", events)
# --- Verify consolidation: no contradictory Beijing/Shanghai pair --------
console.print("[yellow]核对 —— 记忆库应当保持一致，而不是同时保留北京与上海[/yellow]")
memories = await asyncio.to_thread(agent.get_all_memories, user_id)
console.print(f"[bold]用户 {user_id} 当前全部记忆（{len(memories)} 条）：[/bold]")
for i, mem in enumerate(memories, 1):
console.print(f" {i}. {mem.get('memory', mem.get('text', 'N/A'))}")
console.print()
console.print("[dim]提示：观察居住地记忆是否已被 UPDATE 为“上海”，且没有残留矛盾的“北京”条目。[/dim]")
async def multi_session_example():
"""Example showing memory persistence across sessions."""
console.print("\n[bold cyan]Multi-Session Memory Example[/bold cyan]\n")
# Initialize agent
config = Config.from_env()
agent = Mem0Agent(config)
user_id = "persistent_user"
# First session
console.print("[yellow]Starting Session 1...[/yellow]")
session1_id = "session_001"
context1 = agent.create_context(
agent_id="agent_001",
user_id=user_id,
session_id=session1_id
)
# First session conversation
response1, _ = await agent.process_turn_async(
session1_id,
"Hi! I'm working on a project about renewable energy, specifically solar panels."
)
console.print(f"[cyan]Session 1 Response:[/cyan] {response1}\n")
response2, _ = await agent.process_turn_async(
session1_id,
"I need to analyze efficiency data from different manufacturers."
)
console.print(f"[cyan]Session 1 Response:[/cyan] {response2}\n")
# Second session (different session, same user)
console.print("[yellow]Starting Session 2 (after some time)...[/yellow]")
session2_id = "session_002"
context2 = agent.create_context(
agent_id="agent_001",
user_id=user_id,
session_id=session2_id
)
# Second session should remember context from first session
response3, _ = await agent.process_turn_async(
session2_id,
"What was I working on last time we talked?"
)
console.print(f"[cyan]Session 2 Response:[/cyan] {response3}\n")
response4, _ = await agent.process_turn_async(
session2_id,
"Can you help me continue with that project?"
)
console.print(f"[cyan]Session 2 Response:[/cyan] {response4}\n")
# Show all memories
console.print("[bold]All Memories for User:[/bold]")
memories = agent.memory.get_all(user_id=user_id)
for memory in memories:
console.print(f"- {memory.get('memory', memory.get('text', 'N/A'))}")
async def multi_agent_example():
"""Example with multiple agents collaborating."""
console.print("\n[bold cyan]Multi-Agent Collaboration Example[/bold cyan]\n")
# Initialize agent
config = Config.from_env()
agent = Mem0Agent(config)
user_id = "collaboration_user"
session_id = "collab_session"
# Create contexts for multiple agents
agents = ["researcher", "analyst", "advisor"]
contexts = {}
for agent_id in agents:
contexts[agent_id] = agent.create_context(
agent_id=agent_id,
user_id=user_id,
session_id=f"{session_id}_{agent_id}"
)
console.print(f"[green]Created context for {agent_id}[/green]")
# Collaborative conversation
console.print("\n[yellow]Starting collaborative discussion...[/yellow]\n")
# Researcher starts
response1, _ = await agent.process_turn_async(
f"{session_id}_researcher",
"I've found some interesting data on climate change impacts on agriculture."
)
console.print(f"[cyan]Researcher:[/cyan] {response1}\n")
# Analyst responds
response2, _ = await agent.process_turn_async(
f"{session_id}_analyst",
"Based on what the researcher mentioned, what are the key metrics we should analyze?"
)
console.print(f"[cyan]Analyst:[/cyan] {response2}\n")
# Advisor provides guidance
response3, _ = await agent.process_turn_async(
f"{session_id}_advisor",
"Considering both the research and analysis perspectives, what recommendations can we make?"
)
console.print(f"[cyan]Advisor:[/cyan] {response3}\n")
# Show metrics for all agents
console.print("[bold]Performance Metrics:[/bold]")
for agent_id in agents:
console.print(f"\n[yellow]{agent_id.capitalize()}:[/yellow]")
summary = agent.get_performance_summary(f"{session_id}_{agent_id}")
for key, value in summary.items():
if isinstance(value, float):
console.print(f" {key}: {value:.3f}")
else:
console.print(f" {key}: {value}")
async def main():
"""Run all examples."""
console.print(Panel.fit(
"[bold]Mem0 Agent Quickstart Examples[/bold]\n"
"Demonstrating various capabilities of the Mem0 agent with Kimi K3",
title="Welcome"
))
# Check for API key
if not os.getenv("KIMI_API_KEY"):
console.print("[red]Error: KIMI_API_KEY not found in environment[/red]")
console.print("Please set your Kimi API key in the .env file")
return
try:
# Run examples
await memory_pipeline_example()
await asyncio.sleep(1)
await basic_example()
await asyncio.sleep(1)
await multi_session_example()
await asyncio.sleep(1)
await multi_agent_example()
console.print("\n[green]All examples completed successfully![/green]")
except Exception as e:
console.print(f"[red]Error running examples: {e}[/red]")
import traceback
traceback.print_exc()
if __name__ == "__main__":
from rich.panel import Panel
asyncio.run(main())
test_simple.py¶
"""Simple test to verify the Mem0 agent setup."""
import os
import asyncio
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def test_imports():
"""Test that all modules can be imported."""
print("Testing imports...")
try:
from agent import Mem0Agent, AgentContext, KimiK3Client
print("✓ Agent module imported successfully")
from config import Config, KimiConfig, Mem0Config, LOCOMOConfig
print("✓ Config module imported successfully")
from experiment import LOCOMOBenchmark
print("✓ Experiment module imported successfully")
import mem0
print("✓ Mem0 library available")
import openai
print("✓ OpenAI library available")
return True
except ImportError as e:
print(f"✗ Import error: {e}")
return False
def test_configuration():
"""Test configuration loading."""
print("\nTesting configuration...")
try:
from config import Config
config = Config.from_env()
# Check if API key is set
if not config.kimi.api_key:
print("⚠ Warning: KIMI_API_KEY not set in environment")
print(" Please create a .env file from env.example and add your API key")
return False
else:
print(f"✓ Kimi API key configured (length: {len(config.kimi.api_key)})")
print(f"✓ Model: {config.kimi.model_name}")
print(f"✓ Max tokens: {config.kimi.max_tokens}")
print(f"✓ Memory backend: {config.mem0.backend}")
print(f"✓ LOCOMO data path: {config.locomo.data_path}")
# Validate configuration
config.validate()
print("✓ Configuration validation passed")
return True
except Exception as e:
print(f"✗ Configuration error: {e}")
return False
async def test_agent_initialization():
"""Test agent initialization."""
print("\nTesting agent initialization...")
try:
from agent import Mem0Agent
from config import Config
config = Config.from_env()
# Skip if no API key
if not config.kimi.api_key:
print("⚠ Skipping agent test (no API key)")
return False
agent = Mem0Agent(config)
print("✓ Agent initialized successfully")
# Create a test context
context = agent.create_context(
agent_id="test_agent",
user_id="test_user",
session_id="test_session"
)
print(f"✓ Context created: {context.session_id}")
return True
except Exception as e:
print(f"✗ Agent initialization error: {e}")
import traceback
traceback.print_exc()
return False
async def test_memory_system():
"""Test memory system initialization."""
print("\nTesting memory system...")
try:
from mem0 import Memory
from config import Config
config = Config.from_env()
if config.mem0.backend == "local":
# Test local memory initialization
mem0_config = {
"vector_store": {
"provider": "chroma",
"config": {
"collection_name": "test_collection",
"path": "./data/test_chroma_db"
}
}
}
memory = Memory.from_config(mem0_config)
print("✓ Local memory system initialized")
# Clean up test database
import shutil
if os.path.exists("./data/test_chroma_db"):
shutil.rmtree("./data/test_chroma_db")
print("✓ Test database cleaned up")
else:
print("✓ Cloud memory backend configured")
return True
except Exception as e:
print(f"✗ Memory system error: {e}")
return False
def main():
"""Run all tests."""
print("=" * 50)
print("Mem0 Agent Setup Test")
print("=" * 50)
all_passed = True
# Run tests
if not test_imports():
all_passed = False
if not test_configuration():
all_passed = False
# Run async tests
loop = asyncio.get_event_loop()
if not loop.run_until_complete(test_agent_initialization()):
all_passed = False
if not loop.run_until_complete(test_memory_system()):
all_passed = False
# Summary
print("\n" + "=" * 50)
if all_passed:
print("✅ All tests passed! The agent is ready to use.")
print("\nNext steps:")
print("1. Run quickstart examples: python quickstart.py")
print("2. Try interactive mode: python main.py")
print("3. Run LOCOMO benchmark: python experiment.py")
else:
print("⚠️ Some tests failed. Please check the configuration.")
print("\nTroubleshooting:")
print("1. Create .env file from env.example")
print("2. Add your KIMI_API_KEY to .env")
print("3. Install all requirements: pip install -r requirements.txt")
print("=" * 50)
if __name__ == "__main__":
main()