Merge pull request #1004 from MervinPraison/claude/issue-1003-20250718-1604

Fix: Agent.start() now auto-consumes generator for better UX

Author: MervinPraison

src/praisonai-agents/basic-agents.py

@@ -5,4 +5,5 @@
     llm="gpt-4o-mini"
 )
 
+# The start() method now automatically consumes the generator and displays the output
 agent.start("Why sky is Blue?")

src/praisonai-agents/praisonaiagents/__init__.py

@@ -34,7 +34,13 @@
 from .agents.autoagents import AutoAgents
 from .knowledge.knowledge import Knowledge
 from .knowledge.chunking import Chunking
-from .mcp.mcp import MCP
+# MCP support (optional)
+try:
+    from .mcp.mcp import MCP
+    _mcp_available = True
+except ImportError:
+    _mcp_available = False
+    MCP = None
 from .session import Session
 from .memory.memory import Memory
 from .guardrails import GuardrailResult, LLMGuardrail
@@ -124,7 +130,6 @@ def disable_telemetry():
     'async_display_callbacks',
     'Knowledge',
     'Chunking',
-    'MCP',
     'GuardrailResult',
     'LLMGuardrail',
     'Handoff',
@@ -137,4 +142,8 @@ def disable_telemetry():
     'disable_telemetry',
     'MinimalTelemetry',
     'TelemetryCollector'
-] 
+]
+
+# Add MCP to __all__ if available
+if _mcp_available:
+    __all__.append('MCP') 

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

@@ -1981,12 +1981,41 @@ def _cleanup_telemetry(self):
             logging.debug(f"Error cleaning up telemetry: {e}")
 
     def start(self, prompt: str, **kwargs):
-        """Start the agent with a prompt. This is a convenience method that wraps chat()."""
+        """Start the agent with a prompt. This is a convenience method that wraps chat().
+        
+        This method is designed to be convenient for simple use cases. By default, it will
+        automatically consume the generator when streaming is enabled and return the final
+        response, while still displaying the output to the user.
+        
+        For advanced use cases that need the raw generator (e.g., for custom streaming
+        handling), use return_generator=True.
+        
+        Args:
+            prompt: The prompt to send to the agent
+            **kwargs: Additional arguments to pass to the chat method
+                     - stream: If explicitly set to False, disables streaming
+                     - return_generator: If True, returns the raw generator for custom handling
+        
+        Returns:
+            The final response from the agent (default), or a generator if return_generator=True
+        """
         try:
+            # Check if user explicitly wants the raw generator for custom handling
+            return_generator = kwargs.pop('return_generator', False)
+            
             # Check if streaming is enabled and user wants streaming chunks
             if self.stream and kwargs.get('stream', True):
                 result = self._start_stream(prompt, **kwargs)
-                return result
+                
+                if return_generator:
+                    # Return the raw generator for advanced users
+                    return result
+                else:
+                    # Auto-consume the generator for convenience while preserving display
+                    final_response = None
+                    for chunk in result:
+                        final_response = chunk  # Last chunk is typically the final response
+                    return final_response
             else:
                 result = self.chat(prompt, **kwargs)
                 return result

src/praisonai-agents/realtime-streaming.py

@@ -8,5 +8,6 @@
     stream=True
 )
 
-for chunk in agent.start("Write a report on about the history of the world"):
+# Use return_generator=True to get the raw generator for custom streaming handling
+for chunk in agent.start("Write a report on about the history of the world", return_generator=True):
     print(chunk, end="", flush=True)