from praisonaiagents import Agent
from praisonaiagents.context import ManagerConfig

# Configure output reserve via context param
agent = Agent(
    instructions="You are helpful.",
    context=ManagerConfig(
        output_reserve=16000,  # Reserve 16k tokens for output
    ),
)

# Access budget info
budget = agent.context_manager.get_budget()
print(f"Usable context: {budget.usable:,} tokens")

from praisonaiagents.context import ContextBudgeter, get_model_limit

# Create budgeter for your model
budgeter = ContextBudgeter(model="gpt-4o-mini")
budget = budgeter.allocate()

print(f"Model limit: {budget.model_limit:,} tokens")
print(f"Output reserve: {budget.output_reserve:,} tokens")
print(f"Usable context: {budget.usable:,} tokens")

from praisonaiagents.context import get_model_limit, get_output_reserve

limit = get_model_limit("gpt-4o-mini")  # 128000
reserve = get_output_reserve("gpt-4o-mini")  # 16384

budgeter = ContextBudgeter(
    model="gpt-4o",
    system_prompt_budget=3000,
    rules_budget=1000,
    skills_budget=500,
    memory_budget=5000,
    tools_schema_budget=3000,
    tool_outputs_budget=30000,
    buffer_budget=2000,
)
budget = budgeter.allocate()

budgeter = ContextBudgeter(model="gpt-4o-mini")

# Check if current usage exceeds budget
current_tokens = 100000
is_overflow = budgeter.check_overflow(current_tokens)

# Get utilization percentage
utilization = budgeter.get_utilization(current_tokens)
print(f"Utilization: {utilization:.1%}")

# Get remaining capacity
remaining = budgeter.get_remaining(current_tokens)
print(f"Remaining: {remaining:,} tokens")

budgeter = ContextBudgeter(model="gpt-4o-mini")
budget = budgeter.allocate()

# Trigger optimization at 80% utilization
threshold = 0.8
trigger_at = int(budget.usable * threshold)

current_tokens = 95000
if current_tokens > trigger_at:
    print("Time to optimize!")

# Set output reserve
praisonai chat --context-output-reserve 10000

# Set optimization threshold
praisonai chat --context-threshold 0.8

PRAISONAI_CONTEXT_OUTPUT_RESERVE=8000
PRAISONAI_CONTEXT_THRESHOLD=0.8

budgeter = ContextBudgeter(model="gpt-4o-mini")
budget_dict = budgeter.to_dict()

# Returns:
# {
#     'model': 'gpt-4o-mini',
#     'model_limit': 128000,
#     'output_reserve': 16384,
#     'usable': 111616,
#     'allocation': {...}
# }