scripts-thinking-budget-manager

#!/usr/bin/env python3 """

title: "Thinking Budget Manager" component_type: script version: "1.0.0" audience: contributor status: stable summary: "Thinking Budget Manager" keywords: ['budget', 'manager', 'optimization', 'review', 'security'] tokens: ~500 created: 2025-12-22 updated: 2025-12-22 script_name: "thinking-budget-manager.py" language: python executable: true usage: "python3 scripts/thinking-budget-manager.py [options]" python_version: "3.10+" dependencies: [] modifies_files: false network_access: false requires_auth: false

Thinking Budget Manager

Manages thinking token budgets for AI agents based on task complexity, agent capabilities, and cost optimization goals. Tracks usage and provides cost estimates for different thinking levels.

Part of CODITECT Framework Enhancement Inspired by Auto-Claude patterns Created: 2025-12-22

Usage: # Get recommended budget for task python3 scripts/thinking-budget-manager.py --task code_review --complexity standard

# Show usage statistics
python3 scripts/thinking-budget-manager.py --stats

# Estimate cost for custom task
python3 scripts/thinking-budget-manager.py --estimate "architecture review for auth system"

# Get budget for specific agent
python3 scripts/thinking-budget-manager.py --agent orchestrator --task system_design

# Update usage tracking
python3 scripts/thinking-budget-manager.py --track --task code_review --tokens-used 8192

