11_context_optimization.py
1 #!/usr/bin/env python3 2 """ 3 Context Optimization Example 4 5 Demonstrates how to use context optimization with Agent via context= param, 6 and also shows low-level optimizer usage for advanced scenarios. 7 8 Agent-Centric Quick Start: 9 from praisonaiagents import Agent 10 from praisonaiagents.context import ManagerConfig 11 12 agent = Agent( 13 instructions="You are helpful.", 14 context=ManagerConfig( 15 auto_compact=True, 16 strategy="smart", # truncate, sliding_window, summarize, smart 17 ), 18 ) 19 """ 20 21 from praisonaiagents import Agent 22 from praisonaiagents.context import ( 23 ManagerConfig, 24 get_optimizer, 25 OptimizerStrategy, 26 TruncateOptimizer, 27 SlidingWindowOptimizer, 28 PruneToolsOptimizer, 29 SmartOptimizer, 30 estimate_messages_tokens, 31 ) 32 33 34 def agent_centric_example(): 35 """Agent-centric usage - recommended approach.""" 36 print("=" * 60) 37 print("Agent-Centric Context Optimization") 38 print("=" * 60) 39 40 # Enable auto-compaction with smart strategy 41 agent = Agent( 42 instructions="You are a helpful assistant.", 43 context=ManagerConfig( 44 auto_compact=True, 45 compact_threshold=0.8, 46 strategy="smart", 47 ), 48 ) 49 print("Agent with auto-compaction enabled") 50 print(f"Strategy: smart") 51 print(f"Threshold: 0.8 (80% of context budget)") 52 print(f"Context manager: {agent.context_manager is not None}") 53 54 55 def main(): 56 print("=" * 60) 57 print("Context Optimization Example") 58 print("=" * 60) 59 60 # Create sample conversation with tool calls 61 messages = [ 62 {"role": "system", "content": "You are a helpful AI assistant with access to tools."}, 63 {"role": "user", "content": "What's the weather in Paris?"}, 64 {"role": "assistant", "content": "", "tool_calls": [{"id": "call_1", "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}}]}, 65 {"role": "tool", "tool_call_id": "call_1", "content": "Weather in Paris: 15°C, partly cloudy, humidity 65%, wind 10 km/h from the west. Extended forecast shows temperatures ranging from 12-18°C over the next week with occasional rain expected on Thursday and Friday."}, 66 {"role": "assistant", "content": "The weather in Paris is currently 15°C and partly cloudy."}, 67 {"role": "user", "content": "What about London?"}, 68 {"role": "assistant", "content": "", "tool_calls": [{"id": "call_2", "function": {"name": "get_weather", "arguments": '{"city": "London"}'}}]}, 69 {"role": "tool", "tool_call_id": "call_2", "content": "Weather in London: 12°C, overcast with light drizzle, humidity 80%, wind 15 km/h from the southwest. Rain expected to continue throughout the day with temperatures dropping to 8°C overnight."}, 70 {"role": "assistant", "content": "London is 12°C with light drizzle."}, 71 ] 72 73 # Add more history to simulate a long conversation 74 for i in range(10): 75 messages.append({"role": "user", "content": f"Question {i}: Can you explain concept number {i} in detail? " * 20}) 76 messages.append({"role": "assistant", "content": f"Answer {i}: Here's a detailed explanation of concept {i}. " * 30}) 77 78 original_tokens = estimate_messages_tokens(messages) 79 print(f"\nOriginal conversation: {len(messages)} messages, ~{original_tokens} tokens") 80 81 # Example 1: Truncate Optimizer 82 print("\n1. Truncate Optimizer") 83 print("-" * 40) 84 85 optimizer = TruncateOptimizer() 86 result, stats = optimizer.optimize(messages, target_tokens=2000) 87 result_tokens = estimate_messages_tokens(result) 88 89 print(f"Strategy: Remove oldest messages first") 90 print(f"Result: {len(messages)} -> {len(result)} messages") 91 print(f"Tokens: {original_tokens} -> {result_tokens}") 92 93 # Example 2: Sliding Window Optimizer 94 print("\n2. Sliding Window Optimizer") 95 print("-" * 40) 96 97 optimizer = SlidingWindowOptimizer() 98 result, stats = optimizer.optimize(messages, target_tokens=2000) 99 result_tokens = estimate_messages_tokens(result) 100 101 print(f"Strategy: Keep most recent messages within window") 102 print(f"Result: {len(messages)} -> {len(result)} messages") 103 print(f"Tokens: {original_tokens} -> {result_tokens}") 104 105 # Example 3: Prune Tools Optimizer 106 print("\n3. Prune Tools Optimizer") 107 print("-" * 40) 108 109 optimizer = PruneToolsOptimizer() 110 result, stats = optimizer.optimize(messages, target_tokens=2000) 111 result_tokens = estimate_messages_tokens(result) 112 113 print(f"Strategy: Truncate old tool outputs, preserve recent") 114 print(f"Result: {len(messages)} -> {len(result)} messages") 115 print(f"Tokens: {original_tokens} -> {result_tokens}") 116 117 # Example 4: Smart Optimizer (combines strategies) 118 print("\n4. Smart Optimizer (Recommended)") 119 print("-" * 40) 120 121 optimizer = SmartOptimizer() 122 result, stats = optimizer.optimize(messages, target_tokens=2000) 123 result_tokens = estimate_messages_tokens(result) 124 125 print(f"Strategy: Intelligent combination of all strategies") 126 print(f"Result: {len(messages)} -> {len(result)} messages") 127 print(f"Tokens: {original_tokens} -> {result_tokens}") 128 print(f"Stats: {stats}") 129 130 # Example 5: Using the factory function 131 print("\n5. Optimizer Factory") 132 print("-" * 40) 133 134 strategies = [ 135 OptimizerStrategy.TRUNCATE, 136 OptimizerStrategy.SLIDING_WINDOW, 137 OptimizerStrategy.PRUNE_TOOLS, 138 OptimizerStrategy.SMART, 139 ] 140 141 for strategy in strategies: 142 optimizer = get_optimizer(strategy) 143 result, _ = optimizer.optimize(messages, target_tokens=3000) 144 result_tokens = estimate_messages_tokens(result) 145 print(f"{strategy.value:15s}: {len(messages)} -> {len(result)} msgs, ~{result_tokens} tokens") 146 147 # Example 6: Tool call/result pair preservation 148 print("\n6. Tool Call/Result Pair Preservation") 149 print("-" * 40) 150 151 tool_messages = [ 152 {"role": "system", "content": "You are helpful."}, 153 {"role": "user", "content": "Get data"}, 154 {"role": "assistant", "content": "", "tool_calls": [{"id": "tc1"}]}, 155 {"role": "tool", "tool_call_id": "tc1", "content": "Data result"}, 156 {"role": "assistant", "content": "Here's the data."}, 157 ] 158 159 optimizer = SmartOptimizer() 160 result, _ = optimizer.optimize(tool_messages, target_tokens=100) 161 162 # Check if tool pairs are preserved 163 tool_calls = [m for m in result if m.get("tool_calls")] 164 tool_results = [m for m in result if m.get("role") == "tool"] 165 166 print(f"Original: {len(tool_messages)} messages") 167 print(f"Optimized: {len(result)} messages") 168 print(f"Tool calls preserved: {len(tool_calls)}") 169 print(f"Tool results preserved: {len(tool_results)}") 170 171 print("\n" + "=" * 60) 172 print("✓ Context optimization examples complete!") 173 print("=" * 60) 174 175 176 if __name__ == "__main__": 177 main()