first commit
This commit is contained in:
DigiJ
7
agent_dhal/agentdhal_agentchat/ui/__init__.py
Normal file
7
agent_dhal/agentdhal_agentchat/ui/__init__.py
Normal file
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
This module implements utility classes for formatting/printing agent messages.
|
||||
"""
|
||||
|
||||
from ._console import Console, UserInputManager
|
||||
|
||||
__all__ = ["Console", "UserInputManager"]
|
||||
204
agent_dhal/agentdhal_agentchat/ui/_console.py
Normal file
204
agent_dhal/agentdhal_agentchat/ui/_console.py
Normal file
@@ -0,0 +1,204 @@
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from inspect import iscoroutinefunction
|
||||
from typing import AsyncGenerator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union, cast
|
||||
|
||||
from agentdhal_core import CancellationToken
|
||||
from agentdhal_core.models import RequestUsage
|
||||
|
||||
from agentdhal_agentchat.agents import UserProxyAgent
|
||||
from agentdhal_agentchat.base import Response, TaskResult
|
||||
from agentdhal_agentchat.messages import (
|
||||
BaseAgentEvent,
|
||||
BaseChatMessage,
|
||||
ModelClientStreamingChunkEvent,
|
||||
MultiModalMessage,
|
||||
UserInputRequestedEvent,
|
||||
)
|
||||
|
||||
|
||||
def _is_running_in_iterm() -> bool:
|
||||
return os.getenv("TERM_PROGRAM") == "iTerm.app"
|
||||
|
||||
|
||||
def _is_output_a_tty() -> bool:
|
||||
return sys.stdout.isatty()
|
||||
|
||||
|
||||
SyncInputFunc = Callable[[str], str]
|
||||
AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]]
|
||||
InputFuncType = Union[SyncInputFunc, AsyncInputFunc]
|
||||
|
||||
T = TypeVar("T", bound=TaskResult | Response)
|
||||
|
||||
|
||||
class UserInputManager:
|
||||
def __init__(self, callback: InputFuncType):
|
||||
self.input_events: Dict[str, asyncio.Event] = {}
|
||||
self.callback = callback
|
||||
|
||||
def get_wrapped_callback(self) -> AsyncInputFunc:
|
||||
async def user_input_func_wrapper(prompt: str, cancellation_token: Optional[CancellationToken]) -> str:
|
||||
# Lookup the event for the prompt, if it exists wait for it.
|
||||
# If it doesn't exist, create it and store it.
|
||||
# Get request ID:
|
||||
request_id = UserProxyAgent.InputRequestContext.request_id()
|
||||
if request_id in self.input_events:
|
||||
event = self.input_events[request_id]
|
||||
else:
|
||||
event = asyncio.Event()
|
||||
self.input_events[request_id] = event
|
||||
|
||||
await event.wait()
|
||||
|
||||
del self.input_events[request_id]
|
||||
|
||||
if iscoroutinefunction(self.callback):
|
||||
# Cast to AsyncInputFunc for proper typing
|
||||
async_func = cast(AsyncInputFunc, self.callback)
|
||||
return await async_func(prompt, cancellation_token)
|
||||
else:
|
||||
# Cast to SyncInputFunc for proper typing
|
||||
sync_func = cast(SyncInputFunc, self.callback)
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, sync_func, prompt)
|
||||
|
||||
return user_input_func_wrapper
|
||||
|
||||
def notify_event_received(self, request_id: str) -> None:
|
||||
if request_id in self.input_events:
|
||||
self.input_events[request_id].set()
|
||||
else:
|
||||
event = asyncio.Event()
|
||||
self.input_events[request_id] = event
|
||||
|
||||
|
||||
def aprint(output: str, end: str = "\n", flush: bool = False) -> Awaitable[None]:
|
||||
return asyncio.to_thread(print, output, end=end, flush=flush)
|
||||
|
||||
|
||||
async def Console(
|
||||
stream: AsyncGenerator[BaseAgentEvent | BaseChatMessage | T, None],
|
||||
*,
|
||||
no_inline_images: bool = False,
|
||||
output_stats: bool = False,
|
||||
user_input_manager: UserInputManager | None = None,
|
||||
) -> T:
|
||||
"""
|
||||
Consumes the message stream from :meth:`~agentdhal_agentchat.base.TaskRunner.run_stream`
|
||||
or :meth:`~agentdhal_agentchat.base.ChatAgent.on_messages_stream` and renders the messages to the console.
|
||||
Returns the last processed TaskResult or Response.
|
||||
|
||||
.. note::
|
||||
|
||||
`output_stats` is experimental and the stats may not be accurate.
|
||||
It will be improved in future releases.
|
||||
|
||||
Args:
|
||||
stream (AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None] | AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]): Message stream to render.
|
||||
This can be from :meth:`~agentdhal_agentchat.base.TaskRunner.run_stream` or :meth:`~agentdhal_agentchat.base.ChatAgent.on_messages_stream`.
|
||||
no_inline_images (bool, optional): If terminal is iTerm2 will render images inline. Use this to disable this behavior. Defaults to False.
|
||||
output_stats (bool, optional): (Experimental) If True, will output a summary of the messages and inline token usage info. Defaults to False.
|
||||
|
||||
Returns:
|
||||
last_processed: A :class:`~agentdhal_agentchat.base.TaskResult` if the stream is from :meth:`~agentdhal_agentchat.base.TaskRunner.run_stream`
|
||||
or a :class:`~agentdhal_agentchat.base.Response` if the stream is from :meth:`~agentdhal_agentchat.base.ChatAgent.on_messages_stream`.
|
||||
"""
|
||||
render_image_iterm = _is_running_in_iterm() and _is_output_a_tty() and not no_inline_images
|
||||
start_time = time.time()
|
||||
total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)
|
||||
|
||||
last_processed: Optional[T] = None
|
||||
|
||||
streaming_chunks: List[str] = []
|
||||
|
||||
async for message in stream:
|
||||
if isinstance(message, TaskResult):
|
||||
duration = time.time() - start_time
|
||||
if output_stats:
|
||||
output = (
|
||||
f"{'-' * 10} Summary {'-' * 10}\n"
|
||||
f"Number of messages: {len(message.messages)}\n"
|
||||
f"Finish reason: {message.stop_reason}\n"
|
||||
f"Total prompt tokens: {total_usage.prompt_tokens}\n"
|
||||
f"Total completion tokens: {total_usage.completion_tokens}\n"
|
||||
f"Duration: {duration:.2f} seconds\n"
|
||||
)
|
||||
await aprint(output, end="", flush=True)
|
||||
|
||||
# mypy ignore
|
||||
last_processed = message # type: ignore
|
||||
|
||||
elif isinstance(message, Response):
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Print final response.
|
||||
if isinstance(message.chat_message, MultiModalMessage):
|
||||
final_content = message.chat_message.to_text(iterm=render_image_iterm)
|
||||
else:
|
||||
final_content = message.chat_message.to_text()
|
||||
output = f"{'-' * 10} {message.chat_message.source} {'-' * 10}\n{final_content}\n"
|
||||
if message.chat_message.models_usage:
|
||||
if output_stats:
|
||||
output += f"[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]\n"
|
||||
total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens
|
||||
total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens
|
||||
await aprint(output, end="", flush=True)
|
||||
|
||||
# Print summary.
|
||||
if output_stats:
|
||||
if message.inner_messages is not None:
|
||||
num_inner_messages = len(message.inner_messages)
|
||||
else:
|
||||
num_inner_messages = 0
|
||||
output = (
|
||||
f"{'-' * 10} Summary {'-' * 10}\n"
|
||||
f"Number of inner messages: {num_inner_messages}\n"
|
||||
f"Total prompt tokens: {total_usage.prompt_tokens}\n"
|
||||
f"Total completion tokens: {total_usage.completion_tokens}\n"
|
||||
f"Duration: {duration:.2f} seconds\n"
|
||||
)
|
||||
await aprint(output, end="", flush=True)
|
||||
|
||||
# mypy ignore
|
||||
last_processed = message # type: ignore
|
||||
# We don't want to print UserInputRequestedEvent messages, we just use them to signal the user input event.
|
||||
elif isinstance(message, UserInputRequestedEvent):
|
||||
if user_input_manager is not None:
|
||||
user_input_manager.notify_event_received(message.request_id)
|
||||
else:
|
||||
# Cast required for mypy to be happy
|
||||
message = cast(BaseAgentEvent | BaseChatMessage, message) # type: ignore
|
||||
if not streaming_chunks:
|
||||
# Print message sender.
|
||||
await aprint(
|
||||
f"{'-' * 10} {message.__class__.__name__} ({message.source}) {'-' * 10}", end="\n", flush=True
|
||||
)
|
||||
if isinstance(message, ModelClientStreamingChunkEvent):
|
||||
await aprint(message.to_text(), end="", flush=True)
|
||||
streaming_chunks.append(message.content)
|
||||
else:
|
||||
if streaming_chunks:
|
||||
streaming_chunks.clear()
|
||||
# Chunked messages are already printed, so we just print a newline.
|
||||
await aprint("", end="\n", flush=True)
|
||||
elif isinstance(message, MultiModalMessage):
|
||||
await aprint(message.to_text(iterm=render_image_iterm), end="\n", flush=True)
|
||||
else:
|
||||
await aprint(message.to_text(), end="\n", flush=True)
|
||||
if message.models_usage:
|
||||
if output_stats:
|
||||
await aprint(
|
||||
f"[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]",
|
||||
end="\n",
|
||||
flush=True,
|
||||
)
|
||||
total_usage.completion_tokens += message.models_usage.completion_tokens
|
||||
total_usage.prompt_tokens += message.models_usage.prompt_tokens
|
||||
|
||||
if last_processed is None:
|
||||
raise ValueError("No TaskResult or Response was processed.")
|
||||
|
||||
return last_processed
|
||||
Reference in New Issue
Block a user