memobase

第3章 · 用户记忆和知识库 · 配套项目 chapter3/memobase

项目说明

Memobase Agent with Kimi K3 for LOCOMO Benchmark

An advanced AI agent implementation featuring sophisticated memory management, powered by the Kimi K3 model, and evaluated on the LOCOMO (Long Context and Memory Optimization) benchmark.

Two tracks in this folder — don't confuse them:

1. Real Memobase framework demo (profile_demo.py) — uses the actual open-source Memobase SDK (pip install memobase, package memobase>=0.0.27) talking to a running Memobase server. This is the canonical demonstration of Memobase's Profile (structured user attributes) + Event Memory (timeline) split described in book chapter 3. See Memobase Profile + Event Demo.
2. Hand-rolled memory agent (agent.py / main.py) — a self-contained, Memobase-inspired MemoryStore (episodic/semantic/procedural/working memory, pickle-persisted) that calls Kimi K3 directly. It needs no Memobase server — only a KIMI_API_KEY. The --mode commands below (interactive / benchmark / demo / task) all drive this agent.

Features

🧠 Advanced Memory Management

• Multiple Memory Types:
• Episodic Memory: Stores specific task experiences and interactions
• Semantic Memory: Maintains general knowledge and facts
• Procedural Memory: Learns and stores problem-solving patterns

• Working Memory: Manages short-term task context

• Memory Operations:

• Automatic memory compression when thresholds are exceeded
• Memory consolidation and pattern extraction
• Importance-based decay and retention
• Clustering of related memories
• Efficient retrieval based on relevance and recency

🚀 Kimi K3 Model Integration

• Utilizes Kimi K3, a ~2.8-trillion parameter Mixture-of-Experts model
• 32 billion active parameters per forward pass
• Optimized for agentic capabilities including:
• Advanced tool use
• Multi-step reasoning
• Code synthesis
• Long-context understanding (128k tokens)

📊 LOCOMO Benchmark

Comprehensive evaluation across multiple task categories: - Multi-turn Reasoning: Complex problem-solving over multiple interactions - Long Context Q&A: Information extraction from extensive documents - Task Planning: Project and resource planning capabilities - Knowledge Integration: Cross-domain synthesis and analysis - Tool Usage: Effective utilization of external tools and APIs

Installation

1. Clone the repository:

   cd projects/week2/memobase
   

2. Install dependencies:

   pip install -r requirements.txt
   

3. Set up environment variables:

   cp env.example .env
   # Edit .env and add your Kimi API key
   

Configuration

Edit config.py to customize: - Model parameters (temperature, max tokens, etc.) - Memory settings (thresholds, retention policies) - Benchmark configurations - Logging levels

Usage

Interactive Mode

Engage in conversation with the agent while it builds and uses memories:

python main.py --mode interactive

Available commands in interactive mode: - /help - Show available commands - /memory - Display memory statistics - /clear - Clear working memory - /reset - Reset conversation (preserves long-term memories) - /learn - Trigger memory consolidation - /exit - Exit interactive mode

Benchmark Mode

Run the LOCOMO benchmark evaluation:

# Run full benchmark
python main.py --mode benchmark

# Run specific category
python main.py --mode benchmark --category multi_turn_reasoning

# Run limited number of tasks
python main.py --mode benchmark --num-tasks 5

Demo Mode

Run pre-configured demonstration scenarios:

python main.py --mode demo

Demonstrations include: - Memory retention across conversations - Learning from experience - Context management in long interactions

Single Task Mode

Execute a specific task:

python main.py --mode task --task "Plan a 7-day trip to Japan with a $3000 budget"

Additional Options

• --api-key KEY - Override the API key from environment
• --no-memory - Start with empty memory store
• --verbose - Enable detailed logging

Memobase Profile + Event Demo (real SDK)

profile_demo.py uses the real open-source Memobase SDK to show its two memory structures: a Profile (topic → sub-topic → content, e.g. basic_info→城市, work→职位, interest→游戏偏好) and Event Memory (a timeline for "when did we discuss the budget?" style questions). It follows Memobase's buffered pipeline: insert (write to the user buffer) → flush (trigger one LLM extraction) → profile / event / context (recall).

Prerequisites: a running Memobase server + an extraction LLM

Memobase does its memory extraction server-side, so the demo needs a reachable Memobase service (the client does not call the LLM itself):

• Self-hosted: follow memodb-io/memobase (docker compose). Default endpoint http://localhost:8019, default token secret. The extraction model is configured in the server's .env / config.yaml, not by this client (--model is informational only).
• Cloud: get a project_url + api_key from https://www.memobase.io.

Point the client at the server via --project-url / --api-key or the MEMOBASE_PROJECT_URL / MEMOBASE_API_KEY environment variables (see env.example).

Running it

pip install -r requirements.txt          # installs the memobase SDK

# End-to-end demo (built-in sample conversation) against a running server:
python profile_demo.py

# Offline preview — shows the conversation + pipeline without a server,
# does not go online and does not fabricate results:
python profile_demo.py --dry-run

# Individual ops once memories exist:
python profile_demo.py --op profile      # recall structured user profile
python profile_demo.py --op event        # recall the event timeline
python profile_demo.py --op context      # assembled memory context string
python profile_demo.py --input chat.json --output result.json

If no server is reachable the demo exits with an actionable message (pointing to --dry-run or a --project-url) rather than inventing memory output.

Architecture

Memory Store (agent.py)

The MemoryStore class manages different memory types with: - Persistent storage using pickle - Memory compression and clustering - Importance-based retrieval - Time-based decay mechanisms

Agent Core (agent.py)

The MemobaseAgent class provides: - Message processing with memory context - Learning from task outcomes - Memory-aware response generation - Performance metrics tracking

Benchmark System (locomo_benchmark.py)

The LOCOMOBenchmark class implements: - Task generation and management - Response evaluation metrics - Category-specific scoring - Results persistence and analysis

Memory Management Strategies

Compression

When memory exceeds thresholds: 1. Sort memories by importance and recency 2. Keep high-importance memories 3. Compress low-importance memories into clusters 4. Store cluster summaries as compressed memories

Consolidation

During idle or triggered consolidation: 1. Apply decay to all memories 2. Remove very low importance memories 3. Extract patterns from episodic memories 4. Create procedural memories from patterns

Retrieval

When processing queries: 1. Search for relevant memories by content 2. Retrieve recent episodic memories 3. Include applicable procedural knowledge 4. Format memories for context inclusion

Benchmark Results

Results are saved in benchmark_results/ with: - Task-level scores and timings - Category-wise performance metrics - Memory usage statistics - Detailed error logs

Development

Adding Custom Tools

Extend the agent with custom tools by modifying the agent class:

def add_tool(self, tool_name, tool_function):
    # Add tool registration logic
    pass

Custom Memory Types

Add new memory types in config.py:

MEMOBASE_CONFIG["memory_types"].append("custom_type")

Extending Benchmark

Add custom benchmark tasks in locomo_benchmark.py:

self.tasks.append(BenchmarkTask(
    id="custom_001",
    category="custom_category",
    query="Your custom task query",
    expected_capabilities=["capability1", "capability2"]
))

Performance Optimization

Memory Efficiency

• Batch memory operations for better performance
• Use compression for large memory stores
• Implement periodic consolidation
• Clear working memory after task completion

Response Latency

• Pre-fetch likely memories
• Cache embedding computations
• Use parallel processing for memory search
• Optimize context window usage

Troubleshooting

Common Issues

1. API Key Error:
2. Ensure KIMI_API_KEY is set in .env

3. Verify API key validity

4. Memory Overflow:

5. Adjust MAX_MEMORY_ENTRIES in config
6. Enable more aggressive compression

7. Trigger manual consolidation

8. Slow Performance:

9. Reduce MODEL_MAX_TOKENS
10. Enable caching in config
11. Use category-specific benchmarks

Contributing

Contributions are welcome! Areas for improvement: - Additional memory compression algorithms - Enhanced retrieval mechanisms - New benchmark categories - Performance optimizations - Tool integrations

License

MIT License - See LICENSE file for details

Acknowledgments

• Kimi K3 model by Moonshot AI
• Memobase framework concepts
• LOCOMO benchmark design inspirations

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_KEY is 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 to openai/gpt-5.6-luna. Set OPENROUTER_MODEL to 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.

源代码

agent.py

"""
Memobase Agent Implementation with Kimi K3 Model
Advanced memory management for LOCOMO benchmark
"""

import json
import logging
import time
import hashlib
import pickle
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from collections import defaultdict
from openai import OpenAI

from config import (
    KIMI_API_KEY, KIMI_BASE_URL, KIMI_MODEL,
    MODEL_TEMPERATURE, MODEL_MAX_TOKENS, MODEL_TOP_P,
    MEMOBASE_CONFIG, MEMORY_DB_PATH, AGENT_CONFIG,
    MAX_MEMORY_ENTRIES, MEMORY_COMPRESSION_THRESHOLD,
    LOG_LEVEL, LOG_FORMAT
)


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


# Configure logging
logging.basicConfig(level=LOG_LEVEL, format=LOG_FORMAT)
logger = logging.getLogger(__name__)


@dataclass
class Memory:
    """Represents a single memory entry"""
    id: str
    type: str  # episodic, semantic, procedural, working
    content: Any
    embedding: Optional[List[float]] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    accessed_at: datetime = field(default_factory=datetime.now)
    access_count: int = 0
    importance_score: float = 1.0
    decay_rate: float = 0.1

    def __post_init__(self):
        if not self.id:
            # Generate unique ID based on content
            content_str = json.dumps(self.content, sort_keys=True)
            self.id = hashlib.md5(content_str.encode()).hexdigest()[:12]

    def access(self):
        """Update access statistics"""
        self.accessed_at = datetime.now()
        self.access_count += 1
        # Increase importance with access
        self.importance_score = min(10.0, self.importance_score * 1.1)

    def decay(self):
        """Apply time-based decay to importance"""
        time_since_access = (datetime.now() - self.accessed_at).total_seconds() / 3600
        self.importance_score *= (1 - self.decay_rate * min(1, time_since_access / 24))
        self.importance_score = max(0.1, self.importance_score)


