agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] simplify ag-ui run logic, fix mcp bugs, fix anthropic client issues in ag-ui (#3322)

Browse Source 
* Refactor ag-ui to simplify flow

* Refactoring

* Fix backend tool

* Update tests

* Improvements

* Fix mypy

* Fixes

* Fix json serialize errors
This commit is contained in:
Evan Mattson
2026-01-23 14:10:46 +09:00
committed by GitHub
Unverified
parent 9f893a32a6
commit 5436354a83
42 changed files with 2789 additions and 5395 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/ag-ui/agent_framework_ag_ui/__init__.py
-12
View File
@@ -6,13 +6,6 @@ import importlib.metadata
from ._agent import AgentFrameworkAgent
from ._client import AGUIChatClient
from ._confirmation_strategies import (
ConfirmationStrategy,
DefaultConfirmationStrategy,
DocumentWriterConfirmationStrategy,
RecipeConfirmationStrategy,
TaskPlannerConfirmationStrategy,
)
from ._endpoint import add_agent_framework_fastapi_endpoint
from ._event_converters import AGUIEventConverter
from ._http_service import AGUIHttpService
@@ -35,13 +28,8 @@ __all__ = [
"AGUIHttpService",
"AGUIRequest",
"AgentState",
"ConfirmationStrategy",
"DefaultConfirmationStrategy",
"PredictStateConfig",
"RunMetadata",
"TaskPlannerConfirmationStrategy",
"RecipeConfirmationStrategy",
"DocumentWriterConfirmationStrategy",
"DEFAULT_TAGS",
"__version__",
]
python/packages/ag-ui/agent_framework_ag_ui/_agent.py
+12 -88
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
"""AgentFrameworkAgent wrapper for AG-UI protocol - Clean Architecture."""
"""AgentFrameworkAgent wrapper for AG-UI protocol."""
from collections.abc import AsyncGenerator
from typing import Any, cast
@@ -8,13 +8,7 @@ from typing import Any, cast
from ag_ui.core import BaseEvent
from agent_framework import AgentProtocol
from ._confirmation_strategies import ConfirmationStrategy, DefaultConfirmationStrategy
from ._orchestrators import (
DefaultOrchestrator,
ExecutionContext,
HumanInTheLoopOrchestrator,
Orchestrator,
)
from ._run import run_agent_stream
class AgentConfig:
@@ -33,7 +27,7 @@ class AgentConfig:
state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
predict_state_config: Configuration for predictive state updates
use_service_thread: Whether the agent thread is service-managed
require_confirmation: Whether predictive updates require confirmation
require_confirmation: Whether predictive updates require user confirmation before applying
"""
self.state_schema = self._normalize_state_schema(state_schema)
self.predict_state_config = predict_state_config or {}
@@ -58,12 +52,12 @@ class AgentConfig:
base_model_type = None
if base_model_type is not None and isinstance(state_schema, base_model_type):
schema_dict = state_schema.__class__.model_json_schema()
schema_dict = state_schema.__class__.model_json_schema() # type: ignore[union-attr]
return schema_dict.get("properties", {}) or {}
if base_model_type is not None and isinstance(state_schema, type) and issubclass(state_schema, base_model_type):
schema_dict = state_schema.model_json_schema()
return schema_dict.get("properties", {}) or {}
schema_dict = state_schema.model_json_schema() # type: ignore[union-attr]
return schema_dict.get("properties", {}) or {} # type: ignore
return {}
@@ -72,12 +66,7 @@ class AgentFrameworkAgent:
"""Wraps Agent Framework agents for AG-UI protocol compatibility.
Translates between Agent Framework's AgentProtocol and AG-UI's event-based
protocol. Uses orchestrators to handle different execution flows (standard
execution, human-in-the-loop, etc.). Orchestrators are checked in order;
the first matching orchestrator handles the request.
Supports predictive state updates for agentic generative UI, with optional
confirmation requirements configurable per use case.
protocol. Follows a simple linear flow: RunStarted -> content events -> RunFinished.
"""
def __init__(
@@ -88,9 +77,7 @@ class AgentFrameworkAgent:
state_schema: Any | None = None,
predict_state_config: dict[str, dict[str, str]] | None = None,
require_confirmation: bool = True,
orchestrators: list[Orchestrator] | None = None,
use_service_thread: bool = False,
confirmation_strategy: ConfirmationStrategy | None = None,
):
"""Initialize the AG-UI compatible agent wrapper.
@@ -99,15 +86,9 @@ class AgentFrameworkAgent:
name: Optional name for the agent
description: Optional description
state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
predict_state_config: Configuration for predictive state updates.
Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
require_confirmation: Whether predictive updates require confirmation.
Set to False for agentic generative UI that updates automatically.
orchestrators: Custom orchestrators (auto-configured if None).
Orchestrators are checked in order; first match handles the request.
use_service_thread: Whether the agent thread is service-managed.
confirmation_strategy: Strategy for generating confirmation messages.
Defaults to DefaultConfirmationStrategy if None.
predict_state_config: Configuration for predictive state updates
require_confirmation: Whether predictive updates require user confirmation before applying
use_service_thread: Whether the agent thread is service-managed
"""
self.agent = agent
self.name = name or getattr(agent, "name", "agent")
@@ -120,74 +101,17 @@ class AgentFrameworkAgent:
require_confirmation=require_confirmation,
)
# Configure orchestrators
if orchestrators is None:
self.orchestrators = self._default_orchestrators()
else:
self.orchestrators = orchestrators
# Configure confirmation strategy
if confirmation_strategy is None:
self.confirmation_strategy: ConfirmationStrategy = DefaultConfirmationStrategy()
else:
self.confirmation_strategy = confirmation_strategy
def _default_orchestrators(self) -> list[Orchestrator]:
"""Create default orchestrator chain.
Returns:
List of orchestrators in priority order. First matching orchestrator
handles the request, so order matters.
"""
return [
HumanInTheLoopOrchestrator(), # Handle tool approval responses
# Add more specialized orchestrators here as needed
DefaultOrchestrator(), # Fallback: standard agent execution
]
async def run_agent(
self,
input_data: dict[str, Any],
) -> AsyncGenerator[BaseEvent, None]:
"""Run the agent and yield AG-UI events.
This is the ONLY public method - much simpler than the original 376-line
implementation. All orchestration logic has been extracted into dedicated
Orchestrator classes.
The method creates an ExecutionContext with all needed data, then finds
the first orchestrator that can handle the request and delegates to it.
Args:
input_data: The AG-UI run input containing messages, state, etc.
Yields:
AG-UI events
Raises:
RuntimeError: If no orchestrator matches (should never happen if
DefaultOrchestrator is last in the chain)
"""
# Create execution context with all needed data
context = ExecutionContext(
input_data=input_data,
agent=self.agent,
config=self.config,
confirmation_strategy=self.confirmation_strategy,
)
# Find matching orchestrator and execute
for orchestrator in self.orchestrators:
if orchestrator.can_handle(context):
async for event in orchestrator.run(context):
yield event
return
# Should never reach here if DefaultOrchestrator is last
raise RuntimeError("No orchestrator matched - check configuration")
__all__ = [
"AgentFrameworkAgent",
"AgentConfig",
]
async for event in run_agent_stream(input_data, self.agent, self.config):
yield event
python/packages/ag-ui/agent_framework_ag_ui/_client.py
+2 -2
View File
@@ -74,7 +74,7 @@ def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseCha
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 def streaming_wrapper(self: Any, *args: Any, **kwargs: Any) -> AsyncIterable[ChatResponseUpdate]:
async for update in original_get_streaming_response(self, *args, **kwargs):
_unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
yield update
@@ -84,7 +84,7 @@ def _apply_server_function_call_unwrap(chat_client: TBaseChatClient) -> TBaseCha
original_get_response = chat_client.get_response
@wraps(original_get_response)
async def response_wrapper(self, *args: Any, **kwargs: Any) -> ChatResponse:
async def response_wrapper(self: Any, *args: Any, **kwargs: Any) -> ChatResponse:
response = await original_get_response(self, *args, **kwargs)
if response.messages:
for message in response.messages:
python/packages/ag-ui/agent_framework_ag_ui/_confirmation_strategies.py
-217
View File
@@ -1,217 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Confirmation strategies for human-in-the-loop approval flows.
Each agent can provide a custom confirmation strategy to generate domain-specific
messages when users approve or reject changes/actions.
"""
from abc import ABC, abstractmethod
from typing import Any
class ConfirmationStrategy(ABC):
"""Strategy for generating confirmation messages during human-in-the-loop flows.
Subclasses must define the message properties. The methods use those properties
by default, but can be overridden for complete customization.
"""
@property
@abstractmethod
def approval_header(self) -> str:
"""Header for approval accepted message. Must be overridden."""
...
@property
@abstractmethod
def approval_footer(self) -> str:
"""Footer for approval accepted message. Must be overridden."""
...
@property
@abstractmethod
def rejection_message(self) -> str:
"""Message when user rejects. Must be overridden."""
...
@property
@abstractmethod
def state_confirmed_message(self) -> str:
"""Message when state is confirmed. Must be overridden."""
...
@property
@abstractmethod
def state_rejected_message(self) -> str:
"""Message when state is rejected. Must be overridden."""
...
def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
"""Generate message when user approves function execution.
Default implementation uses header/footer properties.
Override for complete customization.
Args:
steps: List of approved steps with 'description', 'status', etc.
Returns:
Message to display to user
"""
enabled_steps = [s for s in steps if s.get("status") == "enabled"]
message_parts = [self.approval_header.format(count=len(enabled_steps))]
for i, step in enumerate(enabled_steps, 1):
message_parts.append(f"{i}. {step['description']}\n")
message_parts.append(self.approval_footer)
return "".join(message_parts)
def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
"""Generate message when user rejects function execution.
Args:
steps: List of rejected steps
Returns:
Message to display to user
"""
return self.rejection_message
def on_state_confirmed(self) -> str:
"""Generate message when user confirms predictive state changes.
Returns:
Message to display to user
"""
return self.state_confirmed_message
def on_state_rejected(self) -> str:
"""Generate message when user rejects predictive state changes.
Returns:
Message to display to user
"""
return self.state_rejected_message
class DefaultConfirmationStrategy(ConfirmationStrategy):
"""Generic confirmation messages suitable for most agents."""
@property
def approval_header(self) -> str:
return "Executing {count} approved steps:\n\n"
@property
def approval_footer(self) -> str:
return "\nAll steps completed successfully!"
@property
def rejection_message(self) -> str:
return "No problem! What would you like me to change about the plan?"
@property
def state_confirmed_message(self) -> str:
return "Changes confirmed and applied successfully!"
@property
def state_rejected_message(self) -> str:
return "No problem! What would you like me to change?"
class TaskPlannerConfirmationStrategy(ConfirmationStrategy):
"""Domain-specific confirmation messages for task planning agents."""
@property
def approval_header(self) -> str:
return "Executing your requested tasks:\n\n"
@property
def approval_footer(self) -> str:
return "\nAll tasks completed successfully!"
@property
def rejection_message(self) -> str:
return "No problem! Let me revise the plan. What would you like me to change?"
@property
def state_confirmed_message(self) -> str:
return "Tasks confirmed and ready to execute!"
@property
def state_rejected_message(self) -> str:
return "No problem! How should I adjust the task list?"
class RecipeConfirmationStrategy(ConfirmationStrategy):
"""Domain-specific confirmation messages for recipe agents."""
@property
def approval_header(self) -> str:
return "Updating your recipe:\n\n"
@property
def approval_footer(self) -> str:
return "\nRecipe updated successfully!"
@property
def rejection_message(self) -> str:
return "No problem! What ingredients or steps should I change?"
@property
def state_confirmed_message(self) -> str:
return "Recipe changes applied successfully!"
@property
def state_rejected_message(self) -> str:
return "No problem! What would you like me to adjust in the recipe?"
class DocumentWriterConfirmationStrategy(ConfirmationStrategy):
"""Domain-specific confirmation messages for document writing agents."""
@property
def approval_header(self) -> str:
return "Applying your edits:\n\n"
@property
def approval_footer(self) -> str:
return "\nDocument updated successfully!"
@property
def rejection_message(self) -> str:
return "No problem! Which changes should I keep or modify?"
@property
def state_confirmed_message(self) -> str:
return "Document edits applied!"
@property
def state_rejected_message(self) -> str:
return "No problem! What should I change about the document?"
def apply_confirmation_strategy(
strategy: ConfirmationStrategy | None,
accepted: bool,
steps: list[dict[str, Any]],
) -> str:
"""Apply a confirmation strategy to generate a message.
This helper consolidates the pattern used in multiple orchestrators.
Args:
strategy: Strategy to use, or None for default
accepted: Whether the user approved
steps: List of steps (may be empty for state confirmations)
Returns:
Generated message string
"""
if strategy is None:
strategy = DefaultConfirmationStrategy()
if not steps:
# State confirmation (no steps)
return strategy.on_state_confirmed() if accepted else strategy.on_state_rejected()
# Step-based approval
return strategy.on_approval_accepted(steps) if accepted else strategy.on_approval_rejected(steps)
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py
+12 -10
View File
@@ -4,7 +4,7 @@
import copy
import logging
from collections.abc import Sequence
from collections.abc import AsyncGenerator, Sequence
from typing import Any
from ag_ui.encoder import EventEncoder
@@ -56,8 +56,8 @@ def add_agent_framework_fastapi_endpoint(
else:
wrapped_agent = agent
@app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies) # type: ignore[arg-type]
async def agent_endpoint(request_body: AGUIRequest): # type: ignore[misc]
@app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies, response_model=None) # type: ignore[arg-type]
async def agent_endpoint(request_body: AGUIRequest) -> StreamingResponse | dict[str, str]:
"""Handle AG-UI agent requests.
Note: Function is accessed via FastAPI's decorator registration,
@@ -77,17 +77,19 @@ def add_agent_framework_fastapi_endpoint(
)
logger.info(f"Received request at {path}: {input_data.get('run_id', 'no-run-id')}")
async def event_generator():
async def event_generator() -> AsyncGenerator[str, None]:
encoder = EventEncoder()
event_count = 0
async for event in wrapped_agent.run_agent(input_data):
event_count += 1
logger.debug(f"[{path}] Event {event_count}: {type(event).__name__}")
# Log event payload for debugging
if hasattr(event, "model_dump"):
event_data = event.model_dump(exclude_none=True)
logger.debug(f"[{path}] Event payload: {event_data}")
event_type_name = getattr(event, "type", type(event).__name__)
# Log important events at INFO level
if "TOOL_CALL" in str(event_type_name) or "RUN" in str(event_type_name):
if hasattr(event, "model_dump"):
event_data = event.model_dump(exclude_none=True)
logger.info(f"[{path}] Event {event_count}: {event_type_name} - {event_data}")
else:
logger.info(f"[{path}] Event {event_count}: {event_type_name}")
encoded = encoder.encode(event)
logger.debug(
python/packages/ag-ui/agent_framework_ag_ui/_events.py
-589
View File
@@ -1,589 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Event bridge for converting Agent Framework events to AG-UI protocol."""
import json
import logging
import re
from copy import deepcopy
from typing import Any
from ag_ui.core import (
BaseEvent,
CustomEvent,
RunFinishedEvent,
RunStartedEvent,
StateDeltaEvent,
StateSnapshotEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework import (
AgentResponseUpdate,
Content,
prepare_function_call_results,
)
from ._utils import extract_state_from_tool_args, generate_event_id, make_json_safe, safe_json_parse
logger = logging.getLogger(__name__)
class AgentFrameworkEventBridge:
"""Converts Agent Framework responses to AG-UI events."""
def __init__(
self,
run_id: str,
thread_id: str,
predict_state_config: dict[str, dict[str, str]] | None = None,
current_state: dict[str, Any] | None = None,
skip_text_content: bool = False,
require_confirmation: bool = True,
approval_tool_name: str | None = None,
) -> None:
"""
Initialize the event bridge.
Args:
run_id: The run identifier.
thread_id: The thread identifier.
predict_state_config: Configuration for predictive state updates.
Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
current_state: Reference to the current state dict for tracking updates.
skip_text_content: If True, skip emitting TextMessageContentEvents (for structured outputs).
require_confirmation: Whether predictive state updates require user confirmation.
"""
self.run_id = run_id
self.thread_id = thread_id
self.current_message_id: str | None = None
self.current_tool_call_id: str | None = None
self.current_tool_call_name: str | None = None # Track the tool name across streaming chunks
self.predict_state_config = predict_state_config or {}
self.current_state = current_state or {}
self.pending_state_updates: dict[str, Any] = {} # Track updates from tool calls
self.skip_text_content = skip_text_content
self.require_confirmation = require_confirmation
self.approval_tool_name = approval_tool_name
# For predictive state updates: accumulate streaming arguments
self.streaming_tool_args: str = "" # Accumulated JSON string
self.last_emitted_state: dict[str, Any] = {} # Track last emitted state to avoid duplicates
self.state_delta_count: int = 0 # Counter for sampling log output
self.should_stop_after_confirm: bool = False # Flag to stop run after confirm_changes
self.suppressed_summary: str = "" # Store LLM summary to show after confirmation
async def from_agent_run_update(self, update: AgentResponseUpdate) -> list[BaseEvent]:
"""
Convert an AgentResponseUpdate to AG-UI events.
Args:
update: The agent run update to convert.
Returns:
List of AG-UI events.
"""
events: list[BaseEvent] = []
logger.info(f"Processing AgentRunUpdate with {len(update.contents)} content items")
for idx, content in enumerate(update.contents):
logger.info(f" Content {idx}: type={type(content).__name__}")
match content.type:
case "text":
events.extend(self._handle_text_content(content))
case "function_call":
events.extend(self._handle_function_call_content(content))
case "function_result":
events.extend(self._handle_function_result_content(content))
case "function_approval_request":
events.extend(self._handle_function_approval_request_content(content))
case _:
logger.warning(f" Unsupported content type: {content.type}, skipping.")
return events
def _handle_text_content(self, content: Content) -> list[BaseEvent]:
events: list[BaseEvent] = []
logger.info(f" TextContent found: length={len(content.text)}") # type: ignore[arg-type]
logger.info(
" Flags: skip_text_content=%s, should_stop_after_confirm=%s",
self.skip_text_content,
self.should_stop_after_confirm,
)
if self.skip_text_content:
logger.info(" SKIPPING TextContent: skip_text_content is True")
return events
if self.should_stop_after_confirm:
logger.info(" SKIPPING TextContent: waiting for confirm_changes response")
self.suppressed_summary += content.text # type: ignore[operator]
logger.info(f" Suppressed summary length={len(self.suppressed_summary)}")
return events
# Skip empty text chunks to avoid emitting
# TextMessageContentEvent with an empty `delta` which fails
# Pydantic validation (AG-UI requires non-empty strings).
if not content.text:
logger.info(" SKIPPING TextContent: empty chunk")
return events
if not self.current_message_id:
self.current_message_id = generate_event_id()
start_event = TextMessageStartEvent(
message_id=self.current_message_id,
role="assistant",
)
logger.info(f" EMITTING TextMessageStartEvent with message_id={self.current_message_id}")
events.append(start_event)
event = TextMessageContentEvent(
message_id=self.current_message_id,
delta=content.text,
)
logger.info(f" EMITTING TextMessageContentEvent with text_len={len(content.text)}")
events.append(event)
return events
def _handle_function_call_content(self, content: Content) -> list[BaseEvent]:
events: list[BaseEvent] = []
if content.name:
logger.debug(f"Tool call: {content.name} (call_id: {content.call_id})")
if not content.name and not content.call_id and not self.current_tool_call_name:
args_length = len(str(content.arguments)) if content.arguments else 0
logger.warning(f"Content missing name and call_id. args_length={args_length}")
tool_call_id = self._coalesce_tool_call_id(content)
# Only emit ToolCallStartEvent once per tool call (when it's a new tool call)
if content.name and tool_call_id != self.current_tool_call_id:
self.streaming_tool_args = ""
self.state_delta_count = 0
self.current_tool_call_id = tool_call_id
self.current_tool_call_name = content.name
tool_start_event = ToolCallStartEvent(
tool_call_id=tool_call_id,
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}'")
events.append(tool_start_event)
elif tool_call_id:
self.current_tool_call_id = tool_call_id
if content.arguments:
delta_str = (
content.arguments
if isinstance(content.arguments, str)
else json.dumps(make_json_safe(content.arguments))
)
logger.info(f"Emitting ToolCallArgsEvent with delta_length={len(delta_str)}, id='{tool_call_id}'")
args_event = ToolCallArgsEvent(
tool_call_id=tool_call_id,
delta=delta_str,
)
events.append(args_event)
events.extend(self._emit_predictive_state_deltas(delta_str))
return events
def _coalesce_tool_call_id(self, content: Content) -> str:
if content.call_id:
return content.call_id
if self.current_tool_call_id:
return self.current_tool_call_id
return generate_event_id()
def _emit_predictive_state_deltas(self, argument_chunk: str) -> list[BaseEvent]:
events: list[BaseEvent] = []
if not self.current_tool_call_name or not self.predict_state_config:
return events
self.streaming_tool_args += argument_chunk
logger.debug(
"Predictive state: accumulated %s chars for tool '%s'",
len(self.streaming_tool_args),
self.current_tool_call_name,
)
parsed_args = safe_json_parse(self.streaming_tool_args)
if parsed_args is None:
for state_key, config in self.predict_state_config.items():
if config["tool"] != self.current_tool_call_name:
continue
tool_arg_name = config["tool_argument"]
pattern = rf'"{re.escape(tool_arg_name)}":\s*"([^"]*)'
match = re.search(pattern, self.streaming_tool_args)
if match:
partial_value = match.group(1).replace("\\n", "\n").replace('\\"', '"').replace("\\\\", "\\")
if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != partial_value:
state_delta_event = StateDeltaEvent(
delta=[
{
"op": "replace",
"path": f"/{state_key}",
"value": partial_value,
}
],
)
self.state_delta_count += 1
if self.state_delta_count % 10 == 1:
logger.info(
"StateDeltaEvent #%s for '%s': op=replace, path=/%s, value_length=%s",
self.state_delta_count,
state_key,
state_key,
len(str(partial_value)),
)
elif self.state_delta_count % 100 == 0:
logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
events.append(state_delta_event)
self.last_emitted_state[state_key] = partial_value
self.pending_state_updates[state_key] = partial_value
if parsed_args:
for state_key, config in self.predict_state_config.items():
if config["tool"] != self.current_tool_call_name:
continue
tool_arg_name = config["tool_argument"]
state_value = extract_state_from_tool_args(parsed_args, tool_arg_name)
if state_value is None:
continue
if state_key not in self.last_emitted_state or self.last_emitted_state[state_key] != state_value:
state_delta_event = StateDeltaEvent(
delta=[
{
"op": "replace",
"path": f"/{state_key}",
"value": state_value,
}
],
)
self.state_delta_count += 1
if self.state_delta_count % 10 == 1:
logger.info(
"StateDeltaEvent #%s for '%s': op=replace, path=/%s, value_length=%s",
self.state_delta_count,
state_key,
state_key,
len(str(state_value)),
)
elif self.state_delta_count % 100 == 0:
logger.info(f"StateDeltaEvent #{self.state_delta_count} emitted")
events.append(state_delta_event)
self.last_emitted_state[state_key] = state_value
self.pending_state_updates[state_key] = state_value
return events
def _handle_function_result_content(self, content: Content) -> list[BaseEvent]:
events: list[BaseEvent] = []
if content.call_id:
end_event = ToolCallEndEvent(
tool_call_id=content.call_id,
)
logger.info(f"Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
events.append(end_event)
if self.state_delta_count > 0:
logger.info(
"Tool call '%s' complete: emitted %s StateDeltaEvents total",
content.call_id,
self.state_delta_count,
)
self.streaming_tool_args = ""
self.state_delta_count = 0
result_message_id = generate_event_id()
result_content = prepare_function_call_results(content.result)
result_event = ToolCallResultEvent(
message_id=result_message_id,
tool_call_id=content.call_id, # type: ignore[arg-type]
content=result_content,
role="tool",
)
events.append(result_event)
events.extend(self._emit_state_snapshot_and_confirmation())
return events
def _emit_state_snapshot_and_confirmation(self) -> list[BaseEvent]:
events: list[BaseEvent] = []
if self.pending_state_updates:
for key, value in self.pending_state_updates.items():
self.current_state[key] = value
logger.info(f"Emitting StateSnapshotEvent with keys: {list(self.current_state.keys())}")
if "recipe" in self.current_state:
recipe = self.current_state["recipe"]
logger.info(
"Recipe fields: title=%s, skill_level=%s, ingredients_count=%s, instructions_count=%s",
recipe.get("title"),
recipe.get("skill_level"),
len(recipe.get("ingredients", [])),
len(recipe.get("instructions", [])),
)
state_snapshot_event = StateSnapshotEvent(
snapshot=self.current_state,
)
events.append(state_snapshot_event)
tool_was_predictive = False
logger.debug(
"Checking predictive state: current_tool='%s', predict_config=%s",
self.current_tool_call_name,
list(self.predict_state_config.keys()) if self.predict_state_config else "None",
)
for state_key, config in self.predict_state_config.items():
if self.current_tool_call_name and config["tool"] == self.current_tool_call_name:
logger.info(
"Tool '%s' matches predictive config for state key '%s'",
self.current_tool_call_name,
state_key,
)
tool_was_predictive = True
break
if tool_was_predictive and self.require_confirmation:
events.extend(self._emit_confirm_changes_tool_call())
elif tool_was_predictive:
logger.info("Skipping confirm_changes - require_confirmation is False")
self.pending_state_updates.clear()
self.last_emitted_state = deepcopy(self.current_state)
self.current_tool_call_name = None
return events
def _emit_confirm_changes_tool_call(self, function_call: Content | None = None) -> list[BaseEvent]:
"""Emit a confirm_changes tool call for Dojo UI compatibility.
Args:
function_call: Optional function call that needs confirmation.
If provided, includes function info in the confirm_changes args
so Dojo UI can display what's being confirmed.
"""
events: list[BaseEvent] = []
confirm_call_id = generate_event_id()
logger.info("Emitting confirm_changes tool call for predictive update")
confirm_start = ToolCallStartEvent(
tool_call_id=confirm_call_id,
tool_call_name="confirm_changes",
parent_message_id=self.current_message_id,
)
events.append(confirm_start)
# Include function info if this is for a function approval
# This helps Dojo UI display meaningful confirmation info
if function_call:
args_dict = {
"function_name": function_call.name,
"function_call_id": function_call.call_id,
"function_arguments": make_json_safe(function_call.parse_arguments() or {}),
"steps": [
{
"description": f"Execute {function_call.name}",
"status": "enabled",
}
],
}
args_json = json.dumps(args_dict)
else:
args_json = "{}"
confirm_args = ToolCallArgsEvent(
tool_call_id=confirm_call_id,
delta=args_json,
)
events.append(confirm_args)
confirm_end = ToolCallEndEvent(
tool_call_id=confirm_call_id,
)
events.append(confirm_end)
self.should_stop_after_confirm = True
logger.info("Set flag to stop run after confirm_changes")
return events
def _emit_function_approval_tool_call(self, function_call: Content) -> list[BaseEvent]:
"""Emit a tool call that can drive UI approval for function requests."""
tool_call_name = "confirm_changes"
if self.approval_tool_name and self.approval_tool_name != function_call.name:
tool_call_name = self.approval_tool_name
tool_call_id = generate_event_id()
tool_start = ToolCallStartEvent(
tool_call_id=tool_call_id,
tool_call_name=tool_call_name,
parent_message_id=self.current_message_id,
)
events: list[BaseEvent] = [tool_start]
args_dict = {
"function_name": function_call.name,
"function_call_id": function_call.call_id,
"function_arguments": make_json_safe(function_call.parse_arguments() or {}),
"steps": [
{
"description": f"Execute {function_call.name}",
"status": "enabled",
}
],
}
args_json = json.dumps(args_dict)
events.append(
ToolCallArgsEvent(
tool_call_id=tool_call_id,
delta=args_json,
)
)
events.append(
ToolCallEndEvent(
tool_call_id=tool_call_id,
)
)
self.should_stop_after_confirm = True
logger.info("Set flag to stop run after confirm_changes")
return events
def _handle_function_approval_request_content(self, content: Content) -> list[BaseEvent]:
events: list[BaseEvent] = []
logger.info("=== FUNCTION APPROVAL REQUEST ===")
logger.info(f" Function: {content.function_call.name}") # type: ignore[union-attr]
logger.info(f" Call ID: {content.function_call.call_id}") # type: ignore[union-attr]
parsed_args = content.function_call.parse_arguments() # type: ignore[union-attr]
parsed_arg_keys = list(parsed_args.keys()) if parsed_args else "None"
logger.info(f" Parsed args keys: {parsed_arg_keys}")
if parsed_args and self.predict_state_config:
logger.info(
" Checking predict_state_config keys: %s",
list(self.predict_state_config.keys()) if self.predict_state_config else "None",
)
for state_key, config in self.predict_state_config.items():
if config["tool"] != content.function_call.name: # type: ignore[union-attr]
continue
tool_arg_name = config["tool_argument"]
logger.info(
" MATCHED tool '%s' for state key '%s', arg='%s'",
content.function_call.name, # type: ignore[union-attr]
state_key,
tool_arg_name,
)
state_value = extract_state_from_tool_args(parsed_args, tool_arg_name)
if state_value is None:
logger.warning(f" Tool argument '{tool_arg_name}' not found in parsed args")
continue
self.current_state[state_key] = state_value
logger.info("Emitting StateSnapshotEvent for key '%s', value type: %s", state_key, type(state_value)) # type: ignore
state_snapshot = StateSnapshotEvent(
snapshot=self.current_state,
)
events.append(state_snapshot)
if content.function_call.call_id: # type: ignore[union-attr]
end_event = ToolCallEndEvent(
tool_call_id=content.function_call.call_id, # type: ignore[union-attr]
)
logger.info(f"Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'") # type: ignore[union-attr]
events.append(end_event)
# Emit the function_approval_request custom event for UI implementations that support it
approval_event = CustomEvent(
name="function_approval_request",
value={
"id": content.id,
"function_call": {
"call_id": content.function_call.call_id, # type: ignore[union-attr]
"name": content.function_call.name, # type: ignore[union-attr]
"arguments": content.function_call.parse_arguments(), # type: ignore[union-attr]
},
},
)
logger.info(f"Emitting function_approval_request custom event for '{content.function_call.name}'") # type: ignore[union-attr]
events.append(approval_event)
# Emit a UI-friendly approval tool call for function approvals.
if self.require_confirmation:
events.extend(self._emit_function_approval_tool_call(content.function_call)) # type: ignore[arg-type]
# Signal orchestrator to stop the run and wait for user approval response
self.should_stop_after_confirm = True
logger.info("Set flag to stop run - waiting for function approval response")
return events
def create_run_started_event(self) -> RunStartedEvent:
"""Create a run started event."""
return RunStartedEvent(
run_id=self.run_id,
thread_id=self.thread_id,
)
def create_run_finished_event(self, result: Any = None) -> RunFinishedEvent:
"""Create a run finished event."""
return RunFinishedEvent(
run_id=self.run_id,
thread_id=self.thread_id,
result=result,
)
def create_message_start_event(self, message_id: str, role: str = "assistant") -> TextMessageStartEvent:
"""Create a message start event."""
return TextMessageStartEvent(
message_id=message_id,
role=role, # type: ignore
)
def create_message_end_event(self, message_id: str) -> TextMessageEndEvent:
"""Create a message end event."""
return TextMessageEndEvent(
message_id=message_id,
)
def create_state_snapshot_event(self, state: dict[str, Any]) -> StateSnapshotEvent:
"""Create a state snapshot event.
Args:
state: The complete state snapshot.
Returns:
StateSnapshotEvent.
"""
return StateSnapshotEvent(
snapshot=state,
)
def create_state_delta_event(self, delta: list[dict[str, Any]]) -> StateDeltaEvent:
"""Create a state delta event using JSON Patch format (RFC 6902).
Args:
delta: List of JSON Patch operations.
Returns:
StateDeltaEvent.
"""
return StateDeltaEvent(
delta=delta,
)
python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
+24 -46
View File
@@ -17,7 +17,6 @@ from ._utils import (
AGUI_TO_FRAMEWORK_ROLE,
FRAMEWORK_TO_AGUI_ROLE,
get_role_value,
make_json_safe,
normalize_agui_role,
safe_json_parse,
)
@@ -253,22 +252,19 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
tool_calls = raw_msg.get("tool_calls") or raw_msg.get("toolCalls")
if not isinstance(tool_calls, list):
continue
tool_calls_list = cast(list[Any], tool_calls)
for tool_call in tool_calls_list:
for tool_call in tool_calls:
if not isinstance(tool_call, dict):
continue
tool_call_dict = cast(dict[str, Any], tool_call)
if str(tool_call_dict.get("id", "")) != tool_call_id:
if str(tool_call.get("id", "")) != tool_call_id:
continue
function_payload = tool_call_dict.get("function")
function_payload = tool_call.get("function")
if not isinstance(function_payload, dict):
return
function_payload_dict = cast(dict[str, Any], function_payload)
existing_args = function_payload_dict.get("arguments")
existing_args = function_payload.get("arguments")
if isinstance(existing_args, str):
function_payload_dict["arguments"] = json.dumps(make_json_safe(modified_args))
function_payload["arguments"] = json.dumps(modified_args)
else:
function_payload_dict["arguments"] = modified_args
function_payload["arguments"] = modified_args
return
def _find_matching_func_call(call_id: str) -> Content | None:
@@ -378,9 +374,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
# a proper FunctionApprovalResponseContent. This enables the agent framework
# to execute the approved tool (fix for GitHub issue #3034).
accepted = parsed.get("accepted", False) if parsed is not None else False
approval_payload_text = (
result_content if isinstance(result_content, str) else json.dumps(make_json_safe(parsed))
)
approval_payload_text = result_content if isinstance(result_content, str) else json.dumps(parsed)
# Log the full approval payload to debug modified arguments
import logging
@@ -436,8 +430,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
if desc:
approved_by_description[str(desc)] = step_item_dict
merged_steps: list[Any] = []
original_steps_list = cast(list[Any], original_steps)
for orig_step in original_steps_list:
for orig_step in original_steps:
if not isinstance(orig_step, dict):
merged_steps.append(orig_step)
continue
@@ -457,9 +450,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
# Keep the original tool call and AG-UI snapshot in sync with approved args.
updated_args = (
json.dumps(make_json_safe(merged_args))
if isinstance(matching_func_call.arguments, str)
else merged_args
json.dumps(merged_args) if isinstance(matching_func_call.arguments, str) else merged_args
)
matching_func_call.arguments = updated_args
_update_tool_call_arguments(messages, str(approval_call_id), merged_args)
@@ -467,7 +458,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
func_call_for_approval = Content.from_function_call(
call_id=matching_func_call.call_id, # type: ignore[arg-type]
name=matching_func_call.name, # type: ignore[arg-type]
arguments=json.dumps(make_json_safe(filtered_args)),
arguments=json.dumps(filtered_args),
)
logger.info(f"Using modified arguments from approval: {filtered_args}")
else:
@@ -503,9 +494,9 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
if isinstance(result_content, str):
func_result = result_content
elif isinstance(result_content, dict):
func_result = cast(dict[str, Any], result_content)
func_result = result_content
elif isinstance(result_content, list):
func_result = cast(list[Any], result_content)
func_result = result_content
else:
func_result = str(result_content)
chat_msg = ChatMessage(
@@ -739,40 +730,35 @@ def agui_messages_to_snapshot_format(messages: list[dict[str, Any]]) -> list[dic
if isinstance(content, list):
# Convert content array format to simple string
text_parts: list[str] = []
content_list = cast(list[Any], content)
for item in content_list:
for item in content:
if isinstance(item, dict):
item_dict = cast(dict[str, Any], item)
# Convert 'input_text' to 'text' type
if item_dict.get("type") == "input_text":
text_parts.append(str(item_dict.get("text", "")))
elif item_dict.get("type") == "text":
text_parts.append(str(item_dict.get("text", "")))
if item.get("type") == "input_text":
text_parts.append(str(item.get("text", "")))
elif item.get("type") == "text":
text_parts.append(str(item.get("text", "")))
else:
# Other types - just extract text field if present
text_parts.append(str(item_dict.get("text", "")))
text_parts.append(str(item.get("text", "")))
normalized_msg["content"] = "".join(text_parts)
elif content is None:
normalized_msg["content"] = ""
tool_calls = normalized_msg.get("tool_calls") or normalized_msg.get("toolCalls")
if isinstance(tool_calls, list):
tool_calls_list = cast(list[Any], tool_calls)
for tool_call in tool_calls_list:
for tool_call in tool_calls:
if not isinstance(tool_call, dict):
continue
tool_call_dict = cast(dict[str, Any], tool_call)
function_payload = tool_call_dict.get("function")
function_payload = tool_call.get("function")
if not isinstance(function_payload, dict):
continue
function_payload_dict = cast(dict[str, Any], function_payload)
if "arguments" not in function_payload_dict:
if "arguments" not in function_payload:
continue
arguments = function_payload_dict.get("arguments")
arguments = function_payload.get("arguments")
if arguments is None:
function_payload_dict["arguments"] = ""
function_payload["arguments"] = ""
elif not isinstance(arguments, str):
function_payload_dict["arguments"] = json.dumps(make_json_safe(arguments))
function_payload["arguments"] = json.dumps(arguments)
# Normalize tool_call_id to toolCallId for tool messages
normalized_msg["role"] = normalize_agui_role(normalized_msg.get("role"))
@@ -786,11 +772,3 @@ def agui_messages_to_snapshot_format(messages: list[dict[str, Any]]) -> list[dic
result.append(normalized_msg)
return result
__all__ = [
"agui_messages_to_agent_framework",
"agent_framework_messages_to_agui",
"agui_messages_to_snapshot_format",
"extract_text_from_contents",
]
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py
+11 -151
View File
@@ -1,22 +1,21 @@
# Copyright (c) Microsoft. All rights reserved.
"""Helper functions for orchestration logic."""
"""Helper functions for orchestration logic.
Most orchestration helpers have been moved inline to _run.py.
This module retains utilities that may be useful for testing or extensions.
"""
import json
import logging
from typing import TYPE_CHECKING, Any
from typing import Any
from ag_ui.core import StateSnapshotEvent
from agent_framework import (
ChatMessage,
Content,
)
from .._utils import get_role_value, make_json_safe, safe_json_parse
if TYPE_CHECKING:
from .._events import AgentFrameworkEventBridge
from ._state_manager import StateManager
from .._utils import get_role_value
logger = logging.getLogger(__name__)
@@ -111,53 +110,6 @@ def tool_name_for_call_id(
return str(name) if name else None
def tool_calls_match_state(
provider_messages: list[ChatMessage],
state_manager: "StateManager",
) -> bool:
"""Check if tool calls in messages match current state.
Args:
provider_messages: Messages to check
state_manager: State manager with config and current state
Returns:
True if tool calls match state configuration
"""
if not state_manager.predict_state_config or not state_manager.current_state:
return False
for state_key, config in state_manager.predict_state_config.items():
tool_name = config["tool"]
tool_arg_name = config["tool_argument"]
tool_args: dict[str, Any] | None = None
for msg in reversed(provider_messages):
if get_role_value(msg) != "assistant":
continue
for content in msg.contents:
if content.type == "function_call" and content.name == tool_name:
tool_args = safe_json_parse(content.arguments)
break
if tool_args is not None:
break
if not tool_args:
return False
if tool_arg_name == "*":
state_value = tool_args
elif tool_arg_name in tool_args:
state_value = tool_args[tool_arg_name]
else:
return False
if state_manager.current_state.get(state_key) != state_value:
return False
return True
def schema_has_steps(schema: Any) -> bool:
"""Check if a schema has a steps array property.
@@ -202,45 +154,10 @@ def select_approval_tool_name(client_tools: list[Any] | None) -> str | None:
return None
def select_messages_to_run(
provider_messages: list[ChatMessage],
state_manager: "StateManager",
) -> list[ChatMessage]:
"""Select and prepare messages for agent execution.
Injects state context message when appropriate.
Args:
provider_messages: Original messages from client
state_manager: State manager instance
Returns:
Messages ready for agent execution
"""
if not provider_messages:
return []
is_new_user_turn = get_role_value(provider_messages[-1]) == "user"
conversation_has_tool_calls = tool_calls_match_state(provider_messages, state_manager)
state_context_msg = state_manager.state_context_message(
is_new_user_turn=is_new_user_turn, conversation_has_tool_calls=conversation_has_tool_calls
)
if not state_context_msg:
return list(provider_messages)
messages_to_run = [msg for msg in provider_messages if not is_state_context_message(msg)]
if pending_tool_call_ids(messages_to_run):
return messages_to_run
insert_index = len(messages_to_run) - 1 if is_new_user_turn else len(messages_to_run)
if insert_index < 0:
insert_index = 0
messages_to_run.insert(insert_index, state_context_msg)
return messages_to_run
def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]:
"""Build metadata dict with truncated string values.
"""Build metadata dict with truncated string values for Azure compatibility.
Azure has a 512 character limit per metadata value.
Args:
thread_metadata: Raw metadata dict
@@ -252,70 +169,13 @@ def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any
return {}
safe_metadata: dict[str, Any] = {}
for key, value in thread_metadata.items():
value_str = value if isinstance(value, str) else json.dumps(make_json_safe(value))
value_str = value if isinstance(value, str) else json.dumps(value)
if len(value_str) > 512:
value_str = value_str[:512]
safe_metadata[key] = value_str
return safe_metadata
def collect_approved_state_snapshots(
provider_messages: list[ChatMessage],
predict_state_config: dict[str, dict[str, str]] | None,
current_state: dict[str, Any],
event_bridge: "AgentFrameworkEventBridge",
) -> list[StateSnapshotEvent]:
"""Collect state snapshots from approved function calls.
Args:
provider_messages: Messages containing approvals
predict_state_config: Predictive state configuration
current_state: Current state dict (will be mutated)
event_bridge: Event bridge for creating events
Returns:
List of state snapshot events
"""
if not predict_state_config:
return []
events: list[StateSnapshotEvent] = []
for msg in provider_messages:
if get_role_value(msg) != "user":
continue
for content in msg.contents:
if content.type == "function_approval_response":
if not content.function_call or not content.approved:
continue
parsed_args = content.function_call.parse_arguments()
state_args = None
if content.additional_properties:
state_args = content.additional_properties.get("ag_ui_state_args")
if not isinstance(state_args, dict):
state_args = parsed_args
if not state_args:
continue
for state_key, config in predict_state_config.items():
if config["tool"] != content.function_call.name:
continue
tool_arg_name = config["tool_argument"]
if tool_arg_name == "*":
state_value = state_args
elif isinstance(state_args, dict) and tool_arg_name in state_args:
state_value = state_args[tool_arg_name]
else:
continue
current_state[state_key] = state_value
event_bridge.current_state[state_key] = state_value
logger.info(
f"Emitting StateSnapshotEvent for approved state key '{state_key}' "
f"with {len(state_value) if isinstance(state_value, list) else 'N/A'} items"
)
events.append(StateSnapshotEvent(snapshot=current_state))
break
return events
def latest_approval_response(messages: list[ChatMessage]) -> Content | None:
"""Get the latest approval response from messages.
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_state_manager.py
-108
View File
@@ -1,108 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""State orchestration utilities."""
import json
from typing import Any
from ag_ui.core import CustomEvent, EventType
from agent_framework import ChatMessage, Content
from .._utils import make_json_safe
class StateManager:
"""Coordinates state defaults, snapshots, and structured updates."""
def __init__(
self,
state_schema: dict[str, Any] | None,
predict_state_config: dict[str, dict[str, str]] | None,
require_confirmation: bool,
) -> None:
self.state_schema = state_schema or {}
self.predict_state_config = predict_state_config or {}
self.require_confirmation = require_confirmation
self.current_state: dict[str, Any] = {}
self._state_from_input: bool = False
def initialize(self, initial_state: dict[str, Any] | None) -> dict[str, Any]:
"""Initialize state with schema defaults."""
self._state_from_input = initial_state is not None
self.current_state = (initial_state or {}).copy()
self._apply_schema_defaults()
return self.current_state
def predict_state_event(self) -> CustomEvent | None:
"""Create predict-state custom event when configured."""
if not self.predict_state_config:
return None
predict_state_value = [
{
"state_key": state_key,
"tool": config["tool"],
"tool_argument": config["tool_argument"],
}
for state_key, config in self.predict_state_config.items()
]
return CustomEvent(
type=EventType.CUSTOM,
name="PredictState",
value=predict_state_value,
)
def initial_snapshot_event(self, event_bridge: Any) -> Any:
"""Emit initial snapshot when schema and state present."""
if not self.state_schema:
return None
self._apply_schema_defaults()
return event_bridge.create_state_snapshot_event(self.current_state)
def state_context_message(self, is_new_user_turn: bool, conversation_has_tool_calls: bool) -> ChatMessage | None:
"""Inject state context only when starting a new user turn."""
if not self.current_state or not self.state_schema:
return None
if not is_new_user_turn:
return None
if conversation_has_tool_calls and not self._state_from_input:
return None
state_json = json.dumps(make_json_safe(self.current_state), indent=2)
return ChatMessage(
role="system",
contents=[
Content.from_text(
text=(
"Current state of the application:\n"
f"{state_json}\n\n"
"When modifying state, you MUST include ALL existing data plus your changes.\n"
"For example, if adding one new item to a list, include ALL existing items PLUS the one new item.\n"
"Never replace existing data - always preserve and append or merge."
)
)
],
)
def extract_state_updates(self, response_dict: dict[str, Any]) -> dict[str, Any]:
"""Extract state updates from structured response payloads."""
if self.state_schema:
return {key: response_dict[key] for key in self.state_schema.keys() if key in response_dict}
return {k: v for k, v in response_dict.items() if k != "message"}
def apply_state_updates(self, updates: dict[str, Any]) -> None:
"""Merge state updates into current state."""
if not updates:
return
self.current_state.update(updates)
def _apply_schema_defaults(self) -> None:
"""Fill missing state fields based on schema hints."""
for key, schema in self.state_schema.items():
if key in self.current_state:
continue
if isinstance(schema, dict) and schema.get("type") == "array": # type: ignore
self.current_state[key] = []
else:
self.current_state[key] = {}
python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py
+25 -1
View File
@@ -84,9 +84,26 @@ def register_additional_client_tools(agent: "AgentProtocol", client_tools: list[
logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
def _has_approval_tools(tools: list[Any]) -> bool:
"""Check if any tools require approval."""
return any(getattr(tool, "approval_mode", None) == "always_require" for tool in tools)
def merge_tools(server_tools: list[Any], client_tools: list[Any] | None) -> list[Any] | None:
"""Combine server and client tools without overriding server metadata."""
"""Combine server and client tools without overriding server metadata.
IMPORTANT: When server tools have approval_mode="always_require", we MUST return
them so they get passed to the streaming response handler. Otherwise, the approval
check in _try_execute_function_calls won't find the tool and won't trigger approval.
"""
if not client_tools:
# Even without client tools, we must pass server tools if any require approval
if server_tools and _has_approval_tools(server_tools):
logger.info(
f"[TOOLS] No client tools but server has approval tools - "
f"passing {len(server_tools)} server tools for approval mode"
)
return server_tools
logger.info("[TOOLS] No client tools - not passing tools= parameter (using agent's configured tools)")
return None
@@ -94,6 +111,13 @@ def merge_tools(server_tools: list[Any], client_tools: list[Any] | None) -> list
unique_client_tools = [tool for tool in client_tools if getattr(tool, "name", None) not in server_tool_names]
if not unique_client_tools:
# Same check: must pass server tools if any require approval
if server_tools and _has_approval_tools(server_tools):
logger.info(
f"[TOOLS] Client tools duplicate server but server has approval tools - "
f"passing {len(server_tools)} server tools for approval mode"
)
return server_tools
logger.info("[TOOLS] All client tools duplicate server tools - not passing tools= parameter")
return None
python/packages/ag-ui/agent_framework_ag_ui/_orchestrators.py
-802
View File
@@ -1,802 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Orchestrators for multi-turn agent flows."""
import json
import logging
import uuid
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Sequence
from typing import TYPE_CHECKING, Any
from ag_ui.core import (
BaseEvent,
MessagesSnapshotEvent,
RunErrorEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework import (
AgentProtocol,
AgentThread,
ChatAgent,
Content,
FunctionInvocationConfiguration,
)
from agent_framework._middleware import extract_and_merge_function_middleware
from agent_framework._tools import (
_collect_approval_responses, # type: ignore
_replace_approval_contents_with_results, # type: ignore
_try_execute_function_calls, # type: ignore
)
from ._orchestration._helpers import (
approval_steps,
build_safe_metadata,
collect_approved_state_snapshots,
ensure_tool_call_entry,
is_step_based_approval,
latest_approval_response,
select_approval_tool_name,
select_messages_to_run,
tool_name_for_call_id,
)
from ._orchestration._tooling import (
collect_server_tools,
merge_tools,
register_additional_client_tools,
)
from ._utils import (
convert_agui_tools_to_agent_framework,
generate_event_id,
get_conversation_id_from_update,
get_role_value,
)
if TYPE_CHECKING:
from ._agent import AgentConfig
from ._confirmation_strategies import ConfirmationStrategy
from ._events import AgentFrameworkEventBridge
from ._orchestration._state_manager import StateManager
logger = logging.getLogger(__name__)
class ExecutionContext:
"""Shared context for orchestrators."""
def __init__(
self,
input_data: dict[str, Any],
agent: AgentProtocol,
config: "AgentConfig", # noqa: F821
confirmation_strategy: "ConfirmationStrategy | None" = None, # noqa: F821
):
"""Initialize execution context.
Args:
input_data: AG-UI run input containing messages, state, etc.
agent: The Agent Framework agent to execute
config: Agent configuration
confirmation_strategy: Strategy for generating confirmation messages
"""
self.input_data = input_data
self.agent = agent
self.config = config
self.confirmation_strategy = confirmation_strategy
# Lazy-loaded properties
self._messages = None
self._snapshot_messages = None
self._last_message = None
self._run_id: str | None = None
self._thread_id: str | None = None
self._supplied_run_id: str | None = None
self._supplied_thread_id: str | None = None
@property
def messages(self):
"""Get converted Agent Framework messages (lazy loaded)."""
if self._messages is None:
from ._message_adapters import normalize_agui_input_messages
raw = self.input_data.get("messages", [])
if not isinstance(raw, list):
raw = []
self._messages, self._snapshot_messages = normalize_agui_input_messages(raw)
return self._messages
@property
def snapshot_messages(self) -> list[dict[str, Any]]:
"""Get normalized AG-UI snapshot messages (lazy loaded)."""
if self._snapshot_messages is None:
if self._messages is None:
_ = self.messages
else:
from ._message_adapters import agent_framework_messages_to_agui, agui_messages_to_snapshot_format
raw_snapshot = agent_framework_messages_to_agui(self._messages)
self._snapshot_messages = agui_messages_to_snapshot_format(raw_snapshot)
return self._snapshot_messages or []
@property
def last_message(self):
"""Get the last message in the conversation (lazy loaded)."""
if self._last_message is None and self.messages:
self._last_message = self.messages[-1]
return self._last_message
@property
def supplied_run_id(self) -> str | None:
"""Get the supplied run ID, if any."""
if self._supplied_run_id is None:
self._supplied_run_id = self.input_data.get("run_id") or self.input_data.get("runId")
return self._supplied_run_id
@property
def run_id(self) -> str:
"""Get supplied run ID or generate a new run ID."""
if self._run_id:
return self._run_id
if self.supplied_run_id:
self._run_id = self.supplied_run_id
if self._run_id is None:
self._run_id = str(uuid.uuid4())
return self._run_id
@property
def supplied_thread_id(self) -> str | None:
"""Get the supplied thread ID, if any."""
if self._supplied_thread_id is None:
self._supplied_thread_id = self.input_data.get("thread_id") or self.input_data.get("threadId")
return self._supplied_thread_id
@property
def thread_id(self) -> str:
"""Get supplied thread ID or generate a new thread ID."""
if self._thread_id:
return self._thread_id
if self.supplied_thread_id:
self._thread_id = self.supplied_thread_id
if self._thread_id is None:
self._thread_id = str(uuid.uuid4())
return self._thread_id
def update_run_id(self, new_run_id: str) -> None:
"""Update the run ID in the context.
Args:
new_run_id: The new run ID to set
"""
self._supplied_run_id = new_run_id
self._run_id = new_run_id
def update_thread_id(self, new_thread_id: str) -> None:
"""Update the thread ID in the context.
Args:
new_thread_id: The new thread ID to set
"""
self._supplied_thread_id = new_thread_id
self._thread_id = new_thread_id
class Orchestrator(ABC):
"""Base orchestrator for agent execution flows."""
@abstractmethod
def can_handle(self, context: ExecutionContext) -> bool:
"""Determine if this orchestrator handles the current request.
Args:
context: Execution context with input data and agent
Returns:
True if this orchestrator should handle the request
"""
...
@abstractmethod
async def run(
self,
context: ExecutionContext,
) -> AsyncGenerator[BaseEvent, None]:
"""Execute the orchestration and yield events.
Args:
context: Execution context
Yields:
AG-UI events
"""
# This is never executed - just satisfies mypy's requirement for async generators
if False: # pragma: no cover
yield
raise NotImplementedError
class HumanInTheLoopOrchestrator(Orchestrator):
"""Handles tool approval responses from user."""
def can_handle(self, context: ExecutionContext) -> bool:
"""Check if last message is a tool approval response.
Args:
context: Execution context
Returns:
True if last message is a tool result
"""
msg = context.last_message
if not msg:
return False
return bool((msg.additional_properties or {}).get("is_tool_result", False))
async def run(
self,
context: ExecutionContext,
) -> AsyncGenerator[BaseEvent, None]:
"""Process approval response and generate confirmation events.
This implementation is extracted from the legacy _agent.py lines 144-244.
Args:
context: Execution context
Yields:
AG-UI events (TextMessage, RunFinished)
"""
from ._confirmation_strategies import DefaultConfirmationStrategy
from ._events import AgentFrameworkEventBridge
logger.info("=== TOOL RESULT DETECTED (HumanInTheLoopOrchestrator) ===")
# Create event bridge for run events
event_bridge = AgentFrameworkEventBridge(
run_id=context.run_id,
thread_id=context.thread_id,
)
# CRITICAL: Every AG-UI run must start with RunStartedEvent
yield event_bridge.create_run_started_event()
# Get confirmation strategy (use default if none provided)
strategy = context.confirmation_strategy
if strategy is None:
strategy = DefaultConfirmationStrategy()
# Parse the tool result content
tool_content_text = ""
last_message = context.last_message
if last_message:
for content in last_message.contents:
if content.type == "text":
tool_content_text = content.text
break
try:
tool_result = json.loads(tool_content_text) # type: ignore[arg-type]
accepted = tool_result.get("accepted", False)
steps = tool_result.get("steps", [])
logger.info(f" Accepted: {accepted}")
logger.info(f" Steps count: {len(steps)}")
# Emit a text message confirming execution
message_id = generate_event_id()
yield TextMessageStartEvent(message_id=message_id, role="assistant")
# Check if this is confirm_changes (no steps) or function approval (has steps)
if not steps:
# This is confirm_changes for predictive state updates
if accepted:
confirmation_message = strategy.on_state_confirmed()
else:
confirmation_message = strategy.on_state_rejected()
elif accepted:
# User approved - execute the enabled steps (function approval flow)
confirmation_message = strategy.on_approval_accepted(steps)
else:
# User rejected
confirmation_message = strategy.on_approval_rejected(steps)
yield TextMessageContentEvent(
message_id=message_id,
delta=confirmation_message,
)
yield TextMessageEndEvent(message_id=message_id)
# Emit run finished
yield event_bridge.create_run_finished_event()
except json.JSONDecodeError:
logger.error(f"Failed to parse tool result: {tool_content_text}")
yield RunErrorEvent(message=f"Invalid tool result format: {tool_content_text[:100]}") # type: ignore[index]
yield event_bridge.create_run_finished_event()
class DefaultOrchestrator(Orchestrator):
"""Standard agent execution (no special handling)."""
def can_handle(self, context: ExecutionContext) -> bool:
"""Always returns True as this is the fallback orchestrator.
Args:
context: Execution context
Returns:
Always True
"""
return True
def _create_initial_events(
self, event_bridge: "AgentFrameworkEventBridge", state_manager: "StateManager"
) -> Sequence[BaseEvent]:
"""Generate initial events for the run.
Args:
event_bridge: Event bridge for creating events
Returns:
Initial AG-UI events
"""
events: list[BaseEvent] = [event_bridge.create_run_started_event()]
predict_event = state_manager.predict_state_event()
if predict_event:
events.append(predict_event)
snapshot_event = state_manager.initial_snapshot_event(event_bridge)
if snapshot_event:
events.append(snapshot_event)
return events
async def run(
self,
context: ExecutionContext,
) -> AsyncGenerator[BaseEvent, None]:
"""Standard agent run with event translation.
This implements the default agent execution flow using the event bridge
to translate Agent Framework events to AG-UI events.
Args:
context: Execution context
Yields:
AG-UI events
"""
from ._events import AgentFrameworkEventBridge
from ._orchestration._state_manager import StateManager
logger.info(f"Starting default agent run for thread_id={context.thread_id}, run_id={context.run_id}")
response_format = None
if isinstance(context.agent, ChatAgent):
response_format = (context.agent.default_options or {}).get("response_format")
skip_text_content = response_format is not None
client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
approval_tool_name = select_approval_tool_name(client_tools)
state_manager = StateManager(
state_schema=context.config.state_schema,
predict_state_config=context.config.predict_state_config,
require_confirmation=context.config.require_confirmation,
)
current_state = state_manager.initialize(context.input_data.get("state"))
event_bridge = AgentFrameworkEventBridge(
run_id=context.run_id,
thread_id=context.thread_id,
predict_state_config=context.config.predict_state_config,
current_state=current_state,
skip_text_content=skip_text_content,
require_confirmation=context.config.require_confirmation,
approval_tool_name=approval_tool_name,
)
if context.config.use_service_thread:
thread = AgentThread(service_thread_id=context.supplied_thread_id)
else:
thread = AgentThread()
thread.metadata = { # type: ignore[attr-defined]
"ag_ui_thread_id": context.thread_id,
"ag_ui_run_id": context.run_id,
}
if current_state:
thread.metadata["current_state"] = current_state # type: ignore[attr-defined]
provider_messages = context.messages or []
snapshot_messages = context.snapshot_messages
if not provider_messages:
for event in self._create_initial_events(event_bridge, state_manager):
yield event
logger.warning("No messages provided in AG-UI input")
yield event_bridge.create_run_finished_event()
return
logger.info(f"Received {len(provider_messages)} provider messages from client")
for i, msg in enumerate(provider_messages):
role = get_role_value(msg)
msg_id = getattr(msg, "message_id", None)
logger.info(f" Message {i}: role={role}, id={msg_id}")
if hasattr(msg, "contents") and msg.contents:
for j, content in enumerate(msg.contents):
if content.type == "text":
logger.debug(" Content %s: %s - text_length=%s", j, content.type, len(content.text)) # type: ignore[arg-type]
elif content.type == "function_call":
arg_length = len(str(content.arguments)) if content.arguments else 0
logger.debug(
" Content %s: %s - %s args_length=%s", j, content.type, content.name, arg_length
)
elif content.type == "function_result":
result_preview = type(content.result).__name__ if content.result is not None else "None"
logger.debug(
" Content %s: %s - call_id=%s, result_type=%s",
j,
content.type,
content.call_id,
result_preview,
)
else:
logger.debug(f" Content {j}: {content.type}")
pending_tool_calls: list[dict[str, Any]] = []
tool_calls_by_id: dict[str, dict[str, Any]] = {}
tool_results: list[dict[str, Any]] = []
tool_calls_ended: set[str] = set()
messages_snapshot_emitted = False
accumulated_text_content = ""
active_message_id: str | None = None
# Check for FunctionApprovalResponseContent and emit updated state snapshot
# This ensures the UI shows the approved state (e.g., 2 steps) not the original (3 steps)
for snapshot_evt in collect_approved_state_snapshots(
provider_messages,
context.config.predict_state_config,
current_state,
event_bridge,
):
yield snapshot_evt
messages_to_run = select_messages_to_run(provider_messages, state_manager)
logger.info(f"[TOOLS] Client sent {len(client_tools) if client_tools else 0} tools")
if client_tools:
for tool in client_tools:
tool_name = getattr(tool, "name", "unknown")
declaration_only = getattr(tool, "declaration_only", None)
logger.info(f"[TOOLS] - Client tool: {tool_name}, declaration_only={declaration_only}")
server_tools = collect_server_tools(context.agent)
register_additional_client_tools(context.agent, client_tools)
tools_param = merge_tools(server_tools, client_tools)
collect_updates = response_format is not None
all_updates: list[Any] | None = [] if collect_updates else None
update_count = 0
# Prepare metadata for chat client (Azure requires string values)
safe_metadata = build_safe_metadata(getattr(thread, "metadata", None))
run_kwargs: dict[str, Any] = {
"thread": thread,
"tools": tools_param,
"options": {"metadata": safe_metadata},
}
if safe_metadata:
run_kwargs["options"]["store"] = True
async def _resolve_approval_responses(
messages: list[Any],
tools_for_execution: list[Any],
) -> None:
fcc_todo = _collect_approval_responses(messages)
if not fcc_todo:
return
approved_responses = [resp for resp in fcc_todo.values() if resp.approved]
approved_function_results: list[Any] = []
if approved_responses and tools_for_execution:
chat_client = getattr(context.agent, "chat_client", None)
config = (
getattr(chat_client, "function_invocation_configuration", None) or FunctionInvocationConfiguration()
)
middleware_pipeline = extract_and_merge_function_middleware(chat_client, run_kwargs)
try:
results, _ = await _try_execute_function_calls(
custom_args=run_kwargs,
attempt_idx=0,
function_calls=approved_responses,
tools=tools_for_execution,
middleware_pipeline=middleware_pipeline,
config=config,
)
approved_function_results = list(results)
except Exception:
logger.error("Failed to execute approved tool calls; injecting error results.")
approved_function_results = []
normalized_results: list[Content] = []
for idx, approval in enumerate(approved_responses):
if idx < len(approved_function_results) and approved_function_results[idx].type == "function_result":
normalized_results.append(approved_function_results[idx])
continue
call_id = approval.function_call.call_id or approval.id # type: ignore[union-attr]
normalized_results.append(
Content.from_function_result(call_id=call_id, result="Error: Tool call invocation failed.") # type: ignore[arg-type]
)
_replace_approval_contents_with_results(messages, fcc_todo, normalized_results) # type: ignore
def _should_emit_tool_snapshot(tool_name: str | None) -> bool:
if not pending_tool_calls or not tool_results:
return False
if tool_name and context.config.predict_state_config and not context.config.require_confirmation:
for config in context.config.predict_state_config.values():
if config["tool"] == tool_name:
logger.info(
f"Skipping intermediate MessagesSnapshotEvent for predictive tool '{tool_name}' "
" - delaying until summary"
)
return False
return True
def _build_messages_snapshot(tool_message_id: str | None = None) -> MessagesSnapshotEvent:
has_text_content = bool(accumulated_text_content)
all_messages = snapshot_messages.copy()
if pending_tool_calls:
if tool_message_id and not has_text_content:
tool_call_message_id = tool_message_id
else:
tool_call_message_id = (
active_message_id if not has_text_content and active_message_id else generate_event_id()
)
tool_call_message = {
"id": tool_call_message_id,
"role": "assistant",
"tool_calls": pending_tool_calls.copy(),
}
all_messages.append(tool_call_message)
all_messages.extend(tool_results)
if has_text_content and active_message_id:
assistant_text_message = {
"id": active_message_id,
"role": "assistant",
"content": accumulated_text_content,
}
all_messages.append(assistant_text_message)
return MessagesSnapshotEvent(
messages=all_messages, # type: ignore[arg-type]
)
# Use tools_param if available (includes client tools), otherwise fall back to server_tools
# This ensures both server tools AND client tools can be executed after approval
tools_for_approval = tools_param if tools_param is not None else server_tools
latest_approval = latest_approval_response(messages_to_run)
await _resolve_approval_responses(messages_to_run, tools_for_approval)
if latest_approval and is_step_based_approval(latest_approval, context.config.predict_state_config):
from ._confirmation_strategies import DefaultConfirmationStrategy
strategy = context.confirmation_strategy
if strategy is None:
strategy = DefaultConfirmationStrategy()
steps = approval_steps(latest_approval)
if steps:
if latest_approval.approved:
confirmation_message = strategy.on_approval_accepted(steps)
else:
confirmation_message = strategy.on_approval_rejected(steps)
else:
if latest_approval.approved:
confirmation_message = strategy.on_state_confirmed()
else:
confirmation_message = strategy.on_state_rejected()
message_id = generate_event_id()
for event in self._create_initial_events(event_bridge, state_manager):
yield event
yield TextMessageStartEvent(message_id=message_id, role="assistant")
yield TextMessageContentEvent(message_id=message_id, delta=confirmation_message)
yield TextMessageEndEvent(message_id=message_id)
yield event_bridge.create_run_finished_event()
return
should_recreate_event_bridge = False
async for update in context.agent.run_stream(messages_to_run, **run_kwargs):
conv_id = get_conversation_id_from_update(update)
if conv_id and conv_id != context.thread_id:
context.update_thread_id(conv_id)
should_recreate_event_bridge = True
if update.response_id and update.response_id != context.run_id:
context.update_run_id(update.response_id)
should_recreate_event_bridge = True
if should_recreate_event_bridge:
event_bridge = AgentFrameworkEventBridge(
run_id=context.run_id,
thread_id=context.thread_id,
predict_state_config=context.config.predict_state_config,
current_state=current_state,
skip_text_content=skip_text_content,
require_confirmation=context.config.require_confirmation,
approval_tool_name=approval_tool_name,
)
should_recreate_event_bridge = False
if update_count == 0:
for event in self._create_initial_events(event_bridge, state_manager):
yield event
update_count += 1
logger.info(f"[STREAM] Received update #{update_count} from agent")
if all_updates is not None:
all_updates.append(update)
if event_bridge.current_message_id is None and update.contents:
has_tool_call = any(content.type == "function_call" for content in update.contents)
has_text = any(content.type == "text" for content in update.contents)
if has_tool_call and not has_text:
tool_message_id = generate_event_id()
event_bridge.current_message_id = tool_message_id
active_message_id = tool_message_id
accumulated_text_content = ""
logger.info(
"[STREAM] Emitting TextMessageStartEvent for tool-only response message_id=%s",
tool_message_id,
)
yield TextMessageStartEvent(message_id=tool_message_id, role="assistant")
events = await event_bridge.from_agent_run_update(update)
logger.info(f"[STREAM] Update #{update_count} produced {len(events)} events")
for event in events:
if isinstance(event, TextMessageStartEvent):
active_message_id = event.message_id
accumulated_text_content = ""
elif isinstance(event, TextMessageContentEvent):
accumulated_text_content += event.delta
elif isinstance(event, ToolCallStartEvent):
tool_call_entry = ensure_tool_call_entry(event.tool_call_id, tool_calls_by_id, pending_tool_calls)
tool_call_entry["function"]["name"] = event.tool_call_name
elif isinstance(event, ToolCallArgsEvent):
tool_call_entry = ensure_tool_call_entry(event.tool_call_id, tool_calls_by_id, pending_tool_calls)
tool_call_entry["function"]["arguments"] += event.delta
elif isinstance(event, ToolCallEndEvent):
tool_calls_ended.add(event.tool_call_id)
elif isinstance(event, ToolCallResultEvent):
tool_results.append(
{
"id": event.message_id,
"role": "tool",
"toolCallId": event.tool_call_id,
"content": event.content,
}
)
logger.info(f"[STREAM] Yielding event: {type(event).__name__}")
yield event
if isinstance(event, ToolCallResultEvent):
tool_name = tool_name_for_call_id(tool_calls_by_id, event.tool_call_id)
if _should_emit_tool_snapshot(tool_name):
messages_snapshot_emitted = True
messages_snapshot = _build_messages_snapshot()
logger.info(f"[STREAM] Yielding event: {type(messages_snapshot).__name__}")
yield messages_snapshot
elif isinstance(event, ToolCallEndEvent):
tool_name = tool_name_for_call_id(tool_calls_by_id, event.tool_call_id)
if tool_name == "confirm_changes":
messages_snapshot_emitted = True
messages_snapshot = _build_messages_snapshot()
logger.info(f"[STREAM] Yielding event: {type(messages_snapshot).__name__}")
yield messages_snapshot
logger.info(f"[STREAM] Agent stream completed. Total updates: {update_count}")
if event_bridge.should_stop_after_confirm:
logger.info("Stopping run - waiting for user approval/confirmation response")
if event_bridge.current_message_id:
logger.info(f"[CONFIRM] Emitting TextMessageEndEvent for message_id={event_bridge.current_message_id}")
yield event_bridge.create_message_end_event(event_bridge.current_message_id)
event_bridge.current_message_id = None
yield event_bridge.create_run_finished_event()
return
if pending_tool_calls:
pending_without_end = [tc for tc in pending_tool_calls if tc.get("id") not in tool_calls_ended]
if pending_without_end:
logger.info(
"Found %s pending tool calls without end event - emitting ToolCallEndEvent",
len(pending_without_end),
)
for tool_call in pending_without_end:
tool_call_id = tool_call.get("id")
if tool_call_id:
end_event = ToolCallEndEvent(tool_call_id=tool_call_id)
logger.info(f"Emitting ToolCallEndEvent for declaration-only tool call '{tool_call_id}'")
yield end_event
if response_format and all_updates:
from agent_framework import AgentResponse
from pydantic import BaseModel
logger.info(f"Processing structured output, update count: {len(all_updates)}")
final_response = AgentResponse.from_agent_run_response_updates(
all_updates, output_format_type=response_format
)
if final_response.value and isinstance(final_response.value, BaseModel):
response_dict = final_response.value.model_dump(mode="json", exclude_none=True)
logger.info(f"Received structured output keys: {list(response_dict.keys())}")
state_updates = state_manager.extract_state_updates(response_dict)
if state_updates:
state_manager.apply_state_updates(state_updates)
state_snapshot = event_bridge.create_state_snapshot_event(current_state)
yield state_snapshot
logger.info(f"Emitted StateSnapshotEvent with updates: {list(state_updates.keys())}")
if "message" in response_dict and response_dict["message"]:
message_id = generate_event_id()
yield TextMessageStartEvent(message_id=message_id, role="assistant")
yield TextMessageContentEvent(message_id=message_id, delta=response_dict["message"])
yield TextMessageEndEvent(message_id=message_id)
logger.info(f"Emitted conversational message with length={len(response_dict['message'])}")
if all_updates is not None and len(all_updates) == 0:
logger.info("No updates received from agent - emitting initial events")
for event in self._create_initial_events(event_bridge, state_manager):
yield event
logger.info(f"[FINALIZE] Checking for unclosed message. current_message_id={event_bridge.current_message_id}")
if event_bridge.current_message_id:
logger.info(f"[FINALIZE] Emitting TextMessageEndEvent for message_id={event_bridge.current_message_id}")
yield event_bridge.create_message_end_event(event_bridge.current_message_id)
messages_snapshot = _build_messages_snapshot(tool_message_id=event_bridge.current_message_id)
messages_snapshot_emitted = True
logger.info(
f"[FINALIZE] Emitting MessagesSnapshotEvent with {len(messages_snapshot.messages)} messages "
f"(text content length: {len(accumulated_text_content)})"
)
yield messages_snapshot
else:
logger.info("[FINALIZE] No current_message_id - skipping TextMessageEndEvent")
if not messages_snapshot_emitted and (pending_tool_calls or tool_results):
messages_snapshot = _build_messages_snapshot()
messages_snapshot_emitted = True
logger.info(
f"[FINALIZE] Emitting MessagesSnapshotEvent with {len(messages_snapshot.messages)} messages"
)
yield messages_snapshot
logger.info("[FINALIZE] Emitting RUN_FINISHED event")
yield event_bridge.create_run_finished_event()
logger.info(f"Completed agent run for thread_id={context.thread_id}, run_id={context.run_id}")
__all__ = [
"Orchestrator",
"ExecutionContext",
"HumanInTheLoopOrchestrator",
"DefaultOrchestrator",
]
python/packages/ag-ui/agent_framework_ag_ui/_run.py
+963
View File
@@ -0,0 +1,963 @@
# Copyright (c) Microsoft. All rights reserved.
"""Simplified AG-UI orchestration - single linear flow."""
import json
import logging
import uuid
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from ag_ui.core import (
BaseEvent,
CustomEvent,
MessagesSnapshotEvent,
RunFinishedEvent,
RunStartedEvent,
StateSnapshotEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework import (
AgentProtocol,
AgentThread,
ChatMessage,
Content,
prepare_function_call_results,
)
from agent_framework._middleware import extract_and_merge_function_middleware
from agent_framework._tools import (
FunctionInvocationConfiguration,
_collect_approval_responses, # type: ignore
_replace_approval_contents_with_results, # type: ignore
_try_execute_function_calls, # type: ignore
)
from ._message_adapters import normalize_agui_input_messages
from ._orchestration._predictive_state import PredictiveStateHandler
from ._orchestration._tooling import collect_server_tools, merge_tools, register_additional_client_tools
from ._utils import (
convert_agui_tools_to_agent_framework,
generate_event_id,
get_conversation_id_from_update,
make_json_safe,
)
if TYPE_CHECKING:
from collections.abc import AsyncGenerator
from ._agent import AgentConfig
logger = logging.getLogger(__name__)
# Keys that are internal to AG-UI orchestration and should not be passed to chat clients
AG_UI_INTERNAL_METADATA_KEYS = {"ag_ui_thread_id", "ag_ui_run_id", "current_state"}
def _build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]:
"""Build metadata dict with truncated string values for Azure compatibility.
Azure has a 512 character limit per metadata value.
Args:
thread_metadata: Raw metadata dict
Returns:
Metadata with string values truncated to 512 chars
"""
if not thread_metadata:
return {}
safe_metadata: dict[str, Any] = {}
for key, value in thread_metadata.items():
value_str = value if isinstance(value, str) else json.dumps(value)
if len(value_str) > 512:
value_str = value_str[:512]
safe_metadata[key] = value_str
return safe_metadata
def _has_only_tool_calls(contents: list[Any]) -> bool:
"""Check if contents have only tool calls (no text).
Args:
contents: List of content items
Returns:
True if there are tool calls but no text content
"""
has_tool_call = any(getattr(c, "type", None) == "function_call" for c in contents)
has_text = any(getattr(c, "type", None) == "text" and getattr(c, "text", None) for c in contents)
return has_tool_call and not has_text
def _should_suppress_intermediate_snapshot(
tool_name: str | None,
predict_state_config: dict[str, dict[str, str]] | None,
require_confirmation: bool,
) -> bool:
"""Check if intermediate MessagesSnapshotEvent should be suppressed for this tool.
For predictive tools without confirmation, we delay the snapshot until the end.
Args:
tool_name: Name of the tool that just completed
predict_state_config: Predictive state configuration
require_confirmation: Whether confirmation is required
Returns:
True if snapshot should be suppressed
"""
if not tool_name or not predict_state_config:
return False
# Only suppress when confirmation is disabled
if require_confirmation:
return False
# Check if this tool is a predictive tool
for config in predict_state_config.values():
if config["tool"] == tool_name:
logger.info(f"Suppressing intermediate MessagesSnapshotEvent for predictive tool '{tool_name}'")
return True
return False
def _extract_approved_state_updates(
messages: list[Any],
predictive_handler: PredictiveStateHandler | None,
) -> dict[str, Any]:
"""Extract state updates from function_approval_response content.
This emits StateSnapshotEvent for approved state-changing tools before running agent.
Args:
messages: List of messages to scan
predictive_handler: Predictive state handler
Returns:
Dict of state updates to apply
"""
if not predictive_handler:
return {}
updates: dict[str, Any] = {}
for msg in messages:
for content in msg.contents:
if getattr(content, "type", None) != "function_approval_response":
continue
if not getattr(content, "approved", False) or not getattr(content, "function_call", None):
continue
parsed_args = content.function_call.parse_arguments()
result = predictive_handler.extract_state_value(content.function_call.name, parsed_args)
if result:
state_key, state_value = result
updates[state_key] = state_value
logger.info(f"Found approved state update for key '{state_key}'")
return updates
@dataclass
class FlowState:
"""Minimal explicit state for a single AG-UI run."""
message_id: str | None = None # Current text message being streamed
tool_call_id: str | None = None # Current tool call being streamed
tool_call_name: str | None = None # Name of current tool call
waiting_for_approval: bool = False # Stop after approval request
current_state: dict[str, Any] = field(default_factory=dict) # Shared state
accumulated_text: str = "" # For MessagesSnapshotEvent
pending_tool_calls: list[dict[str, Any]] = field(default_factory=list) # For MessagesSnapshotEvent
tool_calls_by_id: dict[str, dict[str, Any]] = field(default_factory=dict)
tool_results: list[dict[str, Any]] = field(default_factory=list)
tool_calls_ended: set[str] = field(default_factory=set) # Track which tool calls have been ended
def get_tool_name(self, call_id: str | None) -> str | None:
"""Get tool name by call ID."""
if not call_id or call_id not in self.tool_calls_by_id:
return None
name = self.tool_calls_by_id[call_id]["function"].get("name")
return str(name) if name else None
def get_pending_without_end(self) -> list[dict[str, Any]]:
"""Get tool calls that started but never received an end event (declaration-only)."""
return [tc for tc in self.pending_tool_calls if tc.get("id") not in self.tool_calls_ended]
def _create_state_context_message(
current_state: dict[str, Any],
state_schema: dict[str, Any],
) -> ChatMessage | None:
"""Create a system message with current state context.
This injects the current state into the conversation so the model
knows what state exists and can make informed updates.
Args:
current_state: The current state to inject
state_schema: The state schema (used to determine if injection is needed)
Returns:
ChatMessage with state context, or None if not needed
"""
if not current_state or not state_schema:
return None
state_json = json.dumps(current_state, indent=2)
return ChatMessage(
role="system",
contents=[
Content.from_text(
text=(
"Current state of the application:\n"
f"{state_json}\n\n"
"When modifying state, you MUST include ALL existing data plus your changes.\n"
"For example, if adding one new item to a list, include ALL existing items PLUS the new item.\n"
"Never replace existing data - always preserve and append or merge."
)
)
],
)
def _inject_state_context(
messages: list[ChatMessage],
current_state: dict[str, Any],
state_schema: dict[str, Any],
) -> list[ChatMessage]:
"""Inject state context message into messages if appropriate.
The state context is injected before the last user message to give
the model visibility into the current application state.
Args:
messages: The messages to potentially inject into
current_state: The current state
state_schema: The state schema
Returns:
Messages with state context injected if appropriate
"""
state_msg = _create_state_context_message(current_state, state_schema)
if not state_msg:
return messages
# Check if the last message is from a user (new user turn)
if not messages:
return messages
from ._utils import get_role_value
last_role = get_role_value(messages[-1])
if last_role != "user":
return messages
# Always inject state context if state is provided
# This ensures UI state changes are visible to the model
# Insert state context before the last user message
result = list(messages[:-1])
result.append(state_msg)
result.append(messages[-1])
return result
def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]:
"""Emit TextMessage events for TextContent."""
if not content.text:
return []
# Skip if we're in structured output mode or waiting for approval
if skip_text or flow.waiting_for_approval:
return []
events: list[BaseEvent] = []
if not flow.message_id:
flow.message_id = generate_event_id()
events.append(TextMessageStartEvent(message_id=flow.message_id, role="assistant"))
events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
flow.accumulated_text += content.text
return events
def _emit_tool_call(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
) -> list[BaseEvent]:
"""Emit ToolCall events for FunctionCallContent."""
events: list[BaseEvent] = []
tool_call_id = content.call_id or flow.tool_call_id or generate_event_id()
# Emit start event when we have a new tool call
if content.name and tool_call_id != flow.tool_call_id:
flow.tool_call_id = tool_call_id
flow.tool_call_name = content.name
if predictive_handler:
predictive_handler.reset_streaming()
events.append(
ToolCallStartEvent(
tool_call_id=tool_call_id,
tool_call_name=content.name,
parent_message_id=flow.message_id,
)
)
# Track for MessagesSnapshotEvent
tool_entry = {
"id": tool_call_id,
"type": "function",
"function": {"name": content.name, "arguments": ""},
}
flow.pending_tool_calls.append(tool_entry)
flow.tool_calls_by_id[tool_call_id] = tool_entry
elif tool_call_id:
flow.tool_call_id = tool_call_id
# Emit args if present
if content.arguments:
delta = (
content.arguments if isinstance(content.arguments, str) else json.dumps(make_json_safe(content.arguments))
)
events.append(ToolCallArgsEvent(tool_call_id=tool_call_id, delta=delta))
# Track args for MessagesSnapshotEvent
if tool_call_id in flow.tool_calls_by_id:
flow.tool_calls_by_id[tool_call_id]["function"]["arguments"] += delta
# Emit predictive state deltas
if predictive_handler and flow.tool_call_name:
delta_events = predictive_handler.emit_streaming_deltas(flow.tool_call_name, delta)
events.extend(delta_events)
return events
def _emit_tool_result(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
) -> list[BaseEvent]:
"""Emit ToolCallResult events for FunctionResultContent."""
events: list[BaseEvent] = []
# Cannot emit tool result without a call_id to associate it with
if not content.call_id:
return events
events.append(ToolCallEndEvent(tool_call_id=content.call_id))
flow.tool_calls_ended.add(content.call_id) # Track ended tool calls
result_content = prepare_function_call_results(content.result)
message_id = generate_event_id()
events.append(
ToolCallResultEvent(
message_id=message_id,
tool_call_id=content.call_id,
content=result_content,
role="tool",
)
)
# Track for MessagesSnapshotEvent
flow.tool_results.append(
{
"id": message_id,
"role": "tool",
"toolCallId": content.call_id,
"content": result_content,
}
)
# Apply predictive state updates and emit snapshot
if predictive_handler:
predictive_handler.apply_pending_updates()
if flow.current_state:
events.append(StateSnapshotEvent(snapshot=flow.current_state))
# Reset tool tracking and message context
# After tool result, any subsequent text should start a new message
flow.tool_call_id = None
flow.tool_call_name = None
flow.message_id = None # Reset so next text content starts a new message
return events
def _emit_approval_request(
content: Content,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
require_confirmation: bool = True,
) -> list[BaseEvent]:
"""Emit events for function approval request."""
events: list[BaseEvent] = []
# function_call is required for approval requests - skip if missing
func_call = content.function_call
if not func_call:
logger.warning("Approval request content missing function_call, skipping")
return events
func_name = func_call.name or ""
func_call_id = func_call.call_id
# Extract state from function arguments if predictive
if predictive_handler and func_name:
parsed_args = func_call.parse_arguments()
result = predictive_handler.extract_state_value(func_name, parsed_args)
if result:
state_key, state_value = result
flow.current_state[state_key] = state_value
events.append(StateSnapshotEvent(snapshot=flow.current_state))
# End the original tool call
if func_call_id:
events.append(ToolCallEndEvent(tool_call_id=func_call_id))
flow.tool_calls_ended.add(func_call_id) # Track ended tool calls
# Emit custom event for UI
events.append(
CustomEvent(
name="function_approval_request",
value={
"id": content.id,
"function_call": {
"call_id": func_call_id,
"name": func_name,
"arguments": make_json_safe(func_call.parse_arguments()),
},
},
)
)
# Emit confirm_changes tool call for UI compatibility
# The complete sequence (Start -> Args -> End) signals the UI to show the confirmation dialog
if require_confirmation:
confirm_id = generate_event_id()
events.append(
ToolCallStartEvent(
tool_call_id=confirm_id,
tool_call_name="confirm_changes",
parent_message_id=flow.message_id,
)
)
args = {
"function_name": func_name,
"function_call_id": func_call_id,
"function_arguments": make_json_safe(func_call.parse_arguments()) or {},
"steps": [{"description": f"Execute {func_name}", "status": "enabled"}],
}
events.append(ToolCallArgsEvent(tool_call_id=confirm_id, delta=json.dumps(args)))
events.append(ToolCallEndEvent(tool_call_id=confirm_id))
flow.waiting_for_approval = True
return events
def _emit_content(
content: Any,
flow: FlowState,
predictive_handler: PredictiveStateHandler | None = None,
skip_text: bool = False,
require_confirmation: bool = True,
) -> list[BaseEvent]:
"""Emit appropriate events for any content type."""
content_type = getattr(content, "type", None)
if content_type == "text":
return _emit_text(content, flow, skip_text)
elif content_type == "function_call":
return _emit_tool_call(content, flow, predictive_handler)
elif content_type == "function_result":
return _emit_tool_result(content, flow, predictive_handler)
elif content_type == "function_approval_request":
return _emit_approval_request(content, flow, predictive_handler, require_confirmation)
return []
def _is_confirm_changes_response(messages: list[Any]) -> bool:
"""Check if the last message is a confirm_changes tool result (state confirmation flow).
This returns True for confirm_changes flows where we emit a confirmation message
and stop. The key indicator is the presence of a 'steps' key in the tool result
(even if empty), combined with 'accepted' boolean.
"""
if not messages:
return False
last = messages[-1]
if not last.additional_properties.get("is_tool_result", False):
return False
# Parse the content to check if it has the confirm_changes structure
for content in last.contents:
if getattr(content, "type", None) == "text":
try:
result = json.loads(content.text)
# confirm_changes results have 'accepted' and 'steps' keys
if "accepted" in result and "steps" in result:
return True
except json.JSONDecodeError:
# Content is not valid JSON; continue checking other content items
logger.debug("Failed to parse confirm_changes tool result as JSON; treating as non-confirmation.")
return False
def _handle_step_based_approval(messages: list[Any]) -> list[BaseEvent]:
"""Handle step-based approval response and emit confirmation message."""
events: list[BaseEvent] = []
last = messages[-1]
# Parse the approval content
approval_text = ""
for content in last.contents:
if getattr(content, "type", None) == "text":
approval_text = content.text
break
try:
result = json.loads(approval_text)
accepted = result.get("accepted", False)
steps = result.get("steps", [])
if accepted:
# Generate acceptance message with step descriptions
enabled_steps = [s for s in steps if s.get("status") == "enabled"]
if enabled_steps:
message_parts = [f"Executing {len(enabled_steps)} approved steps:\n\n"]
for i, step in enumerate(enabled_steps, 1):
message_parts.append(f"{i}. {step.get('description', 'Step')}\n")
message_parts.append("\nAll steps completed successfully!")
message = "".join(message_parts)
else:
message = "Changes confirmed and applied successfully!"
else:
# Rejection message
message = "No problem! What would you like me to change about the plan?"
except json.JSONDecodeError:
message = "Acknowledged."
message_id = generate_event_id()
events.append(TextMessageStartEvent(message_id=message_id, role="assistant"))
events.append(TextMessageContentEvent(message_id=message_id, delta=message))
events.append(TextMessageEndEvent(message_id=message_id))
return events
async def _resolve_approval_responses(
messages: list[Any],
tools: list[Any],
agent: AgentProtocol,
run_kwargs: dict[str, Any],
) -> None:
"""Execute approved function calls and replace approval content with results.
This modifies the messages list in place, replacing FunctionApprovalResponseContent
with FunctionResultContent containing the actual tool execution result.
Args:
messages: List of messages (will be modified in place)
tools: List of available tools
agent: The agent instance (to get chat_client and config)
run_kwargs: Kwargs for tool execution
"""
fcc_todo = _collect_approval_responses(messages)
if not fcc_todo:
return
approved_responses = [resp for resp in fcc_todo.values() if resp.approved]
rejected_responses = [resp for resp in fcc_todo.values() if not resp.approved]
approved_function_results: list[Any] = []
# Execute approved tool calls
if approved_responses and tools:
chat_client = getattr(agent, "chat_client", None)
config = getattr(chat_client, "function_invocation_configuration", None) or FunctionInvocationConfiguration()
middleware_pipeline = extract_and_merge_function_middleware(chat_client, run_kwargs)
# Filter out AG-UI-specific kwargs that should not be passed to tool execution
tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"}
try:
results, _ = await _try_execute_function_calls(
custom_args=tool_kwargs,
attempt_idx=0,
function_calls=approved_responses,
tools=tools,
middleware_pipeline=middleware_pipeline,
config=config,
)
approved_function_results = list(results)
except Exception as e:
logger.exception("Failed to execute approved tool calls; injecting error results: %s", e)
approved_function_results = []
# Build normalized results for approved responses
normalized_results: list[Content] = []
for idx, approval in enumerate(approved_responses):
if (
idx < len(approved_function_results)
and getattr(approved_function_results[idx], "type", None) == "function_result"
):
normalized_results.append(approved_function_results[idx])
continue
# Get call_id from function_call if present, otherwise use approval.id
func_call = approval.function_call
call_id = (func_call.call_id if func_call else None) or approval.id or ""
normalized_results.append(
Content.from_function_result(call_id=call_id, result="Error: Tool call invocation failed.")
)
# Build rejection results
for rejection in rejected_responses:
func_call = rejection.function_call
call_id = (func_call.call_id if func_call else None) or rejection.id or ""
normalized_results.append(
Content.from_function_result(call_id=call_id, result="Error: Tool call invocation was rejected by user.")
)
_replace_approval_contents_with_results(messages, fcc_todo, normalized_results) # type: ignore
def _build_messages_snapshot(
flow: FlowState,
snapshot_messages: list[dict[str, Any]],
) -> MessagesSnapshotEvent:
"""Build MessagesSnapshotEvent from current flow state."""
all_messages = list(snapshot_messages)
# Add assistant message with tool calls
if flow.pending_tool_calls:
tool_call_message = {
"id": flow.message_id or generate_event_id(),
"role": "assistant",
"tool_calls": flow.pending_tool_calls.copy(),
}
if flow.accumulated_text:
tool_call_message["content"] = flow.accumulated_text
all_messages.append(tool_call_message)
# Add tool results
all_messages.extend(flow.tool_results)
# Add text-only assistant message if no tool calls
if flow.accumulated_text and not flow.pending_tool_calls:
all_messages.append(
{
"id": flow.message_id or generate_event_id(),
"role": "assistant",
"content": flow.accumulated_text,
}
)
return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type]
async def run_agent_stream(
input_data: dict[str, Any],
agent: AgentProtocol,
config: "AgentConfig",
) -> "AsyncGenerator[BaseEvent, None]":
"""Run agent and yield AG-UI events.
This is the single entry point for all AG-UI agent runs. It follows a simple
linear flow: RunStarted -> content events -> RunFinished.
Args:
input_data: AG-UI request data with messages, state, tools, etc.
agent: The Agent Framework agent to run
config: Agent configuration
Yields:
AG-UI events
"""
# Parse IDs
thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4())
run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4())
# Initialize flow state with schema defaults
flow = FlowState()
if input_data.get("state"):
flow.current_state = dict(input_data["state"])
# Apply schema defaults for missing state keys
if config.state_schema:
for key, schema in config.state_schema.items():
if key in flow.current_state:
continue
if isinstance(schema, dict) and schema.get("type") == "array":
flow.current_state[key] = []
else:
flow.current_state[key] = {}
# Initialize predictive state handler if configured
predictive_handler: PredictiveStateHandler | None = None
if config.predict_state_config:
predictive_handler = PredictiveStateHandler(
predict_state_config=config.predict_state_config,
current_state=flow.current_state,
)
# Normalize messages
raw_messages = input_data.get("messages", [])
messages, snapshot_messages = normalize_agui_input_messages(raw_messages)
# Check for structured output mode (skip text content)
skip_text = False
response_format = None
from agent_framework import ChatAgent
if isinstance(agent, ChatAgent):
response_format = agent.default_options.get("response_format")
skip_text = response_format is not None
# Handle empty messages (emit RunStarted immediately since no agent response)
if not messages:
logger.warning("No messages provided in AG-UI input")
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
yield RunFinishedEvent(run_id=run_id, thread_id=thread_id)
return
# Prepare tools
client_tools = convert_agui_tools_to_agent_framework(input_data.get("tools"))
server_tools = collect_server_tools(agent)
register_additional_client_tools(agent, client_tools)
tools = merge_tools(server_tools, client_tools)
# Create thread (with service thread support)
if config.use_service_thread:
supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId")
thread = AgentThread(service_thread_id=supplied_thread_id)
else:
thread = AgentThread()
# Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation)
base_metadata: dict[str, Any] = {
"ag_ui_thread_id": thread_id,
"ag_ui_run_id": run_id,
}
if flow.current_state:
base_metadata["current_state"] = flow.current_state
thread.metadata = _build_safe_metadata(base_metadata) # type: ignore[attr-defined]
# Build run kwargs (Feature #6: Azure store flag when metadata present)
run_kwargs: dict[str, Any] = {"thread": thread}
if tools:
run_kwargs["tools"] = tools
# Filter out AG-UI internal metadata keys before passing to chat client
# These are used internally for orchestration and should not be sent to the LLM provider
client_metadata = {
k: v for k, v in (getattr(thread, "metadata", None) or {}).items() if k not in AG_UI_INTERNAL_METADATA_KEYS
}
safe_metadata = _build_safe_metadata(client_metadata) if client_metadata else {}
if safe_metadata:
run_kwargs["options"] = {"metadata": safe_metadata, "store": True}
# Resolve approval responses (execute approved tools, replace approvals with results)
# This must happen before running the agent so it sees the tool results
tools_for_execution = tools if tools is not None else server_tools
await _resolve_approval_responses(messages, tools_for_execution, agent, run_kwargs)
# Feature #3: Emit StateSnapshotEvent for approved state-changing tools before agent runs
approved_state_updates = _extract_approved_state_updates(messages, predictive_handler)
approved_state_snapshot_emitted = False
if approved_state_updates:
flow.current_state.update(approved_state_updates)
approved_state_snapshot_emitted = True
# Handle confirm_changes response (state confirmation flow - emit confirmation and stop)
if _is_confirm_changes_response(messages):
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
# Emit approved state snapshot before confirmation message
if approved_state_snapshot_emitted:
yield StateSnapshotEvent(snapshot=flow.current_state)
for event in _handle_step_based_approval(messages):
yield event
yield RunFinishedEvent(run_id=run_id, thread_id=thread_id)
return
# Inject state context message so the model knows current application state
# This is critical for shared state scenarios where the UI state needs to be visible
if config.state_schema and flow.current_state:
messages = _inject_state_context(messages, flow.current_state, config.state_schema)
# Stream from agent - emit RunStarted after first update to get service IDs
run_started_emitted = False
all_updates: list[Any] = [] # Collect for structured output processing
async for update in agent.run_stream(messages, **run_kwargs):
# Collect updates for structured output processing
if response_format is not None:
all_updates.append(update)
# Update IDs from service response on first update and emit RunStarted
if not run_started_emitted:
conv_id = get_conversation_id_from_update(update)
if conv_id:
thread_id = conv_id
if update.response_id:
run_id = update.response_id
# NOW emit RunStarted with proper IDs
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
# Emit PredictState custom event if configured
if config.predict_state_config:
predict_state_value = [
{
"state_key": state_key,
"tool": cfg["tool"],
"tool_argument": cfg["tool_argument"],
}
for state_key, cfg in config.predict_state_config.items()
]
yield CustomEvent(name="PredictState", value=predict_state_value)
# Emit initial state snapshot only if we have both state_schema and state
if config.state_schema and flow.current_state:
yield StateSnapshotEvent(snapshot=flow.current_state)
run_started_emitted = True
# Feature #4: Detect tool-only messages (no text content)
# Emit TextMessageStartEvent to create message context for tool calls
if not flow.message_id and _has_only_tool_calls(update.contents):
flow.message_id = generate_event_id()
logger.info(f"Tool-only response detected, creating message_id={flow.message_id}")
yield TextMessageStartEvent(message_id=flow.message_id, role="assistant")
# Emit events for each content item
for content in update.contents:
for event in _emit_content(
content,
flow,
predictive_handler,
skip_text,
config.require_confirmation,
):
yield event
# Stop if waiting for approval
if flow.waiting_for_approval:
break
# If no updates at all, still emit RunStarted
if not run_started_emitted:
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
if config.predict_state_config:
predict_state_value = [
{
"state_key": state_key,
"tool": cfg["tool"],
"tool_argument": cfg["tool_argument"],
}
for state_key, cfg in config.predict_state_config.items()
]
yield CustomEvent(name="PredictState", value=predict_state_value)
if config.state_schema and flow.current_state:
yield StateSnapshotEvent(snapshot=flow.current_state)
# Process structured output if response_format is set
if response_format is not None and all_updates:
from agent_framework import AgentResponse
from pydantic import BaseModel
logger.info(f"Processing structured output, update count: {len(all_updates)}")
final_response = AgentResponse.from_agent_run_response_updates(all_updates, output_format_type=response_format)
if final_response.value and isinstance(final_response.value, BaseModel):
response_dict = final_response.value.model_dump(mode="json", exclude_none=True)
logger.info(f"Received structured output keys: {list(response_dict.keys())}")
# Extract state updates - if no state_schema, all non-message fields are state
state_keys = (
set(config.state_schema.keys()) if config.state_schema else set(response_dict.keys()) - {"message"}
)
state_updates = {k: v for k, v in response_dict.items() if k in state_keys}
if state_updates:
flow.current_state.update(state_updates)
yield StateSnapshotEvent(snapshot=flow.current_state)
logger.info(f"Emitted StateSnapshotEvent with updates: {list(state_updates.keys())}")
# Emit message field as text if present
if "message" in response_dict and response_dict["message"]:
message_id = generate_event_id()
yield TextMessageStartEvent(message_id=message_id, role="assistant")
yield TextMessageContentEvent(message_id=message_id, delta=response_dict["message"])
yield TextMessageEndEvent(message_id=message_id)
logger.info(f"Emitted conversational message with length={len(response_dict['message'])}")
# Feature #1: Emit ToolCallEndEvent for declaration-only tools (tools without results)
pending_without_end = flow.get_pending_without_end()
if pending_without_end:
logger.info(f"Found {len(pending_without_end)} pending tool calls without end event")
for tool_call in pending_without_end:
tool_call_id = tool_call.get("id")
tool_name = tool_call.get("function", {}).get("name")
if tool_call_id:
logger.info(f"Emitting ToolCallEndEvent for declaration-only tool '{tool_call_id}'")
yield ToolCallEndEvent(tool_call_id=tool_call_id)
# For predictive tools with require_confirmation, emit confirm_changes
if config.require_confirmation and config.predict_state_config and tool_name:
is_predictive_tool = any(cfg["tool"] == tool_name for cfg in config.predict_state_config.values())
if is_predictive_tool:
logger.info(f"Emitting confirm_changes for predictive tool '{tool_name}'")
# Extract state value from tool arguments for StateSnapshot
if predictive_handler:
try:
args_str = tool_call.get("function", {}).get("arguments", "{}")
args = json.loads(args_str) if isinstance(args_str, str) else args_str
result = predictive_handler.extract_state_value(tool_name, args)
if result:
state_key, state_value = result
flow.current_state[state_key] = state_value
yield StateSnapshotEvent(snapshot=flow.current_state)
except json.JSONDecodeError:
# Ignore malformed JSON in tool arguments for predictive state;
# predictive updates are best-effort and should not break the flow.
logger.warning(
"Failed to decode JSON arguments for predictive tool '%s' (tool_call_id=%s).",
tool_name,
tool_call_id,
)
# Emit confirm_changes tool call
confirm_id = generate_event_id()
yield ToolCallStartEvent(
tool_call_id=confirm_id,
tool_call_name="confirm_changes",
parent_message_id=flow.message_id,
)
confirm_args = {
"function_name": tool_name,
"function_call_id": tool_call_id,
"function_arguments": json.loads(tool_call.get("function", {}).get("arguments", "{}")),
"steps": [{"description": f"Execute {tool_name}", "status": "enabled"}],
}
yield ToolCallArgsEvent(tool_call_id=confirm_id, delta=json.dumps(confirm_args))
yield ToolCallEndEvent(tool_call_id=confirm_id)
flow.waiting_for_approval = True
# Close any open message
if flow.message_id:
yield TextMessageEndEvent(message_id=flow.message_id)
# Emit MessagesSnapshotEvent if we have tool calls or results
# Feature #5: Suppress intermediate snapshots for predictive tools without confirmation
should_emit_snapshot = flow.pending_tool_calls or flow.tool_results or flow.accumulated_text
if should_emit_snapshot:
# Check if we should suppress for predictive tool
last_tool_name = None
if flow.tool_results:
last_result = flow.tool_results[-1]
last_call_id = last_result.get("toolCallId")
last_tool_name = flow.get_tool_name(last_call_id)
if not _should_suppress_intermediate_snapshot(
last_tool_name, config.predict_state_config, config.require_confirmation
):
yield _build_messages_snapshot(flow, snapshot_messages)
# Always emit RunFinished - confirm_changes tool call is complete (Start -> Args -> End)
# The UI will show confirmation dialog and send a new request when user responds
yield RunFinishedEvent(run_id=run_id, thread_id=thread_id)
python/packages/ag-ui/agent_framework_ag_ui/_types.py
+16 -7
View File
@@ -13,13 +13,6 @@ if sys.version_info >= (3, 13):
else:
from typing_extensions import TypeVar
__all__ = [
"AGUIChatOptions",
"AgentState",
"PredictStateConfig",
"RunMetadata",
]
class PredictStateConfig(TypedDict):
"""Configuration for predictive state updates."""
@@ -62,6 +55,22 @@ class AGUIRequest(BaseModel):
None,
description="Optional shared state for agentic generative UI",
)
tools: list[dict[str, Any]] | None = Field(
None,
description="Client-side tools to advertise to the LLM",
)
context: list[dict[str, Any]] | None = Field(
None,
description="List of context objects provided to the agent",
)
forwarded_props: dict[str, Any] | None = Field(
None,
description="Additional properties forwarded to the agent",
)
parent_run_id: str | None = Field(
None,
description="ID of the run that spawned this run",
)
# region AG-UI Chat Options TypedDict