mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: AG-UI protocol support (#1826)
* Add AG-UI integration * Fix tests. PR feedback * Cleanup * PR Feedback * Improve README and getting started experience * Fix links
This commit is contained in:
Evan Mattson
committed by
GitHub
Unverified
GitHub
parent
0c862e97a6
commit
35a8565495
@@ -0,0 +1,31 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AG-UI protocol integration for Agent Framework."""
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from ._agent import AgentFrameworkAgent
|
||||
from ._confirmation_strategies import (
|
||||
ConfirmationStrategy,
|
||||
DefaultConfirmationStrategy,
|
||||
DocumentWriterConfirmationStrategy,
|
||||
RecipeConfirmationStrategy,
|
||||
TaskPlannerConfirmationStrategy,
|
||||
)
|
||||
from ._endpoint import add_agent_framework_fastapi_endpoint
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0"
|
||||
|
||||
__all__ = [
|
||||
"AgentFrameworkAgent",
|
||||
"add_agent_framework_fastapi_endpoint",
|
||||
"ConfirmationStrategy",
|
||||
"DefaultConfirmationStrategy",
|
||||
"TaskPlannerConfirmationStrategy",
|
||||
"RecipeConfirmationStrategy",
|
||||
"DocumentWriterConfirmationStrategy",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,160 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""AgentFrameworkAgent wrapper for AG-UI protocol - Clean Architecture."""
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import BaseEvent
|
||||
from agent_framework import AgentProtocol
|
||||
|
||||
from ._confirmation_strategies import ConfirmationStrategy, DefaultConfirmationStrategy
|
||||
from ._orchestrators import (
|
||||
DefaultOrchestrator,
|
||||
ExecutionContext,
|
||||
HumanInTheLoopOrchestrator,
|
||||
Orchestrator,
|
||||
)
|
||||
|
||||
|
||||
class AgentConfig:
|
||||
"""Configuration for agent wrapper."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
state_schema: dict[str, Any] | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
require_confirmation: bool = True,
|
||||
):
|
||||
"""Initialize agent configuration.
|
||||
|
||||
Args:
|
||||
state_schema: Optional state schema for state management
|
||||
predict_state_config: Configuration for predictive state updates
|
||||
require_confirmation: Whether predictive updates require confirmation
|
||||
"""
|
||||
self.state_schema = state_schema or {}
|
||||
self.predict_state_config = predict_state_config or {}
|
||||
self.require_confirmation = require_confirmation
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: AgentProtocol,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
state_schema: dict[str, Any] | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
require_confirmation: bool = True,
|
||||
orchestrators: list[Orchestrator] | None = None,
|
||||
confirmation_strategy: ConfirmationStrategy | None = None,
|
||||
):
|
||||
"""Initialize the AG-UI compatible agent wrapper.
|
||||
|
||||
Args:
|
||||
agent: The Agent Framework agent to wrap
|
||||
name: Optional name for the agent
|
||||
description: Optional description
|
||||
state_schema: Optional state schema for state management
|
||||
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.
|
||||
confirmation_strategy: Strategy for generating confirmation messages.
|
||||
Defaults to DefaultConfirmationStrategy if None.
|
||||
"""
|
||||
self.agent = agent
|
||||
self.name = name or getattr(agent, "name", "agent")
|
||||
self.description = description or getattr(agent, "description", "")
|
||||
|
||||
self.config = AgentConfig(
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
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",
|
||||
]
|
||||
@@ -0,0 +1,175 @@
|
||||
# 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."""
|
||||
|
||||
@abstractmethod
|
||||
def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate message when user approves function execution.
|
||||
|
||||
Args:
|
||||
steps: List of approved steps with 'description', 'status', etc.
|
||||
|
||||
Returns:
|
||||
Message to display to user
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
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
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def on_state_confirmed(self) -> str:
|
||||
"""Generate message when user confirms predictive state changes.
|
||||
|
||||
Returns:
|
||||
Message to display to user
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def on_state_rejected(self) -> str:
|
||||
"""Generate message when user rejects predictive state changes.
|
||||
|
||||
Returns:
|
||||
Message to display to user
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class DefaultConfirmationStrategy(ConfirmationStrategy):
|
||||
"""Generic confirmation messages suitable for most agents.
|
||||
|
||||
This preserves the original behavior from v1.
|
||||
"""
|
||||
|
||||
def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate generic approval message with step list."""
|
||||
enabled_steps = [s for s in steps if s.get("status") == "enabled"]
|
||||
|
||||
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['description']}\n")
|
||||
|
||||
message_parts.append("\nAll steps completed successfully!")
|
||||
|
||||
return "".join(message_parts)
|
||||
|
||||
def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate generic rejection message."""
|
||||
return "No problem! What would you like me to change about the plan?"
|
||||
|
||||
def on_state_confirmed(self) -> str:
|
||||
"""Generate generic state confirmation message."""
|
||||
return "Changes confirmed and applied successfully!"
|
||||
|
||||
def on_state_rejected(self) -> str:
|
||||
"""Generate generic state rejection message."""
|
||||
return "No problem! What would you like me to change?"
|
||||
|
||||
|
||||
class TaskPlannerConfirmationStrategy(ConfirmationStrategy):
|
||||
"""Domain-specific confirmation messages for task planning agents."""
|
||||
|
||||
def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate task-specific approval message."""
|
||||
enabled_steps = [s for s in steps if s.get("status") == "enabled"]
|
||||
|
||||
message_parts = ["Executing your requested tasks:\n\n"]
|
||||
|
||||
for i, step in enumerate(enabled_steps, 1):
|
||||
message_parts.append(f"{i}. {step['description']}\n")
|
||||
|
||||
message_parts.append("\nAll tasks completed successfully!")
|
||||
|
||||
return "".join(message_parts)
|
||||
|
||||
def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate task-specific rejection message."""
|
||||
return "No problem! Let me revise the plan. What would you like me to change?"
|
||||
|
||||
def on_state_confirmed(self) -> str:
|
||||
"""Task planners typically don't use state confirmation."""
|
||||
return "Tasks confirmed and ready to execute!"
|
||||
|
||||
def on_state_rejected(self) -> str:
|
||||
"""Task planners typically don't use state confirmation."""
|
||||
return "No problem! How should I adjust the task list?"
|
||||
|
||||
|
||||
class RecipeConfirmationStrategy(ConfirmationStrategy):
|
||||
"""Domain-specific confirmation messages for recipe agents."""
|
||||
|
||||
def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate recipe-specific approval message."""
|
||||
enabled_steps = [s for s in steps if s.get("status") == "enabled"]
|
||||
|
||||
message_parts = ["Updating your recipe:\n\n"]
|
||||
|
||||
for i, step in enumerate(enabled_steps, 1):
|
||||
message_parts.append(f"{i}. {step['description']}\n")
|
||||
|
||||
message_parts.append("\nRecipe updated successfully!")
|
||||
|
||||
return "".join(message_parts)
|
||||
|
||||
def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate recipe-specific rejection message."""
|
||||
return "No problem! What ingredients or steps should I change?"
|
||||
|
||||
def on_state_confirmed(self) -> str:
|
||||
"""Generate recipe-specific state confirmation message."""
|
||||
return "Recipe changes applied successfully!"
|
||||
|
||||
def on_state_rejected(self) -> str:
|
||||
"""Generate recipe-specific state rejection message."""
|
||||
return "No problem! What would you like me to adjust in the recipe?"
|
||||
|
||||
|
||||
class DocumentWriterConfirmationStrategy(ConfirmationStrategy):
|
||||
"""Domain-specific confirmation messages for document writing agents."""
|
||||
|
||||
def on_approval_accepted(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate document-specific approval message."""
|
||||
enabled_steps = [s for s in steps if s.get("status") == "enabled"]
|
||||
|
||||
message_parts = ["Applying your edits:\n\n"]
|
||||
|
||||
for i, step in enumerate(enabled_steps, 1):
|
||||
message_parts.append(f"{i}. {step['description']}\n")
|
||||
|
||||
message_parts.append("\nDocument updated successfully!")
|
||||
|
||||
return "".join(message_parts)
|
||||
|
||||
def on_approval_rejected(self, steps: list[dict[str, Any]]) -> str:
|
||||
"""Generate document-specific rejection message."""
|
||||
return "No problem! Which changes should I keep or modify?"
|
||||
|
||||
def on_state_confirmed(self) -> str:
|
||||
"""Generate document-specific state confirmation message."""
|
||||
return "Document edits applied!"
|
||||
|
||||
def on_state_rejected(self) -> str:
|
||||
"""Generate document-specific state rejection message."""
|
||||
return "No problem! What should I change about the document?"
|
||||
@@ -0,0 +1,94 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""FastAPI endpoint creation for AG-UI agents."""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.encoder import EventEncoder
|
||||
from agent_framework import AgentProtocol
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from ._agent import AgentFrameworkAgent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def add_agent_framework_fastapi_endpoint(
|
||||
app: FastAPI,
|
||||
agent: AgentProtocol | AgentFrameworkAgent,
|
||||
path: str = "/",
|
||||
state_schema: dict[str, Any] | None = None,
|
||||
predict_state_config: dict[str, dict[str, str]] | None = None,
|
||||
allow_origins: list[str] | None = None,
|
||||
) -> None:
|
||||
"""Add an AG-UI endpoint to a FastAPI app.
|
||||
|
||||
Args:
|
||||
app: The FastAPI application
|
||||
agent: The agent to expose (can be raw AgentProtocol or wrapped)
|
||||
path: The endpoint path
|
||||
state_schema: Optional state schema for shared state management
|
||||
predict_state_config: Optional predictive state update configuration.
|
||||
Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
|
||||
allow_origins: CORS origins (not yet implemented)
|
||||
"""
|
||||
if isinstance(agent, AgentProtocol):
|
||||
wrapped_agent = AgentFrameworkAgent(
|
||||
agent=agent,
|
||||
state_schema=state_schema,
|
||||
predict_state_config=predict_state_config,
|
||||
)
|
||||
else:
|
||||
wrapped_agent = agent
|
||||
|
||||
@app.post(path)
|
||||
async def agent_endpoint(request: Request): # type: ignore[misc]
|
||||
"""Handle AG-UI agent requests.
|
||||
|
||||
Note: Function is accessed via FastAPI's decorator registration,
|
||||
despite appearing unused to static analysis.
|
||||
"""
|
||||
try:
|
||||
input_data = await request.json()
|
||||
logger.debug(
|
||||
f"[{path}] Received request - Run ID: {input_data.get('run_id', 'no-run-id')}, "
|
||||
f"Thread ID: {input_data.get('thread_id', 'no-thread-id')}, "
|
||||
f"Messages: {len(input_data.get('messages', []))}"
|
||||
)
|
||||
logger.info(f"Received request at {path}: {input_data.get('run_id', 'no-run-id')}")
|
||||
|
||||
async def event_generator():
|
||||
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}")
|
||||
|
||||
encoded = encoder.encode(event)
|
||||
logger.debug(
|
||||
f"[{path}] Encoded as: {encoded[:200]}..."
|
||||
if len(encoded) > 200
|
||||
else f"[{path}] Encoded as: {encoded}"
|
||||
)
|
||||
yield encoded
|
||||
logger.info(f"[{path}] Completed streaming {event_count} events")
|
||||
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in agent endpoint: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -0,0 +1,675 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Event bridge for converting Agent Framework events to AG-UI protocol."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import (
|
||||
BaseEvent,
|
||||
CustomEvent,
|
||||
EventType,
|
||||
MessagesSnapshotEvent,
|
||||
RunFinishedEvent,
|
||||
RunStartedEvent,
|
||||
StateDeltaEvent,
|
||||
StateSnapshotEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
ToolCallStartEvent,
|
||||
)
|
||||
from agent_framework import (
|
||||
AgentRunResponseUpdate,
|
||||
FunctionApprovalRequestContent,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
from ._utils import generate_event_id
|
||||
|
||||
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,
|
||||
input_messages: list[Any] | None = None,
|
||||
require_confirmation: bool = True,
|
||||
) -> 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).
|
||||
input_messages: The input messages from the conversation history.
|
||||
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
|
||||
|
||||
# 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
|
||||
|
||||
# For MessagesSnapshotEvent: track tool calls and results
|
||||
self.input_messages = input_messages or []
|
||||
self.pending_tool_calls: list[dict[str, Any]] = [] # Track tool calls for assistant message
|
||||
self.tool_results: list[dict[str, Any]] = [] # Track tool results
|
||||
|
||||
async def from_agent_run_update(self, update: AgentRunResponseUpdate) -> list[BaseEvent]:
|
||||
"""
|
||||
Convert an AgentRunResponseUpdate to AG-UI events.
|
||||
|
||||
Args:
|
||||
update: The agent run update to convert.
|
||||
|
||||
Returns:
|
||||
List of AG-UI events.
|
||||
"""
|
||||
events: list[BaseEvent] = []
|
||||
|
||||
for content in update.contents:
|
||||
if isinstance(content, TextContent):
|
||||
# Skip text content if using structured outputs (it's just the JSON)
|
||||
if self.skip_text_content:
|
||||
continue
|
||||
|
||||
# Skip text content if we're about to emit confirm_changes
|
||||
# The summary should only appear after user confirms
|
||||
if self.should_stop_after_confirm:
|
||||
logger.debug(" >>> Skipping text content - waiting for confirm_changes response")
|
||||
# Save the summary text to show after confirmation
|
||||
self.suppressed_summary += content.text
|
||||
continue
|
||||
|
||||
if not self.current_message_id:
|
||||
self.current_message_id = generate_event_id()
|
||||
start_event = TextMessageStartEvent(
|
||||
message_id=self.current_message_id,
|
||||
role="assistant",
|
||||
)
|
||||
events.append(start_event)
|
||||
|
||||
event = TextMessageContentEvent(
|
||||
message_id=self.current_message_id,
|
||||
delta=content.text,
|
||||
)
|
||||
events.append(event)
|
||||
|
||||
elif isinstance(content, FunctionCallContent):
|
||||
# Log tool calls for debugging
|
||||
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_preview = str(content.arguments)[:50] if content.arguments else "None"
|
||||
logger.warning(f"FunctionCallContent missing name and call_id. Args: {args_preview}")
|
||||
|
||||
# Get or use existing tool call ID - all chunks of same tool call share the same call_id
|
||||
# Important: the first chunk might have name but no call_id yet
|
||||
if content.call_id:
|
||||
tool_call_id = content.call_id
|
||||
elif self.current_tool_call_id:
|
||||
tool_call_id = self.current_tool_call_id
|
||||
else:
|
||||
# Generate a new ID for this tool call
|
||||
tool_call_id = (
|
||||
generate_event_id()
|
||||
) # Handle streaming tool calls - name comes in first chunk, arguments in subsequent chunks
|
||||
if content.name:
|
||||
# This is a new tool call or the first chunk with the name
|
||||
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)
|
||||
|
||||
# Track tool call for MessagesSnapshotEvent
|
||||
# Initialize a new tool call entry
|
||||
self.pending_tool_calls.append(
|
||||
{
|
||||
"id": tool_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": content.name,
|
||||
"arguments": "", # Will accumulate as we get argument chunks
|
||||
},
|
||||
}
|
||||
)
|
||||
else:
|
||||
# Subsequent chunk without name - update our tracked ID if needed
|
||||
if tool_call_id:
|
||||
self.current_tool_call_id = tool_call_id
|
||||
|
||||
# Emit arguments if present
|
||||
if content.arguments:
|
||||
# content.arguments is already a JSON string from the LLM for streaming calls
|
||||
# For non-streaming it could be a dict, so we need to handle both
|
||||
if isinstance(content.arguments, str):
|
||||
delta_str = content.arguments
|
||||
else:
|
||||
# If it's a dict, convert to JSON
|
||||
delta_str = json.dumps(content.arguments)
|
||||
|
||||
logger.info(f" >>> Emitting ToolCallArgsEvent with delta: {delta_str!r}..., id='{tool_call_id}'")
|
||||
args_event = ToolCallArgsEvent(
|
||||
tool_call_id=tool_call_id,
|
||||
delta=delta_str,
|
||||
)
|
||||
events.append(args_event)
|
||||
|
||||
# Accumulate arguments for MessagesSnapshotEvent
|
||||
if self.pending_tool_calls:
|
||||
# Find the matching tool call and append the delta
|
||||
for tool_call in self.pending_tool_calls:
|
||||
if tool_call["id"] == tool_call_id:
|
||||
tool_call["function"]["arguments"] += delta_str
|
||||
break
|
||||
|
||||
# Predictive state updates - accumulate streaming arguments and emit deltas
|
||||
# Use current_tool_call_name since content.name is only present on first chunk
|
||||
if self.current_tool_call_name and self.predict_state_config:
|
||||
# Accumulate the argument string
|
||||
if isinstance(content.arguments, str):
|
||||
self.streaming_tool_args += content.arguments
|
||||
else:
|
||||
self.streaming_tool_args += json.dumps(content.arguments)
|
||||
|
||||
logger.debug(
|
||||
f" >>> Predictive state: accumulated {len(self.streaming_tool_args)} chars for tool '{self.current_tool_call_name}'"
|
||||
)
|
||||
|
||||
# Try to parse accumulated arguments (may be incomplete JSON)
|
||||
# We use a lenient approach: try standard parsing first, then try to extract partial values
|
||||
parsed_args = None
|
||||
try:
|
||||
parsed_args = json.loads(self.streaming_tool_args)
|
||||
except json.JSONDecodeError:
|
||||
# JSON is incomplete - try to extract partial string values
|
||||
# For streaming "document" field, we can extract: {"document": "text...
|
||||
# Look for pattern: {"field": "value (incomplete)
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
if config["tool"] == self.current_tool_call_name:
|
||||
tool_arg_name = config["tool_argument"]
|
||||
|
||||
# Try to extract partial string value for this argument
|
||||
# Pattern: "argument_name": "partial text
|
||||
pattern = rf'"{re.escape(tool_arg_name)}":\s*"([^"]*)'
|
||||
match = re.search(pattern, self.streaming_tool_args)
|
||||
|
||||
if match:
|
||||
partial_value = match.group(1)
|
||||
# Unescape common sequences
|
||||
partial_value = (
|
||||
partial_value.replace("\\n", "\n").replace('\\"', '"').replace("\\\\", "\\")
|
||||
)
|
||||
|
||||
# Emit delta if we have new content
|
||||
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:
|
||||
value_preview = (
|
||||
str(partial_value)[:100] + "..."
|
||||
if len(str(partial_value)) > 100
|
||||
else str(partial_value)
|
||||
)
|
||||
logger.info(
|
||||
f" >>> StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
|
||||
f"op=replace, path=/{state_key}, value={value_preview}"
|
||||
)
|
||||
elif self.state_delta_count % 100 == 0:
|
||||
logger.info(f" >>> StateDeltaEvent #{self.state_delta_count} emitted")
|
||||
|
||||
events.append(state_delta_event)
|
||||
self.last_emitted_state[state_key] = partial_value
|
||||
self.pending_state_updates[state_key] = partial_value
|
||||
|
||||
# If we successfully parsed complete JSON, process it
|
||||
if parsed_args:
|
||||
# Check if this tool matches any predictive state config
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
if config["tool"] == self.current_tool_call_name:
|
||||
tool_arg_name = config["tool_argument"]
|
||||
|
||||
# Extract the state value
|
||||
if tool_arg_name == "*":
|
||||
state_value = parsed_args
|
||||
elif tool_arg_name in parsed_args:
|
||||
state_value = parsed_args[tool_arg_name]
|
||||
else:
|
||||
continue
|
||||
|
||||
# Only emit if state has changed from last emission
|
||||
if (
|
||||
state_key not in self.last_emitted_state
|
||||
or self.last_emitted_state[state_key] != state_value
|
||||
):
|
||||
# Emit StateDeltaEvent for real-time UI updates (JSON Patch format)
|
||||
state_delta_event = StateDeltaEvent(
|
||||
delta=[
|
||||
{
|
||||
"op": "replace", # Use replace since field exists in schema
|
||||
"path": f"/{state_key}", # JSON Pointer path with leading slash
|
||||
"value": state_value,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
# Increment counter and log every 10th emission with sample data
|
||||
self.state_delta_count += 1
|
||||
if self.state_delta_count % 10 == 1: # Log 1st, 11th, 21st, etc.
|
||||
value_preview = (
|
||||
str(state_value)[:100] + "..."
|
||||
if len(str(state_value)) > 100
|
||||
else str(state_value)
|
||||
)
|
||||
logger.info(
|
||||
f" >>> StateDeltaEvent #{self.state_delta_count} for '{state_key}': "
|
||||
f"op=replace, path=/{state_key}, value={value_preview}"
|
||||
)
|
||||
elif self.state_delta_count % 100 == 0: # Also log every 100th
|
||||
logger.info(f" >>> StateDeltaEvent #{self.state_delta_count} emitted")
|
||||
|
||||
events.append(state_delta_event)
|
||||
|
||||
# Track what we emitted
|
||||
self.last_emitted_state[state_key] = state_value
|
||||
self.pending_state_updates[state_key] = state_value
|
||||
|
||||
# Legacy predictive state check (for when arguments are complete)
|
||||
if content.name and content.arguments:
|
||||
parsed_args = content.parse_arguments()
|
||||
|
||||
if parsed_args:
|
||||
logger.info(f"Checking predict_state_config: {self.predict_state_config}")
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
logger.info(f"Checking state_key='{state_key}', config={config}")
|
||||
if config["tool"] == content.name:
|
||||
tool_arg_name = config["tool_argument"]
|
||||
logger.info(
|
||||
f"MATCHED tool '{content.name}' for state key '{state_key}', arg='{tool_arg_name}'"
|
||||
)
|
||||
|
||||
# If tool_argument is "*", use all arguments as the state value
|
||||
if tool_arg_name == "*":
|
||||
state_value = parsed_args
|
||||
logger.info(f"Using all args as state value, keys: {list(state_value.keys())}")
|
||||
elif tool_arg_name in parsed_args:
|
||||
state_value = parsed_args[tool_arg_name]
|
||||
logger.info(f"Using specific arg '{tool_arg_name}' as state value")
|
||||
else:
|
||||
logger.warning(f"Tool argument '{tool_arg_name}' not found in parsed args")
|
||||
continue
|
||||
|
||||
# Emit predictive delta (JSON Patch format)
|
||||
state_delta_event = StateDeltaEvent(
|
||||
delta=[
|
||||
{
|
||||
"op": "replace", # Use replace since field exists in schema
|
||||
"path": f"/{state_key}", # JSON Pointer path with leading slash
|
||||
"value": state_value,
|
||||
}
|
||||
],
|
||||
)
|
||||
logger.info(
|
||||
f" >>> Emitting StateDeltaEvent for key '{state_key}', value type: {type(state_value)}"
|
||||
)
|
||||
events.append(state_delta_event)
|
||||
|
||||
# Track pending update for later snapshot
|
||||
self.pending_state_updates[state_key] = state_value
|
||||
|
||||
# Note: ToolCallEndEvent is emitted when we receive FunctionResultContent,
|
||||
# not here during streaming, since we don't know when the stream is complete
|
||||
|
||||
elif isinstance(content, FunctionResultContent):
|
||||
# First emit ToolCallEndEvent to close the tool call
|
||||
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)
|
||||
|
||||
# Log total StateDeltaEvent count for this tool call
|
||||
if self.state_delta_count > 0:
|
||||
logger.info(
|
||||
f" >>> Tool call '{content.call_id}' complete: emitted {self.state_delta_count} StateDeltaEvents total"
|
||||
)
|
||||
|
||||
# Reset streaming accumulator and counter for next tool call
|
||||
self.streaming_tool_args = ""
|
||||
self.state_delta_count = 0
|
||||
|
||||
# Tool result - emit ToolCallResultEvent
|
||||
result_message_id = generate_event_id()
|
||||
|
||||
# Preserve structured data for backend tool rendering
|
||||
# Serialize dicts to JSON string, otherwise convert to string
|
||||
if isinstance(content.result, dict):
|
||||
result_content = json.dumps(content.result) # type: ignore[arg-type]
|
||||
elif content.result is not None:
|
||||
result_content = str(content.result)
|
||||
else:
|
||||
result_content = ""
|
||||
|
||||
result_event = ToolCallResultEvent(
|
||||
message_id=result_message_id,
|
||||
tool_call_id=content.call_id,
|
||||
content=result_content,
|
||||
role="tool",
|
||||
)
|
||||
events.append(result_event)
|
||||
|
||||
# Track tool result for MessagesSnapshotEvent
|
||||
self.tool_results.append(
|
||||
{
|
||||
"id": result_message_id,
|
||||
"role": "tool",
|
||||
"tool_call_id": content.call_id,
|
||||
"content": result_content,
|
||||
}
|
||||
)
|
||||
|
||||
# Emit MessagesSnapshotEvent with the complete conversation including tool calls and results
|
||||
# This is required for CopilotKit's useCopilotAction to detect tool result
|
||||
if self.pending_tool_calls and self.tool_results:
|
||||
# Build assistant message with tool_calls
|
||||
assistant_message = {
|
||||
"id": generate_event_id(),
|
||||
"role": "assistant",
|
||||
"tool_calls": self.pending_tool_calls.copy(), # Copy the accumulated tool calls
|
||||
}
|
||||
|
||||
# Build complete messages array: input messages + assistant message + tool results
|
||||
all_messages = list(self.input_messages) + [assistant_message] + self.tool_results.copy()
|
||||
|
||||
# Emit MessagesSnapshotEvent using the proper event type
|
||||
messages_snapshot_event = MessagesSnapshotEvent(
|
||||
type=EventType.MESSAGES_SNAPSHOT, messages=all_messages
|
||||
)
|
||||
logger.info(f" >>> Emitting MessagesSnapshotEvent with {len(all_messages)} messages")
|
||||
events.append(messages_snapshot_event)
|
||||
|
||||
# After tool execution, emit StateSnapshotEvent if we have pending state updates
|
||||
if self.pending_state_updates:
|
||||
# Update the current state with pending updates
|
||||
for key, value in self.pending_state_updates.items():
|
||||
self.current_state[key] = value
|
||||
|
||||
# Log the state structure for debugging
|
||||
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(
|
||||
f"Recipe fields: title={recipe.get('title')}, "
|
||||
f"skill_level={recipe.get('skill_level')}, "
|
||||
f"ingredients_count={len(recipe.get('ingredients', []))}, "
|
||||
f"instructions_count={len(recipe.get('instructions', []))}"
|
||||
)
|
||||
|
||||
# Emit complete state snapshot
|
||||
state_snapshot_event = StateSnapshotEvent(
|
||||
snapshot=self.current_state,
|
||||
)
|
||||
events.append(state_snapshot_event)
|
||||
|
||||
# Check if this was a predictive state update tool (e.g., write_document_local)
|
||||
# If so, emit a confirm_changes tool call for the UI modal
|
||||
tool_was_predictive = False
|
||||
logger.debug(
|
||||
f" >>> Checking predictive state: current_tool='{self.current_tool_call_name}', "
|
||||
f"predict_config={list(self.predict_state_config.keys()) if self.predict_state_config else 'None'}"
|
||||
)
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
# Check if this tool call matches a predictive config
|
||||
# We need to match against self.current_tool_call_name
|
||||
if self.current_tool_call_name and config["tool"] == self.current_tool_call_name:
|
||||
logger.info(
|
||||
f" >>> Tool '{self.current_tool_call_name}' matches predictive config for state key '{state_key}'"
|
||||
)
|
||||
tool_was_predictive = True
|
||||
break
|
||||
|
||||
if tool_was_predictive and self.require_confirmation:
|
||||
# Emit confirm_changes tool call sequence
|
||||
confirm_call_id = generate_event_id()
|
||||
|
||||
logger.info(" >>> Emitting confirm_changes tool call for predictive update")
|
||||
|
||||
# Track confirm_changes tool call for MessagesSnapshotEvent (so it persists after RUN_FINISHED)
|
||||
self.pending_tool_calls.append(
|
||||
{
|
||||
"id": confirm_call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "confirm_changes",
|
||||
"arguments": "{}",
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
# Start the confirm_changes tool call
|
||||
confirm_start = ToolCallStartEvent(
|
||||
tool_call_id=confirm_call_id,
|
||||
tool_call_name="confirm_changes",
|
||||
)
|
||||
events.append(confirm_start)
|
||||
|
||||
# Empty args for confirm_changes
|
||||
confirm_args = ToolCallArgsEvent(
|
||||
tool_call_id=confirm_call_id,
|
||||
delta="{}",
|
||||
)
|
||||
events.append(confirm_args)
|
||||
|
||||
# End the confirm_changes tool call
|
||||
confirm_end = ToolCallEndEvent(
|
||||
tool_call_id=confirm_call_id,
|
||||
)
|
||||
events.append(confirm_end)
|
||||
|
||||
# Emit MessagesSnapshotEvent so confirm_changes persists after RUN_FINISHED
|
||||
# Build assistant message with pending confirm_changes tool call
|
||||
assistant_message = {
|
||||
"id": generate_event_id(),
|
||||
"role": "assistant",
|
||||
"tool_calls": self.pending_tool_calls.copy(), # Includes confirm_changes
|
||||
}
|
||||
|
||||
# Build complete messages array: input messages + assistant message + any tool results
|
||||
all_messages = list(self.input_messages) + [assistant_message] + self.tool_results.copy()
|
||||
|
||||
# Emit MessagesSnapshotEvent
|
||||
messages_snapshot_event = MessagesSnapshotEvent(
|
||||
type=EventType.MESSAGES_SNAPSHOT, messages=all_messages
|
||||
)
|
||||
logger.info(
|
||||
f" >>> Emitting MessagesSnapshotEvent for confirm_changes with {len(all_messages)} messages"
|
||||
)
|
||||
events.append(messages_snapshot_event)
|
||||
|
||||
# Set flag to stop the run after this - we're waiting for user response
|
||||
self.should_stop_after_confirm = True
|
||||
logger.info(" >>> Set flag to stop run after confirm_changes")
|
||||
elif tool_was_predictive:
|
||||
logger.info(" >>> Skipping confirm_changes - require_confirmation is False")
|
||||
|
||||
# Clear pending updates and reset tool name tracker
|
||||
self.pending_state_updates.clear()
|
||||
self.last_emitted_state.clear()
|
||||
self.current_tool_call_name = None # Reset for next tool call
|
||||
|
||||
elif isinstance(content, FunctionApprovalRequestContent):
|
||||
# Human in the loop - function approval request
|
||||
logger.info("=== FUNCTION APPROVAL REQUEST ===")
|
||||
logger.info(f" Function: {content.function_call.name}")
|
||||
logger.info(f" Call ID: {content.function_call.call_id}")
|
||||
|
||||
# Parse the arguments to extract state for predictive UI updates
|
||||
parsed_args = content.function_call.parse_arguments()
|
||||
logger.info(f" Parsed args keys: {list(parsed_args.keys()) if parsed_args else 'None'}")
|
||||
|
||||
# Check if this matches our predict_state_config and emit state
|
||||
if parsed_args and self.predict_state_config:
|
||||
logger.info(f" Checking predict_state_config: {self.predict_state_config}")
|
||||
for state_key, config in self.predict_state_config.items():
|
||||
if config["tool"] == content.function_call.name:
|
||||
tool_arg_name = config["tool_argument"]
|
||||
logger.info(
|
||||
f" MATCHED tool '{content.function_call.name}' for state key '{state_key}', arg='{tool_arg_name}'"
|
||||
)
|
||||
|
||||
# Extract the state value
|
||||
if tool_arg_name == "*":
|
||||
state_value = parsed_args
|
||||
elif tool_arg_name in parsed_args:
|
||||
state_value = parsed_args[tool_arg_name]
|
||||
else:
|
||||
logger.warning(f" Tool argument '{tool_arg_name}' not found in parsed args")
|
||||
continue
|
||||
|
||||
# Update current state
|
||||
self.current_state[state_key] = state_value
|
||||
logger.info(
|
||||
f" >>> Emitting StateSnapshotEvent for key '{state_key}', value type: {type(state_value)}"
|
||||
)
|
||||
|
||||
# Emit state snapshot
|
||||
state_snapshot = StateSnapshotEvent(
|
||||
snapshot=self.current_state,
|
||||
)
|
||||
events.append(state_snapshot)
|
||||
|
||||
# The tool call has been streamed already (Start/Args events)
|
||||
# Now we need to close it with an End event before the agent waits for approval
|
||||
if content.function_call.call_id:
|
||||
end_event = ToolCallEndEvent(
|
||||
tool_call_id=content.function_call.call_id,
|
||||
)
|
||||
logger.info(
|
||||
f" >>> Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'"
|
||||
)
|
||||
events.append(end_event)
|
||||
|
||||
# Emit custom event for approval request
|
||||
# Note: In AG-UI protocol, the frontend handles interrupts automatically
|
||||
# when it sees a tool call with the configured name (via predict_state_config)
|
||||
# This custom event is for additional metadata if needed
|
||||
approval_event = CustomEvent(
|
||||
name="function_approval_request",
|
||||
value={
|
||||
"id": content.id,
|
||||
"function_call": {
|
||||
"call_id": content.function_call.call_id,
|
||||
"name": content.function_call.name,
|
||||
"arguments": content.function_call.parse_arguments(),
|
||||
},
|
||||
},
|
||||
)
|
||||
logger.info(f" >>> Emitting function_approval_request custom event for '{content.function_call.name}'")
|
||||
events.append(approval_event)
|
||||
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,218 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Message format conversion between AG-UI and Agent Framework."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import (
|
||||
ChatMessage,
|
||||
FunctionApprovalResponseContent,
|
||||
FunctionCallContent,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
|
||||
# Role mapping constants
|
||||
_AGUI_TO_FRAMEWORK_ROLE = {
|
||||
"user": Role.USER,
|
||||
"assistant": Role.ASSISTANT,
|
||||
"system": Role.SYSTEM,
|
||||
}
|
||||
|
||||
_FRAMEWORK_TO_AGUI_ROLE = {
|
||||
Role.USER: "user",
|
||||
Role.ASSISTANT: "assistant",
|
||||
Role.SYSTEM: "system",
|
||||
}
|
||||
|
||||
|
||||
def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[ChatMessage]:
|
||||
"""Convert AG-UI messages to Agent Framework format.
|
||||
|
||||
Args:
|
||||
messages: List of AG-UI messages
|
||||
|
||||
Returns:
|
||||
List of Agent Framework ChatMessage objects
|
||||
"""
|
||||
result: list[ChatMessage] = []
|
||||
for msg in messages:
|
||||
# Check for backend tool rendering results FIRST (may not have role field)
|
||||
if "actionExecutionId" in msg or "actionName" in msg:
|
||||
# Backend tool rendering - convert to FunctionResultContent
|
||||
from agent_framework import FunctionResultContent
|
||||
|
||||
tool_call_id = msg.get("actionExecutionId", "")
|
||||
result_content = msg.get("result", msg.get("content", ""))
|
||||
|
||||
chat_msg = ChatMessage(
|
||||
role=Role.ASSISTANT, # Tool results are assistant messages
|
||||
contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
|
||||
)
|
||||
|
||||
if "id" in msg:
|
||||
chat_msg.message_id = msg["id"]
|
||||
|
||||
result.append(chat_msg)
|
||||
continue
|
||||
|
||||
role_str = msg.get("role", "user")
|
||||
|
||||
# Handle tool result messages (with role="tool")
|
||||
if role_str == "tool":
|
||||
# Check if this is a standard tool result (has tool_call_id or toolCallId)
|
||||
tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
|
||||
result_content = msg.get("content", "")
|
||||
|
||||
# Distinguish between backend tool results and approval responses
|
||||
# Approval responses have {"accepted": ...} structure
|
||||
is_approval = False
|
||||
if result_content:
|
||||
import json
|
||||
|
||||
try:
|
||||
parsed_content = json.loads(result_content)
|
||||
is_approval = "accepted" in parsed_content
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
is_approval = False
|
||||
|
||||
# Backend tool results have non-empty content WITHOUT "accepted" field
|
||||
if tool_call_id and result_content and not is_approval:
|
||||
# Backend tool execution - convert to FunctionResultContent
|
||||
from agent_framework import FunctionResultContent
|
||||
|
||||
chat_msg = ChatMessage(
|
||||
role=Role.ASSISTANT, # Tool results are assistant messages
|
||||
contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
|
||||
)
|
||||
|
||||
if "id" in msg:
|
||||
chat_msg.message_id = msg["id"]
|
||||
|
||||
result.append(chat_msg)
|
||||
continue
|
||||
else:
|
||||
# Human-in-the-loop approval response - mark for special handling
|
||||
content = msg.get("content", "")
|
||||
chat_msg = ChatMessage(
|
||||
role=Role.USER, # Approval responses are user messages
|
||||
contents=[TextContent(text=content)],
|
||||
)
|
||||
# Mark this as a tool result so we can detect it later
|
||||
chat_msg.metadata = {"is_tool_result": True, "tool_call_id": msg.get("toolCallId", "")} # type: ignore[attr-defined]
|
||||
|
||||
if "id" in msg:
|
||||
chat_msg.message_id = msg["id"]
|
||||
|
||||
result.append(chat_msg)
|
||||
continue
|
||||
|
||||
role = _AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)
|
||||
|
||||
# Check if this message contains function approvals
|
||||
if "function_approvals" in msg and msg["function_approvals"]:
|
||||
# Convert function approvals to FunctionApprovalResponseContent
|
||||
contents: list[Any] = []
|
||||
for approval in msg["function_approvals"]:
|
||||
# Create FunctionCallContent with the modified arguments
|
||||
func_call = FunctionCallContent(
|
||||
call_id=approval.get("call_id", ""),
|
||||
name=approval.get("name", ""),
|
||||
arguments=approval.get("arguments", {}),
|
||||
)
|
||||
|
||||
# Create the approval response
|
||||
approval_response = FunctionApprovalResponseContent(
|
||||
approved=approval.get("approved", True),
|
||||
id=approval.get("id", ""),
|
||||
function_call=func_call,
|
||||
)
|
||||
contents.append(approval_response)
|
||||
|
||||
chat_msg = ChatMessage(role=role, contents=contents) # type: ignore[arg-type]
|
||||
else:
|
||||
# Regular text message
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
chat_msg = ChatMessage(role=role, contents=[TextContent(text=content)])
|
||||
else:
|
||||
chat_msg = ChatMessage(role=role, contents=[TextContent(text=str(content))])
|
||||
|
||||
if "id" in msg:
|
||||
chat_msg.message_id = msg["id"]
|
||||
|
||||
result.append(chat_msg)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def agent_framework_messages_to_agui(messages: list[ChatMessage]) -> list[dict[str, Any]]:
|
||||
"""Convert Agent Framework messages to AG-UI format.
|
||||
|
||||
Args:
|
||||
messages: List of Agent Framework ChatMessage objects
|
||||
|
||||
Returns:
|
||||
List of AG-UI message dictionaries
|
||||
"""
|
||||
result: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
role = _FRAMEWORK_TO_AGUI_ROLE.get(msg.role, "user")
|
||||
|
||||
content_text = ""
|
||||
tool_calls: list[dict[str, Any]] = []
|
||||
|
||||
for content in msg.contents:
|
||||
if isinstance(content, TextContent):
|
||||
content_text += content.text
|
||||
elif isinstance(content, FunctionCallContent):
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": content.call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": content.name,
|
||||
"arguments": content.arguments,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
agui_msg: dict[str, Any] = {
|
||||
"role": role,
|
||||
"content": content_text,
|
||||
}
|
||||
|
||||
if msg.message_id:
|
||||
agui_msg["id"] = msg.message_id
|
||||
|
||||
if tool_calls:
|
||||
agui_msg["tool_calls"] = tool_calls
|
||||
|
||||
result.append(agui_msg)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def extract_text_from_contents(contents: list[Any]) -> str:
|
||||
"""Extract text from Agent Framework contents.
|
||||
|
||||
Args:
|
||||
contents: List of content objects
|
||||
|
||||
Returns:
|
||||
Concatenated text
|
||||
"""
|
||||
text_parts: list[str] = []
|
||||
for content in contents:
|
||||
if isinstance(content, TextContent):
|
||||
text_parts.append(content.text)
|
||||
elif hasattr(content, "text"):
|
||||
text_parts.append(content.text)
|
||||
return "".join(text_parts)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"agui_messages_to_agent_framework",
|
||||
"agent_framework_messages_to_agui",
|
||||
"extract_text_from_contents",
|
||||
]
|
||||
@@ -0,0 +1,439 @@
|
||||
# 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
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ag_ui.core import (
|
||||
BaseEvent,
|
||||
RunErrorEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
TextMessageStartEvent,
|
||||
)
|
||||
from agent_framework import AgentProtocol, AgentThread, TextContent
|
||||
|
||||
from ._utils import generate_event_id
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._agent import AgentConfig
|
||||
from ._confirmation_strategies import ConfirmationStrategy
|
||||
|
||||
|
||||
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._last_message = None
|
||||
self._run_id: str | None = None
|
||||
self._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 agui_messages_to_agent_framework
|
||||
|
||||
raw = self.input_data.get("messages", [])
|
||||
self._messages = agui_messages_to_agent_framework(raw)
|
||||
return self._messages
|
||||
|
||||
@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 run_id(self) -> str:
|
||||
"""Get or generate run ID."""
|
||||
if self._run_id is None:
|
||||
self._run_id = self.input_data.get("run_id") or str(uuid.uuid4())
|
||||
# This should never be None after the if block above, but satisfy type checkers
|
||||
if self._run_id is None: # pragma: no cover
|
||||
raise RuntimeError("Failed to initialize run_id")
|
||||
return self._run_id
|
||||
|
||||
@property
|
||||
def thread_id(self) -> str:
|
||||
"""Get or generate thread ID."""
|
||||
if self._thread_id is None:
|
||||
self._thread_id = self.input_data.get("thread_id") or str(uuid.uuid4())
|
||||
# This should never be None after the if block above, but satisfy type checkers
|
||||
if self._thread_id is None: # pragma: no cover
|
||||
raise RuntimeError("Failed to initialize thread_id")
|
||||
return self._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 or not hasattr(msg, "metadata"):
|
||||
return False
|
||||
|
||||
metadata = getattr(msg, "metadata", None)
|
||||
if not metadata:
|
||||
return False
|
||||
|
||||
return bool(metadata.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 isinstance(content, TextContent):
|
||||
tool_content_text = content.text
|
||||
break
|
||||
|
||||
try:
|
||||
tool_result = json.loads(tool_content_text)
|
||||
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]}")
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
logger.info(f"Starting default agent run for thread_id={context.thread_id}, run_id={context.run_id}")
|
||||
|
||||
# Initialize state tracking
|
||||
initial_state = context.input_data.get("state", {})
|
||||
current_state: dict[str, Any] = initial_state.copy() if initial_state else {}
|
||||
|
||||
# Check if agent uses structured outputs (response_format)
|
||||
chat_options = getattr(context.agent, "chat_options", None)
|
||||
response_format = getattr(chat_options, "response_format", None) if chat_options else None
|
||||
skip_text_content = response_format is not None
|
||||
|
||||
# Create event bridge
|
||||
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,
|
||||
input_messages=context.input_data.get("messages", []),
|
||||
require_confirmation=context.config.require_confirmation,
|
||||
)
|
||||
|
||||
yield event_bridge.create_run_started_event()
|
||||
|
||||
# Emit PredictState custom event if we have predictive state config
|
||||
if context.config.predict_state_config:
|
||||
from ag_ui.core import CustomEvent, EventType
|
||||
|
||||
predict_state_value = [
|
||||
{
|
||||
"state_key": state_key,
|
||||
"tool": config["tool"],
|
||||
"tool_argument": config["tool_argument"],
|
||||
}
|
||||
for state_key, config in context.config.predict_state_config.items()
|
||||
]
|
||||
|
||||
yield CustomEvent(
|
||||
type=EventType.CUSTOM,
|
||||
name="PredictState",
|
||||
value=predict_state_value,
|
||||
)
|
||||
|
||||
# If we have a state schema, ensure we emit initial state snapshot
|
||||
if context.config.state_schema:
|
||||
# Initialize missing state fields with appropriate empty values based on schema type
|
||||
for key, schema in context.config.state_schema.items():
|
||||
if key not in current_state:
|
||||
# Default to empty object; use empty array if schema specifies "array" type
|
||||
current_state[key] = [] if isinstance(schema, dict) and schema.get("type") == "array" else {} # type: ignore
|
||||
yield event_bridge.create_state_snapshot_event(current_state)
|
||||
|
||||
# Create thread for context tracking
|
||||
thread = AgentThread()
|
||||
thread.metadata = { # type: ignore[attr-defined]
|
||||
"ag_ui_thread_id": context.thread_id,
|
||||
"ag_ui_run_id": context.run_id,
|
||||
}
|
||||
|
||||
# Inject current state into thread metadata so agent can access it
|
||||
if current_state:
|
||||
thread.metadata["current_state"] = current_state # type: ignore[attr-defined]
|
||||
|
||||
# Add incoming AG-UI messages to the thread history
|
||||
if context.messages:
|
||||
await thread.on_new_messages(context.messages)
|
||||
|
||||
# Get the last message as the new input
|
||||
new_message = context.last_message
|
||||
if not new_message:
|
||||
logger.warning("No messages provided in AG-UI input")
|
||||
yield event_bridge.create_run_finished_event()
|
||||
return
|
||||
|
||||
# Inject current state as system message context if we have state
|
||||
messages_to_run: list[Any] = []
|
||||
if current_state and context.config.state_schema:
|
||||
state_json = json.dumps(current_state, indent=2)
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
state_context_msg = ChatMessage(
|
||||
role="system",
|
||||
contents=[
|
||||
TextContent(
|
||||
text=f"""Current state of the application:
|
||||
{state_json}
|
||||
|
||||
When modifying state, you MUST include ALL existing data plus your changes.
|
||||
For example, if adding a new ingredient, include all existing ingredients PLUS the new one.
|
||||
Never replace existing data - always append or merge."""
|
||||
)
|
||||
],
|
||||
)
|
||||
messages_to_run.append(state_context_msg)
|
||||
|
||||
messages_to_run.append(new_message)
|
||||
|
||||
# Collect all updates to get the final structured output
|
||||
all_updates: list[Any] = []
|
||||
async for update in context.agent.run_stream(messages_to_run, thread=thread):
|
||||
all_updates.append(update)
|
||||
events = await event_bridge.from_agent_run_update(update)
|
||||
for event in events:
|
||||
yield event
|
||||
|
||||
# After agent completes, check if we should stop (waiting for user to confirm changes)
|
||||
if event_bridge.should_stop_after_confirm:
|
||||
logger.info(" >>> Stopping run after confirm_changes - waiting for user response")
|
||||
yield event_bridge.create_run_finished_event()
|
||||
return
|
||||
|
||||
# After streaming completes, check if agent has response_format and extract structured output
|
||||
if all_updates and response_format:
|
||||
from agent_framework import AgentRunResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger.info(f"Processing structured output, update count: {len(all_updates)}")
|
||||
|
||||
# Convert streaming updates to final response to get the structured output
|
||||
final_response = AgentRunResponse.from_agent_run_response_updates(
|
||||
all_updates, output_format_type=response_format
|
||||
)
|
||||
|
||||
if final_response.value and isinstance(final_response.value, BaseModel):
|
||||
# Convert Pydantic model to dict
|
||||
response_dict = final_response.value.model_dump(mode="json", exclude_none=True)
|
||||
logger.info(f"Received structured output: {list(response_dict.keys())}")
|
||||
|
||||
# Extract state fields based on state_schema
|
||||
state_updates: dict[str, Any] = {}
|
||||
|
||||
if context.config.state_schema:
|
||||
# Use state_schema to determine which fields are state
|
||||
for state_key in context.config.state_schema.keys():
|
||||
if state_key in response_dict:
|
||||
state_updates[state_key] = response_dict[state_key]
|
||||
else:
|
||||
# No schema: treat all non-message fields as state
|
||||
state_updates = {k: v for k, v in response_dict.items() if k != "message"}
|
||||
|
||||
# Apply state updates if any found
|
||||
if state_updates:
|
||||
current_state.update(state_updates)
|
||||
|
||||
# Emit StateSnapshotEvent with the updated state
|
||||
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 there's a message field, emit it as chat text
|
||||
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: {response_dict['message'][:100]}...")
|
||||
|
||||
if event_bridge.current_message_id:
|
||||
yield event_bridge.create_message_end_event(event_bridge.current_message_id)
|
||||
|
||||
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",
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Type definitions for AG-UI integration."""
|
||||
|
||||
from typing import Any, TypedDict
|
||||
|
||||
|
||||
class PredictStateConfig(TypedDict):
|
||||
"""Configuration for predictive state updates."""
|
||||
|
||||
state_key: str
|
||||
tool: str
|
||||
tool_argument: str | None
|
||||
|
||||
|
||||
class RunMetadata(TypedDict):
|
||||
"""Metadata for agent run."""
|
||||
|
||||
run_id: str
|
||||
thread_id: str
|
||||
predict_state: list[PredictStateConfig] | None
|
||||
|
||||
|
||||
class AgentState(TypedDict):
|
||||
"""Base state for AG-UI agents."""
|
||||
|
||||
messages: list[Any] | None
|
||||
@@ -0,0 +1,57 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Utility functions for AG-UI integration."""
|
||||
|
||||
import copy
|
||||
import uuid
|
||||
from dataclasses import asdict, is_dataclass
|
||||
from datetime import date, datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
def generate_event_id() -> str:
|
||||
"""Generate a unique event ID."""
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
def merge_state(current: dict[str, Any], update: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Merge state updates.
|
||||
|
||||
Args:
|
||||
current: Current state dictionary
|
||||
update: Update to apply
|
||||
|
||||
Returns:
|
||||
Merged state
|
||||
"""
|
||||
result = copy.deepcopy(current)
|
||||
result.update(update)
|
||||
return result
|
||||
|
||||
|
||||
def make_json_safe(obj: Any) -> Any: # noqa: ANN401
|
||||
"""Make an object JSON serializable.
|
||||
|
||||
Args:
|
||||
obj: Object to make JSON safe
|
||||
|
||||
Returns:
|
||||
JSON-serializable version of the object
|
||||
"""
|
||||
if obj is None or isinstance(obj, (str, int, float, bool)):
|
||||
return obj
|
||||
if isinstance(obj, (datetime, date)):
|
||||
return obj.isoformat()
|
||||
if is_dataclass(obj):
|
||||
return asdict(obj) # type: ignore[arg-type]
|
||||
if hasattr(obj, "model_dump"):
|
||||
return obj.model_dump() # type: ignore[no-any-return]
|
||||
if hasattr(obj, "dict"):
|
||||
return obj.dict() # type: ignore[no-any-return]
|
||||
if hasattr(obj, "__dict__"):
|
||||
return {key: make_json_safe(value) for key, value in vars(obj).items()} # type: ignore[misc]
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [make_json_safe(item) for item in obj] # type: ignore[misc]
|
||||
if isinstance(obj, dict):
|
||||
return {key: make_json_safe(value) for key, value in obj.items()} # type: ignore[misc]
|
||||
return str(obj)
|
||||
@@ -0,0 +1 @@
|
||||
# Marker file for PEP 561
|
||||
Reference in New Issue
Block a user