liqiang
b119135836
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
581 lines
22 KiB
Python
581 lines
22 KiB
Python
"""
|
|
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
|
|
}
|