Skip to content

Commit 2e05f2c

Browse files
praisonai-triage-agent[bot]CopilotMervinPraison
authored
fix: resolve BudgetExceededError CLI traceback issues (fixes #1627) (#1635)
* fix: resolve BudgetExceededError CLI traceback issues (fixes #1627) - Add clean error handling for BudgetExceededError in CLI * Wrap agent execution calls with graceful error handling * Show user-friendly error messages instead of raw tracebacks * Exit with code 1 and provide clear remediation guidance - Add max_budget parameter to Agent constructor * Accept max_budget directly as convenience parameter * Automatically create/update ExecutionConfig with budget setting * Maintain full backward compatibility * Add proper parameter documentation Both fixes follow architecture guidelines and provide intuitive UX. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: praisonai-triage-agent[bot] <praisonai-triage-agent[bot]@users.noreply.github.com> * fix: address review feedback for budget handling and import safety Agent-Logs-Url: https://github.com/MervinPraison/PraisonAI/sessions/16ae1b3d-1c3c-4051-9465-a718a6746a90 Co-authored-by: MervinPraison <454862+MervinPraison@users.noreply.github.com> * fix: resolve print shadowing and simplify budget error message - Import rich print as rich_print to avoid shadowing builtin - Consolidate multi-line error output into single actionable message - Update guidance to mention new max_budget convenience parameter 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com> * fix: resolve print shadowing and simplify budget error message - Remove dangerous ImportError fallback in budget error handling - Use dataclasses.replace() instead of mutating shared ExecutionConfig - Simplify budget error to single actionable message - Address CodeRabbit, Greptile, and Copilot review feedback Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com> --------- Co-authored-by: praisonai-triage-agent[bot] <272766704+praisonai-triage-agent[bot]@users.noreply.github.com> Co-authored-by: praisonai-triage-agent[bot] <praisonai-triage-agent[bot]@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: MervinPraison <454862+MervinPraison@users.noreply.github.com> Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>
1 parent 4985415 commit 2e05f2c

2 files changed

Lines changed: 101 additions & 29 deletions

File tree

  • src
    • praisonai-agents/praisonaiagents/agent
    • praisonai/praisonai/cli

‎src/praisonai-agents/praisonaiagents/agent/agent.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,7 @@ def __init__(
544544
verification_hooks: Optional[List[Any]] = None, # Deprecated: use autonomy=AutonomyConfig(verification_hooks=[...])
545545
output: Optional[Union[bool, str, Dict[str, Any], 'OutputConfig']] = None,
546546
execution: Optional[Union[bool, str, Dict[str, Any], 'ExecutionConfig']] = None,
547+
max_budget: Optional[float] = None, # Budget limit in USD - convenience alias for ExecutionConfig(max_budget=...)
547548
templates: Optional[Union[Dict[str, Any], 'TemplateConfig']] = None,
548549
caching: Optional[Union[bool, str, Dict[str, Any], 'CachingConfig']] = None,
549550
hooks: Optional[Union[List[Any], Dict[str, Any], 'HooksConfig']] = None,
@@ -622,6 +623,9 @@ def __init__(
622623
- Dict[str, Any]: Config overrides (e.g. {"max_iter": 10, "max_rpm": 60})
623624
- ExecutionConfig: Custom configuration
624625
Controls: max_iter, max_rpm, max_execution_time, max_retry_limit
626+
max_budget: Budget limit in USD (convenience alias). Creates ExecutionConfig(max_budget=value).
627+
If both execution.max_budget and max_budget are provided, max_budget takes precedence.
628+
Use execution=ExecutionConfig(...) for full execution control.
625629
templates: Template configuration. Accepts:
626630
- Dict[str, Any]: Template fields (e.g. {"system": "...", "prompt": "..."})
627631
- TemplateConfig: Custom configuration
@@ -733,6 +737,42 @@ def __init__(
733737
web = apply_config_defaults("web", web, WebConfig)
734738
if output is None:
735739
output = apply_config_defaults("output", output, OutputConfig)
740+
741+
# Handle max_budget convenience parameter
742+
if max_budget is not None:
743+
from ..config.feature_configs import ExecutionConfig, resolve_execution
744+
if execution is None:
745+
execution = ExecutionConfig(max_budget=max_budget)
746+
else:
747+
# If execution config is already provided, merge max_budget into it
748+
resolved_exec = resolve_execution(execution)
749+
if resolved_exec is None:
750+
execution = ExecutionConfig(max_budget=max_budget)
751+
else:
752+
# Update existing config with max_budget
753+
if hasattr(resolved_exec, 'max_budget'):
754+
if (
755+
resolved_exec.max_budget is not None
756+
and resolved_exec.max_budget != max_budget
757+
):
758+
import warnings
759+
warnings.warn(
760+
(
761+
f"Both execution.max_budget={resolved_exec.max_budget} and "
762+
f"max_budget={max_budget} were provided; "
763+
"using max_budget."
764+
),
765+
UserWarning,
766+
stacklevel=3,
767+
)
768+
# Use dataclasses.replace to avoid mutating shared state
769+
import dataclasses
770+
execution = dataclasses.replace(resolved_exec, max_budget=max_budget)
771+
elif not hasattr(resolved_exec, 'max_budget'):
772+
# Fallback for older config objects
773+
import dataclasses
774+
execution = dataclasses.replace(resolved_exec, max_budget=max_budget)
775+
736776
if execution is None:
737777
execution = apply_config_defaults("execution", execution, ExecutionConfig)
738778
if caching is None:

‎src/praisonai/praisonai/cli/main.py‎

Lines changed: 61 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -4065,6 +4065,38 @@ def _run_inline_workflow(self, initial_prompt):
40654065

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

4068+
def _execute_agent_with_budget_handling(self, agent, method_name, *args, **kwargs):
4069+
"""
4070+
Execute an agent method with graceful BudgetExceededError handling.
4071+
4072+
Args:
4073+
agent: The agent instance
4074+
method_name: Name of the method to call ('start' or 'chat')
4075+
*args, **kwargs: Arguments to pass to the method
4076+
4077+
Returns:
4078+
The result of the agent method call
4079+
4080+
Raises:
4081+
SystemExit: On BudgetExceededError with clean error message
4082+
Exception: Re-raises any other exceptions
4083+
"""
4084+
try:
4085+
method = getattr(agent, method_name)
4086+
return method(*args, **kwargs)
4087+
except Exception as e:
4088+
# Handle BudgetExceededError gracefully
4089+
from praisonaiagents.errors import BudgetExceededError
4090+
4091+
if isinstance(e, BudgetExceededError):
4092+
from rich import print as rich_print
4093+
# Single-line error message with actionable guidance
4094+
rich_print(f"[red]Budget limit exceeded: {e!s} - Set max_budget parameter (e.g., Agent(max_budget=1.00))[/red]")
4095+
sys.exit(1)
4096+
else:
4097+
# Re-raise other exceptions
4098+
raise
4099+
40684100
def _extract_cli_config_for_yaml(self):
40694101
"""
40704102
Extract CLI configuration that should be passed to YAML processing.
@@ -4611,26 +4643,26 @@ def level_based_approve(function_name, arguments, risk_level):
46114643
from rich.panel import Panel
46124644

46134645
with Live(Panel(Spinner("dots", text="Generating..."), border_style="cyan"), refresh_per_second=10, transient=True):
4614-
result = auto_rag.chat(prompt)
4646+
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
46154647
else:
4616-
result = auto_rag.chat(prompt)
4648+
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
46174649
else:
46184650
# Resolve display mode from CLI flags
46194651
display_mode = self._resolve_display_mode()
46204652

46214653
if display_mode == 'silent':
46224654
# -qq: No output at all, exit code only
46234655
if hasattr(agent, 'start'):
4624-
result = agent.start(prompt)
4656+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46254657
else:
4626-
result = agent.chat(prompt)
4658+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46274659

46284660
elif display_mode == 'quiet':
46294661
# -q: Result only, no spinners or status
46304662
if hasattr(agent, 'start'):
4631-
result = agent.start(prompt)
4663+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46324664
else:
4633-
result = agent.chat(prompt)
4665+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46344666
if result is not None:
46354667
output = getattr(result, 'output', None) or (str(result) if result else None)
46364668
if output:
@@ -4642,31 +4674,31 @@ def level_based_approve(function_name, arguments, risk_level):
46424674
from praisonaiagents.output.status import enable_status_output, disable_status_output
46434675
enable_status_output(show_timestamps=True, show_metrics=True)
46444676
if hasattr(agent, 'start'):
4645-
result = agent.start(prompt)
4677+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46464678
else:
4647-
result = agent.chat(prompt)
4679+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46484680
disable_status_output()
46494681
except ImportError:
46504682
if hasattr(agent, 'start'):
4651-
result = agent.start(prompt)
4683+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46524684
else:
4653-
result = agent.chat(prompt)
4685+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46544686

46554687
elif display_mode == 'debug':
46564688
# -vv: SDK TraceOutput with markdown rendering
46574689
try:
46584690
from praisonaiagents.output.trace import enable_trace_output, disable_trace_output
46594691
enable_trace_output(use_markdown=True)
46604692
if hasattr(agent, 'start'):
4661-
result = agent.start(prompt)
4693+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46624694
else:
4663-
result = agent.chat(prompt)
4695+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46644696
disable_trace_output()
46654697
except ImportError:
46664698
if hasattr(agent, 'start'):
4667-
result = agent.start(prompt)
4699+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46684700
else:
4669-
result = agent.chat(prompt)
4701+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46704702

46714703
elif display_mode == 'jsonl':
46724704
# --output jsonl: JSONL structured output for CI/CD
@@ -4688,9 +4720,9 @@ def level_based_approve(function_name, arguments, risk_level):
46884720

46894721
start_time = time.time()
46904722
if hasattr(agent, 'start'):
4691-
result = agent.start(prompt)
4723+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46924724
else:
4693-
result = agent.chat(prompt)
4725+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46944726

46954727
# Emit final result
46964728
reason = getattr(result, 'completion_reason', None) if hasattr(result, 'completion_reason') else 'complete'
@@ -4709,9 +4741,9 @@ def level_based_approve(function_name, arguments, risk_level):
47094741
import json as json_mod
47104742
start_time = time.time()
47114743
if hasattr(agent, 'start'):
4712-
result = agent.start(prompt)
4744+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47134745
else:
4714-
result = agent.chat(prompt)
4746+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47154747

47164748
output = result.output if hasattr(result, 'output') else str(result)
47174749
envelope = {
@@ -4732,15 +4764,15 @@ def level_based_approve(function_name, arguments, risk_level):
47324764
flow = track_workflow()
47334765
flow.start()
47344766
if hasattr(agent, 'start'):
4735-
result = agent.start(prompt)
4767+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47364768
else:
4737-
result = agent.chat(prompt)
4769+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47384770
flow.stop()
47394771
except ImportError:
47404772
if hasattr(agent, 'start'):
4741-
result = agent.start(prompt)
4773+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47424774
else:
4743-
result = agent.chat(prompt)
4775+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47444776

47454777
elif display_mode == 'editor':
47464778
# --output editor: User-friendly step-by-step format
@@ -4750,9 +4782,9 @@ def level_based_approve(function_name, arguments, risk_level):
47504782

47514783
# Run agent
47524784
if hasattr(agent, 'start'):
4753-
result = agent.start(prompt)
4785+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47544786
else:
4755-
result = agent.chat(prompt)
4787+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47564788

47574789
# SDK callbacks (interaction, llm_content) handle display —
47584790
# no explicit editor.output() needed here.
@@ -4774,15 +4806,15 @@ def level_based_approve(function_name, arguments, risk_level):
47744806
from praisonaiagents.output.status import enable_status_output, disable_status_output
47754807
enable_status_output(show_timestamps=False, show_metrics=False)
47764808
if hasattr(agent, 'start'):
4777-
result = agent.start(prompt)
4809+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47784810
else:
4779-
result = agent.chat(prompt)
4811+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47804812
disable_status_output()
47814813
except ImportError:
47824814
if hasattr(agent, 'start'):
4783-
result = agent.start(prompt)
4815+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47844816
else:
4785-
result = agent.chat(prompt)
4817+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47864818

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

@@ -4860,7 +4892,7 @@ def level_based_approve(function_name, arguments, risk_level):
48604892
Now, {final_instruction.lower()}:"""
48614893

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

48664898
# Save output if --save is enabled

0 commit comments

Comments
 (0)