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

chapter2/context-compression/main.py

@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+"""
+Interactive demo for context compression strategies
+"""
+
+import os
+import sys
+import argparse
+from colorama import init, Fore, Style
+
+from config import Config
+from agent import ResearchAgent
+from compression_strategies import CompressionStrategy
+
+# Initialize colorama
+init(autoreset=True)
+
+
+# Short CLI aliases -> compression strategy (order matches the book's 实验 2-10)
+STRATEGY_CHOICES = {
+    "no_compression": CompressionStrategy.NO_COMPRESSION,
+    "individual": CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL,
+    "combined": CompressionStrategy.NON_CONTEXT_AWARE_COMBINED,
+    "context_aware": CompressionStrategy.CONTEXT_AWARE,
+    "citations": CompressionStrategy.CONTEXT_AWARE_CITATIONS,
+    "windowed": CompressionStrategy.WINDOWED_CONTEXT,
+}
+
+
+def print_banner():
+    """Print demo banner"""
+    print(f"\n{Fore.CYAN}{'='*70}")
+    print(f"{Fore.CYAN}CONTEXT COMPRESSION RESEARCH AGENT - INTERACTIVE DEMO")
+    print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}")
+    print("\nThis demo allows you to test different compression strategies")
+    print("for researching OpenAI co-founders' current affiliations.\n")
+
+
+def select_strategy() -> CompressionStrategy:
+    """Let user select a compression strategy"""
+    print(f"{Fore.YELLOW}Available Compression Strategies:{Style.RESET_ALL}")
+    print("1. No Compression (expected to fail with large contexts)")
+    print("2. Non-Context-Aware: Individual Summaries (summarize each page, then concatenate)")
+    print("3. Non-Context-Aware: Combined Summary (concatenate all pages, then summarize once)")
+    print("4. Context-Aware Summarization")
+    print("5. Context-Aware with Citations")
+    print("6. Windowed Context (only compress when approaching context limit)")
+    
+    while True:
+        try:
+            choice = input(f"\n{Fore.GREEN}Select strategy (1-6): {Style.RESET_ALL}")
+            strategies = [
+                CompressionStrategy.NO_COMPRESSION,
+                CompressionStrategy.NON_CONTEXT_AWARE_INDIVIDUAL,
+                CompressionStrategy.NON_CONTEXT_AWARE_COMBINED,
+                CompressionStrategy.CONTEXT_AWARE,
+                CompressionStrategy.CONTEXT_AWARE_CITATIONS,
+                CompressionStrategy.WINDOWED_CONTEXT
+            ]
+            return strategies[int(choice) - 1]
+        except (ValueError, IndexError):
+            print(f"{Fore.RED}Invalid choice. Please enter 1-6.{Style.RESET_ALL}")
+
+
+def run_demo(enable_streaming=True, strategy: CompressionStrategy = None):
+    """Run the interactive demo
+
+    Args:
+        enable_streaming: Whether to enable streaming output (default: True)
+        strategy: Preselected compression strategy; if None, prompt the user interactively
+    """
+    print_banner()
+
+    # Check configuration
+    if not Config.validate():
+        print(f"\n{Fore.RED}Configuration validation failed!{Style.RESET_ALL}")
+        print("\nPlease set up your .env file with:")
+        print("  DASHSCOPE_API_KEY=your_api_key_here (for LLM_PROVIDER=dashscope/qwen/bailian)")
+        print("  MOONSHOT_API_KEY=your_api_key_here")
+        print("  SERPER_API_KEY=your_api_key_here (optional, will use mock data)")
+        sys.exit(1)
+
+    # Select strategy (interactively unless one was passed on the command line)
+    if strategy is None:
+        strategy = select_strategy()
+    
+    print(f"\n{Fore.CYAN}Selected: {strategy.value}{Style.RESET_ALL}")
+    
+    # Display streaming status
+    streaming_status = "ENABLED" if enable_streaming else "DISABLED"
+    print(f"{Fore.YELLOW}Streaming output: {streaming_status}{Style.RESET_ALL}")
+    
+    # Create agent
+    print(f"\n{Fore.YELLOW}Initializing agent...{Style.RESET_ALL}")
+    agent = ResearchAgent(
+        api_key=Config.resolve_llm()[0],
+        compression_strategy=strategy,
+        verbose=False,
+        enable_streaming=enable_streaming
+    )
+    
+    print(f"\n{Fore.CYAN}Starting research task...{Style.RESET_ALL}")
+    print("Task: Find current affiliations of all OpenAI co-founders\n")
+    print("-" * 70)
+    
+    try:
+        # Execute research
+        result = agent.execute_research(max_iterations=Config.MAX_ITERATIONS)
+        
+        # Print results
+        print("\n" + "="*70)
+        print(f"{Fore.GREEN}RESEARCH COMPLETE{Style.RESET_ALL}")
+        print("="*70)
+        
+        if result.get('success'):
+            print(f"\n{Fore.GREEN}✅ Success!{Style.RESET_ALL}")
+            print(f"\nFinal Answer:\n{result.get('final_answer', 'No answer found')}")
+        else:
+            print(f"\n{Fore.RED}❌ Failed{Style.RESET_ALL}")
+            if result.get('error'):
+                print(f"Error: {result['error']}")
+        
+        # Print statistics
+        trajectory = result.get('trajectory')
+        if trajectory:
+            print(f"\n{Fore.CYAN}📊 Statistics:{Style.RESET_ALL}")
+            print(f"  Tool Calls: {len(trajectory.tool_calls)}")
+            print(f"  Context Overflows: {trajectory.context_overflows}")
+            print(f"  Execution Time: {result.get('execution_time', 0):.2f}s")
+            print(f"  Total Tokens Used: {trajectory.total_tokens_used:,}")
+            print(f"    - Prompt Tokens: {trajectory.prompt_tokens_used:,}")
+            print(f"    - Completion Tokens: {trajectory.completion_tokens_used:,}")
+            
+            # Calculate compression stats
+            if trajectory.tool_calls:
+                total_original = 0
+                total_compressed = 0
+                
+                for call in trajectory.tool_calls:
+                    if call.compressed_result:
+                        total_original += call.compressed_result.original_length
+                        total_compressed += call.compressed_result.compressed_length
+                
+                if total_original > 0:
+                    ratio = total_compressed / total_original
+                    print(f"  Compression Ratio: {ratio:.1%}")
+                    print(f"  Space Saved: {total_original - total_compressed:,} chars")
+        
+        # Follow-up question demo (for citation strategy)
+        if strategy == CompressionStrategy.CONTEXT_AWARE_CITATIONS and result.get('success'):
+            print(f"\n{Fore.YELLOW}This strategy supports follow-up questions!{Style.RESET_ALL}")
+            follow_up = input("\nAsk a follow-up question (or press Enter to skip): ")
+            
+            if follow_up:
+                print(f"\n{Fore.CYAN}Processing follow-up...{Style.RESET_ALL}")
+                # Add follow-up to conversation
+                agent.conversation_history.append({"role": "user", "content": follow_up})
+                
+                # Get response (simplified for demo)
+                messages = agent.conversation_history.copy()
+                
+                if enable_streaming:
+                    message = agent._stream_response(messages)
+                else:
+                    message = agent._non_streaming_response(messages)
+                
+                if message.get('content'):
+                    print(f"\n{Fore.GREEN}Follow-up Answer:{Style.RESET_ALL}")
+                    print(message['content'])
+    
+    except KeyboardInterrupt:
+        print(f"\n\n{Fore.YELLOW}Demo interrupted by user{Style.RESET_ALL}")
+    except Exception as e:
+        print(f"\n{Fore.RED}Error: {str(e)}{Style.RESET_ALL}")
+
+
+def main():
+    """Main entry point"""
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(
+        prog="main.py",
+        description="上下文压缩策略交互式演示：针对“追踪 OpenAI 联合创始人现状”这一研究任务，"
+                    "单独运行某一种压缩策略并实时观察其执行与压缩过程。",
+        epilog="示例：\n"
+               "  python main.py                         # 交互式选择策略\n"
+               "  python main.py -s citations            # 直接运行“带引用的上下文感知”策略\n"
+               "  python main.py -s windowed --no-streaming\n"
+               "如需批量对比全部策略并生成对比表，请使用 experiment.py。",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        '-s', '--strategy', choices=list(STRATEGY_CHOICES.keys()), metavar="NAME",
+        help="直接指定压缩策略（跳过交互式选择）。可选值：" + ", ".join(STRATEGY_CHOICES.keys()),
+    )
+    parser.add_argument(
+        '-m', '--model', default=None,
+        help=f"覆盖使用的模型名称（默认读取环境变量 MODEL_NAME，当前为 {Config.MODEL_NAME}）",
+    )
+    parser.add_argument(
+        '--no-streaming',
+        action='store_true',
+        help='关闭流式输出（默认开启流式）'
+    )
+    args = parser.parse_args()
+
+    if args.model:
+        Config.MODEL_NAME = args.model
+
+    # Determine streaming preference
+    enable_streaming = not args.no_streaming
+    preset_strategy = STRATEGY_CHOICES[args.strategy] if args.strategy else None
+
+    try:
+        run_demo(enable_streaming=enable_streaming, strategy=preset_strategy)
+
+        # Ask if user wants to try another strategy
+        while True:
+            again = input(f"\n{Fore.GREEN}Try another strategy? (y/n): {Style.RESET_ALL}")
+            if again.lower() == 'y':
+                run_demo(enable_streaming=enable_streaming)
+            else:
+                print(f"\n{Fore.CYAN}Thank you for using the demo!{Style.RESET_ALL}")
+                break
+
+    except KeyboardInterrupt:
+        print(f"\n\n{Fore.YELLOW}Goodbye!{Style.RESET_ALL}")
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()