Revert "fix: resolve BudgetExceededError CLI traceback issues (fixes #1627) (…"

This reverts commit 2e05f2c.

Author: MervinPraison

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

@@ -544,7 +544,6 @@ 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,
@@ -623,9 +622,6 @@ 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
@@ -737,42 +733,6 @@ 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:

src/praisonai/praisonai/cli/main.py

@@ -4065,38 +4065,6 @@ 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.
@@ -4643,26 +4611,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 = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
+                        result = auto_rag.chat(prompt)
                 else:
-                    result = self._execute_agent_with_budget_handling(auto_rag, 'chat', prompt)
+                    result = 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 = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                        result = agent.start(prompt)
                     else:
-                        result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                        result = agent.chat(prompt)
 
                 elif display_mode == 'quiet':
                     # -q: Result only, no spinners or status
                     if hasattr(agent, 'start'):
-                        result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                        result = agent.start(prompt)
                     else:
-                        result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                        result = agent.chat(prompt)
                     if result is not None:
                         output = getattr(result, 'output', None) or (str(result) if result else None)
                         if output:
@@ -4674,31 +4642,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 = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
                         disable_status_output()
                     except ImportError:
                         if hasattr(agent, 'start'):
-                            result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = 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 = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
                         disable_trace_output()
                     except ImportError:
                         if hasattr(agent, 'start'):
-                            result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
 
                 elif display_mode == 'jsonl':
                     # --output jsonl: JSONL structured output for CI/CD
@@ -4720,9 +4688,9 @@ def level_based_approve(function_name, arguments, risk_level):
 
                     start_time = time.time()
                     if hasattr(agent, 'start'):
-                        result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                        result = agent.start(prompt)
                     else:
-                        result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                        result = agent.chat(prompt)
 
                     # Emit final result
                     reason = getattr(result, 'completion_reason', None) if hasattr(result, 'completion_reason') else 'complete'
@@ -4741,9 +4709,9 @@ def level_based_approve(function_name, arguments, risk_level):
                     import json as json_mod
                     start_time = time.time()
                     if hasattr(agent, 'start'):
-                        result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                        result = agent.start(prompt)
                     else:
-                        result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                        result = agent.chat(prompt)
 
                     output = result.output if hasattr(result, 'output') else str(result)
                     envelope = {
@@ -4764,15 +4732,15 @@ def level_based_approve(function_name, arguments, risk_level):
                         flow = track_workflow()
                         flow.start()
                         if hasattr(agent, 'start'):
-                            result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
                         flow.stop()
                     except ImportError:
                         if hasattr(agent, 'start'):
-                            result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
 
                 elif display_mode == 'editor':
                     # --output editor: User-friendly step-by-step format
@@ -4782,9 +4750,9 @@ def level_based_approve(function_name, arguments, risk_level):
 
                     # Run agent
                     if hasattr(agent, 'start'):
-                        result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                        result = agent.start(prompt)
                     else:
-                        result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                        result = agent.chat(prompt)
 
                     # SDK callbacks (interaction, llm_content) handle display —
                     # no explicit editor.output() needed here.
@@ -4806,15 +4774,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 = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
                         disable_status_output()
                     except ImportError:
                         if hasattr(agent, 'start'):
-                            result = self._execute_agent_with_budget_handling(agent, 'start', prompt)
+                            result = agent.start(prompt)
                         else:
-                            result = self._execute_agent_with_budget_handling(agent, 'chat', prompt)
+                            result = agent.chat(prompt)
 
             # ===== POST-PROCESSING WITH NEW FEATURES =====
 
@@ -4892,7 +4860,7 @@ def level_based_approve(function_name, arguments, risk_level):
 Now, {final_instruction.lower()}:"""
 
                 final_agent = PraisonAgent(**final_agent_config)
-                result = self._execute_agent_with_budget_handling(final_agent, 'start', final_prompt)
+                result = final_agent.start(final_prompt)
                 print(f"\n[bold green]✅ Final agent processing complete[/bold green]\n")
 
             # Save output if --save is enabled