diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 7b93c64a5c6..841d14e201c 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -239,12 +239,14 @@ Review focus: whether these names are the right stable contract for Python users | Surface | Public exports | | --- | --- | -| `agent_framework.ag_ui` facade | `AgentFrameworkAgent`, `AgentFrameworkWorkflow`, `AGUIChatClient`, `AGUIEventConverter`, `AGUIHttpService`, `AGUIThreadSnapshot`, `AGUIThreadSnapshotStore`, `InMemoryAGUIThreadSnapshotStore`, `SnapshotScopeResolver`, `add_agent_framework_fastapi_endpoint`, `state_update`, `__version__` | +| `agent_framework.ag_ui` facade | `AgentFrameworkAgent`, `AgentFrameworkWorkflow`, `AGUIChatClient`, `AGUIEventConverter`, `AGUIHttpService`, `AGUIThreadSnapshot`, `AGUIThreadSnapshotStore`, `InMemoryAGUIThreadSnapshotStore`, `SnapshotScopeResolver`, `add_agent_framework_fastapi_endpoint`, `state_carrier`, `state_update`, `__version__` | | Direct `agent_framework_ag_ui` package | Facade exports plus `AGUIChatOptions`, `AGUIRequest`, `AGUIThreadID`, `AgentState`, `DEFAULT_MAX_THREAD_SNAPSHOTS`, `DEFAULT_TAGS`, `PredictStateConfig`, `RunMetadata`, `SnapshotScope`, `WorkflowFactory` | | AG-UI protocol package (`ag_ui.core`) | `Interrupt`, `ResumeEntry`, `RunFinishedInterruptOutcome`, and related run outcome models | Interrupt support is protocol data rather than a separate Agent Framework Python class. Requests accept canonical `availableInterrupts`/`available_interrupts` and `resume` values; `AGUIChatClient` and `AGUIHttpService.post_run(...)` forward those fields with AG-UI wire aliases; agent approval and workflow `request_info` pauses emit `RUN_FINISHED.outcome.interrupts`; `AGUIEventConverter` preserves canonical interrupt outcome metadata on the final `ChatResponseUpdate`; and thread snapshot hydration replays the canonical interrupt outcome when a scoped snapshot stores an unresolved pause. +Use `state_carrier(...)` to mark JSON content that should be sent in the AG-UI request's `state` field rather than as a model-visible document. The client removes explicitly marked carriers from all client-controlled history and uses the most recent carrier. Ordinary `application/json` content remains a document. For migration, pass `allow_legacy_state_carrier=True` in `AGUIChatOptions` to recognize the deprecated final single-content base64 JSON convention; mixed text and document messages remain model-visible. This client-only option emits a `DeprecationWarning` when the legacy convention is used and is not sent to the remote server. + ## Features This integration supports all 7 AG-UI features: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py index 8df80c25782..80c721b26fa 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/__init__.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/__init__.py @@ -19,7 +19,7 @@ SnapshotScope, SnapshotScopeResolver, ) -from ._state import state_update +from ._state import state_carrier, state_update from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata from ._workflow import AgentFrameworkWorkflow, WorkflowFactory @@ -55,6 +55,7 @@ "SnapshotScopeResolver", "DEFAULT_MAX_THREAD_SNAPSHOTS", "DEFAULT_TAGS", + "state_carrier", "state_update", "__version__", # A2UI (lazy — require ag-ui-a2ui-toolkit) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index f7edb458277..7aa43062427 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -4,10 +4,12 @@ from __future__ import annotations +import base64 import json import logging import sys import uuid +import warnings from binascii import Error as BinasciiError from collections.abc import AsyncIterable, Awaitable, Mapping, MutableSequence, Sequence from functools import wraps @@ -32,6 +34,7 @@ from ._feature_usage import FeatureIndex from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume from ._message_adapters import agent_framework_messages_to_agui +from ._state import STATE_CARRIER_KEY from ._utils import convert_tools_to_agui_format if sys.version_info >= (3, 13): @@ -72,6 +75,54 @@ def _unwrap_server_function_call_contents(contents: MutableSequence[Content | di ) +def _is_state_carrier_message(message: Message) -> bool: + """Return whether a message is a dedicated, explicitly marked state carrier.""" + if len(message.contents) != 1: + return False + content = message.contents[0] + return isinstance(content, Content) and (content.additional_properties or {}).get(STATE_CARRIER_KEY) is True + + +def _decode_json_state(content: Content) -> dict[str, Any] | None: + """Decode a base64 JSON state content, returning None for invalid input.""" + if content.type != "data" or content.media_type != "application/json": + return None + + try: + uri = content.uri + prefix, _, encoded_data = uri.partition(",") # type: ignore[union-attr] + if not prefix.startswith("data:"): + return None + + media_type, *parameters = prefix[5:].split(";") + if media_type != "application/json" or "base64" not in parameters: + return None + + decoded_bytes = base64.b64decode(encoded_data, validate=True) + state = json.loads(decoded_bytes.decode("utf-8")) + if not isinstance(state, dict): + logger.warning("AG-UI state carrier JSON must decode to an object") + return None + return state + except (BinasciiError, UnicodeDecodeError, json.JSONDecodeError, ValueError, TypeError, AttributeError) as e: + logger.warning(f"Failed to extract state from message: {e}") + return None + + +def _extract_legacy_json_state(message: Message) -> dict[str, Any] | None: + """Extract the historical implicit state convention from a final state-only message.""" + if len(message.contents) != 1: + return None + + content = message.contents[0] + if not isinstance(content, Content): + return None + if (content.additional_properties or {}).get(STATE_CARRIER_KEY) is True: + return None + + return _decode_json_state(content) + + def _apply_server_function_call_unwrap(client: BaseChatClientT) -> BaseChatClientT: """Class decorator that unwraps server-side function calls after tool handling.""" @@ -279,38 +330,45 @@ def _register_server_tool_placeholder(self, tool_name: str) -> None: self._registered_server_tools = registered logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}") - def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]: - """Extract state from last message if present. + def _extract_state_from_messages( + self, + messages: Sequence[Message], + *, + allow_legacy_state_carrier: bool = False, + ) -> tuple[list[Message], dict[str, Any] | None]: + """Extract explicitly marked state from client-controlled message history. Args: messages: List of chat messages + allow_legacy_state_carrier: Whether to recognize the deprecated implicit + final base64 JSON state convention. Returns: Tuple of (messages_without_state, state_dict) """ - if not messages: - return list(messages), None - - last_message = messages[-1] - - for content in last_message.contents: - if isinstance(content, Content) and content.type == "data" and content.media_type == "application/json": - try: - uri = content.uri - prefix, _, encoded_data = uri.partition(",") # type: ignore[union-attr] - media_type, *parameters = prefix[5:].split(";") - if prefix.startswith("data:") and media_type == "application/json" and "base64" in parameters: - import base64 - - decoded_bytes = base64.b64decode(encoded_data, validate=True) - state = json.loads(decoded_bytes.decode("utf-8")) - - messages_without_state = list(messages[:-1]) if len(messages) > 1 else [] - return messages_without_state, state - except (BinasciiError, json.JSONDecodeError, ValueError, KeyError) as e: - logger.warning(f"Failed to extract state from message: {e}") + messages_to_send: list[Message] = [] + state: dict[str, Any] | None = None + + for message in messages: + if _is_state_carrier_message(message): + content = cast(Content, message.contents[0]) + if (extracted_state := _decode_json_state(content)) is not None: + state = extracted_state + continue + messages_to_send.append(message) + + if allow_legacy_state_carrier and messages and not _is_state_carrier_message(messages[-1]): + legacy_state = _extract_legacy_json_state(messages[-1]) + if legacy_state is not None: + messages_to_send.pop() + state = legacy_state + warnings.warn( + "Implicit AG-UI JSON state extraction is deprecated; use state_carrier() instead.", + DeprecationWarning, + stacklevel=3, + ) - return list(messages), None + return messages_to_send, state def _convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]: """Convert Agent Framework messages to AG-UI format. @@ -401,7 +459,10 @@ async def _streaming_impl( ChatResponseUpdate objects """ mark_feature_used(FeatureIndex.AG_UI) - messages_to_send, state = self._extract_state_from_messages(messages) + messages_to_send, state = self._extract_state_from_messages( + messages, + allow_legacy_state_carrier=options.get("allow_legacy_state_carrier") is True, + ) thread_id = self._get_thread_id(options) run_id = f"run_{uuid.uuid4().hex}" diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index 24624517ebc..31860d05819 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -958,19 +958,55 @@ def _filter_modified_args( return result -def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, Any]]]: - """Encode assistant contents into an AG-UI ``(content, tool_calls)`` pair. +def _convert_framework_content_to_agui(content: Content) -> dict[str, Any] | None: + """Convert Agent Framework media content to an AG-UI input part.""" + if content.type not in {"uri", "data"} or not content.uri: + return None + + media_type = content.media_type + media_type_prefix = media_type.lower().split("/", 1)[0] if media_type else "" + part_type = media_type_prefix if media_type_prefix in {"image", "audio", "video"} else "document" + + if content.type == "data": + data_uri_prefix, separator, encoded_data = content.uri.partition(",") + is_base64_data_uri = bool(separator) and any( + parameter.lower() == "base64" for parameter in data_uri_prefix.split(";")[1:] + ) + if is_base64_data_uri: + source: dict[str, Any] = {"type": "data", "value": encoded_data} + else: + source = {"type": "url", "value": content.uri} + else: + source = {"type": "url", "value": content.uri} + + if media_type is not None: + source["mimeType"] = media_type + return {"type": part_type, "source": source} + + +def _encode_agui_segment( + contents: list[Content], role: str +) -> tuple[str | list[dict[str, Any]], list[dict[str, Any]]]: + """Encode a framework content segment into AG-UI message content and tool calls. - Shared by both the single-message path (``agent_framework_messages_to_agui``) and the - split path (``_split_mixed_message_to_agui``) so the text / function_call - serialization lives in one place. A future argument-format or supported-content - change then updates both paths at once instead of drifting between them. + The shared encoder preserves ordered user text and media parts for both the + single-message and function-result split paths. Non-user messages retain AG-UI's + string-content shape. """ text = "" + input_content_parts: list[dict[str, Any]] = [] + has_multimodal_content = False tool_calls: list[dict[str, Any]] = [] for content in contents: if content.type == "text": - text += content.text or "" + text_content = content.text or "" + text += text_content + if role == "user": + input_content_parts.append({"type": "text", "text": text_content}) + elif role == "user" and content.type in {"uri", "data"}: + if input_part := _convert_framework_content_to_agui(content): + input_content_parts.append(input_part) + has_multimodal_content = True elif content.type == "function_call": tool_calls.append( { @@ -982,7 +1018,8 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A }, } ) - return text, tool_calls + message_content: str | list[dict[str, Any]] = input_content_parts if has_multimodal_content else text + return message_content, tool_calls def _split_mixed_message_to_agui(msg: Message, role: str, unresolved_call_ids: set[str]) -> list[dict[str, Any]]: @@ -1045,7 +1082,7 @@ def flush_segment() -> None: nonlocal seg_contents, seg_has_call if not seg_contents: return - seg_text, seg_tool_calls = _encode_agui_segment(seg_contents) + seg_text, seg_tool_calls = _encode_agui_segment(seg_contents, role) seg_contents = [] seg_has_call = False seg_call_ids.clear() @@ -1079,7 +1116,7 @@ def drain_queued() -> None: queued_results.clear() for content in msg.contents: - if content.type in ("text", "function_call"): + if content.type in ("text", "function_call") or (role == "user" and content.type in {"uri", "data"}): seg_contents.append(content) if content.type == "function_call": seg_has_call = True @@ -1184,12 +1221,12 @@ def track_emitted( result.extend(_split_mixed_message_to_agui(msg, role, unresolved_call_ids)) continue - content_text, tool_calls = _encode_agui_segment(msg.contents) + message_content, tool_calls = _encode_agui_segment(msg.contents, role) agui_msg: dict[str, Any] = { "id": msg.message_id if msg.message_id else generate_event_id(), # Always include id "role": role, - "content": content_text, + "content": message_content, } if tool_calls: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_state.py index 607cc4a8c0d..d1e85d0f0a8 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_state.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. -"""Deterministic tool-driven AG-UI state updates and display payloads. +"""AG-UI state carrier and deterministic tool-result state helpers. Tools wired into the :mod:`agent_framework_ag_ui` endpoint can push a deterministic state update or a per-call tool result display payload by @@ -23,9 +23,12 @@ from ._utils import make_json_safe -__all__ = ["TOOL_RESULT_DISPLAY_KEY", "TOOL_RESULT_STATE_KEY", "state_update"] +__all__ = ["STATE_CARRIER_KEY", "TOOL_RESULT_DISPLAY_KEY", "TOOL_RESULT_STATE_KEY", "state_carrier", "state_update"] +STATE_CARRIER_KEY = "__ag_ui_state_carrier__" +"""Reserved ``Content.additional_properties`` key marking an AG-UI request state carrier.""" + TOOL_RESULT_STATE_KEY = "__ag_ui_tool_result_state__" """Reserved ``Content.additional_properties`` key used to carry a tool-driven state snapshot from a tool return value through to the AG-UI emitter.""" @@ -40,6 +43,44 @@ def _serialize_tool_result(value: Any) -> str: # noqa: ANN401 return value if isinstance(value, str) else json.dumps(make_json_safe(value)) +def state_carrier(state: Mapping[str, Any]) -> Content: + """Build a dedicated message carrier for ``AGUIChatClient`` request state. + + Add the returned content as the only content in a user message. The client + recognizes its explicit marker anywhere in client-controlled history, moves + the most recent carrier's JSON object into the AG-UI request's ``state`` + field, and does not send carriers as chat messages. Ordinary + ``application/json`` content without this marker remains a document input. + + Example: + .. code-block:: python + + from agent_framework import Message + from agent_framework_ag_ui import state_carrier + + messages = [ + Message(role="user", contents=["Update the dashboard"]), + Message(role="user", contents=[state_carrier({"selected_tab": "sales"})]), + ] + + Args: + state: JSON-compatible mapping to send as AG-UI shared state. + + Returns: + A JSON ``Content`` marked as an AG-UI request state carrier. + + Raises: + TypeError: If ``state`` is not a mapping. + """ + if not isinstance(state, Mapping): + raise TypeError(f"state_carrier() 'state' must be a Mapping, got {type(state).__name__}") + return Content.from_data( + json.dumps(make_json_safe(dict(state))).encode("utf-8"), + media_type="application/json", + additional_properties={STATE_CARRIER_KEY: True}, + ) + + def state_update( text: str = "", *, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_types.py b/python/packages/ag-ui/agent_framework_ag_ui/_types.py index cc6dddf19ef..b2236959ad2 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_types.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_types.py @@ -149,7 +149,8 @@ class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], tota Extends base ChatOptions for the AG-UI (Agent-UI) protocol. AG-UI is a streaming protocol for connecting AI agents to user interfaces. - Options are forwarded to the remote AG-UI server. + Options are forwarded to the remote AG-UI server unless explicitly + documented as client-only. See: https://github.com/ag-ui/ag-ui-protocol @@ -182,11 +183,15 @@ class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], tota forward_props: Additional properties to forward to the AG-UI server. Useful for passing custom parameters to specific server implementations. context: Shared context/state to send to the server. + allow_legacy_state_carrier: Client-only migration option. When true, + recognize the deprecated implicit final single-content base64 JSON + state convention and emit a deprecation warning. Defaults to false. Note: AG-UI is a protocol bridge - actual option support depends on the remote server implementation. The client sends all options to the - server, which decides how to handle them. + server, which decides how to handle them, except client-only options + consumed by the client itself. Thread ID management: - Pass ``thread_id`` in ``metadata`` to maintain conversation continuity @@ -200,6 +205,9 @@ class AGUIChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], tota context: dict[str, Any] """Shared context/state to send to the server.""" + allow_legacy_state_carrier: bool + """Recognize the deprecated implicit final JSON state convention.""" + available_interrupts: list[Interrupt] """Canonical AG-UI interrupt descriptors available for resumption.""" diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 4c6d178f63e..2f328f1626a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -6,6 +6,8 @@ from collections.abc import AsyncGenerator, Awaitable, MutableSequence from typing import Any, cast +import httpx +import pytest from ag_ui.core import Interrupt, ResumeEntry from agent_framework import ( ChatOptions, @@ -82,20 +84,18 @@ async def test_extract_state_from_messages_no_state(self) -> None: assert state is None async def test_extract_state_from_messages_with_state(self) -> None: - """Test state extraction from last message.""" - import base64 + """A marked state carrier populates the request state.""" + from agent_framework_ag_ui import state_carrier client = StubAGUIChatClient(endpoint="http://localhost:8888/") state_data = {"key": "value", "count": 42} - state_json = json.dumps(state_data) - state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8") messages = [ Message(role="user", contents=["Hello"]), Message( role="user", - contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")], + contents=[state_carrier(state_data)], ), ] @@ -105,10 +105,53 @@ async def test_extract_state_from_messages_with_state(self) -> None: assert result_messages[0].text == "Hello" assert state == state_data + async def test_extract_state_from_messages_removes_historical_carriers_and_uses_latest_state(self) -> None: + """Historical carriers are removed and the most recent carrier supplies state.""" + from agent_framework_ag_ui import state_carrier + + client = StubAGUIChatClient(endpoint="http://localhost:8888/") + old_carrier = Message(role="user", contents=[state_carrier({"version": "old"})]) + new_carrier = Message(role="user", contents=[state_carrier({"version": "new"})]) + messages = [ + Message(role="user", contents=["Initial prompt"]), + old_carrier, + Message(role="assistant", contents=["Initial response"]), + new_carrier, + Message(role="user", contents=["Follow-up prompt"]), + ] + + result_messages, state = client.extract_state_from_messages(messages) + + assert result_messages == [messages[0], messages[2], messages[4]] + assert state == {"version": "new"} + + async def test_explicit_final_carrier_wins_over_legacy_fallback(self) -> None: + """Legacy mode does not reinterpret an earlier document after removing a final carrier.""" + from agent_framework_ag_ui import state_carrier + + client = StubAGUIChatClient(endpoint="http://localhost:8888/") + messages = [ + Message( + role="user", + contents=[Content.from_data(b'{"document":"keep"}', media_type="application/json")], + ), + Message(role="user", contents=[state_carrier({"source": "explicit"})]), + ] + + result_messages, state = client._extract_state_from_messages( + messages, + allow_legacy_state_carrier=True, + ) + + assert result_messages == messages[:1] + assert state == {"source": "explicit"} + async def test_extract_state_from_messages_with_parameterized_data_uri(self) -> None: """Test state extraction from JSON data URIs with media type parameters.""" import base64 + from agent_framework_ag_ui._state import STATE_CARRIER_KEY + client = StubAGUIChatClient(endpoint="http://localhost:8888/") state_data = {"key": "value", "count": 42} @@ -119,7 +162,12 @@ async def test_extract_state_from_messages_with_parameterized_data_uri(self) -> Message(role="user", contents=["Hello"]), Message( role="user", - contents=[Content.from_uri(uri=f"data:application/json;charset=utf-8;base64,{state_b64}")], + contents=[ + Content.from_uri( + uri=f"data:application/json;charset=utf-8;base64,{state_b64}", + additional_properties={STATE_CARRIER_KEY: True}, + ) + ], ), ] @@ -133,6 +181,8 @@ async def test_extract_state_invalid_json(self) -> None: """Test state extraction with invalid JSON.""" import base64 + from agent_framework_ag_ui._state import STATE_CARRIER_KEY + client = StubAGUIChatClient(endpoint="http://localhost:8888/") invalid_json = "not valid json" @@ -141,29 +191,41 @@ async def test_extract_state_invalid_json(self) -> None: messages = [ Message( role="user", - contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")], + contents=[ + Content.from_uri( + uri=f"data:application/json;base64,{state_b64}", + additional_properties={STATE_CARRIER_KEY: True}, + ) + ], ), ] result_messages, state = client.extract_state_from_messages(messages) - assert result_messages == messages + assert result_messages == [] assert state is None async def test_extract_state_invalid_base64(self) -> None: """Test state extraction with invalid base64.""" + from agent_framework_ag_ui._state import STATE_CARRIER_KEY + client = StubAGUIChatClient(endpoint="http://localhost:8888/") messages = [ Message( role="user", - contents=[Content.from_uri(uri="data:application/json;base64,not-valid-base64!")], + contents=[ + Content.from_uri( + uri="data:application/json;base64,not-valid-base64!", + additional_properties={STATE_CARRIER_KEY: True}, + ) + ], ), ] result_messages, state = client.extract_state_from_messages(messages) - assert result_messages == messages + assert result_messages == [] assert state is None async def test_convert_messages_to_agui_format(self) -> None: @@ -183,6 +245,187 @@ async def test_convert_messages_to_agui_format(self) -> None: assert agui_messages[1]["content"] == "Let me check." assert agui_messages[1]["id"] == "msg_123" + async def test_sends_multimodal_messages_in_request(self) -> None: + """The client sends ordered multimodal content in the HTTP request JSON.""" + captured_request: dict[str, Any] = {} + + async def handler(request: httpx.Request) -> httpx.Response: + captured_request.update(json.loads(request.content)) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"type":"RUN_STARTED","threadId":"thread_1","runId":"run_1"}\n\n' + b'data: {"type":"TEXT_MESSAGE_CONTENT","messageId":"msg_1","delta":"ok"}\n\n' + b'data: {"type":"RUN_FINISHED","threadId":"thread_1","runId":"run_1"}\n\n' + ), + ) + + message = Message( + role="user", + contents=[ + Content.from_text("describe this"), + Content.from_uri("https://example.com/cat.png", media_type="image/png"), + Content.from_data(b"abc", media_type="image/png"), + ], + message_id="msg-request", + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = StubAGUIChatClient(endpoint="http://localhost:8888/", http_client=http_client) + response = await client.inner_get_response(messages=[message], options={}) + + assert response is not None + assert captured_request["messages"] == [ + { + "id": "msg-request", + "role": "user", + "content": [ + {"type": "text", "text": "describe this"}, + { + "type": "image", + "source": { + "type": "url", + "value": "https://example.com/cat.png", + "mimeType": "image/png", + }, + }, + { + "type": "image", + "source": {"type": "data", "value": "YWJj", "mimeType": "image/png"}, + }, + ], + } + ] + + async def test_sends_mixed_json_attachment_when_legacy_compatibility_is_enabled(self) -> None: + """Legacy compatibility does not consume a prompt with a JSON attachment.""" + captured_request: dict[str, Any] = {} + + async def handler(request: httpx.Request) -> httpx.Response: + captured_request.update(json.loads(request.content)) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"type":"RUN_STARTED","threadId":"thread_1","runId":"run_1"}\n\n' + b'data: {"type":"RUN_FINISHED","threadId":"thread_1","runId":"run_1"}\n\n' + ), + ) + + message = Message( + role="user", + contents=[ + Content.from_text("summarize the attached JSON document"), + Content.from_data(b'{"document":"keep me"}', media_type="application/json"), + ], + message_id="msg-json-document", + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = StubAGUIChatClient(endpoint="http://localhost:8888/", http_client=http_client) + response = await client.inner_get_response( + messages=[message], + options={"allow_legacy_state_carrier": True}, + ) + + assert response is not None + assert "state" not in captured_request + assert captured_request["messages"] == [ + { + "id": "msg-json-document", + "role": "user", + "content": [ + {"type": "text", "text": "summarize the attached JSON document"}, + { + "type": "document", + "source": { + "type": "data", + "value": "eyJkb2N1bWVudCI6ImtlZXAgbWUifQ==", + "mimeType": "application/json", + }, + }, + ], + } + ] + + async def test_sends_json_attachment_without_state_carrier_in_request(self) -> None: + """An unmarked JSON-only document is sent as AG-UI document input.""" + captured_request: dict[str, Any] = {} + + async def handler(request: httpx.Request) -> httpx.Response: + captured_request.update(json.loads(request.content)) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"type":"RUN_STARTED","threadId":"thread_1","runId":"run_1"}\n\n' + b'data: {"type":"RUN_FINISHED","threadId":"thread_1","runId":"run_1"}\n\n' + ), + ) + + message = Message( + role="user", + contents=[Content.from_data(b'{"document":"keep me"}', media_type="application/json")], + message_id="msg-json-only-document", + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = StubAGUIChatClient(endpoint="http://localhost:8888/", http_client=http_client) + response = await client.inner_get_response(messages=[message], options={}) + + assert response is not None + assert "state" not in captured_request + assert captured_request["messages"] == [ + { + "id": "msg-json-only-document", + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "data", + "value": "eyJkb2N1bWVudCI6ImtlZXAgbWUifQ==", + "mimeType": "application/json", + }, + } + ], + } + ] + + async def test_sends_legacy_json_state_with_compatibility_option(self) -> None: + """The legacy option extracts an unmarked final JSON state during migration.""" + captured_request: dict[str, Any] = {} + + async def handler(request: httpx.Request) -> httpx.Response: + captured_request.update(json.loads(request.content)) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=( + b'data: {"type":"RUN_STARTED","threadId":"thread_1","runId":"run_1"}\n\n' + b'data: {"type":"RUN_FINISHED","threadId":"thread_1","runId":"run_1"}\n\n' + ), + ) + + message = Message( + role="user", + contents=[Content.from_data(b'{"legacy":"keep"}', media_type="application/json")], + message_id="msg-legacy-state", + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http_client: + client = StubAGUIChatClient(endpoint="http://localhost:8888/", http_client=http_client) + with pytest.warns(DeprecationWarning, match="state_carrier"): + response = await client.inner_get_response( + messages=[message], + options={"allow_legacy_state_carrier": True}, + ) + + assert response is not None + assert captured_request["state"] == {"legacy": "keep"} + assert captured_request["messages"] == [] + async def test_get_thread_id_from_metadata(self) -> None: """Test thread ID extraction from metadata.""" client = StubAGUIChatClient(endpoint="http://localhost:8888/") @@ -370,18 +613,16 @@ async def fake_auto_invoke(*args: object, **kwargs: Any) -> None: pass async def test_state_transmission(self, monkeypatch: MonkeyPatch) -> None: - """Test state is properly transmitted to server.""" - import base64 + """A marked state carrier is transmitted through the request state field.""" + from agent_framework_ag_ui import state_carrier state_data = {"user_id": "123", "session": "abc"} - state_json = json.dumps(state_data) - state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8") messages = [ Message(role="user", contents=["Hello"]), Message( role="user", - contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")], + contents=[state_carrier(state_data)], ), ] diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index a756009df6a..73b974dca44 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -51,6 +51,110 @@ def test_agent_framework_to_agui_basic(sample_agent_framework_message): assert messages[0]["id"] == "msg-123" +def test_agent_framework_to_agui_preserves_uri_content(): + """URI content is serialized as an ordered AG-UI input content part.""" + message = Message( + role="user", + contents=[Content.from_uri("https://example.com/cat.png", media_type="image/png")], + message_id="msg-uri", + ) + + assert agent_framework_messages_to_agui([message]) == [ + { + "id": "msg-uri", + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "url", + "value": "https://example.com/cat.png", + "mimeType": "image/png", + }, + } + ], + } + ] + + +def test_agent_framework_to_agui_converts_data_uri_to_inline_data_source(): + """Inline data content is serialized as an AG-UI data source.""" + message = Message( + role="user", + contents=[Content.from_data(b"abc", media_type="image/png")], + message_id="msg-data", + ) + + assert agent_framework_messages_to_agui([message])[0]["content"] == [ + { + "type": "image", + "source": {"type": "data", "value": "YWJj", "mimeType": "image/png"}, + } + ] + + +def test_agent_framework_to_agui_preserves_non_base64_data_uri_as_url_source(): + """Non-base64 data URIs remain complete URL sources for AG-UI.""" + data_uri = "data:text/plain,hello%20world" + message = Message( + role="user", + contents=[Content.from_uri(data_uri, media_type="text/plain")], + message_id="msg-data-uri", + ) + + assert agent_framework_messages_to_agui([message])[0]["content"] == [ + { + "type": "document", + "source": {"type": "url", "value": data_uri, "mimeType": "text/plain"}, + } + ] + + +def test_agent_framework_to_agui_preserves_mixed_content_order(): + """Mixed text and media content remains in its original order.""" + message = Message( + role="user", + contents=[ + Content.from_text("before"), + Content.from_uri("https://example.com/cat.png", media_type="image/png"), + Content.from_text("after"), + Content.from_data(b"abc", media_type="application/pdf"), + ], + message_id="msg-mixed", + ) + + assert agent_framework_messages_to_agui([message])[0]["content"] == [ + {"type": "text", "text": "before"}, + { + "type": "image", + "source": { + "type": "url", + "value": "https://example.com/cat.png", + "mimeType": "image/png", + }, + }, + {"type": "text", "text": "after"}, + { + "type": "document", + "source": {"type": "data", "value": "YWJj", "mimeType": "application/pdf"}, + }, + ] + + +def test_agent_framework_to_agui_keeps_assistant_content_as_text(): + """Assistant messages keep AG-UI's string content shape when media is present.""" + message = Message( + role="assistant", + contents=[ + Content.from_text("answer"), + Content.from_uri("https://example.com/cat.png", media_type="image/png"), + ], + message_id="msg-assistant", + ) + + assert agent_framework_messages_to_agui([message])[0]["content"] == "answer" + + def test_agent_framework_to_agui_normalizes_dict_roles(): """Dict inputs normalize unknown roles for UI compatibility.""" messages = [ diff --git a/python/packages/ag-ui/tests/ag_ui/test_public_exports.py b/python/packages/ag-ui/tests/ag_ui/test_public_exports.py index daa0d8e4c9c..5e76248e888 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_public_exports.py +++ b/python/packages/ag-ui/tests/ag_ui/test_public_exports.py @@ -18,6 +18,7 @@ def test_core_ag_ui_lazy_exports_include_only_stable_api() -> None: assert hasattr(ag_ui, "AgentFrameworkAgent") assert hasattr(ag_ui, "AGUIChatClient") assert hasattr(ag_ui, "add_agent_framework_fastapi_endpoint") + assert hasattr(ag_ui, "state_carrier") assert hasattr(ag_ui, "state_update") assert not hasattr(ag_ui, "WorkflowFactory") @@ -25,10 +26,11 @@ def test_core_ag_ui_lazy_exports_include_only_stable_api() -> None: assert not hasattr(ag_ui, "RunMetadata") -def test_agent_framework_ag_ui_exports_state_update() -> None: - """Runtime package should export the ``state_update`` helper.""" - from agent_framework_ag_ui import state_update +def test_agent_framework_ag_ui_exports_state_helpers() -> None: + """Runtime package should export the AG-UI state helpers.""" + from agent_framework_ag_ui import state_carrier, state_update + assert callable(state_carrier) assert callable(state_update) diff --git a/python/packages/core/agent_framework/ag_ui/__init__.py b/python/packages/core/agent_framework/ag_ui/__init__.py index 580ae153a9a..36165911d73 100644 --- a/python/packages/core/agent_framework/ag_ui/__init__.py +++ b/python/packages/core/agent_framework/ag_ui/__init__.py @@ -16,6 +16,7 @@ - InMemoryAGUIThreadSnapshotStore - SnapshotScopeResolver - add_agent_framework_fastapi_endpoint +- state_carrier - state_update - __version__ """ @@ -36,6 +37,7 @@ "AGUIThreadSnapshotStore", "InMemoryAGUIThreadSnapshotStore", "SnapshotScopeResolver", + "state_carrier", "state_update", "__version__", ] diff --git a/python/packages/core/agent_framework/ag_ui/__init__.pyi b/python/packages/core/agent_framework/ag_ui/__init__.pyi index e57ba45ac62..84188b3b281 100644 --- a/python/packages/core/agent_framework/ag_ui/__init__.pyi +++ b/python/packages/core/agent_framework/ag_ui/__init__.pyi @@ -12,6 +12,7 @@ from agent_framework_ag_ui import ( SnapshotScopeResolver, __version__, add_agent_framework_fastapi_endpoint, + state_carrier, state_update, ) @@ -27,5 +28,6 @@ __all__ = [ "SnapshotScopeResolver", "__version__", "add_agent_framework_fastapi_endpoint", + "state_carrier", "state_update", ]