"""

import sys import json import argparse from pathlib import Path from datetime import datetime, timezone from typing import Dict, Any, Optional, List, Tuple from enum import Enum

class ComplexityLevel(Enum): """Task complexity levels""" SIMPLE = "simple" STANDARD = "standard" COMPLEX = "complex" CRITICAL = "critical"

class BudgetLevel(Enum): """Budget level tiers""" ULTRATHINK = "ultrathink" HIGH = "high" MEDIUM = "medium" LOW = "low" NONE = "none"

class ThinkingBudgetManager: """Manages thinking budgets for AI agents"""

def __init__(self, framework_root: Path):
    """
    Initialize thinking budget manager

    Args:
        framework_root: Root directory of CODITECT framework
    """
    self.framework_root = framework_root
    self.config_file = framework_root / "config" / "thinking-budgets.json"
    self.usage_file = framework_root / ".coditect" / "thinking-usage.json"

    # Load configuration
    self.config = self._load_config()
    self.budget_levels = self.config.get("budget_levels", {})
    self.task_defaults = self.config.get("task_budget_defaults", {})
    self.complexity_adjustments = self.config.get("complexity_adjustments", {})
    self.agent_overrides = self.config.get("agent_overrides", {})
    self.cost_config = self.config.get("cost_estimation", {})

    # Load usage tracking
    self.usage_data = self._load_usage()

def _load_config(self) -> Dict[str, Any]:
    """Load thinking budget configuration"""
    if not self.config_file.exists():
        print(f"❌ Error: Configuration file not found: {self.config_file}")
        sys.exit(1)

    with open(self.config_file, 'r') as f:
        return json.load(f)

def _load_usage(self) -> Dict[str, Any]:
    """Load usage tracking data"""
    if not self.usage_file.exists():
        return {
            "total_sessions": 0,
            "total_thinking_tokens": 0,
            "total_cost": 0.0,
            "by_level": {},
            "by_task": {},
            "by_agent": {},
            "sessions": []
        }

    with open(self.usage_file, 'r') as f:
        return json.load(f)

def _save_usage(self) -> None:
    """Save usage tracking data"""
    self.usage_file.parent.mkdir(parents=True, exist_ok=True)

    with open(self.usage_file, 'w') as f:
        json.dump(self.usage_data, f, indent=2)

def get_budget_tokens(self, level: str) -> Optional[int]:
    """
    Get token count for budget level

    Args:
        level: Budget level name

    Returns:
        Token count or None for "none" level
    """
    if level not in self.budget_levels:
        return None

    return self.budget_levels[level].get("tokens")

def upgrade_level(self, level: str) -> str:
    """
    Upgrade budget level by one tier

    Args:
        level: Current budget level

    Returns:
        Upgraded budget level
    """
    level_order = ["none", "low", "medium", "high", "ultrathink"]

    if level not in level_order:
        return level

    current_idx = level_order.index(level)
    if current_idx < len(level_order) - 1:
        return level_order[current_idx + 1]

    return level

def downgrade_level(self, level: str) -> str:
    """
    Downgrade budget level by one tier

    Args:
        level: Current budget level

    Returns:
        Downgraded budget level
    """
    level_order = ["none", "low", "medium", "high", "ultrathink"]

    if level not in level_order:
        return level

    current_idx = level_order.index(level)
    if current_idx > 0:
        return level_order[current_idx - 1]

    return level

def min_level(self, level1: str, level2: str) -> str:
    """
    Return the lower of two budget levels

    Args:
        level1: First budget level
        level2: Second budget level

    Returns:
        Lower budget level
    """
    level_order = ["none", "low", "medium", "high", "ultrathink"]

    if level1 not in level_order or level2 not in level_order:
        return level1

    idx1 = level_order.index(level1)
    idx2 = level_order.index(level2)

    return level1 if idx1 < idx2 else level2

def get_thinking_budget(
    self,
    task_type: str,
    complexity: ComplexityLevel = ComplexityLevel.STANDARD,
    agent_name: Optional[str] = None
) -> Tuple[str, Optional[int]]:
    """
    Select appropriate thinking budget based on context

    Args:
        task_type: Type of task to perform
        complexity: Task complexity level
        agent_name: Optional agent name for overrides

    Returns:
        Tuple of (budget_level_name, token_count)
    """
    # Start with task default or fallback to medium
    base_level = self.task_defaults.get(task_type, "medium")

    # Adjust for complexity
    complexity_adjustment = self.complexity_adjustments.get(complexity.value, 0)

    adjusted_level = base_level
    for _ in range(abs(complexity_adjustment)):
        if complexity_adjustment > 0:
            adjusted_level = self.upgrade_level(adjusted_level)
        else:
            adjusted_level = self.downgrade_level(adjusted_level)

    # Apply agent overrides if specified
    if agent_name and agent_name in self.agent_overrides:
        agent_config = self.agent_overrides[agent_name]

        # Respect agent's max level
        max_level = agent_config.get("max_level", "ultrathink")
        adjusted_level = self.min_level(adjusted_level, max_level)

        # Apply agent's minimum level
        min_level = agent_config.get("min_level", "none")
        level_order = ["none", "low", "medium", "high", "ultrathink"]

        if (adjusted_level in level_order and min_level in level_order and
            level_order.index(adjusted_level) < level_order.index(min_level)):
            adjusted_level = min_level

    # Get token count for final level
    token_count = self.get_budget_tokens(adjusted_level)

    return adjusted_level, token_count

def estimate_cost(
    self,
    thinking_tokens: int,
    output_tokens: int = 2000
) -> float:
    """
    Estimate cost for given token usage

    Args:
        thinking_tokens: Number of thinking tokens
        output_tokens: Estimated output tokens

    Returns:
        Estimated cost in USD
    """
    thinking_cost_per_1k = self.cost_config.get("base_cost_per_1k_thinking_tokens", 0.015)
    output_cost_per_1k = self.cost_config.get("base_cost_per_1k_output_tokens", 0.003)

    thinking_cost = (thinking_tokens / 1000) * thinking_cost_per_1k
    output_cost = (output_tokens / 1000) * output_cost_per_1k

    return thinking_cost + output_cost

def track_usage(
    self,
    task_type: str,
    budget_level: str,
    tokens_used: int,
    agent_name: Optional[str] = None
) -> None:
    """
    Track thinking token usage

    Args:
        task_type: Type of task performed
        budget_level: Budget level used
        tokens_used: Actual tokens consumed
        agent_name: Optional agent name
    """
    # Update totals
    self.usage_data["total_sessions"] += 1
    self.usage_data["total_thinking_tokens"] += tokens_used

    # Calculate cost
    cost = self.estimate_cost(tokens_used)
    self.usage_data["total_cost"] += cost

    # Track by level
    if budget_level not in self.usage_data["by_level"]:
        self.usage_data["by_level"][budget_level] = {
            "count": 0,
            "total_tokens": 0,
            "total_cost": 0.0
        }

    self.usage_data["by_level"][budget_level]["count"] += 1
    self.usage_data["by_level"][budget_level]["total_tokens"] += tokens_used
    self.usage_data["by_level"][budget_level]["total_cost"] += cost

    # Track by task
    if task_type not in self.usage_data["by_task"]:
        self.usage_data["by_task"][task_type] = {
            "count": 0,
            "total_tokens": 0,
            "total_cost": 0.0
        }

    self.usage_data["by_task"][task_type]["count"] += 1
    self.usage_data["by_task"][task_type]["total_tokens"] += tokens_used
    self.usage_data["by_task"][task_type]["total_cost"] += cost

    # Track by agent if provided
    if agent_name:
        if agent_name not in self.usage_data["by_agent"]:
            self.usage_data["by_agent"][agent_name] = {
                "count": 0,
                "total_tokens": 0,
                "total_cost": 0.0
            }

        self.usage_data["by_agent"][agent_name]["count"] += 1
        self.usage_data["by_agent"][agent_name]["total_tokens"] += tokens_used
        self.usage_data["by_agent"][agent_name]["total_cost"] += cost

    # Add session record
    session = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "task_type": task_type,
        "budget_level": budget_level,
        "tokens_used": tokens_used,
        "cost": cost,
        "agent_name": agent_name
    }

    self.usage_data["sessions"].append(session)

    # Keep only last 1000 sessions
    if len(self.usage_data["sessions"]) > 1000:
        self.usage_data["sessions"] = self.usage_data["sessions"][-1000:]

    # Save updated usage
    self._save_usage()

def get_stats(self) -> Dict[str, Any]:
    """
    Get usage statistics

    Returns:
        Dictionary of usage statistics
    """
    stats = {
        "total_sessions": self.usage_data.get("total_sessions", 0),
        "total_thinking_tokens": self.usage_data.get("total_thinking_tokens", 0),
        "total_cost": self.usage_data.get("total_cost", 0.0),
        "by_level": self.usage_data.get("by_level", {}),
        "by_task": self.usage_data.get("by_task", {}),
        "by_agent": self.usage_data.get("by_agent", {})
    }

    # Calculate averages
    if stats["total_sessions"] > 0:
        stats["avg_tokens_per_session"] = stats["total_thinking_tokens"] / stats["total_sessions"]
        stats["avg_cost_per_session"] = stats["total_cost"] / stats["total_sessions"]
    else:
        stats["avg_tokens_per_session"] = 0
        stats["avg_cost_per_session"] = 0.0

    return stats

def print_stats(self) -> None:
    """Print formatted usage statistics"""
    stats = self.get_stats()

    print("\n" + "=" * 70)
    print("THINKING BUDGET USAGE STATISTICS")
    print("=" * 70)

    print(f"\n📊 Overall Statistics:")
    print(f"   Total Sessions: {stats['total_sessions']:,}")
    print(f"   Total Thinking Tokens: {stats['total_thinking_tokens']:,}")
    print(f"   Total Cost: ${stats['total_cost']:.4f}")
    print(f"   Avg Tokens/Session: {stats['avg_tokens_per_session']:,.0f}")
    print(f"   Avg Cost/Session: ${stats['avg_cost_per_session']:.4f}")

    if stats['by_level']:
        print(f"\n📈 By Budget Level:")
        for level, data in sorted(stats['by_level'].items()):
            print(f"   {level:12} - {data['count']:4} sessions, "
                  f"{data['total_tokens']:8,} tokens, ${data['total_cost']:.4f}")

    if stats['by_task']:
        print(f"\n🎯 Top Tasks by Usage:")
        sorted_tasks = sorted(
            stats['by_task'].items(),
            key=lambda x: x[1]['total_tokens'],
            reverse=True
        )
        for task, data in sorted_tasks[:10]:
            print(f"   {task:25} - {data['count']:4} sessions, "
                  f"{data['total_tokens']:8,} tokens, ${data['total_cost']:.4f}")

    if stats['by_agent']:
        print(f"\n🤖 By Agent:")
        sorted_agents = sorted(
            stats['by_agent'].items(),
            key=lambda x: x[1]['total_tokens'],
            reverse=True
        )
        for agent, data in sorted_agents:
            print(f"   {agent:25} - {data['count']:4} sessions, "
                  f"{data['total_tokens']:8,} tokens, ${data['total_cost']:.4f}")

    print("\n" + "=" * 70 + "\n")

def estimate_task_cost(self, description: str) -> Dict[str, Any]:
    """
    Estimate cost for a task based on description

    Args:
        description: Task description

    Returns:
        Cost estimates for different budget levels
    """
    # Simple keyword-based task classification
    description_lower = description.lower()

    # Determine task type
    task_type = "feature_implementation"  # Default

    if any(word in description_lower for word in ["architecture", "design", "system", "plan"]):
        task_type = "architecture_review"
    elif any(word in description_lower for word in ["review", "qa", "quality", "audit"]):
        task_type = "code_review"
    elif any(word in description_lower for word in ["security", "vulnerability", "penetration"]):
        task_type = "security_audit"
    elif any(word in description_lower for word in ["bug", "fix", "issue"]):
        task_type = "bug_fix"
    elif any(word in description_lower for word in ["format", "style", "lint"]):
        task_type = "formatting"
    elif any(word in description_lower for word in ["doc", "documentation", "readme"]):
        task_type = "documentation"

    # Determine complexity
    complexity = ComplexityLevel.STANDARD

    if any(word in description_lower for word in ["complex", "comprehensive", "multi", "large"]):
        complexity = ComplexityLevel.COMPLEX
    elif any(word in description_lower for word in ["critical", "production", "security"]):
        complexity = ComplexityLevel.CRITICAL
    elif any(word in description_lower for word in ["simple", "quick", "small", "minor"]):
        complexity = ComplexityLevel.SIMPLE

    # Get recommended budget
    level, tokens = self.get_thinking_budget(task_type, complexity)

    # Calculate estimates for all levels
    estimates = {}
    for level_name, level_config in self.budget_levels.items():
        token_count = level_config.get("tokens")
        if token_count:
            cost = self.estimate_cost(token_count)
            estimates[level_name] = {
                "tokens": token_count,
                "cost": cost,
                "description": level_config.get("description", "")
            }

    return {
        "task_description": description,
        "detected_task_type": task_type,
        "detected_complexity": complexity.value,
        "recommended_level": level,
        "recommended_tokens": tokens,
        "recommended_cost": self.estimate_cost(tokens) if tokens else 0.0,
        "all_level_estimates": estimates
    }

def main(): """Main entry point""" parser = argparse.ArgumentParser( description="Thinking Budget Manager for CODITECT Framework", formatter_class=argparse.RawDescriptionHelpFormatter )

# Task recommendation options
parser.add_argument(
    "--task",
    help="Task type to get budget for"
)
parser.add_argument(
    "--complexity",
    choices=["simple", "standard", "complex", "critical"],
    default="standard",
    help="Task complexity level"
)
parser.add_argument(
    "--agent",
    help="Agent name for override rules"
)

# Usage tracking options
parser.add_argument(
    "--track",
    action="store_true",
    help="Track thinking token usage"
)
parser.add_argument(
    "--tokens-used",
    type=int,
    help="Actual tokens used (for tracking)"
)
parser.add_argument(
    "--budget-level",
    help="Budget level used (for tracking)"
)

# Statistics options
parser.add_argument(
    "--stats",
    action="store_true",
    help="Show usage statistics"
)

# Cost estimation
parser.add_argument(
    "--estimate",
    help="Estimate cost for task description"
)

# List options
parser.add_argument(
    "--list-tasks",
    action="store_true",
    help="List all task types"
)
parser.add_argument(
    "--list-levels",
    action="store_true",
    help="List all budget levels"
)

# Output format
parser.add_argument(
    "--json",
    action="store_true",
    help="Output in JSON format"
)

args = parser.parse_args()

# Find framework root
script_dir = Path(__file__).parent
framework_root = script_dir.parent

# Initialize manager
manager = ThinkingBudgetManager(framework_root)

# Handle different operations
if args.stats:
    if args.json:
        stats = manager.get_stats()
        print(json.dumps(stats, indent=2))
    else:
        manager.print_stats()
    return

if args.list_tasks:
    print("\n📋 Available Task Types:")
    for task_type, level in sorted(manager.task_defaults.items()):
        print(f"   {task_type:30} → {level}")
    return

if args.list_levels:
    print("\n📊 Budget Levels:")
    for level_name, level_config in manager.budget_levels.items():
        tokens = level_config.get("tokens")
        desc = level_config.get("description", "")
        token_str = f"{tokens:,}" if tokens else "none"
        print(f"\n   {level_name.upper():12} ({token_str:8} tokens)")
        print(f"   {desc}")
    return

if args.estimate:
    estimates = manager.estimate_task_cost(args.estimate)

    if args.json:
        print(json.dumps(estimates, indent=2))
    else:
        print(f"\n💡 Task Estimation")
        print(f"   Description: {estimates['task_description']}")
        print(f"   Detected Type: {estimates['detected_task_type']}")
        print(f"   Detected Complexity: {estimates['detected_complexity']}")
        print(f"\n✅ Recommended:")
        print(f"   Level: {estimates['recommended_level']}")
        print(f"   Tokens: {estimates['recommended_tokens']:,}" if estimates['recommended_tokens'] else "   Tokens: none")
        print(f"   Est. Cost: ${estimates['recommended_cost']:.4f}")

        print(f"\n📊 All Levels:")
        for level, data in estimates['all_level_estimates'].items():
            print(f"   {level:12} - {data['tokens']:8,} tokens, ${data['cost']:.4f}")
    return

if args.track:
    if not args.task or not args.tokens_used or not args.budget_level:
        print("❌ Error: --track requires --task, --tokens-used, and --budget-level")
        sys.exit(1)

    manager.track_usage(
        args.task,
        args.budget_level,
        args.tokens_used,
        args.agent
    )

    print(f"✅ Usage tracked: {args.task} ({args.budget_level}) - {args.tokens_used:,} tokens")
    return

if args.task:
    complexity = ComplexityLevel(args.complexity)
    level, tokens = manager.get_thinking_budget(args.task, complexity, args.agent)

    if args.json:
        result = {
            "task_type": args.task,
            "complexity": args.complexity,
            "agent": args.agent,
            "budget_level": level,
            "tokens": tokens,
            "estimated_cost": manager.estimate_cost(tokens) if tokens else 0.0
        }
        print(json.dumps(result, indent=2))
    else:
        print(f"\n💡 Thinking Budget Recommendation")
        print(f"   Task: {args.task}")
        print(f"   Complexity: {args.complexity}")
        if args.agent:
            print(f"   Agent: {args.agent}")
        print(f"\n✅ Recommended:")
        print(f"   Budget Level: {level}")
        print(f"   Tokens: {tokens:,}" if tokens else "   Tokens: none")
        if tokens:
            cost = manager.estimate_cost(tokens)
            print(f"   Estimated Cost: ${cost:.4f}")
    return

# No operation specified
parser.print_help()

if name == "main": main()