diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index ab074ae15cd..9f62df10f91 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -51,8 +51,8 @@ from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id from ._message_adapters import normalize_agui_input_messages -from ._orchestration._predictive_state import PredictiveStateHandler -from ._orchestration._tooling import collect_server_tools, merge_tools, register_additional_client_tools +from ._predictive_state import PredictiveStateHandler +from ._tooling import collect_server_tools, merge_tools, register_additional_client_tools from ._run_common import ( FlowState, _approval_interrupt_for_function_call, # type: ignore diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/__init__.py deleted file mode 100644 index 2a50eae8941..00000000000 --- a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py b/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py deleted file mode 100644 index 09f637061ca..00000000000 --- a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_helpers.py +++ /dev/null @@ -1,247 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Helper functions for orchestration logic. - -This module retains utilities that may be useful for testing or extensions. -""" - -from __future__ import annotations - -import json -import logging -from typing import Any - -from agent_framework import ( - Content, - Message, -) - -from .._utils import get_role_value - -logger = logging.getLogger(__name__) - - -def pending_tool_call_ids(messages: list[Message]) -> set[str]: - """Get IDs of tool calls without corresponding results. - - Args: - messages: List of messages to scan - - Returns: - Set of pending tool call IDs - """ - pending_ids: set[str] = set() - resolved_ids: set[str] = set() - for msg in messages: - for content in msg.contents: - if content.type == "function_call" and content.call_id: - pending_ids.add(str(content.call_id)) - elif content.type == "function_result" and content.call_id: - resolved_ids.add(str(content.call_id)) - return pending_ids - resolved_ids - - -def is_state_context_message(message: Message) -> bool: - """Check if a message is a state context system message. - - Args: - message: Message to check - - Returns: - True if this is a state context message - """ - if get_role_value(message) != "system": - return False - for content in message.contents: - if content.type == "text" and content.text.startswith("Current state of the application:"): # type: ignore[union-attr] - return True - return False - - -def ensure_tool_call_entry( - tool_call_id: str, - tool_calls_by_id: dict[str, dict[str, Any]], - pending_tool_calls: list[dict[str, Any]], -) -> dict[str, Any]: - """Get or create a tool call entry in the tracking dicts. - - Args: - tool_call_id: The tool call ID - tool_calls_by_id: Dict mapping IDs to tool call entries - pending_tool_calls: List of pending tool calls - - Returns: - The tool call entry dict - """ - entry = tool_calls_by_id.get(tool_call_id) - if entry is None: - entry = { - "id": tool_call_id, - "type": "function", - "function": { - "name": "", - "arguments": "", - }, - } - tool_calls_by_id[tool_call_id] = entry - pending_tool_calls.append(entry) - return entry - - -def tool_name_for_call_id( - tool_calls_by_id: dict[str, dict[str, Any]], - tool_call_id: str, -) -> str | None: - """Get the tool name for a given call ID. - - Args: - tool_calls_by_id: Dict mapping IDs to tool call entries - tool_call_id: The tool call ID to look up - - Returns: - Tool name or None if not found - """ - entry = tool_calls_by_id.get(tool_call_id) - if not entry: - return None - function = entry.get("function") - if not isinstance(function, dict): - return None - name = function.get("name") - return str(name) if name else None - - -def schema_has_steps(schema: Any) -> bool: - """Check if a schema has a steps array property. - - Args: - schema: JSON schema to check - - Returns: - True if schema has steps array - """ - if not isinstance(schema, dict): - return False - properties = schema.get("properties") - if not isinstance(properties, dict): - return False - steps_schema = properties.get("steps") - if not isinstance(steps_schema, dict): - return False - return steps_schema.get("type") == "array" - - -def select_approval_tool_name(client_tools: list[Any] | None) -> str | None: - """Select appropriate approval tool from client tools. - - Args: - client_tools: List of client tool definitions - - Returns: - Name of approval tool, or None if not found - """ - if not client_tools: - return None - for tool in client_tools: - tool_name = getattr(tool, "name", None) - if not tool_name: - continue - params_fn = getattr(tool, "parameters", None) - if not callable(params_fn): - continue - schema = params_fn() - if schema_has_steps(schema): - return str(tool_name) - return None - - -def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any]: - """Build metadata dict with truncated string values for Azure compatibility. - - Azure has a 512 character limit per metadata value. - - Args: - thread_metadata: Raw metadata dict - - Returns: - Metadata with string values truncated to 512 chars - """ - if not thread_metadata: - return {} - safe_metadata: dict[str, Any] = {} - for key, value in thread_metadata.items(): - value_str = value if isinstance(value, str) else json.dumps(value) - if len(value_str) > 512: - value_str = value_str[:512] - safe_metadata[key] = value_str - return safe_metadata - - -def latest_approval_response(messages: list[Message]) -> Content | None: - """Get the latest approval response from messages. - - Args: - messages: Messages to search - - Returns: - Latest approval response or None - """ - if not messages: - return None - last_message = messages[-1] - for content in last_message.contents: - if content.type == "function_approval_response": - return content - return None - - -def approval_steps(approval: Content) -> list[Any]: - """Extract steps from an approval response. - - Args: - approval: Approval response content - - Returns: - List of steps, or empty list if none - """ - state_args = approval.additional_properties.get("ag_ui_state_args", None) - if isinstance(state_args, dict): - steps = state_args.get("steps") - if isinstance(steps, list): - return steps - - if approval.function_call: - parsed_args = approval.function_call.parse_arguments() - if isinstance(parsed_args, dict): - steps = parsed_args.get("steps") - if isinstance(steps, list): - return steps - - return [] - - -def is_step_based_approval( - approval: Content, - predict_state_config: dict[str, dict[str, str]] | None, -) -> bool: - """Check if an approval is step-based. - - Args: - approval: Approval response to check - predict_state_config: Predictive state configuration - - Returns: - True if this is a step-based approval - """ - steps = approval_steps(approval) - if steps: - return True - if not approval.function_call: - return False - if not predict_state_config: - return False - tool_name = approval.function_call.name - for config in predict_state_config.values(): - if config.get("tool") == tool_name and config.get("tool_argument") == "steps": - return True - return False diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_predictive_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_predictive_state.py similarity index 99% rename from python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_predictive_state.py rename to python/packages/ag-ui/agent_framework_ag_ui/_predictive_state.py index 216b1ca6626..99c18b9e45c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_predictive_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_predictive_state.py @@ -11,7 +11,7 @@ from ag_ui.core import StateDeltaEvent -from .._utils import safe_json_parse +from ._utils import safe_json_parse logger = logging.getLogger(__name__) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index da7cbd7d36f..28095883ae6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -34,7 +34,7 @@ ) from agent_framework import Content -from ._orchestration._predictive_state import PredictiveStateHandler +from ._predictive_state import PredictiveStateHandler from ._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY from ._utils import generate_event_id, make_json_safe, normalize_agui_role diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py b/python/packages/ag-ui/agent_framework_ag_ui/_tooling.py similarity index 100% rename from python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py rename to python/packages/ag-ui/agent_framework_ag_ui/_tooling.py diff --git a/python/packages/ag-ui/tests/ag_ui/test_helpers.py b/python/packages/ag-ui/tests/ag_ui/test_helpers.py deleted file mode 100644 index bfcea08afca..00000000000 --- a/python/packages/ag-ui/tests/ag_ui/test_helpers.py +++ /dev/null @@ -1,504 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Tests for orchestration helper functions.""" - -from typing import Any - -from agent_framework import Content, Message - -from agent_framework_ag_ui._orchestration._helpers import ( - approval_steps, - build_safe_metadata, - ensure_tool_call_entry, - is_state_context_message, - is_step_based_approval, - latest_approval_response, - pending_tool_call_ids, - schema_has_steps, - select_approval_tool_name, - tool_name_for_call_id, -) - - -class TestPendingToolCallIds: - """Tests for pending_tool_call_ids function.""" - - def test_empty_messages(self): - """Returns empty set for empty messages list.""" - result = pending_tool_call_ids([]) - assert result == set() - - def test_no_tool_calls(self): - """Returns empty set when no tool calls in messages.""" - messages = [ - Message(role="user", contents=[Content.from_text("Hello")]), - Message(role="assistant", contents=[Content.from_text("Hi there")]), - ] - result = pending_tool_call_ids(messages) - assert result == set() - - def test_pending_tool_call(self): - """Returns pending tool call ID when no result exists.""" - messages = [ - Message( - role="assistant", - contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")], - ), - ] - result = pending_tool_call_ids(messages) - assert result == {"call_123"} - - def test_resolved_tool_call(self): - """Returns empty set when tool call has result.""" - messages = [ - Message( - role="assistant", - contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")], - ), - Message( - role="tool", - contents=[Content.from_function_result(call_id="call_123", result="sunny")], - ), - ] - result = pending_tool_call_ids(messages) - assert result == set() - - def test_multiple_tool_calls_some_resolved(self): - """Returns only unresolved tool call IDs.""" - messages = [ - Message( - role="assistant", - contents=[ - Content.from_function_call(call_id="call_1", name="tool_a", arguments="{}"), - Content.from_function_call(call_id="call_2", name="tool_b", arguments="{}"), - Content.from_function_call(call_id="call_3", name="tool_c", arguments="{}"), - ], - ), - Message( - role="tool", - contents=[Content.from_function_result(call_id="call_1", result="result_a")], - ), - Message( - role="tool", - contents=[Content.from_function_result(call_id="call_3", result="result_c")], - ), - ] - result = pending_tool_call_ids(messages) - assert result == {"call_2"} - - -class TestIsStateContextMessage: - """Tests for is_state_context_message function.""" - - def test_state_context_message(self): - """Returns True for state context message.""" - message = Message( - role="system", - contents=[Content.from_text("Current state of the application: {}")], - ) - assert is_state_context_message(message) is True - - def test_non_system_message(self): - """Returns False for non-system message.""" - message = Message( - role="user", - contents=[Content.from_text("Current state of the application: {}")], - ) - assert is_state_context_message(message) is False - - def test_system_message_without_state_prefix(self): - """Returns False for system message without state prefix.""" - message = Message( - role="system", - contents=[Content.from_text("You are a helpful assistant.")], - ) - assert is_state_context_message(message) is False - - def test_empty_contents(self): - """Returns False for message with empty contents.""" - message = Message(role="system", contents=[]) - assert is_state_context_message(message) is False - - -class TestEnsureToolCallEntry: - """Tests for ensure_tool_call_entry function.""" - - def test_creates_new_entry(self): - """Creates new entry when ID not found.""" - tool_calls_by_id: dict = {} - pending_tool_calls: list = [] - - entry = ensure_tool_call_entry("call_123", tool_calls_by_id, pending_tool_calls) - - assert entry["id"] == "call_123" - assert entry["type"] == "function" - assert entry["function"]["name"] == "" - assert entry["function"]["arguments"] == "" - assert "call_123" in tool_calls_by_id - assert len(pending_tool_calls) == 1 - - def test_returns_existing_entry(self): - """Returns existing entry when ID found.""" - existing_entry: dict[str, Any] = { - "id": "call_123", - "type": "function", - "function": {"name": "get_weather", "arguments": '{"city": "NYC"}'}, - } - tool_calls_by_id: dict[str, dict[str, Any]] = {"call_123": existing_entry} - pending_tool_calls: list[dict[str, Any]] = [] - - entry = ensure_tool_call_entry("call_123", tool_calls_by_id, pending_tool_calls) - - assert entry is existing_entry - assert entry["function"]["name"] == "get_weather" - assert len(pending_tool_calls) == 0 # Not added again - - -class TestToolNameForCallId: - """Tests for tool_name_for_call_id function.""" - - def test_returns_tool_name(self): - """Returns tool name for valid entry.""" - tool_calls_by_id = { - "call_123": { - "id": "call_123", - "function": {"name": "get_weather", "arguments": "{}"}, - } - } - result = tool_name_for_call_id(tool_calls_by_id, "call_123") - assert result == "get_weather" - - def test_returns_none_for_missing_id(self): - """Returns None when ID not found.""" - tool_calls_by_id: dict = {} - result = tool_name_for_call_id(tool_calls_by_id, "call_123") - assert result is None - - def test_returns_none_for_missing_function(self): - """Returns None when function key missing.""" - tool_calls_by_id = {"call_123": {"id": "call_123"}} - result = tool_name_for_call_id(tool_calls_by_id, "call_123") - assert result is None - - def test_returns_none_for_non_dict_function(self): - """Returns None when function is not a dict.""" - tool_calls_by_id = {"call_123": {"id": "call_123", "function": "not_a_dict"}} - result = tool_name_for_call_id(tool_calls_by_id, "call_123") - assert result is None - - def test_returns_none_for_empty_name(self): - """Returns None when name is empty.""" - tool_calls_by_id = {"call_123": {"id": "call_123", "function": {"name": "", "arguments": "{}"}}} - result = tool_name_for_call_id(tool_calls_by_id, "call_123") - assert result is None - - -class TestSchemaHasSteps: - """Tests for schema_has_steps function.""" - - def test_schema_with_steps_array(self): - """Returns True when schema has steps array property.""" - schema = {"properties": {"steps": {"type": "array"}}} - assert schema_has_steps(schema) is True - - def test_schema_without_steps(self): - """Returns False when schema doesn't have steps.""" - schema = {"properties": {"name": {"type": "string"}}} - assert schema_has_steps(schema) is False - - def test_schema_with_non_array_steps(self): - """Returns False when steps is not array type.""" - schema = {"properties": {"steps": {"type": "string"}}} - assert schema_has_steps(schema) is False - - def test_non_dict_schema(self): - """Returns False for non-dict schema.""" - assert schema_has_steps(None) is False - assert schema_has_steps("not a dict") is False - assert schema_has_steps([]) is False - - def test_missing_properties(self): - """Returns False when properties key is missing.""" - schema = {"type": "object"} - assert schema_has_steps(schema) is False - - def test_non_dict_properties(self): - """Returns False when properties is not a dict.""" - schema = {"properties": "not a dict"} - assert schema_has_steps(schema) is False - - def test_non_dict_steps(self): - """Returns False when steps is not a dict.""" - schema = {"properties": {"steps": "not a dict"}} - assert schema_has_steps(schema) is False - - -class TestSelectApprovalToolName: - """Tests for select_approval_tool_name function.""" - - def test_none_client_tools(self): - """Returns None when client_tools is None.""" - result = select_approval_tool_name(None) - assert result is None - - def test_empty_client_tools(self): - """Returns None when client_tools is empty.""" - result = select_approval_tool_name([]) - assert result is None - - def test_finds_approval_tool(self): - """Returns tool name when tool has steps schema.""" - - class MockTool: - name = "generate_task_steps" - - def parameters(self): - return {"properties": {"steps": {"type": "array"}}} - - result = select_approval_tool_name([MockTool()]) - assert result == "generate_task_steps" - - def test_skips_tool_without_name(self): - """Skips tools without name attribute.""" - - class MockToolNoName: - def parameters(self): - return {"properties": {"steps": {"type": "array"}}} - - result = select_approval_tool_name([MockToolNoName()]) - assert result is None - - def test_skips_tool_without_parameters_method(self): - """Skips tools without callable parameters method.""" - - class MockToolNoParams: - name = "some_tool" - parameters = "not callable" - - result = select_approval_tool_name([MockToolNoParams()]) - assert result is None - - def test_skips_tool_without_steps_schema(self): - """Skips tools that don't have steps in schema.""" - - class MockToolNoSteps: - name = "other_tool" - - def parameters(self): - return {"properties": {"data": {"type": "string"}}} - - result = select_approval_tool_name([MockToolNoSteps()]) - assert result is None - - -class TestBuildSafeMetadata: - """Tests for build_safe_metadata function.""" - - def test_none_metadata(self): - """Returns empty dict for None metadata.""" - result = build_safe_metadata(None) - assert result == {} - - def test_empty_metadata(self): - """Returns empty dict for empty metadata.""" - result = build_safe_metadata({}) - assert result == {} - - def test_string_values_under_limit(self): - """Preserves string values under 512 chars.""" - metadata = {"key1": "short value", "key2": "another value"} - result = build_safe_metadata(metadata) - assert result == metadata - - def test_truncates_long_string_values(self): - """Truncates string values over 512 chars.""" - long_value = "x" * 1000 - metadata = {"key": long_value} - result = build_safe_metadata(metadata) - assert len(result["key"]) == 512 - assert result["key"] == "x" * 512 - - def test_non_string_values_serialized(self): - """Serializes non-string values to JSON.""" - metadata = {"count": 42, "items": ["a", "b"]} - result = build_safe_metadata(metadata) - assert result["count"] == "42" - assert result["items"] == '["a", "b"]' - - def test_truncates_serialized_values(self): - """Truncates serialized JSON values over 512 chars.""" - long_list = list(range(200)) # Will serialize to >512 chars - metadata = {"data": long_list} - result = build_safe_metadata(metadata) - assert len(result["data"]) == 512 - - -class TestLatestApprovalResponse: - """Tests for latest_approval_response function.""" - - def test_empty_messages(self): - """Returns None for empty messages.""" - result = latest_approval_response([]) - assert result is None - - def test_no_approval_response(self): - """Returns None when no approval response in last message.""" - messages = [ - Message(role="assistant", contents=[Content.from_text("Hello")]), - ] - result = latest_approval_response(messages) - assert result is None - - def test_finds_approval_response(self): - """Returns approval response from last message.""" - # Create a function call content first - fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}") - approval_content = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - messages = [ - Message(role="user", contents=[approval_content]), - ] - result = latest_approval_response(messages) - assert result is approval_content - - -class TestApprovalSteps: - """Tests for approval_steps function.""" - - def test_steps_from_ag_ui_state_args(self): - """Extracts steps from ag_ui_state_args.""" - fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - additional_properties={"ag_ui_state_args": {"steps": [{"id": 1}, {"id": 2}]}}, - ) - result = approval_steps(approval) - assert result == [{"id": 1}, {"id": 2}] - - def test_steps_from_function_call(self): - """Extracts steps from function call arguments.""" - fc = Content.from_function_call( - call_id="call_123", - name="test", - arguments='{"steps": [{"step": 1}]}', - ) - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - result = approval_steps(approval) - assert result == [{"step": 1}] - - def test_empty_steps_when_no_state_args(self): - """Returns empty list when no ag_ui_state_args.""" - fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - result = approval_steps(approval) - assert result == [] - - def test_empty_steps_when_state_args_not_dict(self): - """Returns empty list when ag_ui_state_args is not a dict.""" - fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - additional_properties={"ag_ui_state_args": "not a dict"}, - ) - result = approval_steps(approval) - assert result == [] - - def test_empty_steps_when_steps_not_list(self): - """Returns empty list when steps is not a list.""" - fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - additional_properties={"ag_ui_state_args": {"steps": "not a list"}}, - ) - result = approval_steps(approval) - assert result == [] - - -class TestIsStepBasedApproval: - """Tests for is_step_based_approval function.""" - - def test_returns_true_when_has_steps(self): - """Returns True when approval has steps.""" - fc = Content.from_function_call(call_id="call_123", name="test_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - additional_properties={"ag_ui_state_args": {"steps": [{"id": 1}]}}, - ) - result = is_step_based_approval(approval, None) - assert result is True - - def test_returns_false_no_steps_no_function_call(self): - """Returns False when no steps and no function call.""" - # Create content directly to have no function_call - approval = Content( - type="function_approval_response", - function_call=None, - ) - result = is_step_based_approval(approval, None) - assert result is False - - def test_returns_false_no_predict_config(self): - """Returns False when no predict_state_config.""" - fc = Content.from_function_call(call_id="call_123", name="some_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - result = is_step_based_approval(approval, None) - assert result is False - - def test_returns_true_when_tool_matches_config(self): - """Returns True when tool matches predict_state_config with steps.""" - fc = Content.from_function_call(call_id="call_123", name="generate_steps", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - config = {"steps": {"tool": "generate_steps", "tool_argument": "steps"}} - result = is_step_based_approval(approval, config) - assert result is True - - def test_returns_false_when_tool_not_in_config(self): - """Returns False when tool not in predict_state_config.""" - fc = Content.from_function_call(call_id="call_123", name="other_tool", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - config = {"steps": {"tool": "generate_steps", "tool_argument": "steps"}} - result = is_step_based_approval(approval, config) - assert result is False - - def test_returns_false_when_tool_arg_not_steps(self): - """Returns False when tool_argument is not 'steps'.""" - fc = Content.from_function_call(call_id="call_123", name="generate_steps", arguments="{}") - approval = Content.from_function_approval_response( - approved=True, - id="approval_123", - function_call=fc, - ) - config = {"document": {"tool": "generate_steps", "tool_argument": "content"}} - result = is_step_based_approval(approval, config) - assert result is False diff --git a/python/packages/ag-ui/tests/ag_ui/test_predictive_state.py b/python/packages/ag-ui/tests/ag_ui/test_predictive_state.py index 31ad46fc3a7..5d9a5ec3417 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_predictive_state.py +++ b/python/packages/ag-ui/tests/ag_ui/test_predictive_state.py @@ -4,7 +4,7 @@ from ag_ui.core import StateDeltaEvent -from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler +from agent_framework_ag_ui._predictive_state import PredictiveStateHandler class TestPredictiveStateHandlerInit: diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 217983c3d9e..261f400e143 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -938,7 +938,7 @@ def test_extract_approved_state_updates_no_handler(): def test_extract_approved_state_updates_no_approval(): """Test _extract_approved_state_updates returns empty when no approval content.""" from agent_framework_ag_ui._agent_run import _extract_approved_state_updates - from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler + from agent_framework_ag_ui._predictive_state import PredictiveStateHandler handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "content"}}) messages = [Message(role="user", contents=[Content.from_text("Hello")])] diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index 9784c81c3df..fe1f3bc4dce 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -14,7 +14,7 @@ from agent_framework import Content from agent_framework_ag_ui import state_update -from agent_framework_ag_ui._orchestration._predictive_state import PredictiveStateHandler +from agent_framework_ag_ui._predictive_state import PredictiveStateHandler from agent_framework_ag_ui._run_common import ( FlowState, _build_run_finished_event, diff --git a/python/packages/ag-ui/tests/ag_ui/test_tooling.py b/python/packages/ag-ui/tests/ag_ui/test_tooling.py index 26fc57cd650..bcdf285b0f4 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_tooling.py +++ b/python/packages/ag-ui/tests/ag_ui/test_tooling.py @@ -6,7 +6,7 @@ import pytest from agent_framework import Agent, tool -from agent_framework_ag_ui._orchestration._tooling import ( +from agent_framework_ag_ui._tooling import ( collect_server_tools, merge_tools, register_additional_client_tools, @@ -167,7 +167,7 @@ def test_register_additional_client_tools_no_tools() -> None: def test_register_additional_client_tools_no_chat_client() -> None: """register_additional_client_tools does nothing when agent has no client.""" - from agent_framework_ag_ui._orchestration._tooling import register_additional_client_tools + from agent_framework_ag_ui._tooling import register_additional_client_tools class MockAgent: pass