Skip to content

Commit f29326a

Browse files
author
praisonai-bot
committed
fix(cli): graceful BudgetExceededError handling (wrapper-only, no Agent param bloat)
Re-implements the CLI UX half of reverted PR #1635 without the core SDK changes that violated AGENTS.md. ## What changed - Add _execute_agent_with_budget_handling(agent, method_name, *args, **kwargs) helper on PraisonAIApp: calls agent.<method_name>(...) and catches BudgetExceededError to print a single-line actionable CLI message and exit(1) instead of leaking a raw traceback. - Route all 17 agent.start(prompt) / agent.chat(prompt) / auto_rag.chat(prompt) call sites in the direct-prompt dispatch through the helper. Covers every display mode: silent, quiet, verbose, debug, jsonl, json, flow, editor, default. ## What is NOT changed (intentionally) - Zero edits to praisonaiagents (core SDK). - NO new 'max_budget=' parameter on Agent.__init__. - Users continue to configure budgets via the already-consolidated execution=ExecutionConfig(max_budget=...) pattern per AGENTS.md §5.3 'Parameter Consolidation'. ## Why wrapper-only AGENTS.md §4.1 reserves core SDK for protocols/hooks/adapters — CLI UX belongs in the wrapper. AGENTS.md §5.3 explicitly consolidates execution- related knobs (max_iter, max_rpm, max_budget, rate_limiter, ...) into ExecutionConfig precisely to avoid adding more top-level Agent params. ## Error message Points users to the correct API: Budget limit exceeded: <reason>. Hint: set budget via execution=ExecutionConfig(max_budget=1.00) on your Agent. Closes #1627
1 parent b2611e5 commit f29326a

1 file changed

Lines changed: 53 additions & 28 deletions

File tree

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

