Skip to content

fix: resolve BudgetExceededError CLI traceback issues (fixes #1627) #1635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MervinPraison merged 4 commits into from
May 9, 2026
Merged

fix: resolve BudgetExceededError CLI traceback issues (fixes #1627) #1635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9e146ad
fix: resolve BudgetExceededError CLI traceback issues (fixes #1627)
praisonai-triage-agent[bot] May 7, 2026
f4ff710
fix: address review feedback for budget handling and import safety
Copilot May 7, 2026
5ccc652
fix: resolve print shadowing and simplify budget error message
praisonai-triage-agent[bot] May 7, 2026
3bf7ec4
fix: resolve print shadowing and simplify budget error message
praisonai-triage-agent[bot] May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions src/praisonai-agents/praisonaiagents/agent/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -543,6 +543,7 @@ def __init__(
verification_hooks: Optional[List[Any]] = None, # Deprecated: use autonomy=AutonomyConfig(verification_hooks=[...])
output: Optional[Union[bool, str, Dict[str, Any], 'OutputConfig']] = None,
execution: Optional[Union[bool, str, Dict[str, Any], 'ExecutionConfig']] = None,
max_budget: Optional[float] = None, # Budget limit in USD - convenience alias for ExecutionConfig(max_budget=...)
templates: Optional[Union[Dict[str, Any], 'TemplateConfig']] = None,
caching: Optional[Union[bool, str, Dict[str, Any], 'CachingConfig']] = None,
hooks: Optional[Union[List[Any], Dict[str, Any], 'HooksConfig']] = None,
Expand Down Expand Up @@ -621,6 +622,9 @@ def __init__(
- Dict[str, Any]: Config overrides (e.g. {"max_iter": 10, "max_rpm": 60})
- ExecutionConfig: Custom configuration
Controls: max_iter, max_rpm, max_execution_time, max_retry_limit
max_budget: Budget limit in USD (convenience alias). Creates ExecutionConfig(max_budget=value).
If both execution.max_budget and max_budget are provided, max_budget takes precedence.
Use execution=ExecutionConfig(...) for full execution control.
templates: Template configuration. Accepts:
- Dict[str, Any]: Template fields (e.g. {"system": "...", "prompt": "..."})
- TemplateConfig: Custom configuration
Expand Down Expand Up @@ -732,6 +736,42 @@ def __init__(
web = apply_config_defaults("web", web, WebConfig)
if output is None:
output = apply_config_defaults("output", output, OutputConfig)

# Handle max_budget convenience parameter
if max_budget is not None:
from ..config.feature_configs import ExecutionConfig, resolve_execution
if execution is None:
execution = ExecutionConfig(max_budget=max_budget)
else:
# If execution config is already provided, merge max_budget into it
resolved_exec = resolve_execution(execution)
if resolved_exec is None:
execution = ExecutionConfig(max_budget=max_budget)
else:
# Update existing config with max_budget
if hasattr(resolved_exec, 'max_budget'):
if (
resolved_exec.max_budget is not None
and resolved_exec.max_budget != max_budget
):
import warnings
warnings.warn(
(
f"Both execution.max_budget={resolved_exec.max_budget} and "
f"max_budget={max_budget} were provided; "
"using max_budget."
),
UserWarning,
stacklevel=3,
)
# Use dataclasses.replace to avoid mutating shared state
import dataclasses
execution = dataclasses.replace(resolved_exec, max_budget=max_budget)
elif not hasattr(resolved_exec, 'max_budget'):
# Fallback for older config objects
import dataclasses
execution = dataclasses.replace(resolved_exec, max_budget=max_budget)

if execution is None:
execution = apply_config_defaults("execution", execution, ExecutionConfig)
if caching is None:
Expand Down
90 changes: 61 additions & 29 deletions src/praisonai/praisonai/cli/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4063,6 +4063,38 @@ def _run_inline_workflow(self, initial_prompt):

return results[-1].get("output", "") if results else ""

def _execute_agent_with_budget_handling(self, agent, method_name, *args, **kwargs):
"""
Execute an agent method with graceful BudgetExceededError handling.

Args:
agent: The agent instance
method_name: Name of the method to call ('start' or 'chat')
*args, **kwargs: Arguments to pass to the method

Returns:
The result of the agent method call

Raises:
SystemExit: On BudgetExceededError with clean error message
Exception: Re-raises any other exceptions
"""
try:
method = getattr(agent, method_name)
return method(*args, **kwargs)
except Exception as e:
# Handle BudgetExceededError gracefully
from praisonaiagents.errors import BudgetExceededError

if isinstance(e, BudgetExceededError):
from rich import print as rich_print
# Single-line error message with actionable guidance
rich_print(f"[red]Budget limit exceeded: {e!s} - Set max_budget parameter (e.g., Agent(max_budget=1.00))[/red]")
sys.exit(1)
else:
# Re-raise other exceptions
raise

def _extract_cli_config_for_yaml(self):
"""
Extract CLI configuration that should be passed to YAML processing.
Expand Down Expand Up @@ -4609,26 +4641,26 @@ def level_based_approve(function_name, arguments, risk_level):
from rich.panel import Panel

with Live(Panel(Spinner("dots", text="Generating..."), border_style="cyan"), refresh_per_second=10, transient=True):
result = auto_rag.chat(prompt)
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
else:
result = auto_rag.chat(prompt)
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
else:
# Resolve display mode from CLI flags
display_mode = self._resolve_display_mode()

if display_mode == 'silent':
# -qq: No output at all, exit code only
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'quiet':
# -q: Result only, no spinners or status
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
if result is not None:
output = getattr(result, 'output', None) or (str(result) if result else None)
if output:
Expand All @@ -4640,31 +4672,31 @@ def level_based_approve(function_name, arguments, risk_level):
from praisonaiagents.output.status import enable_status_output, disable_status_output
enable_status_output(show_timestamps=True, show_metrics=True)
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
disable_status_output()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'debug':
# -vv: SDK TraceOutput with markdown rendering
try:
from praisonaiagents.output.trace import enable_trace_output, disable_trace_output
enable_trace_output(use_markdown=True)
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
disable_trace_output()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'jsonl':
# --output jsonl: JSONL structured output for CI/CD
Expand All @@ -4686,9 +4718,9 @@ def level_based_approve(function_name, arguments, risk_level):

start_time = time.time()
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

# Emit final result
reason = getattr(result, 'completion_reason', None) if hasattr(result, 'completion_reason') else 'complete'
Expand All @@ -4707,9 +4739,9 @@ def level_based_approve(function_name, arguments, risk_level):
import json as json_mod
start_time = time.time()
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

output = result.output if hasattr(result, 'output') else str(result)
envelope = {
Expand All @@ -4730,15 +4762,15 @@ def level_based_approve(function_name, arguments, risk_level):
flow = track_workflow()
flow.start()
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
flow.stop()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

elif display_mode == 'editor':
# --output editor: User-friendly step-by-step format
Expand All @@ -4748,9 +4780,9 @@ def level_based_approve(function_name, arguments, risk_level):

# Run agent
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

# SDK callbacks (interaction, llm_content) handle display —
# no explicit editor.output() needed here.
Expand All @@ -4772,15 +4804,15 @@ def level_based_approve(function_name, arguments, risk_level):
from praisonaiagents.output.status import enable_status_output, disable_status_output
enable_status_output(show_timestamps=False, show_metrics=False)
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
disable_status_output()
except ImportError:
if hasattr(agent, 'start'):
result = agent.start(prompt)
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
else:
result = agent.chat(prompt)
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)

# ===== POST-PROCESSING WITH NEW FEATURES =====

Expand Down Expand Up @@ -4858,7 +4890,7 @@ def level_based_approve(function_name, arguments, risk_level):
Now, {final_instruction.lower()}:"""

final_agent = PraisonAgent(**final_agent_config)
result = final_agent.start(final_prompt)
result = self._execute_agent_with_budget_handling(final_agent, 'start', final_prompt)
print(f"\n[bold green]✅ Final agent processing complete[/bold green]\n")

# Save output if --save is enabled
Expand Down