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

chapter3/contextual-retrieval/main.py

@@ -0,0 +1,292 @@
+"""Main entry point for Agentic RAG system"""
+
+import os
+import json
+import logging
+import argparse
+from typing import Optional
+from config import Config, KnowledgeBaseType
+from agent import AgenticRAG
+from chunking import DocumentIndexer
+
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+
+def setup_environment():
+    """Setup environment and check requirements"""
+    # Check for required API keys
+    config = Config.from_env()
+    
+    # Check LLM API key
+    try:
+        api_key = config.llm.get_api_key(config.llm.provider)
+        if not api_key:
+            logger.warning(f"No API key found for provider {config.llm.provider}")
+            logger.info("Please set the appropriate environment variable:")
+            logger.info("  - MOONSHOT_API_KEY for Kimi")
+            logger.info("  - ARK_API_KEY for Doubao")
+            logger.info("  - SILICONFLOW_API_KEY for SiliconFlow")
+            logger.info("  - OPENAI_API_KEY for OpenAI")
+            return False
+    except Exception as e:
+        logger.error(f"Error checking API keys: {e}")
+        return False
+    
+    # Check knowledge base setup
+    if config.knowledge_base.type == KnowledgeBaseType.LOCAL:
+        # Check if local retrieval pipeline is running
+        import requests
+        try:
+            response = requests.get(f"{config.knowledge_base.local_base_url}/health", timeout=30)
+            if response.status_code != 200:
+                logger.warning("Local retrieval pipeline not responding")
+                logger.info(f"Please ensure the retrieval pipeline is running at {config.knowledge_base.local_base_url}")
+                logger.info("Run: cd ../retrieval-pipeline && python main.py")
+        except Exception:
+            logger.warning("Cannot connect to local retrieval pipeline")
+            logger.info("Will continue anyway - searches may fail")
+    
+    elif config.knowledge_base.type == KnowledgeBaseType.DIFY:
+        if not config.knowledge_base.dify_api_key:
+            logger.warning("Dify API key not set")
+            logger.info("Please set DIFY_API_KEY environment variable")
+    
+    return True
+
+
+def run_interactive_mode(agent: AgenticRAG, mode: str = "agentic"):
+    """Run interactive query mode"""
+    print(f"\n{'='*60}")
+    print(f"Agentic RAG System - {mode.capitalize()} Mode")
+    print(f"Verbose: {'Enabled' if agent.config.agent.verbose else 'Disabled'} | Top-K: {agent.config.knowledge_base.local_top_k}")
+    print(f"{'='*60}")
+    print("Type 'quit' or 'exit' to stop")
+    print("Type 'clear' to clear conversation history")
+    print("Type 'mode' to switch between agentic/non-agentic modes")
+    print(f"{'='*60}\n")
+    
+    current_mode = mode
+    
+    while True:
+        try:
+            user_input = input("\n[USER] > ").strip()
+            
+            if user_input.lower() in ['quit', 'exit']:
+                print("\nGoodbye!")
+                break
+            
+            if user_input.lower() == 'clear':
+                agent.clear_history()
+                print("Conversation history cleared.")
+                continue
+            
+            if user_input.lower() == 'mode':
+                current_mode = "non-agentic" if current_mode == "agentic" else "agentic"
+                print(f"Switched to {current_mode} mode")
+                continue
+            
+            if not user_input:
+                continue
+            
+            # Process query
+            print(f"\n[ASSISTANT ({current_mode})] > ", end="", flush=True)
+            
+            if current_mode == "agentic":
+                response = agent.query(user_input, stream=True)
+            else:
+                response = agent.query_non_agentic(user_input, stream=True)
+            
+            # Handle streaming response
+            if hasattr(response, '__iter__'):
+                for chunk in response:
+                    print(chunk, end="", flush=True)
+                print()  # New line after response
+            else:
+                print(response)
+                
+        except KeyboardInterrupt:
+            print("\n\nInterrupted. Type 'quit' to exit.")
+        except Exception as e:
+            logger.error(f"Error: {e}")
+            print(f"\nError processing query: {e}")
+
+
+def run_batch_mode(agent: AgenticRAG, queries_file: str, output_file: str, mode: str = "agentic"):
+    """Run batch queries from file"""
+    try:
+        with open(queries_file, 'r', encoding='utf-8') as f:
+            queries = [line.strip() for line in f if line.strip()]
+    except Exception as e:
+        logger.error(f"Error reading queries file: {e}")
+        return
+    
+    results = []
+    
+    for i, query in enumerate(queries, 1):
+        print(f"\n[{i}/{len(queries)}] Processing: {query[:100]}...")
+        
+        try:
+            if mode == "agentic":
+                response = agent.query(query, stream=False)
+            else:
+                response = agent.query_non_agentic(query, stream=False)
+            
+            results.append({
+                "query": query,
+                "response": response,
+                "mode": mode
+            })
+            
+        except Exception as e:
+            logger.error(f"Error processing query: {e}")
+            results.append({
+                "query": query,
+                "response": f"Error: {str(e)}",
+                "mode": mode
+            })
+    
+    # Save results
+    try:
+        with open(output_file, 'w', encoding='utf-8') as f:
+            json.dump(results, f, ensure_ascii=False, indent=2)
+        print(f"\nResults saved to {output_file}")
+    except Exception as e:
+        logger.error(f"Error saving results: {e}")
+
+
+def run_comparison_mode(agent: AgenticRAG, query: str):
+    """Run both modes and compare results"""
+    print(f"\n{'='*60}")
+    print("Comparison Mode - Running both Agentic and Non-Agentic")
+    print(f"{'='*60}")
+    print(f"Query: {query}")
+    print(f"{'='*60}")
+    
+    # Run non-agentic mode
+    print("\n[NON-AGENTIC MODE]")
+    print("-" * 40)
+    non_agentic_response = agent.query_non_agentic(query, stream=False)
+    print(non_agentic_response)
+    
+    # Clear history for fair comparison
+    agent.clear_history()
+    
+    # Run agentic mode
+    print("\n[AGENTIC MODE]")
+    print("-" * 40)
+    agentic_response = agent.query(query, stream=False)
+    print(agentic_response)
+    
+    print(f"\n{'='*60}")
+
+
+def main():
+    """Main function"""
+    parser = argparse.ArgumentParser(description="Agentic RAG System")
+    
+    # Mode selection
+    parser.add_argument("--mode", choices=["agentic", "non-agentic", "compare"], 
+                       default="agentic", help="Query mode")
+    
+    # Query options
+    parser.add_argument("--query", type=str, help="Single query to process")
+    parser.add_argument("--batch", type=str, help="Path to file with queries (one per line)")
+    parser.add_argument("--output", type=str, default="results.json", 
+                       help="Output file for batch results")
+    
+    # Configuration options
+    parser.add_argument("--provider", type=str, help="LLM provider")
+    parser.add_argument("--model", type=str, help="LLM model")
+    parser.add_argument("--kb-type", choices=["local", "dify"], help="Knowledge base type")
+    parser.add_argument("--verbose", action="store_true", help="Verbose output (default: True)")
+    parser.add_argument("--no-verbose", action="store_true", help="Disable verbose output")
+    
+    # Indexing options
+    parser.add_argument("--index", type=str, help="Path to index (file or directory)")
+    parser.add_argument("--chunk-size", type=int, default=2048, help="Chunk size for indexing")
+    
+    args = parser.parse_args()
+    
+    # Setup environment
+    if not setup_environment():
+        logger.warning("Environment setup incomplete, continuing anyway...")
+    
+    # Load or create config
+    config = Config.from_env()
+    
+    # Set verbose mode by default (can be disabled with --no-verbose)
+    config.agent.verbose = True  # Default to verbose mode
+    
+    # Override config with command line args
+    if args.provider:
+        config.llm.provider = args.provider
+    if args.model:
+        config.llm.model = args.model
+    if args.kb_type:
+        config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
+    
+    # Handle verbose mode (default is True, can be disabled with --no-verbose)
+    if args.no_verbose:
+        config.agent.verbose = False
+    elif args.verbose:
+        config.agent.verbose = True  # Explicitly set if --verbose is used
+    
+    # Handle indexing if requested
+    if args.index:
+        print(f"\n{'='*60}")
+        print("Indexing Documents")
+        print(f"{'='*60}")
+        
+        config.chunking.chunk_size = args.chunk_size
+        indexer = DocumentIndexer(config.knowledge_base, config.chunking)
+        
+        from pathlib import Path
+        path = Path(args.index)
+        
+        if path.is_file():
+            result = indexer.index_file(str(path))
+        elif path.is_dir():
+            result = indexer.index_directory(str(path))
+        else:
+            print(f"Path not found: {path}")
+            return
+        
+        print(json.dumps(result, indent=2, ensure_ascii=False))
+        print(f"{'='*60}\n")
+    
+    # Create agent
+    agent = AgenticRAG(config)
+    
+    # Handle different execution modes
+    if args.query and args.mode == "compare":
+        # Comparison mode with single query
+        run_comparison_mode(agent, args.query)
+        
+    elif args.query:
+        # Single query mode
+        print(f"\n[Query] {args.query}")
+        print(f"[Mode] {args.mode}")
+        print(f"[Verbose] {'Enabled' if config.agent.verbose else 'Disabled'}")
+        print(f"[Top-K] {config.knowledge_base.local_top_k}")
+        print("-" * 40)
+        
+        if args.mode == "agentic":
+            response = agent.query(args.query, stream=False)
+        else:
+            response = agent.query_non_agentic(args.query, stream=False)
+        
+        print(response)
+        
+    elif args.batch:
+        # Batch mode
+        run_batch_mode(agent, args.batch, args.output, args.mode)
+        
+    else:
+        # Interactive mode (default)
+        run_interactive_mode(agent, args.mode)
+
+
+if __name__ == "__main__":
+    main()