Lines changed: 53 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4065,6 +4065,31 @@ 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+
"""Run ``agent.<method_name>(*args, **kwargs)`` with a graceful
4070+
BudgetExceededError handler.
4071+
4072+
Wrapper-only fix (no core SDK changes). Users configure budgets via
4073+
``execution=ExecutionConfig(max_budget=...)`` on the Agent — per
4074+
AGENTS.md §5.3 there is NO top-level ``max_budget=`` parameter on
4075+
Agent.__init__ (avoids parameter bloat).
4076+
4077+
When the budget is hit this prints a single-line actionable error
4078+
message and exits with code 1 instead of leaking a raw traceback.
4079+
Any other exception is re-raised unchanged.
4080+
"""
4081+
from praisonaiagents.errors import BudgetExceededError
4082+
try:
4083+
return getattr(agent, method_name)(*args, **kwargs)
4084+
except BudgetExceededError as e:
4085+
from rich import print as rich_print
4086+
rich_print(
4087+
f"[red]Budget limit exceeded: {e!s}. "
4088+
"Hint: set budget via "
4089+
"execution=ExecutionConfig(max_budget=1.00) on your Agent.[/red]"
4090+
)
4091+
sys.exit(1)
4092+
40684093
def _extract_cli_config_for_yaml(self):
40694094
"""
40704095
Extract CLI configuration that should be passed to YAML processing.
@@ -4611,26 +4636,26 @@ def level_based_approve(function_name, arguments, risk_level):
46114636
from rich.panel import Panel
46124637

46134638
with Live(Panel(Spinner("dots", text="Generating..."), border_style="cyan"), refresh_per_second=10, transient=True):
4614-
result = auto_rag.chat(prompt)
4639+
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
46154640
else:
4616-
result = auto_rag.chat(prompt)
4641+
result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
46174642
else:
46184643
# Resolve display mode from CLI flags
46194644
display_mode = self._resolve_display_mode()
46204645

46214646
if display_mode == 'silent':
46224647
# -qq: No output at all, exit code only
46234648
if hasattr(agent, 'start'):
4624-
result = agent.start(prompt)
4649+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46254650
else:
4626-
result = agent.chat(prompt)
4651+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46274652

46284653
elif display_mode == 'quiet':
46294654
# -q: Result only, no spinners or status
46304655
if hasattr(agent, 'start'):
4631-
result = agent.start(prompt)
4656+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46324657
else:
4633-
result = agent.chat(prompt)
4658+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46344659
if result is not None:
46354660
output = getattr(result, 'output', None) or (str(result) if result else None)
46364661
if output:
@@ -4642,31 +4667,31 @@ def level_based_approve(function_name, arguments, risk_level):
46424667
from praisonaiagents.output.status import enable_status_output, disable_status_output
46434668
enable_status_output(show_timestamps=True, show_metrics=True)
46444669
if hasattr(agent, 'start'):
4645-
result = agent.start(prompt)
4670+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46464671
else:
4647-
result = agent.chat(prompt)
4672+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46484673
disable_status_output()
46494674
except ImportError:
46504675
if hasattr(agent, 'start'):
4651-
result = agent.start(prompt)
4676+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46524677
else:
4653-
result = agent.chat(prompt)
4678+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46544679

46554680
elif display_mode == 'debug':
46564681
# -vv: SDK TraceOutput with markdown rendering
46574682
try:
46584683
from praisonaiagents.output.trace import enable_trace_output, disable_trace_output
46594684
enable_trace_output(use_markdown=True)
46604685
if hasattr(agent, 'start'):
4661-
result = agent.start(prompt)
4686+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46624687
else:
4663-
result = agent.chat(prompt)
4688+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46644689
disable_trace_output()
46654690
except ImportError:
46664691
if hasattr(agent, 'start'):
4667-
result = agent.start(prompt)
4692+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46684693
else:
4669-
result = agent.chat(prompt)
4694+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46704695

46714696
elif display_mode == 'jsonl':
46724697
# --output jsonl: JSONL structured output for CI/CD
@@ -4688,9 +4713,9 @@ def level_based_approve(function_name, arguments, risk_level):
46884713

46894714
start_time = time.time()
46904715
if hasattr(agent, 'start'):
4691-
result = agent.start(prompt)
4716+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
46924717
else:
4693-
result = agent.chat(prompt)
4718+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
46944719

46954720
# Emit final result
46964721
reason = getattr(result, 'completion_reason', None) if hasattr(result, 'completion_reason') else 'complete'
@@ -4709,9 +4734,9 @@ def level_based_approve(function_name, arguments, risk_level):
47094734
import json as json_mod
47104735
start_time = time.time()
47114736
if hasattr(agent, 'start'):
4712-
result = agent.start(prompt)
4737+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47134738
else:
4714-
result = agent.chat(prompt)
4739+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47154740

47164741
output = result.output if hasattr(result, 'output') else str(result)
47174742
envelope = {
@@ -4732,15 +4757,15 @@ def level_based_approve(function_name, arguments, risk_level):
47324757
flow = track_workflow()
47334758
flow.start()
47344759
if hasattr(agent, 'start'):
4735-
result = agent.start(prompt)
4760+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47364761
else:
4737-
result = agent.chat(prompt)
4762+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47384763
flow.stop()
47394764
except ImportError:
47404765
if hasattr(agent, 'start'):
4741-
result = agent.start(prompt)
4766+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47424767
else:
4743-
result = agent.chat(prompt)
4768+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47444769

47454770
elif display_mode == 'editor':
47464771
# --output editor: User-friendly step-by-step format
@@ -4750,9 +4775,9 @@ def level_based_approve(function_name, arguments, risk_level):
47504775

47514776
# Run agent
47524777
if hasattr(agent, 'start'):
4753-
result = agent.start(prompt)
4778+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47544779
else:
4755-
result = agent.chat(prompt)
4780+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47564781

47574782
# SDK callbacks (interaction, llm_content) handle display —
47584783
# no explicit editor.output() needed here.
@@ -4774,15 +4799,15 @@ def level_based_approve(function_name, arguments, risk_level):
47744799
from praisonaiagents.output.status import enable_status_output, disable_status_output
47754800
enable_status_output(show_timestamps=False, show_metrics=False)
47764801
if hasattr(agent, 'start'):
4777-
result = agent.start(prompt)
4802+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47784803
else:
4779-
result = agent.chat(prompt)
4804+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47804805
disable_status_output()
47814806
except ImportError:
47824807
if hasattr(agent, 'start'):
4783-
result = agent.start(prompt)
4808+
result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
47844809
else:
4785-
result = agent.chat(prompt)
4810+
result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
47864811

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

0 commit comments

Comments
 (0)