/ examples / context / 09_context_budgeting.py
09_context_budgeting.py
  1  #!/usr/bin/env python3
  2  """
  3  Context Budgeting Example
  4  
  5  Demonstrates how to use context budgeting with Agent via context= param,
  6  and also shows low-level ContextBudgeter 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(output_reserve=16000),  # Custom output reserve
 15      )
 16  """
 17  
 18  from praisonaiagents import Agent
 19  from praisonaiagents.context import (
 20      ManagerConfig,
 21      ContextBudgeter,
 22      BudgetAllocation,
 23      get_model_limit,
 24      get_output_reserve,
 25  )
 26  
 27  
 28  def agent_centric_example():
 29      """Agent-centric usage - recommended approach."""
 30      print("=" * 60)
 31      print("Agent-Centric Context Budgeting")
 32      print("=" * 60)
 33      
 34      # Simple: Enable context management with defaults
 35      agent = Agent(
 36          instructions="You are a helpful assistant.",
 37          context=True,
 38      )
 39      print(f"Agent with context=True created")
 40      print(f"Context manager: {agent.context_manager is not None}")
 41      
 42      # Custom: Specify output reserve
 43      agent2 = Agent(
 44          instructions="You are a code assistant.",
 45          context=ManagerConfig(output_reserve=16000),
 46      )
 47      print(f"Agent with custom output_reserve created")
 48      if agent2.context_manager:
 49          stats = agent2.context_manager.get_stats()
 50          print(f"Output reserve: {stats.get('output_reserve', 'N/A')}")
 51  
 52  
 53  def main():
 54      print("=" * 60)
 55      print("Context Budgeting Example")
 56      print("=" * 60)
 57      
 58      # Example 1: Basic budget allocation for GPT-4o-mini
 59      print("\n1. Basic Budget Allocation (gpt-4o-mini)")
 60      print("-" * 40)
 61      
 62      budgeter = ContextBudgeter(model="gpt-4o-mini")
 63      budget = budgeter.allocate()
 64      
 65      print(f"Model: gpt-4o-mini")
 66      print(f"Model Limit: {budget.model_limit:,} tokens")
 67      print(f"Output Reserve: {budget.output_reserve:,} tokens")
 68      print(f"Usable Context: {budget.usable:,} tokens")
 69      
 70      # Example 2: Different models have different limits
 71      print("\n2. Model Comparison")
 72      print("-" * 40)
 73      
 74      models = ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "claude-3-opus", "gemini-1.5-pro"]
 75      
 76      for model in models:
 77          limit = get_model_limit(model)
 78          reserve = get_output_reserve(model)
 79          print(f"{model:20s}: {limit:>10,} tokens (reserve: {reserve:,})")
 80      
 81      # Example 3: Custom segment budgets
 82      print("\n3. Custom Segment Budgets")
 83      print("-" * 40)
 84      
 85      budgeter = ContextBudgeter(
 86          model="gpt-4o",
 87          system_prompt_budget=2000,
 88          rules_budget=1000,
 89          skills_budget=500,
 90          memory_budget=5000,
 91          tools_schema_budget=3000,
 92      )
 93      budget = budgeter.allocate()
 94      
 95      print(f"System prompt budget: 2,000 tokens")
 96      print(f"Rules budget: 1,000 tokens")
 97      print(f"Skills budget: 500 tokens")
 98      print(f"Memory budget: 5,000 tokens")
 99      print(f"Tools schema budget: 3,000 tokens")
