ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/mem0/config.py

@@ -0,0 +1,218 @@
+"""Configuration module for Mem0 agent with Kimi K3 integration."""
+
+import os
+from pathlib import Path
+from typing import Optional, Dict, Any
+from dataclasses import dataclass, field
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+
+def _reasoning_safe_temperature(model, requested=1.0):
+    """Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
+    Return 1 for those; otherwise the requested value so non-reasoning
+    providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
+    m = str(model or "").lower().replace("/", "-")
+    return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
+
+
+def _openrouter_model_id(model) -> str:
+    """Map a provider-native model name to an OpenRouter model id, used by the
+    universal OpenRouter fallback. An explicit OPENROUTER_MODEL env var wins."""
+    override = os.getenv("OPENROUTER_MODEL")
+    if override:
+        return override
+    m = (model or "").strip()
+    if not m:
+        return "openai/gpt-5.6-luna"
+    if "/" in m:
+        return m  # already an OpenRouter-style id (e.g. openai/gpt-5.6-luna)
+    ml = m.lower()
+    if ml.startswith(("gpt-", "o1", "o3", "o4", "chatgpt")):
+        return "openai/" + m
+    if ml.startswith("claude-"):
+        return "anthropic/claude-opus-4.8"
+    if ml.startswith("kimi"):
+        # kimi-k3 is not on OpenRouter; moonshotai/kimi-k2.6 is the closest hosted id.
+        return "moonshotai/kimi-k2.6"
+    # Provider-native ids (kimi-*/doubao-*/qwen/deepseek-*) not hosted on
+    # OpenRouter under the same name -> a widely-available OpenAI chat model.
+    return "openai/gpt-5.6-luna"
+
+
+@dataclass
+class KimiConfig:
+    """Configuration for Kimi K3 model."""
+
+    api_key: str = field(default_factory=lambda: os.getenv("KIMI_API_KEY", ""))
+    model_name: str = field(default_factory=lambda: os.getenv("MODEL_NAME", "kimi-k3"))
+    max_tokens: int = field(default_factory=lambda: int(os.getenv("MAX_TOKENS", "128000")))
+    temperature: float = field(default_factory=lambda: float(os.getenv("TEMPERATURE", "0.7")))
+    api_base: str = field(default_factory=lambda: os.getenv("KIMI_API_BASE", "https://api.moonshot.cn/v1"))
+
+    def __post_init__(self):
+        """Universal OpenRouter fallback for the chat LLM: when KIMI_API_KEY is
+        absent but OPENROUTER_API_KEY is present, route the chat model (used by
+        KimiK3Client and threaded into mem0's own LLM config) through OpenRouter.
+        NB: mem0's embedder still uses OpenAI embeddings (OpenRouter has no
+        embeddings endpoint), so OPENAI_API_KEY remains needed for memory add."""
+        if not self.api_key and os.getenv("OPENROUTER_API_KEY"):
+            self.api_key = os.getenv("OPENROUTER_API_KEY")
+            self.api_base = "https://openrouter.ai/api/v1"
+            self.model_name = _openrouter_model_id(self.model_name)
+
+    def validate(self) -> bool:
+        """Validate Kimi configuration."""
+        if not self.api_key:
+            raise ValueError("KIMI_API_KEY is required (or set OPENROUTER_API_KEY for the fallback)")
+        if self.max_tokens <= 0 or self.max_tokens > 128000:
+            raise ValueError("MAX_TOKENS must be between 1 and 128000")
+        if self.temperature < 0 or self.temperature > 2:
+            raise ValueError("TEMPERATURE must be between 0 and 2")
+        return True
+
+
+@dataclass
+class Mem0Config:
+    """Configuration for Mem0 memory system."""
+    
+    api_key: Optional[str] = field(default_factory=lambda: os.getenv("MEM0_API_KEY"))
+    backend: str = field(default_factory=lambda: os.getenv("MEMORY_BACKEND", "local"))
+    collection_name: str = field(default_factory=lambda: os.getenv("MEMORY_COLLECTION", "locomo_benchmark"))
+    embedding_model: str = field(default_factory=lambda: os.getenv("MEMORY_EMBEDDING_MODEL", "text-embedding-3-small"))
+    vector_store_config: Dict[str, Any] = field(default_factory=dict)
+    
+    def __post_init__(self):
+        """Initialize vector store configuration based on backend."""
+        if self.backend == "local":
+            # NB: mem0 >=1.0 validates the chroma config against a fixed field
+            # set (collection_name/path/host/port/api_key/tenant/client). The
+            # embedding model belongs to the top-level "embedder" block (set in
+            # agent.py), NOT here — passing embedding_function raises a
+            # MemoryConfig validation error.
+            self.vector_store_config = {
+                "provider": "chroma",
+                "config": {
+                    "collection_name": self.collection_name,
+                    "path": "./data/chroma_db",
+                }
+            }
+        elif self.backend == "cloud":
+            if not self.api_key:
+                raise ValueError("MEM0_API_KEY is required for cloud backend")
+            self.vector_store_config = {
+                "provider": "mem0_cloud",
+                "config": {
+                    "api_key": self.api_key,
+                    "collection_name": self.collection_name
+                }
+            }
+        else:
+            raise ValueError(f"Invalid backend: {self.backend}. Must be 'local' or 'cloud'")
+    
+    def validate(self) -> bool:
+        """Validate Mem0 configuration."""
+        if self.backend not in ["local", "cloud"]:
+            raise ValueError("MEMORY_BACKEND must be 'local' or 'cloud'")
+        if self.backend == "cloud" and not self.api_key:
+            raise ValueError("MEM0_API_KEY is required for cloud backend")
+        return True
+
+
+@dataclass
+class LOCOMOConfig:
+    """Configuration for LOCOMO benchmark."""
+    
+    data_path: Path = field(default_factory=lambda: Path(os.getenv("BENCHMARK_DATA_PATH", "./data/locomo")))
+    max_sessions: int = field(default_factory=lambda: int(os.getenv("MAX_SESSIONS", "100")))
+    max_agents: int = field(default_factory=lambda: int(os.getenv("MAX_AGENTS", "10")))
+    context_window_size: int = field(default_factory=lambda: int(os.getenv("CONTEXT_WINDOW_SIZE", "128000")))
+    evaluation_metrics: list = field(default_factory=lambda: [
+        "consistency_score",
+        "coherence_score",
+        "memory_retention",
+        "context_utilization",
+        "response_relevance"
+    ])
+    
+    def __post_init__(self):
+        """Ensure data path exists."""
+        self.data_path.mkdir(parents=True, exist_ok=True)
+    
+    def validate(self) -> bool:
+        """Validate LOCOMO configuration."""
+        if self.max_sessions <= 0:
+            raise ValueError("MAX_SESSIONS must be positive")
+        if self.max_agents <= 0:
+            raise ValueError("MAX_AGENTS must be positive")
+        if self.context_window_size <= 0:
+            raise ValueError("CONTEXT_WINDOW_SIZE must be positive")
+        return True
+
+
+@dataclass
+class LoggingConfig:
+    """Configuration for logging."""
+    
+    level: str = field(default_factory=lambda: os.getenv("LOG_LEVEL", "INFO"))
+    file_path: Optional[Path] = field(default_factory=lambda: Path(os.getenv("LOG_FILE", "./logs/mem0_agent.log")) if os.getenv("LOG_FILE") else None)
+    
+    def __post_init__(self):
+        """Ensure log directory exists."""
+        if self.file_path:
+            self.file_path.parent.mkdir(parents=True, exist_ok=True)
+
+
+@dataclass
+class Config:
+    """Main configuration class."""
+    
+    kimi: KimiConfig = field(default_factory=KimiConfig)
+    mem0: Mem0Config = field(default_factory=Mem0Config)
+    locomo: LOCOMOConfig = field(default_factory=LOCOMOConfig)
+    logging: LoggingConfig = field(default_factory=LoggingConfig)
+    
+    def validate(self) -> bool:
+        """Validate all configurations."""
+        self.kimi.validate()
+        self.mem0.validate()
+        self.locomo.validate()
+        return True
+    
+    @classmethod
+    def from_env(cls) -> "Config":
+        """Create configuration from environment variables."""
+        return cls()
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert configuration to dictionary."""
+        return {
+            "kimi": {
+                "model_name": self.kimi.model_name,
+                "max_tokens": self.kimi.max_tokens,
+                "temperature": _reasoning_safe_temperature(self.kimi.model_name, self.kimi.temperature),
+                "api_base": self.kimi.api_base
+            },
+            "mem0": {
+                "backend": self.mem0.backend,
+                "collection_name": self.mem0.collection_name,
+                "embedding_model": self.mem0.embedding_model
+            },
+            "locomo": {
+                "data_path": str(self.locomo.data_path),
+                "max_sessions": self.locomo.max_sessions,
+                "max_agents": self.locomo.max_agents,
+                "context_window_size": self.locomo.context_window_size,
+                "evaluation_metrics": self.locomo.evaluation_metrics
+            },
+            "logging": {
+                "level": self.logging.level,
+                "file_path": str(self.logging.file_path) if self.logging.file_path else None
+            }
+        }
+
+
+# Global configuration instance
+config = Config.from_env()