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

chapter3/contextual-retrieval/config.py

@@ -0,0 +1,259 @@
+"""Configuration for Agentic RAG System"""
+
+import os
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any
+from enum import Enum
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def _openrouter_model_id(model: Optional[str]) -> 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"
+
+
+class Provider(str, Enum):
+    """Supported LLM providers"""
+    DASHSCOPE = "dashscope"  # Alibaba Cloud Model Studio / Bailian (Qwen)
+    SILICONFLOW = "siliconflow"
+    DOUBAO = "doubao"
+    KIMI = "kimi"
+    MOONSHOT = "moonshot"
+    OPENROUTER = "openrouter"
+    OPENAI = "openai"
+    GROQ = "groq"
+    TOGETHER = "together"
+    DEEPSEEK = "deepseek"
+
+
+class KnowledgeBaseType(str, Enum):
+    """Knowledge base backend types"""
+    LOCAL = "local"  # Local retrieval pipeline
+    DIFY = "dify"    # Dify knowledge base API
+    RAPTOR = "raptor"  # RAPTOR tree-based index
+    GRAPHRAG = "graphrag"  # GraphRAG graph-based index
+
+
+@dataclass
+class LLMConfig:
+    """LLM configuration"""
+    provider: str = "kimi"  # Default provider
+    model: Optional[str] = None  # Will use provider defaults if not specified
+    api_key: Optional[str] = None  # Will read from env if not provided
+    temperature: float = 0.7
+    max_tokens: int = 1024
+    stream: bool = True
+    
+    # Provider-specific defaults
+    PROVIDER_DEFAULTS = {
+        "dashscope": {
+            "model": "qwen3.7-plus",
+            "base_url": os.getenv(
+                "DASHSCOPE_BASE_URL",
+                "https://dashscope.aliyuncs.com/compatible-mode/v1",
+            ),
+        },
+        "siliconflow": {
+            "model": "Qwen/Qwen3-235B-A22B-Thinking-2507",
+            "base_url": "https://api.siliconflow.cn/v1"
+        },
+        "doubao": {
+            "model": "doubao-seed-1-6-thinking-250715",
+            "base_url": "https://ark.cn-beijing.volces.com/api/v3"
+        },
+        "kimi": {
+            "model": "kimi-k3",
+            "base_url": "https://api.moonshot.cn/v1"
+        },
+        "moonshot": {
+            "model": "kimi-k3",
+            "base_url": "https://api.moonshot.cn/v1"
+        },
+        "openrouter": {
+            "model": "openai/gpt-5.6-luna",
+            "base_url": "https://openrouter.ai/api/v1"
+        },
+        "openai": {
+            "model": "gpt-5.6-luna",
+            "base_url": "https://api.openai.com/v1"
+        },
+        "groq": {
+            "model": "llama-3.3-70b-versatile",
+            "base_url": "https://api.groq.com/openai/v1"
+        },
+        "together": {
+            "model": "meta-llama/Llama-3.3-70B-Instruct-Turbo",
+            "base_url": "https://api.together.xyz"
+        },
+        "deepseek": {
+            "model": "deepseek-reasoner",
+            "base_url": "https://api.deepseek.com/v1"
+        }
+    }
+    
+    @classmethod
+    def get_api_key(cls, provider: str) -> Optional[str]:
+        """Get API key from environment"""
+        env_mappings = {
+            "dashscope": "DASHSCOPE_API_KEY",
+            "qwen": "DASHSCOPE_API_KEY",
+            "bailian": "DASHSCOPE_API_KEY",
+            "siliconflow": "SILICONFLOW_API_KEY",
+            "doubao": "ARK_API_KEY",
+            "kimi": "MOONSHOT_API_KEY",
+            "moonshot": "MOONSHOT_API_KEY",
+            "openrouter": "OPENROUTER_API_KEY",
+            "openai": "OPENAI_API_KEY",
+            "groq": "GROQ_API_KEY",
+            "together": "TOGETHER_API_KEY",
+            "deepseek": "DEEPSEEK_API_KEY"
+        }
+        return os.getenv(env_mappings.get(provider.lower(), ""))
+    
+    def get_client_config(self) -> Dict[str, Any]:
+        """Get OpenAI client configuration"""
+        provider_lower = self.provider.lower()
+        provider_lower = {"qwen": "dashscope", "bailian": "dashscope"}.get(
+            provider_lower, provider_lower
+        )
+        defaults = self.PROVIDER_DEFAULTS.get(provider_lower, {})
+        
+        # Get API key
+        api_key = self.api_key or self.get_api_key(provider_lower)
+
+        # Universal OpenRouter fallback: primary provider key absent but
+        # OPENROUTER_API_KEY present -> route through OpenRouter. Additionally,
+        # gpt-5.x (incl. gpt-5.6*) needs OpenAI org-verification on the direct
+        # API, so prefer OpenRouter for those ids whenever an OR key is present.
+        model_name = self.model or defaults.get("model")
+        openrouter_key = os.getenv("OPENROUTER_API_KEY")
+        prefer_openrouter = bool(openrouter_key) and str(model_name or "").lower().startswith("gpt-5")
+        if (not api_key or prefer_openrouter) and provider_lower != "openrouter" and openrouter_key:
+            model = _openrouter_model_id(model_name)
+            return {
+                "api_key": openrouter_key,
+                "base_url": "https://openrouter.ai/api/v1",
+            }, model
+
+        if not api_key:
+            raise ValueError(
+                f"API key required for provider '{provider_lower}'. Set the "
+                f"provider's key (e.g. MOONSHOT_API_KEY / OPENAI_API_KEY) or "
+                f"OPENROUTER_API_KEY to use the OpenRouter fallback."
+            )
+
+        # Build config
+        config = {
+            "api_key": api_key,
+            "model": self.model or defaults.get("model")
+        }
+
+        # Add base_url if not OpenAI
+        if "base_url" in defaults:
+            config["base_url"] = defaults["base_url"]
+
+        return config, config.pop("model")
+
+
+@dataclass
+class KnowledgeBaseConfig:
+    """Knowledge base configuration"""
+    type: KnowledgeBaseType = KnowledgeBaseType.LOCAL
+    
+    # Local retrieval pipeline config
+    local_base_url: str = "http://localhost:4242"
+    local_top_k: int = 3
+    
+    # Dify config
+    dify_api_key: Optional[str] = field(default_factory=lambda: os.getenv("DIFY_API_KEY"))
+    dify_base_url: str = "https://api.dify.ai/v1"
+    dify_dataset_id: Optional[str] = None
+    dify_top_k: int = 10
+    
+    # RAPTOR tree-based index config
+    raptor_base_url: str = "http://localhost:4242"
+    raptor_top_k: int = 10
+    raptor_search_levels: bool = True  # Search across multiple tree levels
+    
+    # GraphRAG graph-based index config
+    graphrag_base_url: str = "http://localhost:4242"
+    graphrag_top_k: int = 10
+    graphrag_search_type: str = "hybrid"  # entity, community, or hybrid
+    
+    # Document storage
+    document_store_path: str = "document_store.json"
+    
+    
+@dataclass
+class ChunkingConfig:
+    """Document chunking configuration"""
+    chunk_size: int = 2048  # Characters per chunk
+    max_chunk_size: int = 1024  # Max size when respecting paragraph boundaries
+    chunk_overlap: int = 200  # Overlap between chunks
+    respect_paragraph_boundary: bool = True
+    min_chunk_size: int = 100  # Minimum chunk size
+
+
+@dataclass 
+class AgentConfig:
+    """Agent configuration"""
+    max_iterations: int = 10  # Max reasoning iterations
+    enable_reasoning_trace: bool = True
+    enable_citations: bool = True
+    strict_knowledge_base: bool = True  # Only answer from knowledge base
+    conversation_history_limit: int = 20  # Max conversation turns to keep
+    verbose: bool = True
+
+
+@dataclass
+class EvaluationConfig:
+    """Evaluation configuration"""
+    dataset_path: str = "evaluation/legal_qa_dataset.json"
+    results_path: str = "evaluation/results"
+    metrics: list = field(default_factory=lambda: ["accuracy", "relevance", "citation_quality"])
+
+
+@dataclass
+class Config:
+    """Main configuration"""
+    llm: LLMConfig = field(default_factory=LLMConfig)
+    knowledge_base: KnowledgeBaseConfig = field(default_factory=KnowledgeBaseConfig)
+    chunking: ChunkingConfig = field(default_factory=ChunkingConfig)
+    agent: AgentConfig = field(default_factory=AgentConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
+    
+    @classmethod
+    def from_env(cls) -> "Config":
+        """Create config from environment variables"""
+        config = cls()
+        
+        # Override from env
+        if provider := os.getenv("LLM_PROVIDER"):
+            config.llm.provider = provider
+        if model := os.getenv("LLM_MODEL"):
+            config.llm.model = model
+        if kb_type := os.getenv("KB_TYPE"):
+            config.knowledge_base.type = KnowledgeBaseType(kb_type.lower())
+        
+        return config