Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/praisonai-agents/basic-agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -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?")
15 changes: 12 additions & 3 deletions src/praisonai-agents/praisonaiagents/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -124,7 +130,6 @@ def disable_telemetry():
'async_display_callbacks',
'Knowledge',
'Chunking',
'MCP',
'GuardrailResult',
'LLMGuardrail',
'Handoff',
Expand All @@ -137,4 +142,8 @@ def disable_telemetry():
'disable_telemetry',
'MinimalTelemetry',
'TelemetryCollector'
]
]

# Add MCP to __all__ if available
if _mcp_available:
__all__.append('MCP')
33 changes: 31 additions & 2 deletions src/praisonai-agents/praisonaiagents/agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +2015 to +2018
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The current implementation for auto-consuming the generator is incorrect. It only returns the last chunk of the streamed response, not the complete response, because the final_response variable is overwritten in each iteration of the loop. Additionally, this implementation doesn't display the streaming output to the user, which is part of the intended behavior described in the pull request.

To fix this, accumulate all the chunks to form the complete final response and print each chunk as it is received to provide a live streaming display.

Suggested change
final_response = None
for chunk in result:
final_response = chunk # Last chunk is typically the final response
return final_response
final_response_parts = []
for chunk in result:
self.console.print(str(chunk), end="")
final_response_parts.append(str(chunk))
self.console.print() # Add a newline at the end of the stream.
return "".join(final_response_parts)

else:
result = self.chat(prompt, **kwargs)
return result
Expand Down
3 changes: 2 additions & 1 deletion src/praisonai-agents/realtime-streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading