Revert "Merge from main"

This reverts commit b8206a85d7.
This commit is contained in:
Dmytro Struk
2025-11-11 18:44:25 -08:00
Unverified
parent b8206a85d7
commit 85fcd230bf
231 changed files with 4138 additions and 19654 deletions
@@ -5,7 +5,6 @@
import importlib.metadata
from ._agent import AgentFrameworkAgent
from ._client import AGUIChatClient
from ._confirmation_strategies import (
ConfirmationStrategy,
DefaultConfirmationStrategy,
@@ -14,8 +13,6 @@ from ._confirmation_strategies import (
TaskPlannerConfirmationStrategy,
)
from ._endpoint import add_agent_framework_fastapi_endpoint
from ._event_converters import AGUIEventConverter
from ._http_service import AGUIHttpService
try:
__version__ = importlib.metadata.version(__name__)
@@ -25,9 +22,6 @@ except importlib.metadata.PackageNotFoundError:
__all__ = [
"AgentFrameworkAgent",
"add_agent_framework_fastapi_endpoint",
"AGUIChatClient",
"AGUIEventConverter",
"AGUIHttpService",
"ConfirmationStrategy",
"DefaultConfirmationStrategy",
"TaskPlannerConfirmationStrategy",
@@ -1,407 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""AG-UI Chat Client implementation."""
import json
import logging
import uuid
from collections.abc import AsyncIterable, MutableSequence
from functools import wraps
from typing import Any, TypeVar, cast
import httpx
from agent_framework import (
AIFunction,
BaseChatClient,
ChatMessage,
ChatOptions,
ChatResponse,
ChatResponseUpdate,
DataContent,
FunctionCallContent,
)
from agent_framework._middleware import use_chat_middleware
from agent_framework._tools import use_function_invocation
from agent_framework._types import BaseContent, Contents
from agent_framework.observability import use_observability
from ._event_converters import AGUIEventConverter
from ._http_service import AGUIHttpService
from ._message_adapters import agent_framework_messages_to_agui
from ._utils import convert_tools_to_agui_format
logger: logging.Logger = logging.getLogger(__name__)
class ServerFunctionCallContent(BaseContent):
"""Wrapper for server function calls to prevent client re-execution.
All function calls from the remote server are server-side executions.
This wrapper prevents @use_function_invocation from trying to execute them again.
"""
function_call_content: FunctionCallContent
def __init__(self, function_call_content: FunctionCallContent) -> None:
"""Initialize with the function call content."""
super().__init__(type="server_function_call")
self.function_call_content = function_call_content
def _unwrap_server_function_call_contents(contents: MutableSequence[Contents | dict[str, Any]]) -> None:
"""Replace ServerFunctionCallContent instances with their underlying call content."""
for idx, content in enumerate(contents):
if isinstance(content, ServerFunctionCallContent):
contents[idx] = content.function_call_content # type: ignore[assignment]
TBaseChatClient = TypeVar("TBaseChatClient", bound=type[BaseChatClient])
def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseChatClient:
"""Class decorator that unwraps server-side function calls after tool handling."""
original_get_streaming_response = chat_client.get_streaming_response
@wraps(original_get_streaming_response)
async def streaming_wrapper(self, *args: Any, **kwargs: Any) -> AsyncIterable[ChatResponseUpdate]:
async for update in original_get_streaming_response(self, *args, **kwargs):
_unwrap_server_function_call_contents(cast(MutableSequence[Contents | dict[str, Any]], update.contents))
yield update
chat_client.get_streaming_response = streaming_wrapper # type: ignore[assignment]
original_get_response = chat_client.get_response
@wraps(original_get_response)
async def response_wrapper(self, *args: Any, **kwargs: Any) -> ChatResponse:
response = await original_get_response(self, *args, **kwargs)
if response.messages:
for message in response.messages:
_unwrap_server_function_call_contents(
cast(MutableSequence[Contents | dict[str, Any]], message.contents)
)
return response
chat_client.get_response = response_wrapper # type: ignore[assignment]
return chat_client
@_apply_server_function_call_unwrap
@use_function_invocation
@use_observability
@use_chat_middleware
class AGUIChatClient(BaseChatClient):
"""Chat client for communicating with AG-UI compliant servers.
This client implements the BaseChatClient interface and automatically handles:
- Thread ID management for conversation continuity
- State synchronization between client and server
- Server-Sent Events (SSE) streaming
- Event conversion to Agent Framework types
Important: Message History Management
This client sends exactly the messages it receives to the server. It does NOT
automatically maintain conversation history. The server must handle history via thread_id.
For stateless servers: Use ChatAgent wrapper which will send full message history on each
request. However, even with ChatAgent, the server must echo back all context for the
agent to maintain history across turns.
Important: Tool Handling (Hybrid Execution - matches .NET)
1. Client tool metadata sent to server - LLM knows about both client and server tools
2. Server has its own tools that execute server-side
3. When LLM calls a client tool, @use_function_invocation executes it locally
4. Both client and server tools work together (hybrid pattern)
The wrapping ChatAgent's @use_function_invocation handles client tool execution
automatically when the server's LLM decides to call them.
Examples:
Direct usage (server manages thread history):
.. code-block:: python
from agent_framework.ag_ui import AGUIChatClient
client = AGUIChatClient(endpoint="http://localhost:8888/")
# First message - thread ID auto-generated
response = await client.get_response("Hello!")
thread_id = response.additional_properties.get("thread_id")
# Second message - server retrieves history using thread_id
response2 = await client.get_response(
"How are you?",
metadata={"thread_id": thread_id}
)
Recommended usage with ChatAgent (client manages history):
.. code-block:: python
from agent_framework import ChatAgent
from agent_framework.ag_ui import AGUIChatClient
client = AGUIChatClient(endpoint="http://localhost:8888/")
agent = ChatAgent(name="assistant", client=client)
thread = await agent.get_new_thread()
# ChatAgent automatically maintains history and sends full context
response = await agent.run("Hello!", thread=thread)
response2 = await agent.run("How are you?", thread=thread)
Streaming usage:
.. code-block:: python
async for update in client.get_streaming_response("Tell me a story"):
if update.contents:
for content in update.contents:
if hasattr(content, "text"):
print(content.text, end="", flush=True)
Context manager:
.. code-block:: python
async with AGUIChatClient(endpoint="http://localhost:8888/") as client:
response = await client.get_response("Hello!")
print(response.messages[0].text)
"""
OTEL_PROVIDER_NAME = "agui"
def __init__(
self,
*,
endpoint: str,
http_client: httpx.AsyncClient | None = None,
timeout: float = 60.0,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
"""Initialize the AG-UI chat client.
Args:
endpoint: The AG-UI server endpoint URL (e.g., "http://localhost:8888/")
http_client: Optional httpx.AsyncClient instance. If None, one will be created.
timeout: Request timeout in seconds (default: 60.0)
additional_properties: Additional properties to store
**kwargs: Additional arguments passed to BaseChatClient
"""
super().__init__(additional_properties=additional_properties, **kwargs)
self._http_service = AGUIHttpService(
endpoint=endpoint,
http_client=http_client,
timeout=timeout,
)
async def close(self) -> None:
"""Close the HTTP client."""
await self._http_service.close()
async def __aenter__(self) -> "AGUIChatClient":
"""Enter async context manager."""
return self
async def __aexit__(self, *args: Any) -> None:
"""Exit async context manager."""
await self.close()
def _register_server_tool_placeholder(self, tool_name: str) -> None:
"""Register a declaration-only placeholder so function invocation skips execution."""
config = getattr(self, "function_invocation_configuration", None)
if not config:
return
if any(getattr(tool, "name", None) == tool_name for tool in config.additional_tools):
return
placeholder: AIFunction[Any, Any] = AIFunction(
name=tool_name,
description="Server-managed tool placeholder (AG-UI)",
func=None,
)
config.additional_tools = list(config.additional_tools) + [placeholder]
registered: set[str] = getattr(self, "_registered_server_tools", set())
registered.add(tool_name)
self._registered_server_tools = registered # type: ignore[attr-defined]
from agent_framework._logging import get_logger
logger = get_logger()
logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
def _extract_state_from_messages(
self, messages: MutableSequence[ChatMessage]
) -> tuple[list[ChatMessage], dict[str, Any] | None]:
"""Extract state from last message if present.
Args:
messages: List of chat messages
Returns:
Tuple of (messages_without_state, state_dict)
"""
if not messages:
return list(messages), None
last_message = messages[-1]
for content in last_message.contents:
if isinstance(content, DataContent) and content.media_type == "application/json":
try:
uri = content.uri
if uri.startswith("data:application/json;base64,"):
import base64
encoded_data = uri.split(",", 1)[1]
decoded_bytes = base64.b64decode(encoded_data)
state = json.loads(decoded_bytes.decode("utf-8"))
messages_without_state = list(messages[:-1]) if len(messages) > 1 else []
return messages_without_state, state
except (json.JSONDecodeError, ValueError, KeyError) as e:
from agent_framework._logging import get_logger
logger = get_logger()
logger.warning(f"Failed to extract state from message: {e}")
return list(messages), None
def _convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
"""Convert Agent Framework messages to AG-UI format.
Args:
messages: List of ChatMessage objects
Returns:
List of AG-UI formatted message dictionaries
"""
return agent_framework_messages_to_agui(messages)
def _get_thread_id(self, chat_options: ChatOptions) -> str:
"""Get or generate thread ID from chat options.
Args:
chat_options: Chat options containing metadata
Returns:
Thread ID string
"""
thread_id = None
if chat_options.metadata:
thread_id = chat_options.metadata.get("thread_id")
if not thread_id:
thread_id = f"thread_{uuid.uuid4().hex}"
return thread_id
async def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
"""Internal method to get non-streaming response.
Keyword Args:
messages: List of chat messages
chat_options: Chat options for the request
**kwargs: Additional keyword arguments
Returns:
ChatResponse object
"""
return await ChatResponse.from_chat_response_generator(
self._inner_get_streaming_response(
messages=messages,
chat_options=chat_options,
**kwargs,
)
)
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Internal method to get streaming response.
Keyword Args:
messages: List of chat messages
chat_options: Chat options for the request
**kwargs: Additional keyword arguments
Yields:
ChatResponseUpdate objects
"""
messages_to_send, state = self._extract_state_from_messages(messages)
thread_id = self._get_thread_id(chat_options)
run_id = f"run_{uuid.uuid4().hex}"
agui_messages = self._convert_messages_to_agui_format(messages_to_send)
# Send client tools to server so LLM knows about them
# Client tools execute via ChatAgent's @use_function_invocation wrapper
agui_tools = convert_tools_to_agui_format(chat_options.tools)
# Build set of client tool names (matches .NET clientToolSet)
# Used to distinguish client vs server tools in response stream
client_tool_set: set[str] = set()
if chat_options.tools:
for tool in chat_options.tools:
if hasattr(tool, "name"):
client_tool_set.add(tool.name) # type: ignore[arg-type]
self._last_client_tool_set = client_tool_set # type: ignore[attr-defined]
logger.debug(
"[AGUIChatClient] Preparing request",
extra={
"thread_id": thread_id,
"run_id": run_id,
"client_tools": list(client_tool_set),
"messages": [msg.text for msg in messages_to_send if msg.text],
},
)
logger.debug(f"[AGUIChatClient] Client tool set: {client_tool_set}")
converter = AGUIEventConverter()
async for event in self._http_service.post_run(
thread_id=thread_id,
run_id=run_id,
messages=agui_messages,
state=state,
tools=agui_tools,
):
logger.debug(f"[AGUIChatClient] Raw AG-UI event: {event}")
update = converter.convert_event(event)
if update is not None:
logger.debug(
"[AGUIChatClient] Converted update",
extra={"role": update.role, "contents": [type(c).__name__ for c in update.contents]},
)
# Distinguish client vs server tools
for i, content in enumerate(update.contents):
if isinstance(content, FunctionCallContent):
logger.debug(
f"[AGUIChatClient] Function call: {content.name}, in client_tool_set: {content.name in client_tool_set}"
)
if content.name in client_tool_set:
# Client tool - let @use_function_invocation execute it
if not content.additional_properties:
content.additional_properties = {}
content.additional_properties["agui_thread_id"] = thread_id
else:
# Server tool - wrap so @use_function_invocation ignores it
logger.debug(f"[AGUIChatClient] Wrapping server tool: {content.name}")
self._register_server_tool_placeholder(content.name)
update.contents[i] = ServerFunctionCallContent(content) # type: ignore
yield update
@@ -1,209 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Event converter for AG-UI protocol events to Agent Framework types."""
from typing import Any
from agent_framework import (
ChatResponseUpdate,
ErrorContent,
FinishReason,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
)
class AGUIEventConverter:
"""Converter for AG-UI events to Agent Framework types.
Handles conversion of AG-UI protocol events to ChatResponseUpdate objects
while maintaining state, aggregating content, and tracking metadata.
"""
def __init__(self) -> None:
"""Initialize the converter with fresh state."""
self.current_message_id: str | None = None
self.current_tool_call_id: str | None = None
self.current_tool_name: str | None = None
self.accumulated_tool_args: str = ""
self.thread_id: str | None = None
self.run_id: str | None = None
def convert_event(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
"""Convert a single AG-UI event to ChatResponseUpdate.
Args:
event: AG-UI event dictionary
Returns:
ChatResponseUpdate if event produces content, None otherwise
Examples:
RUN_STARTED event:
.. code-block:: python
converter = AGUIEventConverter()
event = {"type": "RUN_STARTED", "threadId": "t1", "runId": "r1"}
update = converter.convert_event(event)
assert update.additional_properties["thread_id"] == "t1"
TEXT_MESSAGE_CONTENT event:
.. code-block:: python
event = {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "Hello"}
update = converter.convert_event(event)
assert update.contents[0].text == "Hello"
"""
event_type = event.get("type", "")
if event_type == "RUN_STARTED":
return self._handle_run_started(event)
elif event_type == "TEXT_MESSAGE_START":
return self._handle_text_message_start(event)
elif event_type == "TEXT_MESSAGE_CONTENT":
return self._handle_text_message_content(event)
elif event_type == "TEXT_MESSAGE_END":
return self._handle_text_message_end(event)
elif event_type == "TOOL_CALL_START":
return self._handle_tool_call_start(event)
elif event_type == "TOOL_CALL_ARGS":
return self._handle_tool_call_args(event)
elif event_type == "TOOL_CALL_END":
return self._handle_tool_call_end(event)
elif event_type == "TOOL_CALL_RESULT":
return self._handle_tool_call_result(event)
elif event_type == "RUN_FINISHED":
return self._handle_run_finished(event)
elif event_type == "RUN_ERROR":
return self._handle_run_error(event)
return None
def _handle_run_started(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle RUN_STARTED event."""
self.thread_id = event.get("threadId")
self.run_id = event.get("runId")
return ChatResponseUpdate(
role=Role.ASSISTANT,
contents=[],
additional_properties={
"thread_id": self.thread_id,
"run_id": self.run_id,
},
)
def _handle_text_message_start(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
"""Handle TEXT_MESSAGE_START event."""
self.current_message_id = event.get("messageId")
return ChatResponseUpdate(
role=Role.ASSISTANT,
message_id=self.current_message_id,
contents=[],
)
def _handle_text_message_content(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle TEXT_MESSAGE_CONTENT event."""
message_id = event.get("messageId")
delta = event.get("delta", "")
if message_id != self.current_message_id:
self.current_message_id = message_id
return ChatResponseUpdate(
role=Role.ASSISTANT,
message_id=self.current_message_id,
contents=[TextContent(text=delta)],
)
def _handle_text_message_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
"""Handle TEXT_MESSAGE_END event."""
return None
def _handle_tool_call_start(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle TOOL_CALL_START event."""
self.current_tool_call_id = event.get("toolCallId")
self.current_tool_name = event.get("toolName") or event.get("toolCallName") or event.get("tool_call_name")
self.accumulated_tool_args = ""
return ChatResponseUpdate(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id=self.current_tool_call_id or "",
name=self.current_tool_name or "",
arguments="",
)
],
)
def _handle_tool_call_args(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle TOOL_CALL_ARGS event."""
delta = event.get("delta", "")
self.accumulated_tool_args += delta
return ChatResponseUpdate(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id=self.current_tool_call_id or "",
name=self.current_tool_name or "",
arguments=delta,
)
],
)
def _handle_tool_call_end(self, event: dict[str, Any]) -> ChatResponseUpdate | None:
"""Handle TOOL_CALL_END event."""
self.accumulated_tool_args = ""
return None
def _handle_tool_call_result(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle TOOL_CALL_RESULT event."""
tool_call_id = event.get("toolCallId", "")
result = event.get("result") if event.get("result") is not None else event.get("content")
return ChatResponseUpdate(
role=Role.TOOL,
contents=[
FunctionResultContent(
call_id=tool_call_id,
result=result,
)
],
)
def _handle_run_finished(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle RUN_FINISHED event."""
return ChatResponseUpdate(
role=Role.ASSISTANT,
finish_reason=FinishReason.STOP,
contents=[],
additional_properties={
"thread_id": self.thread_id,
"run_id": self.run_id,
},
)
def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate:
"""Handle RUN_ERROR event."""
error_message = event.get("message", "Unknown error")
return ChatResponseUpdate(
role=Role.ASSISTANT,
finish_reason=FinishReason.CONTENT_FILTER,
contents=[
ErrorContent(
message=error_message,
error_code="RUN_ERROR",
)
],
additional_properties={
"thread_id": self.thread_id,
"run_id": self.run_id,
},
)
@@ -107,7 +107,7 @@ class AgentFrameworkEventBridge:
# Skip text content if we're about to emit confirm_changes
# The summary should only appear after user confirms
if self.should_stop_after_confirm:
logger.debug("Skipping text content - waiting for confirm_changes response")
logger.debug(" >>> Skipping text content - waiting for confirm_changes response")
# Save the summary text to show after confirmation
self.suppressed_summary += content.text
continue
@@ -156,7 +156,7 @@ class AgentFrameworkEventBridge:
tool_call_name=content.name,
parent_message_id=self.current_message_id,
)
logger.info(f"Emitting ToolCallStartEvent with name='{content.name}', id='{tool_call_id}'")
logger.info(f" >>> Emitting ToolCallStartEvent with name='{content.name}', id='{tool_call_id}'")
events.append(tool_start_event)
# Track tool call for MessagesSnapshotEvent
@@ -186,7 +186,7 @@ class AgentFrameworkEventBridge:
# If it's a dict, convert to JSON
delta_str = json.dumps(content.arguments)
logger.info(f"Emitting ToolCallArgsEvent with delta: {delta_str!r}..., id='{tool_call_id}'")
logger.info(f" >>> Emitting ToolCallArgsEvent with delta: {delta_str!r}..., id='{tool_call_id}'")
args_event = ToolCallArgsEvent(
tool_call_id=tool_call_id,
delta=delta_str,
@@ -211,7 +211,7 @@ class AgentFrameworkEventBridge:
self.streaming_tool_args += json.dumps(content.arguments)
logger.debug(
f"Predictive state: accumulated {len(self.streaming_tool_args)} chars for tool '{self.current_tool_call_name}'"
f" >>> Predictive state: accumulated {len(self.streaming_tool_args)} chars for tool '{self.current_tool_call_name}'"
)
# Try to parse accumulated arguments (may be incomplete JSON)
@@ -262,11 +262,11 @@ class AgentFrameworkEventBridge:
else str(partial_value)
)
logger.info(
f"StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
f" >>> StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
f"op=replace, path=/{state_key}, value={value_preview}"
)
elif self.state_delta_count % 100 == 0:
logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
logger.info(f" >>> StateDeltaEvent #{self.state_delta_count} emitted")
events.append(state_delta_event)
self.last_emitted_state[state_key] = partial_value
@@ -312,11 +312,11 @@ class AgentFrameworkEventBridge:
else str(state_value)
)
logger.info(
f"StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
f" >>> StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
f"op=replace, path=/{state_key}, value={value_preview}"
)
elif self.state_delta_count % 100 == 0: # Also log every 100th
logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
logger.info(f" >>> StateDeltaEvent #{self.state_delta_count} emitted")
events.append(state_delta_event)
@@ -360,7 +360,7 @@ class AgentFrameworkEventBridge:
],
)
logger.info(
f"Emitting StateDeltaEvent for key '{state_key}', value type: {type(state_value)}"
f" >>> Emitting StateDeltaEvent for key '{state_key}', value type: {type(state_value)}"
)
events.append(state_delta_event)
@@ -376,13 +376,13 @@ class AgentFrameworkEventBridge:
end_event = ToolCallEndEvent(
tool_call_id=content.call_id,
)
logger.info(f"Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
logger.info(f" >>> Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
events.append(end_event)
# Log total StateDeltaEvent count for this tool call
if self.state_delta_count > 0:
logger.info(
f"Tool call '{content.call_id}' complete: emitted {self.state_delta_count} StateDeltaEvents total"
f" >>> Tool call '{content.call_id}' complete: emitted {self.state_delta_count} StateDeltaEvents total"
)
# Reset streaming accumulator and counter for next tool call
@@ -410,13 +410,11 @@ class AgentFrameworkEventBridge:
events.append(result_event)
# Track tool result for MessagesSnapshotEvent
# AG-UI protocol expects: { role: "tool", toolCallId: ..., content: ... }
# Use camelCase for Pydantic's alias_generator=to_camel
self.tool_results.append(
{
"id": result_message_id,
"role": "tool",
"toolCallId": content.call_id,
"tool_call_id": content.call_id,
"content": result_content,
}
)
@@ -424,9 +422,6 @@ class AgentFrameworkEventBridge:
# Emit MessagesSnapshotEvent with the complete conversation including tool calls and results
# This is required for CopilotKit's useCopilotAction to detect tool result
if self.pending_tool_calls and self.tool_results:
# Import message adapter
from ._message_adapters import agent_framework_messages_to_agui
# Build assistant message with tool_calls
assistant_message = {
"id": generate_event_id(),
@@ -434,19 +429,14 @@ class AgentFrameworkEventBridge:
"tool_calls": self.pending_tool_calls.copy(), # Copy the accumulated tool calls
}
# Convert Agent Framework messages to AG-UI format (adds required 'id' field)
converted_input_messages = agent_framework_messages_to_agui(self.input_messages)
# Build complete messages array: input messages + assistant message + tool results
all_messages = converted_input_messages + [assistant_message] + self.tool_results.copy()
all_messages = list(self.input_messages) + [assistant_message] + self.tool_results.copy()
# Emit MessagesSnapshotEvent using the proper event type
# Note: messages are dict[str, Any] but Pydantic will validate them as Message types
messages_snapshot_event = MessagesSnapshotEvent(
type=EventType.MESSAGES_SNAPSHOT,
messages=all_messages, # type: ignore[arg-type]
type=EventType.MESSAGES_SNAPSHOT, messages=all_messages
)
logger.info(f"Emitting MessagesSnapshotEvent with {len(all_messages)} messages")
logger.info(f" >>> Emitting MessagesSnapshotEvent with {len(all_messages)} messages")
events.append(messages_snapshot_event)
# After tool execution, emit StateSnapshotEvent if we have pending state updates
@@ -476,7 +466,7 @@ class AgentFrameworkEventBridge:
# If so, emit a confirm_changes tool call for the UI modal
tool_was_predictive = False
logger.debug(
f"Checking predictive state: current_tool='{self.current_tool_call_name}', "
f" >>> Checking predictive state: current_tool='{self.current_tool_call_name}', "
f"predict_config={list(self.predict_state_config.keys()) if self.predict_state_config else 'None'}"
)
for state_key, config in self.predict_state_config.items():
@@ -484,7 +474,7 @@ class AgentFrameworkEventBridge:
# We need to match against self.current_tool_call_name
if self.current_tool_call_name and config["tool"] == self.current_tool_call_name:
logger.info(
f"Tool '{self.current_tool_call_name}' matches predictive config for state key '{state_key}'"
f" >>> Tool '{self.current_tool_call_name}' matches predictive config for state key '{state_key}'"
)
tool_was_predictive = True
break
@@ -493,7 +483,7 @@ class AgentFrameworkEventBridge:
# Emit confirm_changes tool call sequence
confirm_call_id = generate_event_id()
logger.info("Emitting confirm_changes tool call for predictive update")
logger.info(" >>> Emitting confirm_changes tool call for predictive update")
# Track confirm_changes tool call for MessagesSnapshotEvent (so it persists after RUN_FINISHED)
self.pending_tool_calls.append(
@@ -528,9 +518,6 @@ class AgentFrameworkEventBridge:
events.append(confirm_end)
# Emit MessagesSnapshotEvent so confirm_changes persists after RUN_FINISHED
# Import message adapter
from ._message_adapters import agent_framework_messages_to_agui
# Build assistant message with pending confirm_changes tool call
assistant_message = {
"id": generate_event_id(),
@@ -538,28 +525,23 @@ class AgentFrameworkEventBridge:
"tool_calls": self.pending_tool_calls.copy(), # Includes confirm_changes
}
# Convert Agent Framework messages to AG-UI format (adds required 'id' field)
converted_input_messages = agent_framework_messages_to_agui(self.input_messages)
# Build complete messages array: input messages + assistant message + any tool results
all_messages = converted_input_messages + [assistant_message] + self.tool_results.copy()
all_messages = list(self.input_messages) + [assistant_message] + self.tool_results.copy()
# Emit MessagesSnapshotEvent
# Note: messages are dict[str, Any] but Pydantic will validate them as Message types
messages_snapshot_event = MessagesSnapshotEvent(
type=EventType.MESSAGES_SNAPSHOT,
messages=all_messages, # type: ignore[arg-type]
type=EventType.MESSAGES_SNAPSHOT, messages=all_messages
)
logger.info(
f"Emitting MessagesSnapshotEvent for confirm_changes with {len(all_messages)} messages"
f" >>> Emitting MessagesSnapshotEvent for confirm_changes with {len(all_messages)} messages"
)
events.append(messages_snapshot_event)
# Set flag to stop the run after this - we're waiting for user response
self.should_stop_after_confirm = True
logger.info("Set flag to stop run after confirm_changes")
logger.info(" >>> Set flag to stop run after confirm_changes")
elif tool_was_predictive:
logger.info("Skipping confirm_changes - require_confirmation is False")
logger.info(" >>> Skipping confirm_changes - require_confirmation is False")
# Clear pending updates and reset tool name tracker
self.pending_state_updates.clear()
@@ -598,7 +580,7 @@ class AgentFrameworkEventBridge:
# Update current state
self.current_state[state_key] = state_value
logger.info(
f"Emitting StateSnapshotEvent for key '{state_key}', value type: {type(state_value)}"
f" >>> Emitting StateSnapshotEvent for key '{state_key}', value type: {type(state_value)}"
)
# Emit state snapshot
@@ -614,7 +596,7 @@ class AgentFrameworkEventBridge:
tool_call_id=content.function_call.call_id,
)
logger.info(
f"Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'"
f" >>> Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'"
)
events.append(end_event)
@@ -633,7 +615,7 @@ class AgentFrameworkEventBridge:
},
},
)
logger.info(f"Emitting function_approval_request custom event for '{content.function_call.name}'")
logger.info(f" >>> Emitting function_approval_request custom event for '{content.function_call.name}'")
events.append(approval_event)
return events
@@ -1,157 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""HTTP service for AG-UI protocol communication."""
import json
import logging
from collections.abc import AsyncIterable
from typing import Any
import httpx
logger = logging.getLogger(__name__)
class AGUIHttpService:
"""HTTP service for AG-UI protocol communication.
Handles HTTP POST requests and Server-Sent Events (SSE) stream parsing
for the AG-UI protocol.
Examples:
Basic usage:
.. code-block:: python
service = AGUIHttpService("http://localhost:8888/")
async for event in service.post_run(
thread_id="thread_123",
run_id="run_456",
messages=[{"role": "user", "content": "Hello"}]
):
print(event["type"])
With context manager:
.. code-block:: python
async with AGUIHttpService("http://localhost:8888/") as service:
async for event in service.post_run(...):
print(event)
"""
def __init__(
self,
endpoint: str,
http_client: httpx.AsyncClient | None = None,
timeout: float = 60.0,
) -> None:
"""Initialize the HTTP service.
Args:
endpoint: AG-UI server endpoint URL (e.g., "http://localhost:8888/")
http_client: Optional httpx AsyncClient. If None, creates a new one.
timeout: Request timeout in seconds (default: 60.0)
"""
self.endpoint = endpoint.rstrip("/")
self._owns_client = http_client is None
self.http_client = http_client or httpx.AsyncClient(timeout=timeout)
async def post_run(
self,
thread_id: str,
run_id: str,
messages: list[dict[str, Any]],
state: dict[str, Any] | None = None,
tools: list[dict[str, Any]] | None = None,
) -> AsyncIterable[dict[str, Any]]:
"""Post a run request and stream AG-UI events.
Args:
thread_id: Thread identifier for conversation continuity
run_id: Unique run identifier
messages: List of messages in AG-UI format
state: Optional state object to send to server
tools: Optional list of tools available to the agent
Yields:
AG-UI event dictionaries parsed from SSE stream
Raises:
httpx.HTTPStatusError: If the HTTP request fails
ValueError: If SSE parsing encounters invalid data
Examples:
.. code-block:: python
service = AGUIHttpService("http://localhost:8888/")
async for event in service.post_run(
thread_id="thread_abc",
run_id="run_123",
messages=[{"role": "user", "content": "Hello"}],
state={"user_context": {"name": "Alice"}}
):
if event["type"] == "TEXT_MESSAGE_CONTENT":
print(event["delta"])
"""
# Build request payload
request_data: dict[str, Any] = {
"thread_id": thread_id,
"run_id": run_id,
"messages": messages,
}
if state is not None:
request_data["state"] = state
if tools is not None:
request_data["tools"] = tools
logger.debug(
f"Posting run to {self.endpoint}: thread_id={thread_id}, run_id={run_id}, "
f"messages={len(messages)}, has_state={state is not None}, has_tools={tools is not None}"
)
# Stream the response using SSE
async with self.http_client.stream(
"POST",
self.endpoint,
json=request_data,
headers={"Accept": "text/event-stream"},
) as response:
try:
response.raise_for_status()
except httpx.HTTPStatusError as e:
logger.error(f"HTTP request failed: {e.response.status_code} - {e.response.text}")
raise
async for line in response.aiter_lines():
# Parse Server-Sent Events format
if line.startswith("data: "):
data = line[6:] # Remove "data: " prefix
try:
event = json.loads(data)
logger.debug(f"Received event: {event.get('type', 'UNKNOWN')}")
yield event
except json.JSONDecodeError as e:
logger.warning(f"Failed to parse SSE data: {data}. Error: {e}")
# Continue processing other events instead of failing
continue
async def close(self) -> None:
"""Close the HTTP client if owned by this service.
Only closes the client if it was created by this service instance.
If an external client was provided, it remains the caller's
responsibility to close it.
"""
if self._owns_client and self.http_client:
await self.http_client.aclose()
async def __aenter__(self) -> "AGUIHttpService":
"""Enter async context manager."""
return self
async def __aexit__(self, *args: Any) -> None:
"""Exit async context manager and clean up resources."""
await self.close()
@@ -2,13 +2,12 @@
"""Message format conversion between AG-UI and Agent Framework."""
from typing import Any, cast
from typing import Any
from agent_framework import (
ChatMessage,
FunctionApprovalResponseContent,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
)
@@ -47,7 +46,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
result_content = msg.get("result", msg.get("content", ""))
chat_msg = ChatMessage(
role=Role.TOOL, # Tool results must be tool role
role=Role.ASSISTANT, # Tool results are assistant messages
contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
)
@@ -57,42 +56,6 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
result.append(chat_msg)
continue
# If assistant message includes tool calls, convert to FunctionCallContent(s)
tool_calls = msg.get("tool_calls") or msg.get("toolCalls")
if tool_calls:
contents: list[Any] = []
# Include any assistant text content if present
content_text = msg.get("content")
if isinstance(content_text, str) and content_text:
contents.append(TextContent(text=content_text))
# Convert each tool call entry
for tc in tool_calls:
if not isinstance(tc, dict):
continue
# Cast to typed dict for proper type inference
tc_dict = cast(dict[str, Any], tc)
tc_type = tc_dict.get("type")
if tc_type == "function":
func_data = tc_dict.get("function", {})
func_dict = cast(dict[str, Any], func_data) if isinstance(func_data, dict) else {}
call_id = str(tc_dict.get("id", ""))
name = str(func_dict.get("name", ""))
arguments = func_dict.get("arguments")
contents.append(
FunctionCallContent(
call_id=call_id,
name=name,
arguments=arguments,
)
)
chat_msg = ChatMessage(role=Role.ASSISTANT, contents=contents)
if "id" in msg:
chat_msg.message_id = msg["id"]
result.append(chat_msg)
continue
role_str = msg.get("role", "user")
# Handle tool result messages (with role="tool")
@@ -115,11 +78,11 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
# Backend tool results have non-empty content WITHOUT "accepted" field
if tool_call_id and result_content and not is_approval:
# Tool execution result - convert to FunctionResultContent with correct role
# Backend tool execution - convert to FunctionResultContent
from agent_framework import FunctionResultContent
chat_msg = ChatMessage(
role=Role.TOOL,
role=Role.ASSISTANT, # Tool results are assistant messages
contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
)
@@ -134,8 +97,9 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
chat_msg = ChatMessage(
role=Role.USER, # Approval responses are user messages
contents=[TextContent(text=content)],
additional_properties={"is_tool_result": True, "tool_call_id": msg.get("toolCallId", "")},
)
# Mark this as a tool result so we can detect it later
chat_msg.metadata = {"is_tool_result": True, "tool_call_id": msg.get("toolCallId", "")} # type: ignore[attr-defined]
if "id" in msg:
chat_msg.message_id = msg["id"]
@@ -148,7 +112,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
# Check if this message contains function approvals
if "function_approvals" in msg and msg["function_approvals"]:
# Convert function approvals to FunctionApprovalResponseContent
approval_contents: list[Any] = []
contents: list[Any] = []
for approval in msg["function_approvals"]:
# Create FunctionCallContent with the modified arguments
func_call = FunctionCallContent(
@@ -163,9 +127,9 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
id=approval.get("id", ""),
function_call=func_call,
)
approval_contents.append(approval_response)
contents.append(approval_response)
chat_msg = ChatMessage(role=role, contents=approval_contents) # type: ignore[arg-type]
chat_msg = ChatMessage(role=role, contents=contents) # type: ignore[arg-type]
else:
# Regular text message
content = msg.get("content", "")
@@ -182,44 +146,21 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
return result
def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str, Any]]) -> list[dict[str, Any]]:
def agent_framework_messages_to_agui(messages: list[ChatMessage]) -> list[dict[str, Any]]:
"""Convert Agent Framework messages to AG-UI format.
Args:
messages: List of Agent Framework ChatMessage objects or AG-UI dicts (already converted)
messages: List of Agent Framework ChatMessage objects
Returns:
List of AG-UI message dictionaries
"""
from ._utils import generate_event_id
result: list[dict[str, Any]] = []
for msg in messages:
# If already a dict (AG-UI format), ensure it has an ID and normalize keys for Pydantic
if isinstance(msg, dict):
# Always work on a copy to avoid mutating input
normalized_msg = msg.copy()
# Ensure ID exists
if "id" not in normalized_msg:
normalized_msg["id"] = generate_event_id()
# Normalize tool_call_id to toolCallId for Pydantic's alias_generator=to_camel
if normalized_msg.get("role") == "tool":
if "tool_call_id" in normalized_msg:
normalized_msg["toolCallId"] = normalized_msg["tool_call_id"]
del normalized_msg["tool_call_id"]
elif "toolCallId" not in normalized_msg:
# Tool message missing toolCallId - add empty string to satisfy schema
normalized_msg["toolCallId"] = ""
# Always append the normalized copy, not the original
result.append(normalized_msg)
continue
# Convert ChatMessage to AG-UI format
role = _FRAMEWORK_TO_AGUI_ROLE.get(msg.role, "user")
content_text = ""
tool_calls: list[dict[str, Any]] = []
tool_result_call_id: str | None = None
for content in msg.contents:
if isinstance(content, TextContent):
@@ -235,32 +176,18 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
},
}
)
elif isinstance(content, FunctionResultContent):
# Tool result content - extract call_id and result
tool_result_call_id = content.call_id
# Serialize result to string
if isinstance(content.result, dict):
import json
content_text = json.dumps(content.result) # type: ignore
elif content.result is not None:
content_text = str(content.result)
agui_msg: dict[str, Any] = {
"id": msg.message_id if msg.message_id else generate_event_id(), # Always include id
"role": role,
"content": content_text,
}
if msg.message_id:
agui_msg["id"] = msg.message_id
if tool_calls:
agui_msg["tool_calls"] = tool_calls
# If this is a tool result message, add toolCallId (using camelCase for Pydantic)
if tool_result_call_id:
agui_msg["toolCallId"] = tool_result_call_id
# Tool result messages should have role="tool"
agui_msg["role"] = "tool"
result.append(agui_msg)
return result
@@ -16,9 +16,9 @@ from ag_ui.core import (
TextMessageEndEvent,
TextMessageStartEvent,
)
from agent_framework import AgentProtocol, AgentThread, ChatAgent, TextContent
from agent_framework import AgentProtocol, AgentThread, TextContent
from ._utils import convert_agui_tools_to_agent_framework, generate_event_id
from ._utils import generate_event_id
if TYPE_CHECKING:
from ._agent import AgentConfig
@@ -142,10 +142,14 @@ class HumanInTheLoopOrchestrator(Orchestrator):
True if last message is a tool result
"""
msg = context.last_message
if not msg:
if not msg or not hasattr(msg, "metadata"):
return False
return bool(msg.additional_properties.get("is_tool_result", False))
metadata = getattr(msg, "metadata", None)
if not metadata:
return False
return bool(metadata.get("is_tool_result", False))
async def run(
self,
@@ -270,10 +274,8 @@ class DefaultOrchestrator(Orchestrator):
current_state: dict[str, Any] = initial_state.copy() if initial_state else {}
# Check if agent uses structured outputs (response_format)
# Use isinstance to narrow type for proper attribute access
response_format = None
if isinstance(context.agent, ChatAgent):
response_format = context.agent.chat_options.response_format
chat_options = getattr(context.agent, "chat_options", None)
response_format = getattr(chat_options, "response_format", None) if chat_options else None
skip_text_content = response_format is not None
# Create event bridge
@@ -332,8 +334,9 @@ class DefaultOrchestrator(Orchestrator):
if context.messages:
await thread.on_new_messages(context.messages)
# Use the full incoming message batch to preserve tool-call adjacency
if not context.messages:
# Get the last message as the new input
new_message = context.last_message
if not new_message:
logger.warning("No messages provided in AG-UI input")
yield event_bridge.create_run_finished_event()
return
@@ -359,68 +362,11 @@ Never replace existing data - always append or merge."""
)
messages_to_run.append(state_context_msg)
# Preserve order from client to satisfy provider constraints (assistant tool_calls must
# immediately precede tool result messages). Using the full batch avoids reordering.
messages_to_run.extend(context.messages)
# Handle client tools for hybrid execution
# Client sends tool metadata, server merges with its own tools.
# Client tools have func=None (declaration-only), so @use_function_invocation
# will return the function call without executing (passes back to client).
from agent_framework import BaseChatClient
client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
# Extract server tools - use type narrowing when possible
server_tools: list[Any] = []
if isinstance(context.agent, ChatAgent):
server_tools = context.agent.chat_options.tools or []
else:
# AgentProtocol allows duck-typed implementations - fallback to attribute access
# This supports test mocks and custom agent implementations
try:
chat_options_attr = getattr(context.agent, "chat_options", None)
if chat_options_attr is not None:
server_tools = getattr(chat_options_attr, "tools", None) or []
except AttributeError:
pass
# Register client tools as additional (declaration-only) so they are not executed on server
if client_tools:
if isinstance(context.agent, ChatAgent):
# Type-safe path for ChatAgent
chat_client = context.agent.chat_client
if (
isinstance(chat_client, BaseChatClient)
and chat_client.function_invocation_configuration is not None
):
chat_client.function_invocation_configuration.additional_tools = client_tools
logger.debug(
f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
)
else:
# Fallback for AgentProtocol implementations (test mocks, custom agents)
try:
chat_client_attr = getattr(context.agent, "chat_client", None)
if chat_client_attr is not None:
fic = getattr(chat_client_attr, "function_invocation_configuration", None)
if fic is not None:
fic.additional_tools = client_tools # type: ignore[attr-defined]
logger.debug(
f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
)
except AttributeError:
pass
combined_tools: list[Any] = []
if server_tools:
combined_tools.extend(server_tools)
if client_tools:
combined_tools.extend(client_tools)
messages_to_run.append(new_message)
# Collect all updates to get the final structured output
all_updates: list[Any] = []
async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=combined_tools or None):
async for update in context.agent.run_stream(messages_to_run, thread=thread):
all_updates.append(update)
events = await event_bridge.from_agent_run_update(update)
for event in events:
@@ -428,7 +374,7 @@ Never replace existing data - always append or merge."""
# After agent completes, check if we should stop (waiting for user to confirm changes)
if event_bridge.should_stop_after_confirm:
logger.info("Stopping run after confirm_changes - waiting for user response")
logger.info(" >>> Stopping run after confirm_changes - waiting for user response")
yield event_bridge.create_run_finished_event()
return
@@ -4,13 +4,10 @@
import copy
import uuid
from collections.abc import Callable, MutableMapping, Sequence
from dataclasses import asdict, is_dataclass
from datetime import date, datetime
from typing import Any
from agent_framework import AIFunction, ToolProtocol
def generate_event_id() -> str:
"""Generate a unique event ID."""
@@ -58,109 +55,3 @@ def make_json_safe(obj: Any) -> Any: # noqa: ANN401
if isinstance(obj, dict):
return {key: make_json_safe(value) for key, value in obj.items()} # type: ignore[misc]
return str(obj)
def convert_agui_tools_to_agent_framework(
agui_tools: list[dict[str, Any]] | None,
) -> list[AIFunction[Any, Any]] | None:
"""Convert AG-UI tool definitions to Agent Framework AIFunction declarations.
Creates declaration-only AIFunction instances (no executable implementation).
These are used to tell the LLM about available tools. The actual execution
happens on the client side via @use_function_invocation.
CRITICAL: These tools MUST have func=None so that declaration_only returns True.
This prevents the server from trying to execute client-side tools.
Args:
agui_tools: List of AG-UI tool definitions with name, description, parameters
Returns:
List of AIFunction declarations, or None if no tools provided
"""
if not agui_tools:
return None
result: list[AIFunction[Any, Any]] = []
for tool_def in agui_tools:
# Create declaration-only AIFunction (func=None means no implementation)
# When func=None, the declaration_only property returns True,
# which tells @use_function_invocation to return the function call
# without executing it (so it can be sent back to the client)
func: AIFunction[Any, Any] = AIFunction(
name=tool_def.get("name", ""),
description=tool_def.get("description", ""),
func=None, # CRITICAL: Makes declaration_only=True
input_model=tool_def.get("parameters", {}),
)
result.append(func)
return result
def convert_tools_to_agui_format(
tools: (
ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| Sequence[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None
),
) -> list[dict[str, Any]] | None:
"""Convert tools to AG-UI format.
This sends only the metadata (name, description, JSON schema) to the server.
The actual executable implementation stays on the client side.
The @use_function_invocation decorator handles client-side execution when
the server requests a function.
Args:
tools: Tools to convert (single tool or sequence of tools)
Returns:
List of tool specifications in AG-UI format, or None if no tools provided
"""
if not tools:
return None
# Normalize to list
if not isinstance(tools, list):
tool_list: list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]] = [tools] # type: ignore[list-item]
else:
tool_list = tools # type: ignore[assignment]
results: list[dict[str, Any]] = []
for tool in tool_list:
if isinstance(tool, dict):
# Already in dict format, pass through
results.append(tool) # type: ignore[arg-type]
elif isinstance(tool, AIFunction):
# Convert AIFunction to AG-UI tool format
results.append(
{
"name": tool.name,
"description": tool.description,
"parameters": tool.parameters(),
}
)
elif callable(tool):
# Convert callable to AIFunction first, then to AG-UI format
from agent_framework import ai_function
ai_func = ai_function(tool)
results.append(
{
"name": ai_func.name,
"description": ai_func.description,
"parameters": ai_func.parameters(),
}
)
elif isinstance(tool, ToolProtocol):
# Handle other ToolProtocol implementations
# For now, we'll skip non-AIFunction tools as they may not have
# the parameters() method. This matches .NET behavior which only
# converts AIFunctionDeclaration instances.
continue
return results if results else None