diff --git a/src/praisonai-agents/praisonaiagents/agent/agent.py b/src/praisonai-agents/praisonaiagents/agent/agent.py index 7fad33816..5f93ecfaf 100644 --- a/src/praisonai-agents/praisonaiagents/agent/agent.py +++ b/src/praisonai-agents/praisonaiagents/agent/agent.py @@ -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, @@ -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 @@ -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: diff --git a/src/praisonai/praisonai/cli/main.py b/src/praisonai/praisonai/cli/main.py index b6298e1b4..839149b1a 100644 --- a/src/praisonai/praisonai/cli/main.py +++ b/src/praisonai/praisonai/cli/main.py @@ -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. @@ -4609,9 +4641,9 @@ 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() @@ -4619,16 +4651,16 @@ def level_based_approve(function_name, arguments, risk_level): 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: @@ -4640,15 +4672,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=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 @@ -4656,15 +4688,15 @@ def level_based_approve(function_name, arguments, risk_level): 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 @@ -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' @@ -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 = { @@ -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 @@ -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. @@ -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 ===== @@ -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