Skip to content

feat: Stream tool progress messages to Open WebUI #4032

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
sroecker wants to merge 1 commit into from
+80 −4
Closed

feat: Stream tool progress messages to Open WebUI #4032

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
30 changes: 30 additions & 0 deletions gateway/platforms/api_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -380,6 +380,7 @@ def _create_agent(
ephemeral_system_prompt: Optional[str] = None,
session_id: Optional[str] = None,
stream_delta_callback=None,
tool_progress_callback=None,
) -> Any:
"""
Create an AIAgent instance using the gateway's runtime config.
Expand Down Expand Up @@ -412,6 +413,7 @@ def _create_agent(
session_id=session_id,
platform="api_server",
stream_delta_callback=stream_delta_callback,
tool_progress_callback=tool_progress_callback,
)
return agent

Expand Down Expand Up @@ -514,6 +516,16 @@ def _on_delta(delta):
if delta is not None:
_stream_q.put(delta)

def _on_tool_progress(name, preview, args, status=None, duration=None, result=None):
# Stream tool progress messages for visibility in Open WebUI
# The 'preview' already contains the cute formatted message
logger.debug("Tool progress callback: name=%s status=%s", name, status)
if status == "complete":
# Format: add newlines before and after for visibility in the UI
msg = f"\n{preview}\n"
logger.debug("Putting tool message in queue: %s", msg[:50])
_stream_q.put(msg)

# Start agent in background. agent_ref is a mutable container
# so the SSE writer can interrupt the agent on client disconnect.
agent_ref = [None]
Expand All @@ -523,6 +535,7 @@ def _on_delta(delta):
ephemeral_system_prompt=system_prompt,
session_id=session_id,
stream_delta_callback=_on_delta,
tool_progress_callback=_on_tool_progress,
agent_ref=agent_ref,
))

Expand Down Expand Up @@ -647,6 +660,7 @@ async def _write_sse_chat_completion(
if delta is None: # End of stream sentinel
break

logger.debug("SSE writing delta: %s", repr(delta)[:100])
content_chunk = {
"id": completion_id, "object": "chat.completion.chunk",
"created": created, "model": model,
Expand Down Expand Up @@ -692,6 +706,20 @@ async def _write_sse_chat_completion(
except (asyncio.CancelledError, Exception):
pass
logger.info("SSE client disconnected; interrupted agent task %s", completion_id)
except Exception as e:
# Log any other unexpected errors in the streaming loop
logger.error("Unexpected error in SSE streaming for %s: %s", completion_id, e, exc_info=True)
# Try to send an error chunk before closing
try:
error_chunk = {
"id": completion_id, "object": "chat.completion.chunk",
"created": created, "model": model,
"choices": [{"index": 0, "delta": {"content": f"\n[Error: {e}]\n"}, "finish_reason": "stop"}],
}
await response.write(f"data: {json.dumps(error_chunk)}\n\n".encode())
await response.write(b"data: [DONE]\n\n")
except Exception:
pass

return response

Expand Down Expand Up @@ -1194,6 +1222,7 @@ async def _run_agent(
ephemeral_system_prompt: Optional[str] = None,
session_id: Optional[str] = None,
stream_delta_callback=None,
tool_progress_callback=None,
agent_ref: Optional[list] = None,
) -> tuple:
"""
Expand All @@ -1214,6 +1243,7 @@ def _run():
ephemeral_system_prompt=ephemeral_system_prompt,
session_id=session_id,
stream_delta_callback=stream_delta_callback,
tool_progress_callback=tool_progress_callback,
)
if agent_ref is not None:
agent_ref[0] = agent
Expand Down
54 changes: 50 additions & 4 deletions run_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5467,6 +5467,12 @@ def _run_tool(index, tool_call, function_name, function_args):
if self.quiet_mode:
cute_msg = _get_cute_tool_message_impl(name, args, tool_duration, result=function_result)
self._safe_print(f" {cute_msg}")
# Also notify progress callback for API server streaming
if self.tool_progress_callback:
try:
self.tool_progress_callback(name, cute_msg, args, "complete", tool_duration, function_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
elif not self.quiet_mode:
if self.verbose_logging:
print(f" ✅ Tool {i+1} completed in {tool_duration:.2f}s")
Expand Down Expand Up @@ -5597,8 +5603,14 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
store=self._todo_store,
)
tool_duration = time.time() - tool_start_time
cute_msg = _get_cute_tool_message_impl('todo', function_args, tool_duration, result=function_result)
if self.quiet_mode:
self._vprint(f" {_get_cute_tool_message_impl('todo', function_args, tool_duration, result=function_result)}")
self._vprint(f" {cute_msg}")
if self.tool_progress_callback:
try:
self.tool_progress_callback('todo', cute_msg, function_args, "complete", tool_duration, function_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
elif function_name == "session_search":
if not self._session_db:
function_result = json.dumps({"success": False, "error": "Session database not available."})
Expand All @@ -5612,8 +5624,14 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
current_session_id=self.session_id,
)
tool_duration = time.time() - tool_start_time
cute_msg = _get_cute_tool_message_impl('session_search', function_args, tool_duration, result=function_result)
if self.quiet_mode:
self._vprint(f" {_get_cute_tool_message_impl('session_search', function_args, tool_duration, result=function_result)}")
self._vprint(f" {cute_msg}")
if self.tool_progress_callback:
try:
self.tool_progress_callback('session_search', cute_msg, function_args, "complete", tool_duration, function_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
elif function_name == "memory":
target = function_args.get("target", "memory")
from tools.memory_tool import memory_tool as _memory_tool
Expand All @@ -5628,8 +5646,14 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
if self._honcho and target == "user" and function_args.get("action") == "add":
self._honcho_save_user_observation(function_args.get("content", ""))
tool_duration = time.time() - tool_start_time
cute_msg = _get_cute_tool_message_impl('memory', function_args, tool_duration, result=function_result)
if self.quiet_mode:
self._vprint(f" {_get_cute_tool_message_impl('memory', function_args, tool_duration, result=function_result)}")
self._vprint(f" {cute_msg}")
if self.tool_progress_callback:
try:
self.tool_progress_callback('memory', cute_msg, function_args, "complete", tool_duration, function_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
elif function_name == "clarify":
from tools.clarify_tool import clarify_tool as _clarify_tool
function_result = _clarify_tool(
Expand All @@ -5638,8 +5662,14 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
callback=self.clarify_callback,
)
tool_duration = time.time() - tool_start_time
cute_msg = _get_cute_tool_message_impl('clarify', function_args, tool_duration, result=function_result)
if self.quiet_mode:
self._vprint(f" {_get_cute_tool_message_impl('clarify', function_args, tool_duration, result=function_result)}")
self._vprint(f" {cute_msg}")
if self.tool_progress_callback:
try:
self.tool_progress_callback('clarify', cute_msg, function_args, "complete", tool_duration, function_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
elif function_name == "delegate_task":
from tools.delegate_tool import delegate_task as _delegate_task
tasks_arg = function_args.get("tasks")
Expand Down Expand Up @@ -5673,6 +5703,11 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
spinner.stop(cute_msg)
elif self.quiet_mode:
self._vprint(f" {cute_msg}")
if self.tool_progress_callback:
try:
self.tool_progress_callback('delegate_task', cute_msg, function_args, "complete", tool_duration, _delegate_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
elif self.quiet_mode:
spinner = None
if not self.tool_progress_callback:
Expand Down Expand Up @@ -5700,6 +5735,11 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
spinner.stop(cute_msg)
else:
self._vprint(f" {cute_msg}")
if self.tool_progress_callback:
try:
self.tool_progress_callback(function_name, cute_msg, function_args, "complete", tool_duration, _spinner_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
else:
try:
function_result = handle_function_call(
Expand All @@ -5712,6 +5752,12 @@ def _execute_tool_calls_sequential(self, assistant_message, messages: list, effe
function_result = f"Error executing tool '{function_name}': {tool_error}"
logger.error("handle_function_call raised for %s: %s", function_name, tool_error, exc_info=True)
tool_duration = time.time() - tool_start_time
cute_msg = _get_cute_tool_message_impl(function_name, function_args, tool_duration, result=function_result)
if self.tool_progress_callback:
try:
self.tool_progress_callback(function_name, cute_msg, function_args, "complete", tool_duration, function_result)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")

result_preview = function_result if self.verbose_logging else (
function_result[:200] if len(function_result) > 200 else function_result
Expand Down