@dataclass
class MemoryCluster:
    """Represents a cluster of related memories"""
    id: str
    memories: List[Memory]
    summary: Optional[str] = None
    centroid_embedding: Optional[List[float]] = None
    created_at: datetime = field(default_factory=datetime.now)

    def add_memory(self, memory: Memory):
        """Add a memory to the cluster"""
        self.memories.append(memory)
        # TODO: Update centroid embedding

    def compress(self) -> str:
        """Compress cluster into a summary"""
        if not self.summary:
            # Create summary from memories
            contents = [m.content for m in self.memories]
            self.summary = f"Cluster of {len(self.memories)} related memories: {contents[:3]}..."
        return self.summary


class MemoryStore:
    """Manages different types of memories with persistence"""

    def __init__(self, db_path: Path = MEMORY_DB_PATH):
        self.db_path = db_path
        self.memories: Dict[str, List[Memory]] = defaultdict(list)
        self.clusters: List[MemoryCluster] = []
        self.embeddings_cache: Dict[str, List[float]] = {}
        self._load_memories()
        logger.info(f"Initialized MemoryStore at {db_path}")

    def _load_memories(self):
        """Load memories from persistent storage"""
        memory_file = self.db_path / "memories.pkl"
        if memory_file.exists():
            try:
                with open(memory_file, 'rb') as f:
                    data = pickle.load(f)
                    self.memories = data.get('memories', defaultdict(list))
                    self.clusters = data.get('clusters', [])
                    logger.info(f"Loaded {sum(len(m) for m in self.memories.values())} memories")
            except Exception as e:
                logger.error(f"Failed to load memories: {e}")

    def _save_memories(self):
        """Save memories to persistent storage"""
        memory_file = self.db_path / "memories.pkl"
        try:
            with open(memory_file, 'wb') as f:
                pickle.dump({
                    'memories': dict(self.memories),
                    'clusters': self.clusters
                }, f)
            logger.debug("Saved memories to disk")
        except Exception as e:
            logger.error(f"Failed to save memories: {e}")

    def add_memory(self, memory_type: str, content: Any, 
                   metadata: Optional[Dict] = None,
                   importance: float = 1.0) -> Memory:
        """Add a new memory"""
        memory = Memory(
            id="",  # Will be auto-generated
            type=memory_type,
            content=content,
            metadata=metadata or {},
            importance_score=importance
        )

        self.memories[memory_type].append(memory)

        # Apply compression if needed
        if len(self.memories[memory_type]) > MEMORY_COMPRESSION_THRESHOLD:
            self._compress_memories(memory_type)

        self._save_memories()
        logger.debug(f"Added {memory_type} memory: {memory.id}")
        return memory

    def get_memories(self, memory_type: Optional[str] = None,
                     limit: int = 10,
                     min_importance: float = 0.5) -> List[Memory]:
        """Retrieve memories, optionally filtered by type and importance"""
        if memory_type:
            memories = self.memories.get(memory_type, [])
        else:
            memories = [m for mlist in self.memories.values() for m in mlist]

        # Filter by importance and sort by relevance
        memories = [m for m in memories if m.importance_score >= min_importance]
        memories.sort(key=lambda m: (m.importance_score, m.access_count), reverse=True)

        # Update access stats
        for memory in memories[:limit]:
            memory.access()

        return memories[:limit]

    def search_memories(self, query: str, limit: int = 5) -> List[Memory]:
        """Search memories by content similarity"""
        results = []
        query_lower = query.lower()

        for memory_list in self.memories.values():
            for memory in memory_list:
                # Simple text matching (would use embeddings in production)
                content_str = str(memory.content).lower()
                if query_lower in content_str:
                    score = content_str.count(query_lower) * memory.importance_score
                    results.append((score, memory))

        results.sort(key=lambda x: x[0], reverse=True)

        # Update access stats
        for _, memory in results[:limit]:
            memory.access()

        return [m for _, m in results[:limit]]

    def _compress_memories(self, memory_type: str):
        """Compress old memories to save space"""
        memories = self.memories[memory_type]
        if len(memories) <= MEMORY_COMPRESSION_THRESHOLD:
            return

        # Sort by importance and recency
        memories.sort(key=lambda m: (m.importance_score, m.accessed_at.timestamp()))

        # Keep top memories, compress others
        to_keep = memories[-MAX_MEMORY_ENTRIES//2:]
        to_compress = memories[:-MAX_MEMORY_ENTRIES//2]

        if to_compress:
            # Create a cluster from compressed memories
            cluster = MemoryCluster(
                id=hashlib.md5(f"{memory_type}_{datetime.now()}".encode()).hexdigest()[:12],
                memories=to_compress
            )
            cluster.compress()
            self.clusters.append(cluster)

            # Replace with compressed version
            compressed_memory = Memory(
                id=cluster.id,
                type=memory_type,
                content=cluster.summary,
                metadata={"cluster_id": cluster.id, "compressed_count": len(to_compress)},
                importance_score=sum(m.importance_score for m in to_compress) / len(to_compress)
            )

            self.memories[memory_type] = [compressed_memory] + to_keep
            logger.info(f"Compressed {len(to_compress)} {memory_type} memories into cluster {cluster.id}")

    def consolidate_memories(self):
        """Consolidate and reorganize memories for efficiency"""
        for memory_type in self.memories:
            memories = self.memories[memory_type]

            # Apply decay to all memories
            for memory in memories:
                memory.decay()

            # Remove very low importance memories
            self.memories[memory_type] = [
                m for m in memories if m.importance_score > 0.1
            ]

        self._save_memories()
        logger.info("Consolidated memories")

    def clear_working_memory(self):
        """Clear working memory (short-term)"""
        self.memories['working'] = []
        logger.debug("Cleared working memory")


class MemobaseAgent:
    """
    Advanced agent with Memobase memory management for LOCOMO benchmark
    """

    def __init__(self, api_key: str = KIMI_API_KEY):
        """Initialize the Memobase agent"""
        self.client = OpenAI(
            api_key=api_key,
            base_url=KIMI_BASE_URL
        )
        self.model = KIMI_MODEL
        self.memory_store = MemoryStore()
        self.conversation_history = []
        self.current_task = None
        self.task_context = {}

        # Initialize system prompt
        self._init_system_prompt()
        logger.info(f"Initialized MemobaseAgent with model {self.model}")

    def _init_system_prompt(self):
        """Initialize the system prompt with memory capabilities"""
        self.system_prompt = """You are an advanced AI agent with sophisticated memory management capabilities.

You have access to multiple types of memory:
1. **Episodic Memory**: Specific experiences and events from tasks
2. **Semantic Memory**: General knowledge and facts
3. **Procedural Memory**: Learned procedures and problem-solving patterns
4. **Working Memory**: Current task context and temporary information

Memory Management Guidelines:
- Store important information for future reference
- Retrieve relevant memories when solving new problems
- Learn from past experiences to improve performance
- Compress and consolidate memories to maintain efficiency
- Use procedural memories to apply learned strategies

Your goal is to complete tasks efficiently while learning and adapting from experience.
When you encounter similar problems, use your memories to solve them more effectively.

Always think step-by-step and use your memory system strategically."""

    def _store_interaction(self, role: str, content: str, memory_type: str = "episodic"):
        """Store an interaction in memory"""
        self.memory_store.add_memory(
            memory_type=memory_type,
            content={
                "role": role,
                "content": content,
                "task": self.current_task,
                "timestamp": datetime.now().isoformat()
            },
            metadata={
                "task_id": self.current_task,
                "turn": len(self.conversation_history)
            }
        )

    def _retrieve_relevant_memories(self, query: str, limit: int = 5) -> List[Memory]:
        """Retrieve memories relevant to current query"""
        # Search across all memory types
        relevant_memories = []

        # Get recent episodic memories
        episodic = self.memory_store.get_memories("episodic", limit=limit//2)
        relevant_memories.extend(episodic)

        # Search for similar content
        searched = self.memory_store.search_memories(query, limit=limit//2)
        relevant_memories.extend(searched)

        # Get procedural memories if task-related
        if "solve" in query.lower() or "how" in query.lower():
            procedural = self.memory_store.get_memories("procedural", limit=2)
            relevant_memories.extend(procedural)

        # Remove duplicates
        seen = set()
        unique_memories = []
        for memory in relevant_memories:
            if memory.id not in seen:
                seen.add(memory.id)
                unique_memories.append(memory)

        return unique_memories[:limit]

    def _format_memories_for_context(self, memories: List[Memory]) -> str:
        """Format memories for inclusion in context"""
        if not memories:
            return ""

        formatted = "\n=== Relevant Memories ===\n"
        for memory in memories:
            formatted += f"[{memory.type.upper()}] (importance: {memory.importance_score:.2f})\n"
            if isinstance(memory.content, dict):
                formatted += json.dumps(memory.content, indent=2)
            else:
                formatted += str(memory.content)
            formatted += "\n---\n"

        return formatted

    def _learn_from_outcome(self, task: str, approach: str, outcome: str, success: bool):
        """Learn from task outcomes and store procedural knowledge"""
        # Store the learning as procedural memory
        self.memory_store.add_memory(
            memory_type="procedural",
            content={
                "task_pattern": task,
                "approach": approach,
                "outcome": outcome,
                "success": success,
                "learned_at": datetime.now().isoformat()
            },
            importance=2.0 if success else 1.0,
            metadata={"task_id": self.current_task}
        )

        if success:
            logger.info(f"Learned successful approach for task type: {task}")
        else:
            logger.info(f"Learned from failure in task type: {task}")

    def process_message(self, message: str, task_id: Optional[str] = None) -> str:
        """
        Process a message with memory-aware reasoning

        Args:
            message: User message to process
            task_id: Optional task identifier for context

        Returns:
            Agent's response
        """
        self.current_task = task_id or f"task_{int(time.time())}"

        # Store the query in working memory
        self.memory_store.add_memory(
            memory_type="working",
            content=message,
            metadata={"task_id": self.current_task}
        )

        # Retrieve relevant memories
        relevant_memories = self._retrieve_relevant_memories(message)
        memory_context = self._format_memories_for_context(relevant_memories)

        # Build messages with memory context
        messages = [
            {"role": "system", "content": self.system_prompt}
        ]

        if memory_context:
            messages.append({
                "role": "system",
                "content": memory_context
            })

        # Add conversation history
        messages.extend(self.conversation_history)
        messages.append({"role": "user", "content": message})

        try:
            # Call Kimi K3 model
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=_reasoning_safe_temperature(self.model, MODEL_TEMPERATURE),
                max_tokens=MODEL_MAX_TOKENS,
                top_p=MODEL_TOP_P
            )

            assistant_response = response.choices[0].message.content

            # Store the interaction in episodic memory
            self._store_interaction("user", message)
            self._store_interaction("assistant", assistant_response)

            # Update conversation history
            self.conversation_history.append({"role": "user", "content": message})
            self.conversation_history.append({"role": "assistant", "content": assistant_response})

            # Keep conversation history manageable
            if len(self.conversation_history) > 20:
                # Move old conversations to episodic memory and compress
                old_convs = self.conversation_history[:10]
                for conv in old_convs:
                    self.memory_store.add_memory(
                        memory_type="episodic",
                        content=conv,
                        importance=0.5
                    )
                self.conversation_history = self.conversation_history[10:]

            return assistant_response

        except Exception as e:
            logger.error(f"Error processing message: {e}")
            return f"Error: {str(e)}"

    def execute_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute a LOCOMO benchmark task

        Args:
            task: Task dictionary with 'id', 'type', 'query', and optional 'context'

        Returns:
            Result dictionary with 'response', 'memories_used', 'execution_time'
        """
        start_time = time.time()
        task_id = task.get('id', f"task_{int(time.time())}")
        task_type = task.get('type', 'unknown')
        query = task['query']
        context = task.get('context', '')

        self.current_task = task_id

        # Check for similar past tasks in procedural memory
        similar_tasks = self.memory_store.search_memories(f"{task_type} {query[:50]}", limit=3)

        # Build enhanced query with context
        enhanced_query = query
        if context:
            enhanced_query = f"Context: {context}\n\nTask: {query}"

        # Process the task
        response = self.process_message(enhanced_query, task_id)

        # Extract approach and outcome for learning
        approach = f"Used {len(similar_tasks)} similar memories"
        outcome = response[:100]  # First 100 chars as outcome summary

        # Learn from this task
        self._learn_from_outcome(
            task=task_type,
            approach=approach,
            outcome=outcome,
            success=True  # Would be determined by evaluation
        )

        execution_time = time.time() - start_time

        return {
            "task_id": task_id,
            "response": response,
            "memories_used": len(similar_tasks),
            "execution_time": execution_time,
            "memory_stats": {
                "episodic": len(self.memory_store.memories.get('episodic', [])),
                "semantic": len(self.memory_store.memories.get('semantic', [])),
                "procedural": len(self.memory_store.memories.get('procedural', [])),
                "working": len(self.memory_store.memories.get('working', []))
            }
        }

    def consolidate_and_learn(self):
        """Consolidate memories and extract learnings"""
        logger.info("Starting memory consolidation...")

        # Consolidate memories
        self.memory_store.consolidate_memories()

        # Extract patterns from episodic memories
        episodic_memories = self.memory_store.get_memories('episodic', limit=50)

        # Group by task type and extract patterns
        task_patterns = defaultdict(list)
        for memory in episodic_memories:
            if isinstance(memory.content, dict):
                task = memory.content.get('task', 'unknown')
                task_patterns[task].append(memory)

        # Create procedural memories from patterns
        for task_type, memories in task_patterns.items():
            if len(memories) >= 3:  # Need multiple examples to learn
                # Extract common approach
                pattern = {
                    "task_type": task_type,
                    "successful_approaches": [],
                    "common_challenges": [],
                    "learned_from": len(memories)
                }

                self.memory_store.add_memory(
                    memory_type="procedural",
                    content=pattern,
                    importance=2.0,
                    metadata={"consolidation_run": datetime.now().isoformat()}
                )

        # Clear working memory
        self.memory_store.clear_working_memory()

        logger.info("Memory consolidation complete")

    def reset(self, keep_memories: bool = True):
        """Reset the agent state"""
        self.conversation_history = []
        self.current_task = None
        self.task_context = {}

        if not keep_memories:
            self.memory_store = MemoryStore()
        else:
            # Only clear working memory
            self.memory_store.clear_working_memory()

        logger.info(f"Agent reset (memories kept: {keep_memories})")

    def get_performance_metrics(self) -> Dict[str, Any]:
        """Get agent performance metrics"""
        memory_stats = {
            memory_type: len(memories)
            for memory_type, memories in self.memory_store.memories.items()
        }

        total_memories = sum(memory_stats.values())
        cluster_count = len(self.memory_store.clusters)

        return {
            "total_memories": total_memories,
            "memory_distribution": memory_stats,
            "clusters_created": cluster_count,
            "conversation_length": len(self.conversation_history),
            "current_task": self.current_task
        }

config.py

"""
Configuration for Memobase Agent with Kimi K3 Model
"""

import os
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables
load_dotenv()


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
    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"
    return "openai/gpt-5.6-luna"


# Kimi K3 Model Configuration
KIMI_API_KEY = os.getenv("KIMI_API_KEY", "") or os.getenv("MOONSHOT_API_KEY", "")
KIMI_BASE_URL = "https://api.moonshot.cn/v1"
KIMI_MODEL = "kimi-k3"  # Kimi K3 model identifier

# Universal OpenRouter fallback: primary key (KIMI/MOONSHOT) absent but
# OPENROUTER_API_KEY present -> route the chat LLM through OpenRouter.
if not KIMI_API_KEY and os.getenv("OPENROUTER_API_KEY"):
    KIMI_API_KEY = os.getenv("OPENROUTER_API_KEY")
    KIMI_BASE_URL = "https://openrouter.ai/api/v1"
    KIMI_MODEL = _openrouter_model_id(KIMI_MODEL)

# Model Parameters
MODEL_TEMPERATURE = 0.7
MODEL_MAX_TOKENS = 4096
MODEL_TOP_P = 0.95

# Context Window Configuration
CONTEXT_WINDOW_SIZE = 128000  # Experiment context budget (K3 itself supports up to 1M tokens)
MAX_MEMORY_ENTRIES = 100
MEMORY_COMPRESSION_THRESHOLD = 50  # Compress when memory exceeds this

# Memobase Configuration
MEMOBASE_CONFIG = {
    "memory_types": [
        "episodic",      # Task-specific memories
        "semantic",      # General knowledge
        "procedural",    # Learned procedures and patterns
        "working"        # Short-term working memory
    ],
    "retention_policy": "adaptive",  # adaptive, fixed, or decay
    "compression_strategy": "hierarchical",  # hierarchical, summary, or selective
    "storage_backend": "local",  # local, redis, or postgresql
}

# LOCOMO Benchmark Configuration
LOCOMO_CONFIG = {
    "benchmark_path": Path("benchmarks/locomo"),
    "evaluation_metrics": [
        "task_completion",
        "reasoning_accuracy",
        "memory_utilization",
        "context_efficiency",
        "adaptation_score"
    ],
    "task_categories": [
        "multi_turn_reasoning",
        "long_context_qa",
        "task_planning",
        "knowledge_integration",
        "tool_usage"
    ],
    "max_turns": 20,
    "timeout_seconds": 300
}

# Memory Database Configuration
MEMORY_DB_PATH = Path("memory_store")
MEMORY_DB_PATH.mkdir(exist_ok=True)

# Logging Configuration
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
LOG_FILE = Path("logs") / "memobase_agent.log"
LOG_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"

# Agent Configuration
AGENT_CONFIG = {
    "name": "MemobaseAgent",
    "version": "1.0.0",
    "capabilities": [
        "long_term_memory",
        "context_compression",
        "adaptive_learning",
        "tool_calling",
        "multi_turn_reasoning"
    ],
    "max_retries": 3,
    "retry_delay": 1.0,
}

# Tool Configuration
ENABLE_WEB_SEARCH = True
ENABLE_CODE_EXECUTION = True
ENABLE_FILE_OPERATIONS = True
ENABLE_DATABASE_ACCESS = True

# Performance Optimization
BATCH_SIZE = 10
CACHE_ENABLED = True
CACHE_TTL = 3600  # 1 hour
PARALLEL_PROCESSING = True
MAX_WORKERS = 4

# Experimental Features
ENABLE_MEMORY_CONSOLIDATION = True  # Consolidate memories during idle time
ENABLE_PREDICTIVE_CACHING = True    # Pre-fetch likely needed memories
ENABLE_ADAPTIVE_COMPRESSION = True  # Adjust compression based on usage patterns

locomo_benchmark.py

"""
LOCOMO Benchmark Implementation for Memobase Agent
Evaluates agent performance on long-context and memory-intensive tasks
"""

import json
import time
import logging
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from collections import defaultdict
import random

from config import LOCOMO_CONFIG, LOG_LEVEL, LOG_FORMAT

# Configure logging
logging.basicConfig(level=LOG_LEVEL, format=LOG_FORMAT)
logger = logging.getLogger(__name__)


@dataclass
class BenchmarkTask:
    """Represents a single benchmark task"""
    id: str
    category: str  # multi_turn_reasoning, long_context_qa, etc.
    query: str
    context: Optional[str] = None
    expected_capabilities: List[str] = field(default_factory=list)
    ground_truth: Optional[Any] = None
    max_turns: int = 1
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class TaskResult:
    """Result of executing a benchmark task"""
    task_id: str
    response: str
    execution_time: float
    memory_usage: Dict[str, int]
    success: bool
    score: float
    turns_used: int
    errors: List[str] = field(default_factory=list)


class LOCOMOBenchmark:
    """
    LOCOMO (Long Context and Memory Optimization) Benchmark
    Evaluates agent capabilities in handling long contexts and memory management
    """

    def __init__(self, benchmark_path: Optional[Path] = None):
        """Initialize the benchmark suite"""
        self.benchmark_path = benchmark_path or LOCOMO_CONFIG["benchmark_path"]
        self.tasks: List[BenchmarkTask] = []
        self.results: List[TaskResult] = []
        self.metrics: Dict[str, Any] = defaultdict(list)

        # Load or generate benchmark tasks
        self._initialize_tasks()
        logger.info(f"Initialized LOCOMO benchmark with {len(self.tasks)} tasks")

    def _initialize_tasks(self):
        """Initialize benchmark tasks"""
        # Try to load from file
        if self.benchmark_path and (self.benchmark_path / "tasks.json").exists():
            self._load_tasks_from_file()
        else:
            # Generate default benchmark tasks
            self._generate_default_tasks()

    def _load_tasks_from_file(self):
        """Load tasks from JSON file"""
        try:
            with open(self.benchmark_path / "tasks.json", 'r') as f:
                tasks_data = json.load(f)
                for task_dict in tasks_data:
                    self.tasks.append(BenchmarkTask(**task_dict))
            logger.info(f"Loaded {len(self.tasks)} tasks from file")
        except Exception as e:
            logger.error(f"Failed to load tasks: {e}")
            self._generate_default_tasks()

    def _generate_default_tasks(self):
        """Generate default benchmark tasks for each category"""

        # Multi-turn reasoning tasks
        self.tasks.extend([
            BenchmarkTask(
                id="mtr_001",
                category="multi_turn_reasoning",
                query="Let's solve a complex problem step by step. First, calculate the compound interest on $10,000 at 5% annual rate for 3 years.",
                expected_capabilities=["mathematical_reasoning", "multi_step_problem_solving"],
                max_turns=5,
                ground_truth=11576.25
            ),
            BenchmarkTask(
                id="mtr_002",
                category="multi_turn_reasoning",
                query="Plan a detailed itinerary for a 7-day trip to Japan, considering budget constraints of $3000.",
                expected_capabilities=["planning", "constraint_satisfaction", "cultural_knowledge"],
                max_turns=8
            ),
            BenchmarkTask(
                id="mtr_003",
                category="multi_turn_reasoning",
                query="Debug this code issue: A recursive function is causing a stack overflow. Help me identify and fix it step by step.",
                expected_capabilities=["code_analysis", "debugging", "iterative_refinement"],
                max_turns=6
            )
        ])

        # Long context Q&A tasks
        self.tasks.extend([
            BenchmarkTask(
                id="lcqa_001",
                category="long_context_qa",
                query="Based on the provided research papers, summarize the key findings about climate change impacts on ocean acidification.",
                context=self._generate_long_context("climate_research", 50000),
                expected_capabilities=["information_extraction", "summarization", "scientific_reasoning"]
            ),
            BenchmarkTask(
                id="lcqa_002",
                category="long_context_qa",
                query="Analyze the financial statements and identify the top 3 risk factors for the company.",
                context=self._generate_long_context("financial_report", 30000),
                expected_capabilities=["financial_analysis", "risk_assessment", "data_interpretation"]
            ),
            BenchmarkTask(
                id="lcqa_003",
                category="long_context_qa",
                query="Review the legal documents and identify any potential conflicts or inconsistencies.",
                context=self._generate_long_context("legal_document", 40000),
                expected_capabilities=["legal_reasoning", "contradiction_detection", "document_analysis"]
            )
        ])

        # Task planning tasks
        self.tasks.extend([
            BenchmarkTask(
                id="tp_001",
                category="task_planning",
                query="Create a detailed project plan for developing a mobile app, including timeline, resources, and milestones.",
                expected_capabilities=["project_management", "resource_allocation", "timeline_planning"]
            ),
            BenchmarkTask(
                id="tp_002",
                category="task_planning",
                query="Design an optimal study plan for learning machine learning in 3 months with 2 hours daily.",
                expected_capabilities=["curriculum_design", "learning_optimization", "scheduling"]
            )
        ])

        # Knowledge integration tasks
        self.tasks.extend([
            BenchmarkTask(
                id="ki_001",
                category="knowledge_integration",
                query="Combine insights from psychology, neuroscience, and education to explain how humans learn languages.",
                expected_capabilities=["interdisciplinary_thinking", "knowledge_synthesis", "conceptual_integration"]
            ),
            BenchmarkTask(
                id="ki_002",
                category="knowledge_integration",
                query="Integrate historical events, economic theories, and sociological concepts to analyze the 2008 financial crisis.",
                expected_capabilities=["historical_analysis", "economic_reasoning", "systemic_thinking"]
            )
        ])

        # Tool usage tasks
        self.tasks.extend([
            BenchmarkTask(
                id="tu_001",
                category="tool_usage",
                query="Use available tools to gather real-time weather data and create a 5-day forecast analysis.",
                expected_capabilities=["tool_selection", "data_gathering", "predictive_analysis"]
            ),
            BenchmarkTask(
                id="tu_002",
                category="tool_usage",
                query="Research and compare the top 5 programming languages for web development using current data.",
                expected_capabilities=["web_search", "comparative_analysis", "technology_assessment"]
            )
        ])

        logger.info(f"Generated {len(self.tasks)} default benchmark tasks")

    def _generate_long_context(self, context_type: str, char_count: int) -> str:
        """Generate synthetic long context for testing"""
        templates = {
            "climate_research": [
                "Recent studies on ocean acidification show significant changes in pH levels. ",
                "The correlation between CO2 emissions and marine ecosystem degradation is evident. ",
                "Temperature variations in deep ocean currents affect global climate patterns. ",
                "Coral reef bleaching events have increased by 40% in the last decade. ",
                "Phytoplankton populations show remarkable adaptation to changing conditions. "
            ],
            "financial_report": [
                "Revenue increased by 15% year-over-year, driven by strong product demand. ",
                "Operating expenses rose due to increased R&D investments. ",
                "Market volatility poses risks to future earnings projections. ",
                "Cash flow remains strong with $2.3B in liquid assets. ",
                "Debt-to-equity ratio improved to 0.8 from previous 1.2. "
            ],
            "legal_document": [
                "The party of the first part agrees to the terms specified in Section 3.2. ",
                "Notwithstanding the above, exceptions may apply under force majeure. ",
                "Intellectual property rights remain with the original creator as per Article 7. ",
                "Dispute resolution shall follow binding arbitration procedures. ",
                "Confidentiality clauses extend for 5 years post-termination. "
            ]
        }

        sentences = templates.get(context_type, templates["climate_research"])
        context = ""

        while len(context) < char_count:
            context += random.choice(sentences)
            # Add some variation
            if random.random() > 0.7:
                context += f"In {random.randint(2020, 2024)}, researchers found that "

        return context[:char_count]

    def evaluate_response(self, task: BenchmarkTask, response: str,
                         execution_time: float, memory_usage: Dict[str, int]) -> TaskResult:
        """
        Evaluate agent response for a task

        Args:
            task: The benchmark task
            response: Agent's response
            execution_time: Time taken to generate response
            memory_usage: Memory statistics

        Returns:
            TaskResult with evaluation metrics
        """
        score = 0.0
        success = False
        errors = []

        # Basic response validation
        if not response or len(response) < 10:
            errors.append("Response too short or empty")
            score = 0.0
        else:
            # Category-specific evaluation
            if task.category == "multi_turn_reasoning":
                score = self._evaluate_reasoning(task, response)
            elif task.category == "long_context_qa":
                score = self._evaluate_context_qa(task, response)
            elif task.category == "task_planning":
                score = self._evaluate_planning(task, response)
            elif task.category == "knowledge_integration":
                score = self._evaluate_integration(task, response)
            elif task.category == "tool_usage":
                score = self._evaluate_tool_usage(task, response)
            else:
                score = self._evaluate_generic(task, response)

            success = score >= 0.6  # 60% threshold for success

        # Bonus for efficiency
        if execution_time < 5.0:
            score += 0.05
        if sum(memory_usage.values()) < 100:
            score += 0.05

        score = min(1.0, score)  # Cap at 1.0

        return TaskResult(
            task_id=task.id,
            response=response,
            execution_time=execution_time,
            memory_usage=memory_usage,
            success=success,
            score=score,
            turns_used=1,  # Will be updated for multi-turn tasks
            errors=errors
        )

    def _evaluate_reasoning(self, task: BenchmarkTask, response: str) -> float:
        """Evaluate multi-turn reasoning response"""
        score = 0.0

        # Check for step-by-step reasoning
        if any(marker in response.lower() for marker in ["step 1", "first", "then", "finally"]):
            score += 0.3

        # Check for mathematical accuracy (if applicable)
        if task.ground_truth and str(task.ground_truth) in response:
            score += 0.4

        # Check for logical flow
        if len(response.split('\n')) > 3:
            score += 0.2

        # Check for conclusion
        if any(word in response.lower() for word in ["therefore", "conclusion", "result"]):
            score += 0.1

        return score

    def _evaluate_context_qa(self, task: BenchmarkTask, response: str) -> float:
        """Evaluate long context Q&A response"""
        score = 0.0

        # Check for relevant content extraction
        keywords = ["finding", "summary", "key point", "important", "significant"]
        keyword_matches = sum(1 for kw in keywords if kw in response.lower())
        score += min(0.3, keyword_matches * 0.1)

        # Check for structured response
        if any(marker in response for marker in ["1.", "•", "-", "*"]):
            score += 0.2

        # Check response length (should be comprehensive but concise)
        optimal_length = 500
        length_ratio = min(len(response), optimal_length) / optimal_length
        score += length_ratio * 0.3

        # Check for citations or references to context
        if any(phrase in response.lower() for phrase in ["according to", "based on", "the document states"]):
            score += 0.2

        return score

    def _evaluate_planning(self, task: BenchmarkTask, response: str) -> float:
        """Evaluate task planning response"""
        score = 0.0

        # Check for timeline/schedule
        if any(word in response.lower() for word in ["timeline", "schedule", "deadline", "milestone"]):
            score += 0.25

        # Check for resource consideration
        if any(word in response.lower() for word in ["resource", "budget", "cost", "team"]):
            score += 0.25

        # Check for structured plan
        if any(marker in response for marker in ["phase", "stage", "step"]):
            score += 0.25

        # Check for risk/contingency consideration
        if any(word in response.lower() for word in ["risk", "contingency", "backup", "alternative"]):
            score += 0.25

        return score

    def _evaluate_integration(self, task: BenchmarkTask, response: str) -> float:
        """Evaluate knowledge integration response"""
        score = 0.0

        # Check for multiple domain references
        domains = ["psychology", "neuroscience", "education", "history", "economics", "sociology"]
        domain_count = sum(1 for domain in domains if domain in response.lower())
        score += min(0.4, domain_count * 0.2)

        # Check for synthesis language
        synthesis_words = ["integrate", "combine", "together", "relationship", "connection"]
        if any(word in response.lower() for word in synthesis_words):
            score += 0.3

        # Check for depth of analysis
        if len(response) > 300:
            score += 0.3

        return score

    def _evaluate_tool_usage(self, task: BenchmarkTask, response: str) -> float:
        """Evaluate tool usage response"""
        score = 0.0

        # Check for tool mentions
        if any(word in response.lower() for word in ["search", "query", "fetch", "retrieve", "analyze"]):
            score += 0.3

        # Check for data presentation
        if any(marker in response for marker in ["data", "result", "finding", "information"]):
            score += 0.3

        # Check for analysis/interpretation
        if any(word in response.lower() for word in ["analysis", "comparison", "trend", "pattern"]):
            score += 0.4

        return score

    def _evaluate_generic(self, task: BenchmarkTask, response: str) -> float:
        """Generic evaluation for unspecified task types"""
        score = 0.0

        # Response length
        if len(response) > 100:
            score += 0.3

        # Coherence (simple check)
        if response.count('.') > 2:
            score += 0.3

        # Relevance (keyword matching)
        query_words = task.query.lower().split()
        if query_words:
            matching_words = sum(1 for word in query_words if word in response.lower())
            score += min(0.4, matching_words / len(query_words))

        return score

    def run_benchmark(self, agent, tasks: Optional[List[BenchmarkTask]] = None,
                      verbose: bool = True) -> Dict[str, Any]:
        """
        Run benchmark evaluation on an agent

        Args:
            agent: The agent to evaluate (must have execute_task method)
            tasks: Optional list of tasks (uses all if None)
            verbose: Print progress information

        Returns:
            Benchmark results and metrics
        """
        tasks_to_run = tasks or self.tasks
        self.results = []

        if verbose:
            print(f"\n{'='*60}")
            print(f"Running LOCOMO Benchmark: {len(tasks_to_run)} tasks")
            print(f"{'='*60}\n")

        for i, task in enumerate(tasks_to_run, 1):
            if verbose:
                print(f"[{i}/{len(tasks_to_run)}] Running task {task.id} ({task.category})...")

            try:
                # Execute task
                start_time = time.time()
                result = agent.execute_task({
                    "id": task.id,
                    "type": task.category,
                    "query": task.query,
                    "context": task.context
                })
                execution_time = time.time() - start_time

                # Evaluate response
                task_result = self.evaluate_response(
                    task=task,
                    response=result.get("response", ""),
                    execution_time=execution_time,
                    memory_usage=result.get("memory_stats", {})
                )

                self.results.append(task_result)

                if verbose:
                    print(f"    ✓ Score: {task_result.score:.2f} | Time: {execution_time:.2f}s | Success: {task_result.success}")

            except Exception as e:
                logger.error(f"Error running task {task.id}: {e}")
                self.results.append(TaskResult(
                    task_id=task.id,
                    response="",
                    execution_time=0,
                    memory_usage={},
                    success=False,
                    score=0.0,
                    turns_used=0,
                    errors=[str(e)]
                ))

                if verbose:
                    print(f"    ✗ Error: {str(e)}")

        # Calculate metrics
        metrics = self._calculate_metrics()

        if verbose:
            self._print_summary(metrics)

        return metrics

    def _calculate_metrics(self) -> Dict[str, Any]:
        """Calculate benchmark metrics from results"""
        if not self.results:
            return {"error": "No results to calculate metrics"}

        # Overall metrics
        total_tasks = len(self.results)
        successful_tasks = sum(1 for r in self.results if r.success)
        avg_score = sum(r.score for r in self.results) / total_tasks
        avg_time = sum(r.execution_time for r in self.results) / total_tasks

        # Category-specific metrics
        category_scores = defaultdict(list)
        category_times = defaultdict(list)

        for task, result in zip(self.tasks, self.results):
            category_scores[task.category].append(result.score)
            category_times[task.category].append(result.execution_time)

        category_metrics = {}
        for category in category_scores:
            category_metrics[category] = {
                "avg_score": sum(category_scores[category]) / len(category_scores[category]),
                "avg_time": sum(category_times[category]) / len(category_times[category]),
                "task_count": len(category_scores[category])
            }

        # Memory efficiency
        total_memory = {}
        for result in self.results:
            for mem_type, count in result.memory_usage.items():
                total_memory[mem_type] = total_memory.get(mem_type, 0) + count

        return {
            "overall": {
                "total_tasks": total_tasks,
                "successful_tasks": successful_tasks,
                "success_rate": successful_tasks / total_tasks,
                "average_score": avg_score,
                "average_time": avg_time
            },
            "categories": category_metrics,
            "memory_usage": total_memory,
            "detailed_results": [
                {
                    "task_id": r.task_id,
                    "score": r.score,
                    "success": r.success,
                    "time": r.execution_time,
                    "errors": r.errors
                }
                for r in self.results
            ]
        }

    def _print_summary(self, metrics: Dict[str, Any]):
        """Print benchmark summary"""
        print(f"\n{'='*60}")
        print("BENCHMARK SUMMARY")
        print(f"{'='*60}")

        overall = metrics["overall"]
        print(f"\nOverall Performance:")
        print(f"  • Success Rate: {overall['success_rate']:.1%}")
        print(f"  • Average Score: {overall['average_score']:.2f}/1.00")
        print(f"  • Average Time: {overall['average_time']:.2f}s")
        print(f"  • Tasks Completed: {overall['successful_tasks']}/{overall['total_tasks']}")

        print(f"\nCategory Breakdown:")
        for category, cat_metrics in metrics["categories"].items():
            print(f"  {category}:")
            print(f"    - Score: {cat_metrics['avg_score']:.2f}")
            print(f"    - Time: {cat_metrics['avg_time']:.2f}s")
            print(f"    - Tasks: {cat_metrics['task_count']}")

        print(f"\nMemory Usage:")
        for mem_type, count in metrics["memory_usage"].items():
            print(f"  • {mem_type}: {count} entries")

        print(f"\n{'='*60}\n")

    def save_results(self, filepath: Path):
        """Save benchmark results to file"""
        results_data = {
            "timestamp": datetime.now().isoformat(),
            "tasks": [
                {
                    "id": task.id,
                    "category": task.category,
                    "query": task.query[:100]  # Truncate for storage
                }
                for task in self.tasks
            ],
            "results": [
                {
                    "task_id": r.task_id,
                    "score": r.score,
                    "success": r.success,
                    "execution_time": r.execution_time,
                    "memory_usage": r.memory_usage,
                    "turns_used": r.turns_used,
                    "errors": r.errors
                }
                for r in self.results
            ],
            "metrics": self._calculate_metrics()
        }

        with open(filepath, 'w') as f:
            json.dump(results_data, f, indent=2)

        logger.info(f"Saved benchmark results to {filepath}")

main.py

"""
Main entry point for Memobase Agent with LOCOMO Benchmark
"""

import argparse
import logging
import sys
from pathlib import Path
from datetime import datetime
import json
from typing import Optional

from agent import MemobaseAgent
from locomo_benchmark import LOCOMOBenchmark, BenchmarkTask
from config import (
    KIMI_API_KEY, LOG_LEVEL, LOG_FORMAT, 
    MEMORY_DB_PATH, LOCOMO_CONFIG
)

# Configure logging
logging.basicConfig(level=LOG_LEVEL, format=LOG_FORMAT)
logger = logging.getLogger(__name__)


def run_interactive_mode(agent: MemobaseAgent):
    """Run agent in interactive mode"""
    print("\n" + "="*60)
    print("MEMOBASE AGENT - Interactive Mode")
    print("="*60)
    print("\nCommands:")
    print("  /help     - Show this help message")
    print("  /memory   - Show memory statistics")
    print("  /clear    - Clear working memory")
    print("  /reset    - Reset agent (keep long-term memories)")
    print("  /learn    - Trigger memory consolidation")
    print("  /exit     - Exit interactive mode")
    print("\nType your message or command:\n")

    while True:
        try:
            user_input = input("\n👤 You: ").strip()

            if not user_input:
                continue

            # Handle commands
            if user_input.startswith('/'):
                command = user_input.lower()

                if command == '/help':
                    print("\nAvailable commands:")
                    print("  /memory   - Display memory statistics")
                    print("  /clear    - Clear working memory")
                    print("  /reset    - Reset conversation (keep memories)")
                    print("  /learn    - Consolidate and learn from memories")
                    print("  /exit     - Exit interactive mode")

                elif command == '/memory':
                    metrics = agent.get_performance_metrics()
                    print("\n📊 Memory Statistics:")
                    print(f"  Total memories: {metrics['total_memories']}")
                    for mem_type, count in metrics['memory_distribution'].items():
                        print(f"    • {mem_type}: {count}")
                    print(f"  Clusters created: {metrics['clusters_created']}")
                    print(f"  Conversation length: {metrics['conversation_length']}")

                elif command == '/clear':
                    agent.memory_store.clear_working_memory()
                    print("✅ Working memory cleared")

                elif command == '/reset':
                    agent.reset(keep_memories=True)
                    print("✅ Agent reset (memories preserved)")

                elif command == '/learn':
                    print("🧠 Consolidating memories...")
                    agent.consolidate_and_learn()
                    print("✅ Memory consolidation complete")

                elif command == '/exit':
                    print("\n👋 Goodbye!")
                    break

                else:
                    print(f"❌ Unknown command: {command}")
                    print("   Type /help for available commands")

            else:
                # Process regular message
                print("\n🤖 Assistant: ", end="", flush=True)
                response = agent.process_message(user_input)
                print(response)

        except KeyboardInterrupt:
            print("\n\n👋 Interrupted. Goodbye!")
            break
        except Exception as e:
            logger.error(f"Error in interactive mode: {e}")
            print(f"\n❌ Error: {str(e)}")


def run_benchmark_mode(agent: MemobaseAgent, 
                      category: Optional[str] = None,
                      num_tasks: Optional[int] = None):
    """Run LOCOMO benchmark evaluation"""
    print("\n" + "="*60)
    print("LOCOMO BENCHMARK EVALUATION")
    print("="*60)

    # Initialize benchmark
    benchmark = LOCOMOBenchmark()

    # Filter tasks if category specified
    tasks = benchmark.tasks
    if category:
        tasks = [t for t in tasks if t.category == category]
        print(f"\nRunning {category} tasks only")

    # Limit number of tasks if specified
    if num_tasks:
        tasks = tasks[:num_tasks]
        print(f"Limited to {num_tasks} tasks")

    # Run benchmark
    results = benchmark.run_benchmark(agent, tasks=tasks, verbose=True)

    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_dir = Path("benchmark_results")
    results_dir.mkdir(exist_ok=True)
    results_file = results_dir / f"locomo_results_{timestamp}.json"

    benchmark.save_results(results_file)
    print(f"\n📊 Results saved to: {results_file}")

    return results


def run_demo_mode(agent: MemobaseAgent):
    """Run demonstration scenarios"""
    print("\n" + "="*60)
    print("MEMOBASE AGENT - Demo Mode")
    print("="*60)

    demos = [
        {
            "name": "Memory Retention Demo",
            "messages": [
                "Remember that my favorite programming language is Python and I prefer functional programming style.",
                "What's my favorite programming language?",
                "I'm working on a new project. What programming style should I use?"
            ]
        },
        {
            "name": "Learning from Experience Demo",
            "messages": [
                "Help me debug this issue: my recursive function is causing a stack overflow.",
                "I have another recursive function problem. What should I check first?",
                "What are common causes of stack overflow in recursive functions?"
            ]
        },
        {
            "name": "Context Management Demo",
            "messages": [
                "Let's plan a machine learning project. We need data collection, preprocessing, model training, and deployment phases.",
                "What were the phases we discussed for the ML project?",
                "Add evaluation phase between training and deployment.",
                "Now summarize the complete project plan with all phases."
            ]
        }
    ]

    for demo in demos:
        print(f"\n\n{'='*40}")
        print(f"Demo: {demo['name']}")
        print(f"{'='*40}")

        for i, message in enumerate(demo['messages'], 1):
            print(f"\n[{i}/{len(demo['messages'])}] User: {message}")
            response = agent.process_message(message)
            print(f"Assistant: {response}")

            # Show memory stats after each interaction
            metrics = agent.get_performance_metrics()
            print(f"\n📊 Memory: {metrics['total_memories']} total, Distribution: {metrics['memory_distribution']}")

        # Reset for next demo but keep memories
        agent.reset(keep_memories=True)

    # Final memory consolidation
    print("\n\n🧠 Consolidating learnings from all demos...")
    agent.consolidate_and_learn()

    final_metrics = agent.get_performance_metrics()
    print(f"\n📊 Final Memory Statistics:")
    print(f"  Total memories: {final_metrics['total_memories']}")
    for mem_type, count in final_metrics['memory_distribution'].items():
        print(f"    • {mem_type}: {count}")
    print(f"  Clusters created: {final_metrics['clusters_created']}")


def run_single_task(agent: MemobaseAgent, task_query: str):
    """Run a single task"""
    print("\n" + "="*60)
    print("SINGLE TASK EXECUTION")
    print("="*60)

    print(f"\n📝 Task: {task_query}")

    result = agent.execute_task({
        "id": f"custom_{int(datetime.now().timestamp())}",
        "type": "custom",
        "query": task_query
    })

    print(f"\n🤖 Response: {result['response']}")
    print(f"\n📊 Execution Statistics:")
    print(f"  • Execution time: {result['execution_time']:.2f}s")
    print(f"  • Memories used: {result['memories_used']}")
    print(f"  • Memory distribution: {result['memory_stats']}")


def main():
    """Main entry point"""
    parser = argparse.ArgumentParser(
        description="Memobase Agent with LOCOMO Benchmark"
    )

    parser.add_argument(
        "--mode",
        choices=["interactive", "benchmark", "demo", "task"],
        default="interactive",
        help="Execution mode"
    )

    parser.add_argument(
        "--api-key",
        type=str,
        default=KIMI_API_KEY,
        help="Kimi API key"
    )

    parser.add_argument(
        "--category",
        type=str,
        choices=["multi_turn_reasoning", "long_context_qa", "task_planning", 
                "knowledge_integration", "tool_usage"],
        help="Benchmark category to run (benchmark mode only)"
    )

    parser.add_argument(
        "--num-tasks",
        type=int,
        help="Number of benchmark tasks to run"
    )

    parser.add_argument(
        "--task",
        type=str,
        help="Single task query to execute (task mode only)"
    )

    parser.add_argument(
        "--no-memory",
        action="store_true",
        help="Start with empty memory store"
    )

    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose logging"
    )

    args = parser.parse_args()

    # Set logging level
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Initialize agent
    print("\n🚀 Initializing Memobase Agent with Kimi K3...")
    agent = MemobaseAgent(api_key=args.api_key)

    if args.no_memory:
        agent.reset(keep_memories=False)
        print("   Started with empty memory store")

    # Run selected mode
    try:
        if args.mode == "interactive":
            run_interactive_mode(agent)

        elif args.mode == "benchmark":
            results = run_benchmark_mode(
                agent, 
                category=args.category,
                num_tasks=args.num_tasks
            )

        elif args.mode == "demo":
            run_demo_mode(agent)

        elif args.mode == "task":
            if not args.task:
                print("❌ Error: --task argument required for task mode")
                sys.exit(1)
            run_single_task(agent, args.task)

    except KeyboardInterrupt:
        print("\n\n👋 Interrupted. Goodbye!")
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        print(f"\n❌ Fatal error: {str(e)}")
        sys.exit(1)

    # Final cleanup
    print("\n💾 Saving final memory state...")
    agent.memory_store._save_memories()
    print("✅ Memory state saved")


if __name__ == "__main__":
    main()

profile_demo.py

"""
Memobase 用户画像（Profile）+ 事件记忆（Event Memory）演示

对应《深入理解 AI Agent》第 3 章"记忆框架案例"中对 Memobase 的介绍：
Memobase（开源项目 memodb-io/memobase）把用户记忆组织为两部分——
  * 用户画像（Profile）：按"主题—子主题"两级组织的稳定用户属性
    （如 basic_info→姓名、interest→游戏偏好、work→职位），从对话中提取；
  * 事件记忆（Event Memory）：按时间线记录用户经历的事件，
    用于回答"我们上次讨论预算是什么时候"这类与时间有关的问题。
工程上 Memobase 采用"缓冲批处理"：对话先 insert 进缓冲区累积，
flush 时才统一触发一次 LLM 记忆提取，查询侧（profile/event/context）
只读取已整理好的结果，保证低延迟。

本脚本用真实的 memobase Python SDK 演示这条链路：
    insert（写入对话缓冲）→ flush（触发提取）→ profile / event / context（召回）

运行前提：Memobase 需要一个正在运行的服务端 + 用于提取记忆的 LLM。
  * 自托管：见 https://github.com/memodb-io/memobase （docker compose 启动，
    默认服务地址 http://localhost:8019，默认 token 为 secret）；
  * 云服务：在 https://www.memobase.io 申请 project_url 与 api_key。
无服务端时可用 --dry-run 查看示例对话与将要执行的操作（不联网、不伪造结果）。
"""

import argparse
import json
import os
import sys
from pathlib import Path

# 示例多轮对话：内容刻意覆盖三个画像主题，便于观察 Profile 的两级结构
SAMPLE_CONVERSATION = [
    {"role": "user", "content": "你好，我叫李明，今年 28 岁，住在上海。"},
    {"role": "assistant", "content": "你好李明，很高兴认识你！有什么可以帮你的吗？"},
    {"role": "user", "content": "我在一家游戏公司做后端工程师，主要用 Go 语言。"},
    {"role": "assistant", "content": "后端工程师很酷，Go 在高并发服务里很常用。"},
    {"role": "user", "content": "平时最喜欢玩《塞尔达传说》和《艾尔登法环》这类开放世界游戏。"},
    {"role": "assistant", "content": "开放世界游戏的探索感确实很棒。"},
    {"role": "user", "content": "对了，上周我们把新项目的预算定在了 50 万。"},
    {"role": "assistant", "content": "好的，我记住了新项目预算是 50 万。"},
]

# 期望被提取出的画像主题（仅用于 --dry-run 的结构说明，非真实结果）
EXPECTED_TOPICS = {
    "basic_info": ["姓名", "年龄", "城市"],
    "work": ["职位", "公司", "技术栈"],
    "interest": ["游戏偏好"],
}

DEFAULT_PROJECT_URL = os.getenv("MEMOBASE_PROJECT_URL", "http://localhost:8019")
DEFAULT_API_KEY = os.getenv("MEMOBASE_API_KEY", "secret")


def load_conversation(input_path: str | None) -> list[dict]:
    """从 --input 指定的 JSON 文件读取对话，缺省则用内置示例对话。

    文件格式：[{"role": "user"|"assistant", "content": "..."}, ...]
    """
    if not input_path:
        return SAMPLE_CONVERSATION
    data = json.loads(Path(input_path).read_text(encoding="utf-8"))
    if not isinstance(data, list) or not all(
        isinstance(m, dict) and "role" in m and "content" in m for m in data
    ):
        raise ValueError("对话文件应为 [{'role':..., 'content':...}, ...] 形式的 JSON")
    return data


def print_conversation(conversation: list[dict]) -> None:
    print("\n💬 对话内容：")
    for m in conversation:
        who = "👤 用户" if m["role"] == "user" else "🤖 助手"
        print(f"  {who}: {m['content']}")


def build_client(args):
    """构造并连通 Memobase 客户端（联网）。失败时给出可操作的提示。"""
    try:
        from memobase import MemoBaseClient
    except ImportError:
        print("❌ 未安装 memobase SDK，请先运行：pip install -r requirements.txt")
        sys.exit(1)

    client = MemoBaseClient(project_url=args.project_url, api_key=args.api_key)
    try:
        reachable = client.ping()
    except Exception as exc:  # 连接被拒绝、超时、DNS 失败等
        reachable = False
        print(f"❌ 连接 Memobase 服务端时出错：{exc}")
    if not reachable:
        print("❌ 无法连接 Memobase 服务端：", args.project_url)
        print("   请确认服务已启动（自托管见 memodb-io/memobase 的 docker compose），")
        print("   或用 --project-url / --api-key 指定云服务地址与密钥。")
        print("   仅想查看示例与流程（不联网）可加 --dry-run。")
        sys.exit(1)
    return client


def render_profiles(profiles) -> list[dict]:
    """把 SDK 返回的 UserProfile 列表整理成 主题→子主题→内容 的结构。"""
    rows = []
    for p in profiles:
        rows.append({"topic": p.topic, "sub_topic": p.sub_topic, "content": p.content})
    return rows


def render_events(events) -> list[dict]:
    """把 SDK 返回的 UserEventData 列表整理成可读的时间线。"""
    rows = []
    for e in events:
        tip = None
        tags = None
        if e.event_data is not None:
            tip = e.event_data.event_tip
            if e.event_data.event_tags:
                tags = {t.tag: t.value for t in e.event_data.event_tags}
        rows.append(
            {
                "created_at": str(e.created_at),
                "event_tip": tip,
                "event_tags": tags,
            }
        )
    return rows


def op_insert(user, conversation) -> None:
    from memobase import ChatBlob

    print_conversation(conversation)
    bid = user.insert(ChatBlob(messages=conversation))
    print(f"\n✅ 已写入对话缓冲区（blob id: {bid}）")
    print("   提示：写入只进缓冲区，尚未提取记忆；执行 flush 才会触发一次提取。")


def op_flush(user) -> None:
    print("\n🧠 正在 flush 缓冲区，触发记忆提取（由服务端 LLM 完成，可能需要数秒）...")
    ok = user.flush(sync=True)
    print("✅ 记忆提取完成" if ok else "⚠️ flush 返回 False，请检查服务端日志")


def op_profile(user) -> list[dict]:
    profiles = render_profiles(user.profile())
    print("\n📇 用户画像（Profile，主题 → 子主题 → 内容）：")
    if not profiles:
        print("  （暂无画像，请先 insert 对话并 flush 提取）")
    for r in profiles:
        print(f"  • [{r['topic']}] {r['sub_topic']}: {r['content']}")
    return profiles


def op_event(user) -> list[dict]:
    events = render_events(user.event())
    print("\n🗓️  事件记忆（Event Memory，按时间线）：")
    if not events:
        print("  （暂无事件）")
    for r in events:
        print(f"  • {r['created_at']}  {r['event_tip'] or ''}")
        if r["event_tags"]:
            print(f"      标签: {r['event_tags']}")
    return events


def op_context(user) -> str:
    ctx = user.context()
    print("\n🧩 组装好的记忆上下文（context，可直接拼进 LLM 提示词）：")
    print(ctx if ctx else "  （空）")
    return ctx


def run_demo(user, conversation, output_path):
    """端到端演示：从对话构建画像，再召回画像 / 事件 / 上下文。"""
    print("=" * 64)
    print("Memobase 用户画像 + 事件记忆 端到端演示")
    print("=" * 64)
    op_insert(user, conversation)
    op_flush(user)
    profiles = op_profile(user)
    events = op_event(user)
    ctx = op_context(user)
    result = {"profiles": profiles, "events": events, "context": ctx}
    if output_path:
        Path(output_path).write_text(
            json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        print(f"\n💾 结果已写入：{output_path}")
    return result


def run_dry_run(conversation):
    """离线路径：不联网，展示示例对话、缓冲批处理流程与期望画像结构。"""
    print("=" * 64)
    print("Memobase 演示（--dry-run 离线预览，不连接服务端、不伪造结果）")
    print("=" * 64)
    print_conversation(conversation)
    print("\n🔀 将要执行的操作（缓冲批处理流水线）：")
    print("  1) insert  —— 把上述对话写入用户缓冲区（不触发提取）")
    print("  2) flush   —— 统一触发一次 LLM 记忆提取（摊薄调用成本）")
    print("  3) profile —— 读取提取出的结构化用户画像")
    print("  4) event   —— 读取按时间线组织的事件记忆")
    print("  5) context —— 读取组装好的记忆上下文")
    print("\n🧬 期望被提取的画像结构（主题 → 子主题，示意，非真实结果）：")
    for topic, subs in EXPECTED_TOPICS.items():
        print(f"  • {topic}: {', '.join(subs)}")
    print("\n▶️  接入真实服务后，去掉 --dry-run 即可执行上述完整链路。")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Memobase 用户画像（Profile）+ 事件记忆（Event Memory）演示",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "示例：\n"
            "  python profile_demo.py                      # 端到端演示（内置示例对话）\n"
            "  python profile_demo.py --dry-run            # 离线预览流程，不连接服务端\n"
            "  python profile_demo.py --op profile         # 只召回已提取的用户画像\n"
            "  python profile_demo.py --op event           # 只查看事件记忆时间线\n"
            "  python profile_demo.py --input chat.json --output result.json\n"
        ),
    )
    parser.add_argument(
        "--op",
        choices=["demo", "insert", "flush", "profile", "event", "context", "reset"],
        default="demo",
        help="记忆操作：demo=端到端演示(默认)；insert=写入对话缓冲；"
        "flush=触发提取；profile=召回画像；event=事件记忆；"
        "context=记忆上下文；reset=删除该用户重置",
    )
    parser.add_argument(
        "--input",
        type=str,
        default=None,
        help="对话输入文件（JSON：[{'role','content'}, ...]），缺省用内置示例对话",
    )
    parser.add_argument(
        "--user-id",
        type=str,
        default="demo_user_liming",
        help="用户 ID（默认 demo_user_liming）",
    )
    parser.add_argument(
        "--project-url",
        type=str,
        default=DEFAULT_PROJECT_URL,
        help=f"Memobase 服务地址（默认 {DEFAULT_PROJECT_URL}，可用环境变量 MEMOBASE_PROJECT_URL）",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=DEFAULT_API_KEY,
        help="Memobase 访问密钥（默认取环境变量 MEMOBASE_API_KEY，自托管默认 secret）",
    )
    parser.add_argument(
        "--model",
        type=str,
        default=None,
        help="记忆提取所用 LLM（仅作说明：Memobase 由服务端配置提取模型，"
        "客户端不直接指定；如需更换请改服务端 .env）",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="将画像/事件/上下文结果写入指定 JSON 文件",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="离线预览：仅展示示例对话与操作流程，不连接服务端、不伪造结果",
    )
    return parser


def main():
    args = build_parser().parse_args()
    try:
        conversation = load_conversation(args.input)
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        print(f"❌ 读取对话文件失败（{args.input}）：{exc}")
        sys.exit(1)

    if args.dry_run:
        run_dry_run(conversation)
        return

    if args.model:
        print(f"ℹ️  --model={args.model}：Memobase 的提取模型由服务端配置，此处仅作记录。")

    client = build_client(args)
    user = client.get_or_create_user(args.user_id)

    if args.op == "demo":
        run_demo(user, conversation, args.output)
        return

    result = None
    if args.op == "insert":
        op_insert(user, conversation)
    elif args.op == "flush":
        op_flush(user)
    elif args.op == "profile":
        result = {"profiles": op_profile(user)}
    elif args.op == "event":
        result = {"events": op_event(user)}
    elif args.op == "context":
        result = {"context": op_context(user)}
    elif args.op == "reset":
        client.delete_user(args.user_id)
        print(f"🧹 已删除用户 {args.user_id}，画像与事件已清空。")

    if result and args.output:
        Path(args.output).write_text(
            json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        print(f"\n💾 结果已写入：{args.output}")


if __name__ == "__main__":
    main()

quickstart.py

"""
Quick start script for testing Memobase Agent
"""

import os
import sys
from pathlib import Path

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from agent import MemobaseAgent
from locomo_benchmark import LOCOMOBenchmark, BenchmarkTask


def quick_demo():
    """Run a quick demonstration of the agent's capabilities"""

    print("=" * 60)
    print("MEMOBASE AGENT - Quick Start Demo")
    print("=" * 60)

    # Check for API key
    api_key = os.getenv("KIMI_API_KEY", "")
    if not api_key or api_key == "sk-your-kimi-api-key-here":
        print("\n⚠️  Warning: KIMI_API_KEY not set properly!")
        print("Please set your API key in .env file or as environment variable")
        print("\nContinuing with demo setup...")
        api_key = "demo-key"  # Use demo key for structure demonstration

    try:
        # Initialize agent
        print("\n🚀 Initializing Memobase Agent...")
        agent = MemobaseAgent(api_key=api_key)
        print("✅ Agent initialized successfully!")

        # Show initial memory state
        metrics = agent.get_performance_metrics()
        print(f"\n📊 Initial Memory State:")
        print(f"  • Total memories: {metrics['total_memories']}")
        print(f"  • Memory types: {list(metrics['memory_distribution'].keys())}")

        # Demo 1: Simple interaction with memory
        print("\n" + "-" * 40)
        print("Demo 1: Memory Storage and Retrieval")
        print("-" * 40)

        test_messages = [
            "Remember that I prefer Python for data science and JavaScript for web development.",
            "What programming language should I use for data analysis?",
            "What about for building a web application?"
        ]

        for msg in test_messages:
            print(f"\n👤 User: {msg}")
            if api_key != "demo-key":
                response = agent.process_message(msg)
                print(f"🤖 Agent: {response}")
            else:
                print("🤖 Agent: [Demo mode - API key required for actual response]")

            # Show memory growth
            metrics = agent.get_performance_metrics()
            print(f"   📊 Memories: {metrics['total_memories']} total")

        # Demo 2: Learning from experience
        print("\n" + "-" * 40)
        print("Demo 2: Learning from Experience")
        print("-" * 40)

        # Simulate learning
        agent._learn_from_outcome(
            task="debugging",
            approach="Check for infinite loops in recursive functions",
            outcome="Successfully identified stack overflow cause",
            success=True
        )
        print("✅ Learned from debugging experience")

        agent._learn_from_outcome(
            task="optimization",
            approach="Use memoization for recursive algorithms",
            outcome="Reduced computation time by 80%",
            success=True
        )
        print("✅ Learned optimization technique")

        # Show procedural memories
        procedural_memories = agent.memory_store.get_memories("procedural", limit=5)
        print(f"\n📚 Procedural Knowledge Acquired: {len(procedural_memories)} patterns")

        # Demo 3: Memory consolidation
        print("\n" + "-" * 40)
        print("Demo 3: Memory Consolidation")
        print("-" * 40)

        print("🧠 Triggering memory consolidation...")
        agent.consolidate_and_learn()
        print("✅ Consolidation complete")

        # Final memory statistics
        final_metrics = agent.get_performance_metrics()
        print(f"\n📊 Final Memory Statistics:")
        print(f"  • Total memories: {final_metrics['total_memories']}")
        for mem_type, count in final_metrics['memory_distribution'].items():
            print(f"    - {mem_type}: {count}")
        print(f"  • Memory clusters: {final_metrics['clusters_created']}")

        # Demo 4: Mini benchmark
        print("\n" + "-" * 40)
        print("Demo 4: Mini Benchmark Test")
        print("-" * 40)

        # Create a simple benchmark task
        task = BenchmarkTask(
            id="demo_001",
            category="multi_turn_reasoning",
            query="What are the key factors to consider when choosing a database for a web application?",
            expected_capabilities=["technical_knowledge", "comparative_analysis"]
        )

        print(f"📝 Task: {task.query}")

        if api_key != "demo-key":
            # Initialize mini benchmark
            benchmark = LOCOMOBenchmark()
            benchmark.tasks = [task]

            # Run benchmark
            print("🏃 Running benchmark...")
            results = benchmark.run_benchmark(agent, tasks=[task], verbose=False)

            # Show results
            print(f"\n📊 Benchmark Results:")
            print(f"  • Score: {results['overall']['average_score']:.2f}/1.00")
            print(f"  • Time: {results['overall']['average_time']:.2f}s")
            print(f"  • Success: {results['overall']['success_rate']:.0%}")
        else:
            print("📊 [Demo mode - API key required for benchmark execution]")

        print("\n" + "=" * 60)
        print("✅ Quick Start Demo Complete!")
        print("=" * 60)
        print("\nNext steps:")
        print("1. Set your KIMI_API_KEY in .env file")
        print("2. Run 'python main.py --mode interactive' for full interaction")
        print("3. Run 'python main.py --mode benchmark' for complete evaluation")
        print("4. Check README.md for detailed documentation")

    except Exception as e:
        print(f"\n❌ Error during demo: {str(e)}")
        print("\nTroubleshooting:")
        print("1. Ensure all dependencies are installed: pip install -r requirements.txt")
        print("2. Check that KIMI_API_KEY is properly set")
        print("3. Verify network connectivity for API calls")


if __name__ == "__main__":
    quick_demo()

test_evaluate_generic_empty_query.py

"""Regression tests: benchmark evaluation must not raise ZeroDivisionError on
tasks with an empty/whitespace query (e.g. loaded from tasks.json) that fall
through to the generic evaluator."""
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from locomo_benchmark import BenchmarkTask, LOCOMOBenchmark

LONG_RESPONSE = (
    "This is a reasonably long agent response. It has several sentences. "
    "And it keeps going with more detail. Even more detail here."
)


def _suite():
    return LOCOMOBenchmark.__new__(LOCOMOBenchmark)


def test_evaluate_generic_empty_query():
    task = BenchmarkTask(id="t1", category="custom_category", query="")
    score = _suite()._evaluate_generic(task, LONG_RESPONSE)
    assert 0.0 <= score <= 1.0


def test_evaluate_response_whitespace_query_unknown_category():
    task = BenchmarkTask(id="t2", category="custom_category", query="   ")
    result = _suite().evaluate_response(
        task, response=LONG_RESPONSE, execution_time=1.0, memory_usage={}
    )
    assert 0.0 <= result.score <= 1.0


def test_evaluate_generic_normal_query_unchanged():
    task = BenchmarkTask(id="t3", category="custom_category", query="weather forecast analysis")
    score = _suite()._evaluate_generic(task, LONG_RESPONSE.replace("agent", "weather"))
    assert score > 0.0

test_memory.py

"""
Test script for Memobase memory functionality
"""

import os
import sys
import time
from pathlib import Path
from datetime import datetime

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from agent import Memory, MemoryStore, MemoryCluster, MemobaseAgent
from config import MEMORY_DB_PATH


def test_memory_basics():
    """Test basic memory operations"""
    print("\n" + "=" * 50)
    print("Testing Basic Memory Operations")
    print("=" * 50)

    # Create a memory
    memory = Memory(
        id="",
        type="episodic",
        content={"event": "User asked about Python", "response": "Python is great for data science"},
        importance_score=2.0
    )

    print(f"✅ Created memory with ID: {memory.id}")
    print(f"   Type: {memory.type}")
    print(f"   Importance: {memory.importance_score}")

    # Test access
    initial_count = memory.access_count
    memory.access()
    print(f"✅ Memory accessed: count {initial_count} -> {memory.access_count}")

    # Test decay
    initial_importance = memory.importance_score
    time.sleep(0.1)  # Small delay
    memory.decay()
    print(f"✅ Memory decayed: importance {initial_importance:.2f} -> {memory.importance_score:.2f}")

    return True


def test_memory_store():
    """Test memory store operations"""
    print("\n" + "=" * 50)
    print("Testing Memory Store")
    print("=" * 50)

    # Create temporary store
    test_path = Path("test_memory_store")
    test_path.mkdir(exist_ok=True)
    store = MemoryStore(db_path=test_path)

    # Add memories of different types
    memories_added = []

    # Episodic memory
    mem1 = store.add_memory(
        memory_type="episodic",
        content="First interaction with user about machine learning",
        metadata={"session": "test_001"},
        importance=2.0
    )
    memories_added.append(mem1)
    print(f"✅ Added episodic memory: {mem1.id}")

    # Semantic memory
    mem2 = store.add_memory(
        memory_type="semantic",
        content="Python is a programming language used for data science",
        metadata={"domain": "programming"},
        importance=1.5
    )
    memories_added.append(mem2)
    print(f"✅ Added semantic memory: {mem2.id}")

    # Procedural memory
    mem3 = store.add_memory(
        memory_type="procedural",
        content={"pattern": "debugging", "approach": "check logs first"},
        metadata={"learned_from": "experience"},
        importance=3.0
    )
    memories_added.append(mem3)
    print(f"✅ Added procedural memory: {mem3.id}")

    # Working memory
    mem4 = store.add_memory(
        memory_type="working",
        content="Current task: testing memory system",
        importance=1.0
    )
    memories_added.append(mem4)
    print(f"✅ Added working memory: {mem4.id}")

    # Test retrieval
    print("\n📚 Testing Memory Retrieval:")

    # Get episodic memories
    episodic = store.get_memories("episodic", limit=5)
    print(f"   Episodic memories retrieved: {len(episodic)}")

    # Search memories
    search_results = store.search_memories("Python", limit=3)
    print(f"   Search for 'Python': {len(search_results)} results")

    # Test persistence
    store._save_memories()
    print("✅ Memories saved to disk")

    # Create new store and load
    store2 = MemoryStore(db_path=test_path)
    total_memories = sum(len(mems) for mems in store2.memories.values())
    print(f"✅ Memories loaded from disk: {total_memories} total")

    # Test consolidation
    store.consolidate_memories()
    print("✅ Memory consolidation completed")

    # Clear working memory
    store.clear_working_memory()
    working_after = len(store.memories.get('working', []))
    print(f"✅ Working memory cleared: {working_after} remaining")

    # Cleanup
    import shutil
    shutil.rmtree(test_path)
    print("✅ Test data cleaned up")

    return True


def test_memory_compression():
    """Test memory compression functionality"""
    print("\n" + "=" * 50)
    print("Testing Memory Compression")
    print("=" * 50)

    # Create store with low threshold for testing
    test_path = Path("test_compression")
    test_path.mkdir(exist_ok=True)
    store = MemoryStore(db_path=test_path)

    # Add many memories to trigger compression
    print("Adding memories to trigger compression...")
    for i in range(10):
        store.add_memory(
            memory_type="episodic",
            content=f"Memory {i}: Event occurred at time {i}",
            importance=1.0 if i < 5 else 2.0
        )

    # Force compression by adding more with low threshold
    original_threshold = 50  # From config
    test_threshold = 5

    # Manually trigger compression
    if len(store.memories['episodic']) > test_threshold:
        print(f"   Memories before compression: {len(store.memories['episodic'])}")
        store._compress_memories('episodic')
        print(f"   Memories after compression: {len(store.memories['episodic'])}")
        print(f"   Clusters created: {len(store.clusters)}")

    print("✅ Compression test completed")

    # Cleanup
    import shutil
    shutil.rmtree(test_path)

    return True


def test_agent_memory_integration():
    """Test agent integration with memory system"""
    print("\n" + "=" * 50)
    print("Testing Agent Memory Integration")
    print("=" * 50)

    # Check for API key
    api_key = os.getenv("KIMI_API_KEY", "test-key")

    # Initialize agent
    agent = MemobaseAgent(api_key=api_key)
    print("✅ Agent initialized with memory system")

    # Test memory operations through agent
    print("\n📊 Initial state:")
    metrics = agent.get_performance_metrics()
    print(f"   Total memories: {metrics['total_memories']}")
    print(f"   Distribution: {metrics['memory_distribution']}")

    # Simulate learning
    agent._learn_from_outcome(
        task="test_task",
        approach="test_approach",
        outcome="successful",
        success=True
    )
    print("✅ Agent learned from outcome")

    # Check procedural memory
    procedural = agent.memory_store.get_memories("procedural", limit=1)
    if procedural:
        print(f"✅ Procedural memory created: {procedural[0].id}")

    # Test memory retrieval in context
    relevant = agent._retrieve_relevant_memories("test query", limit=3)
    print(f"✅ Retrieved {len(relevant)} relevant memories")

    # Test consolidation
    agent.consolidate_and_learn()
    print("✅ Agent consolidation completed")

    # Test reset with memory preservation
    agent.reset(keep_memories=True)
    metrics_after = agent.get_performance_metrics()
    print(f"✅ Agent reset (memories preserved): {metrics_after['total_memories']} memories retained")

    # Test reset without memory preservation
    agent.reset(keep_memories=False)
    metrics_final = agent.get_performance_metrics()
    print(f"✅ Agent reset (memories cleared): {metrics_final['total_memories']} memories remaining")

    return True


def run_all_tests():
    """Run all memory tests"""
    print("\n" + "=" * 60)
    print("MEMOBASE MEMORY SYSTEM TEST SUITE")
    print("=" * 60)

    tests = [
        ("Basic Memory Operations", test_memory_basics),
        ("Memory Store", test_memory_store),
        ("Memory Compression", test_memory_compression),
        ("Agent Integration", test_agent_memory_integration)
    ]

    results = []
    for test_name, test_func in tests:
        try:
            success = test_func()
            results.append((test_name, success))
        except Exception as e:
            print(f"\n❌ Test '{test_name}' failed: {str(e)}")
            results.append((test_name, False))

    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)

    passed = sum(1 for _, success in results if success)
    total = len(results)

    for test_name, success in results:
        status = "✅ PASSED" if success else "❌ FAILED"
        print(f"{status}: {test_name}")

    print(f"\nTotal: {passed}/{total} tests passed")

    if passed == total:
        print("\n🎉 All tests passed successfully!")
    else:
        print(f"\n⚠️  {total - passed} test(s) failed")

    return passed == total


if __name__ == "__main__":
    success = run_all_tests()
    sys.exit(0 if success else 1)