100      print(f"Total usable: {budget.usable:,} tokens")
101      
102      # Example 4: Check for overflow
103      print("\n4. Overflow Detection")
104      print("-" * 40)
105      
106      budgeter = ContextBudgeter(model="gpt-4o-mini")
107      
108      # Simulate different usage levels
109      test_cases = [
110          ("Low usage", 10000),
111          ("Medium usage", 80000),
112          ("High usage", 100000),
113          ("Near limit", 110000),
114          ("Over limit", 130000),
115      ]
116      
117      for name, tokens in test_cases:
118          overflow = budgeter.check_overflow(tokens)
119          utilization = budgeter.get_utilization(tokens)
120          remaining = budgeter.get_remaining(tokens)
121          status = "⚠️ OVERFLOW" if overflow else "✓ OK"
122          print(f"{name:15s}: {tokens:>7,} tokens | {utilization:>5.1%} | {status}")
123      
124      # Example 5: Budget to dict for serialization
125      print("\n5. Budget Serialization")
126      print("-" * 40)
127      
128      budget_dict = budgeter.to_dict()
129      print(f"Budget as dict: {budget_dict}")
130      
131      # Example 6: Overflow handling strategies overview
132      print("\n6. Overflow Handling Strategies")
133      print("-" * 40)
134      strategies = [
135          ("truncate", "Remove oldest messages first", "Fast, simple", "Loses early context"),
136          ("sliding_window", "Keep N most recent messages", "Preserves recent", "Loses early context"),
137          ("prune_tools", "Truncate old tool outputs", "Keeps messages", "May lose tool details"),
138          ("summarize", "Replace old messages with summary", "Preserves meaning", "Slower, uses API"),
139          ("smart", "Combine strategies intelligently", "Best balance", "More complex"),
140      ]
141      print(f"{'Strategy':<15} | {'Description':<30} | {'Pros':<18} | {'Cons'}")
142      print("-" * 90)
143      for name, desc, pros, cons in strategies:
144          print(f"{name:<15} | {desc:<30} | {pros:<18} | {cons}")
145      
146      # Example 7: Simulated overflow playbook
147      print("\n7. Overflow Playbook (Thresholds → Actions)")
148      print("-" * 40)
149      playbook = [
150          (0.70, "INFO", "Monitor usage, no action needed"),
151          (0.80, "NOTICE", "Consider optimization soon"),
152          (0.90, "WARNING", "Trigger auto-compact if enabled"),
153          (0.95, "CRITICAL", "Aggressive optimization required"),
154          (1.00, "OVERFLOW", "Immediate truncation to prevent API error"),
155      ]
156      print(f"{'Usage':<8} | {'Level':<10} | {'Action'}")
157      print("-" * 60)
158      for threshold, level, action in playbook:
159          print(f"{threshold:>6.0%}   | {level:<10} | {action}")
160      
161      # Example 8: Strategy comparison table
162      print("\n8. Strategy Comparison")
163      print("-" * 40)
164      print("Strategy         | Preserves      | Loses          | Best For")
165      print("-" * 70)
166      print("truncate         | Recent msgs    | Old msgs       | Simple chats")
167      print("sliding_window   | Last N msgs    | Earlier msgs   | Long conversations")
168      print("prune_tools      | All messages   | Tool details   | Tool-heavy agents")
169      print("summarize        | Context meaning| Exact wording  | Important history")
170      print("smart            | Balanced       | Minimal        | Production use")
171      
172      # Example 9: Default strategy for interactive mode
173      print("\n9. Default Strategy: Interactive Mode")
174      print("-" * 40)
175      print("Interactive mode default: context=False (zero overhead)")
176      print("To enable context management in interactive mode:")
177      print("  praisonai chat --context  # Enable with defaults")
178      print("  Or in code:")
179      print("    agent = Agent(instructions='...', context=True)")
180      print("")
181      print("When enabled, defaults are:")
182      print("  - auto_compact: True")
183      print("  - compact_threshold: 0.8 (80% usage)")
184      print("  - strategy: smart")
185      print("  - output_reserve: model-specific (8K-16K)")
186      
187      # Example 10: Default strategy for auto-agents mode
188      print("\n10. Default Strategy: Auto-Agents Mode")
189      print("-" * 40)
190      print("Auto-agents mode default: context=False (zero overhead)")
191      print("To enable for multi-agent workflows:")
192      print("    from praisonaiagents import AgentTeam")
193      print("    agents = AgentTeam(agents=[...], context=True)")
194      print("")
195      print("Recommended config for long tasks:")
196      print("    context=ManagerConfig(")
197      print("        auto_compact=True,")
198      print("        compact_threshold=0.7,  # Earlier trigger for safety")
199      print("        strategy='smart',")
200      print("        output_reserve=16384,   # Larger reserve for complex outputs")
201      print("    )")
202      
203      print("\n" + "=" * 60)
204      print("✓ Context budgeting examples complete!")
205      print("=" * 60)
206  
207  
208  if __name__ == "__main__":
209      main()