#!/usr/bin/env python3 """ Token Optimizer - Simulation Test Demonstrates before/after performance comparison with mock data. """ import json import os import sys import shutil from pathlib import Path from datetime import datetime # Fix Windows encoding if sys.platform == 'win32': sys.stdout.reconfigure(encoding='utf-8', errors='replace') # Add src to path sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) from analyzer import OpenClawAnalyzer from optimizer import TokenOptimizer # ANSI colors class Colors: RED = '\033[91m' GREEN = '\033[92m' YELLOW = '\033[93m' BLUE = '\033[94m' CYAN = '\033[96m' BOLD = '\033[1m' END = '\033[0m' def colorize(text, color): if sys.stdout.isatty(): return f"{color}{text}{Colors.END}" return text # Cost constants COSTS = { 'opus': {'input': 0.015, 'output': 0.075}, 'sonnet': {'input': 0.003, 'output': 0.015}, 'haiku': {'input': 0.00025, 'output': 0.00125}, 'ollama': {'input': 0.0, 'output': 0.0} } def create_mock_environment(test_dir: Path): """Create a mock OpenClaw environment for testing.""" print(colorize("\n=== CREATING MOCK OPENCLAW ENVIRONMENT ===\n", Colors.BOLD + Colors.CYAN)) # Create directory structure openclaw_dir = test_dir / '.openclaw' workspace_dir = openclaw_dir / 'workspace' workspace_dir.mkdir(parents=True, exist_ok=True) # Create UNOPTIMIZED config (typical default setup) unoptimized_config = { "agents": { "defaults": { "model": { "primary": "anthropic/claude-sonnet-4-5" # Expensive default! } } }, "heartbeat": { "every": "1h", "model": "anthropic/claude-sonnet-4-5", # Paid API for heartbeats! "prompt": "Status check" } # No caching, no budgets, no rate limits } config_path = openclaw_dir / 'openclaw.json' with open(config_path, 'w') as f: json.dump(unoptimized_config, f, indent=2) # Create BLOATED workspace files (typical unoptimized setup) # Large SOUL.md (15KB - too big!) soul_content = """# SOUL.md - Agent Configuration ## Identity You are an AI assistant... ## Detailed History """ + "\n".join([f"- Historical entry {i}: Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris." for i in range(100)]) with open(workspace_dir / 'SOUL.md', 'w') as f: f.write(soul_content) # Large MEMORY.md (25KB - loaded every time!) memory_content = """# MEMORY.md - Full History ## All Previous Sessions """ + "\n".join([f"### Session {i}\nUser asked about topic {i}. Assistant responded with detailed explanation about subject {i}. This conversation covered multiple aspects including technical details, examples, and follow-up questions. The user was satisfied with the response.\n" for i in range(200)]) with open(workspace_dir / 'MEMORY.md', 'w') as f: f.write(memory_content) # Large USER.md (10KB) user_content = """# USER.md - User Profile ## Complete User History """ + "\n".join([f"- Preference {i}: User likes detailed explanations with examples and code snippets when relevant." for i in range(150)]) with open(workspace_dir / 'USER.md', 'w') as f: f.write(user_content) print(f" Created mock OpenClaw directory: {openclaw_dir}") print(f" Config: Sonnet default, paid heartbeats, no caching") print(f" Workspace files: ~50KB total (bloated)") return openclaw_dir def calculate_costs(config: dict, workspace_size_kb: float, daily_messages: int = 100, daily_heartbeats: int = 24): """Calculate estimated daily/monthly costs.""" # Determine model model_name = config.get('agents', {}).get('defaults', {}).get('model', {}).get('primary', '') if 'haiku' in model_name.lower(): model = 'haiku' elif 'opus' in model_name.lower(): model = 'opus' else: model = 'sonnet' # Heartbeat model hb_model_name = config.get('heartbeat', {}).get('model', '') if 'ollama' in hb_model_name.lower(): hb_model = 'ollama' elif 'haiku' in hb_model_name.lower(): hb_model = 'haiku' else: hb_model = 'sonnet' # Caching cache_enabled = config.get('agents', {}).get('defaults', {}).get('cache', {}).get('enabled', False) cache_discount = 0.9 if cache_enabled else 0.0 # Calculate costs costs = COSTS[model] hb_costs = COSTS[hb_model] # Context tokens (workspace loaded each message) context_tokens = workspace_size_kb * 250 # ~250 tokens per KB # Message costs (context + average 500 token response) avg_output_tokens = 500 # Per-message cost input_cost = (context_tokens / 1000) * costs['input'] output_cost = (avg_output_tokens / 1000) * costs['output'] # Apply cache discount to input (agent prompt) if cache_enabled: # First message full price, subsequent 90% off cached_input_cost = input_cost * (1 - cache_discount * 0.8) # 80% of messages cached else: cached_input_cost = input_cost message_cost = cached_input_cost + output_cost # Heartbeat cost hb_tokens = 500 # tokens per heartbeat hb_cost = (hb_tokens / 1000) * (hb_costs['input'] + hb_costs['output']) # Daily costs daily_message_cost = message_cost * daily_messages daily_hb_cost = hb_cost * daily_heartbeats daily_total = daily_message_cost + daily_hb_cost return { 'model': model, 'heartbeat_model': hb_model, 'cache_enabled': cache_enabled, 'context_tokens': context_tokens, 'per_message_cost': message_cost, 'daily_message_cost': daily_message_cost, 'daily_heartbeat_cost': daily_hb_cost, 'daily_total': daily_total, 'monthly_total': daily_total * 30, 'yearly_total': daily_total * 365 } def print_cost_report(title: str, costs: dict, color: str): """Print a formatted cost report.""" print(colorize(f"\n{'='*60}", color)) print(colorize(f" {title}", Colors.BOLD + color)) print(colorize(f"{'='*60}", color)) print(f"\n Configuration:") print(f" Model: {costs['model'].upper()}") print(f" Heartbeat: {costs['heartbeat_model'].upper()}") print(f" Caching: {'Enabled' if costs['cache_enabled'] else 'Disabled'}") print(f" Context: {costs['context_tokens']:,.0f} tokens ({costs['context_tokens']/250:.1f}KB)") print(f"\n Per-Message Cost: ${costs['per_message_cost']:.4f}") print(f"\n Daily Costs:") print(f" Messages (100): ${costs['daily_message_cost']:.2f}") print(f" Heartbeats: ${costs['daily_heartbeat_cost']:.2f}") print(colorize(f" TOTAL: ${costs['daily_total']:.2f}", Colors.BOLD)) print(f"\n Projected Costs:") print(colorize(f" Monthly: ${costs['monthly_total']:.2f}", Colors.BOLD)) print(colorize(f" Yearly: ${costs['yearly_total']:.2f}", Colors.BOLD)) def run_simulation(): """Run the full simulation test.""" print(colorize(""" +---------------------------------------------------------------+ | | | TOKEN OPTIMIZER - SIMULATION TEST | | | | Demonstrates before/after performance comparison | | | +---------------------------------------------------------------+ """, Colors.BOLD + Colors.CYAN)) # Setup test directory test_dir = Path(__file__).parent / 'mock_environment' if test_dir.exists(): shutil.rmtree(test_dir) test_dir.mkdir(parents=True) # Create mock environment openclaw_dir = create_mock_environment(test_dir) # ========== BEFORE OPTIMIZATION ========== print(colorize("\n\n" + "="*60, Colors.RED)) print(colorize(" PHASE 1: BEFORE OPTIMIZATION (Typical Default Setup)", Colors.BOLD + Colors.RED)) print(colorize("="*60, Colors.RED)) # Load unoptimized config with open(openclaw_dir / 'openclaw.json') as f: before_config = json.load(f) # Calculate workspace size workspace_dir = openclaw_dir / 'workspace' before_workspace_size = sum(f.stat().st_size for f in workspace_dir.iterdir() if f.is_file()) / 1024 print(f"\n Workspace size: {before_workspace_size:.1f}KB") print(f" Files loaded every message: SOUL.md, MEMORY.md, USER.md") before_costs = calculate_costs(before_config, before_workspace_size) print_cost_report("BEFORE OPTIMIZATION - COST ANALYSIS", before_costs, Colors.RED) # ========== APPLY OPTIMIZATION ========== print(colorize("\n\n" + "="*60, Colors.YELLOW)) print(colorize(" PHASE 2: APPLYING TOKEN OPTIMIZER", Colors.BOLD + Colors.YELLOW)) print(colorize("="*60, Colors.YELLOW)) # Create optimized config optimized_config = { "agents": { "defaults": { "model": { "primary": "anthropic/claude-haiku-4-5" # Cheap default }, "cache": { "enabled": True, "ttl": "5m", "priority": "high" }, "models": { "anthropic/claude-sonnet-4-5": {"alias": "sonnet", "cache": True}, "anthropic/claude-haiku-4-5": {"alias": "haiku", "cache": False}, "anthropic/claude-opus-4-5": {"alias": "opus", "cache": True} } } }, "heartbeat": { "every": "1h", "model": "ollama/llama3.2:3b", # FREE local LLM "prompt": "Status check" }, "rate_limits": { "api_calls": {"min_interval_seconds": 5} }, "budgets": { "daily": 5.00, "monthly": 150.00, "warning_threshold": 0.75 } } # Save optimized config with open(openclaw_dir / 'openclaw.json', 'w') as f: json.dump(optimized_config, f, indent=2) print("\n Optimizations applied:") print(colorize(" [OK] Model routing: Haiku default (92% cheaper)", Colors.GREEN)) print(colorize(" [OK] Heartbeat: Ollama local (100% free)", Colors.GREEN)) print(colorize(" [OK] Prompt caching: Enabled (90% discount)", Colors.GREEN)) print(colorize(" [OK] Budget controls: $5/day, $150/month", Colors.GREEN)) # Create lean workspace files lean_soul = """# SOUL.md ## Core Principles - Efficiency first - Use Haiku for routine tasks - Sonnet only for complex reasoning """ lean_user = """# USER.md - Name: User - Preference: Concise responses """ # Remove bloated files, create lean ones (workspace_dir / 'MEMORY.md').unlink() # Don't auto-load with open(workspace_dir / 'SOUL.md', 'w') as f: f.write(lean_soul) with open(workspace_dir / 'USER.md', 'w') as f: f.write(lean_user) print(colorize(" [OK] Workspace: Reduced from 50KB to 2KB", Colors.GREEN)) print(colorize(" [OK] Memory: On-demand loading only", Colors.GREEN)) # ========== AFTER OPTIMIZATION ========== print(colorize("\n\n" + "="*60, Colors.GREEN)) print(colorize(" PHASE 3: AFTER OPTIMIZATION", Colors.BOLD + Colors.GREEN)) print(colorize("="*60, Colors.GREEN)) # Calculate new workspace size after_workspace_size = sum(f.stat().st_size for f in workspace_dir.iterdir() if f.is_file()) / 1024 print(f"\n Workspace size: {after_workspace_size:.1f}KB") print(f" Files loaded: SOUL.md, USER.md only (lean)") after_costs = calculate_costs(optimized_config, after_workspace_size) print_cost_report("AFTER OPTIMIZATION - COST ANALYSIS", after_costs, Colors.GREEN) # ========== COMPARISON ========== print(colorize("\n\n" + "="*60, Colors.BOLD + Colors.CYAN)) print(colorize(" SAVINGS SUMMARY", Colors.BOLD + Colors.CYAN)) print(colorize("="*60, Colors.BOLD + Colors.CYAN)) daily_savings = before_costs['daily_total'] - after_costs['daily_total'] monthly_savings = before_costs['monthly_total'] - after_costs['monthly_total'] yearly_savings = before_costs['yearly_total'] - after_costs['yearly_total'] savings_percent = (1 - after_costs['monthly_total'] / before_costs['monthly_total']) * 100 print(f"\n {'Metric':<20} {'Before':>12} {'After':>12} {'Savings':>12}") print(f" {'-'*56}") print(f" {'Daily Cost':<20} ${before_costs['daily_total']:>10.2f} ${after_costs['daily_total']:>10.2f} {colorize(f'${daily_savings:>10.2f}', Colors.GREEN)}") print(f" {'Monthly Cost':<20} ${before_costs['monthly_total']:>10.2f} ${after_costs['monthly_total']:>10.2f} {colorize(f'${monthly_savings:>10.2f}', Colors.GREEN)}") print(f" {'Yearly Cost':<20} ${before_costs['yearly_total']:>10.2f} ${after_costs['yearly_total']:>10.2f} {colorize(f'${yearly_savings:>10.2f}', Colors.GREEN)}") print(colorize(f"\n +-------------------------------------------------------+", Colors.BOLD + Colors.GREEN)) print(colorize(f" | TOTAL SAVINGS: {savings_percent:.0f}% |", Colors.BOLD + Colors.GREEN)) print(colorize(f" | ${monthly_savings:.2f}/month = ${yearly_savings:.2f}/year |", Colors.BOLD + Colors.GREEN)) print(colorize(f" +-------------------------------------------------------+", Colors.BOLD + Colors.GREEN)) # Breakdown print(colorize("\n Savings Breakdown:", Colors.BOLD)) model_savings = (COSTS['sonnet']['input'] - COSTS['haiku']['input']) / COSTS['sonnet']['input'] * 100 print(f" - Model Routing (Sonnet->Haiku): {model_savings:.0f}% per token") print(f" - Heartbeat (Paid->Ollama): 100% (free)") print(f" - Context Reduction (50KB->2KB): 96% less tokens") print(f" - Prompt Caching: 90% on repeated content") # Cleanup shutil.rmtree(test_dir) print(colorize("\n\n[OK] Simulation complete! Mock environment cleaned up.", Colors.GREEN)) print(colorize("\nTo apply these optimizations to your real OpenClaw setup:", Colors.CYAN)) print(" python src/optimizer.py --mode full\n") if __name__ == '__main__': run_simulation()