Python: Clean up imports (#2318)

* chore: tidy imports

* Update python/packages/azurefunctions/agent_framework_azurefunctions/_errors.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update python/packages/azurefunctions/agent_framework_azurefunctions/_callbacks.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* chore: revert stub file change

* chore: trigger pre-commit hook, re-add `annotations` import

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

python/packages/ag-ui/README.md

@@ -16,7 +16,7 @@ pip install agent-framework-ag-ui
 from fastapi import FastAPI
 from agent_framework import ChatAgent
 from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 
 # Create your agent
 agent = ChatAgent(
@@ -41,7 +41,7 @@ add_agent_framework_fastapi_endpoint(app, agent, "/")
 ```python
 import asyncio
 from agent_framework import TextContent
-from agent_framework_ag_ui import AGUIChatClient
+from agent_framework.ag_ui import AGUIChatClient
 
 async def main():
     async with AGUIChatClient(endpoint="http://localhost:8000/") as client:

python/packages/ag-ui/agent_framework_ag_ui_examples/README.md

@@ -18,7 +18,7 @@ All example agents are factory functions that accept any `ChatClientProtocol`-co
 from fastapi import FastAPI
 from agent_framework.azure import AzureOpenAIChatClient
 from agent_framework.openai import OpenAIChatClient
-from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework_ag_ui_examples.agents import simple_agent, weather_agent
 
 app = FastAPI()
@@ -40,7 +40,7 @@ add_agent_framework_fastapi_endpoint(app, weather_agent(openai_client), "/weathe
 from fastapi import FastAPI
 from agent_framework import ChatAgent
 from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 
 # Create your agent
 agent = ChatAgent(
@@ -136,7 +136,7 @@ The server exposes endpoints at:
 ```python
 from fastapi import FastAPI
 from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework_ag_ui_examples.agents import (
     simple_agent,
     weather_agent,
@@ -188,8 +188,8 @@ You can create your own agent factories following the same pattern as the exampl
 
 ```python
 from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
-from agent_framework_ag_ui import AgentFrameworkAgent
+from agent_framework import ChatClientProtocol
+from agent_framework.ag_ui import AgentFrameworkAgent
 
 @ai_function
 def my_tool(param: str) -> str:

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/document_writer_agent.py

@@ -2,10 +2,8 @@
 
 """Example agent demonstrating predictive state updates with document writing."""
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
-
-from agent_framework_ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
+from agent_framework.ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
 
 
 @ai_function

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py

@@ -4,8 +4,7 @@
 
 from enum import Enum
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
 from pydantic import BaseModel, Field
 
 

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py

@@ -4,12 +4,10 @@
 
 from enum import Enum
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
+from agent_framework.ag_ui import AgentFrameworkAgent, RecipeConfirmationStrategy
 from pydantic import BaseModel, Field
 
-from agent_framework_ag_ui import AgentFrameworkAgent, RecipeConfirmationStrategy
-
 
 class SkillLevel(str, Enum):
     """The skill level required for the recipe."""

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py

@@ -4,10 +4,8 @@
 
 import asyncio
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
-
-from agent_framework_ag_ui import AgentFrameworkAgent
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
+from agent_framework.ag_ui import AgentFrameworkAgent
 
 
 @ai_function

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py

@@ -2,8 +2,7 @@
 
 """Simple agentic chat example (Feature 1: Agentic Chat)."""
 
-from agent_framework import ChatAgent
-from agent_framework._clients import ChatClientProtocol
+from agent_framework import ChatAgent, ChatClientProtocol
 
 
 def simple_agent(chat_client: ChatClientProtocol) -> ChatAgent:

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py

@@ -2,10 +2,8 @@
 
 """Example agent demonstrating human-in-the-loop with function approvals."""
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
-
-from agent_framework_ag_ui import AgentFrameworkAgent, TaskPlannerConfirmationStrategy
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
+from agent_framework.ag_ui import AgentFrameworkAgent, TaskPlannerConfirmationStrategy
 
 
 @ai_function(approval_mode="always_require")

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py

@@ -18,12 +18,10 @@ from ag_ui.core import (
     TextMessageStartEvent,
     ToolCallStartEvent,
 )
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
+from agent_framework.ag_ui import AgentFrameworkAgent
 from pydantic import BaseModel, Field
 
-from agent_framework_ag_ui import AgentFrameworkAgent
-
 
 class StepStatus(str, Enum):
     """Status of a task step."""

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py

@@ -4,10 +4,8 @@
 
 from typing import Any
 
-from agent_framework import AIFunction, ChatAgent
-from agent_framework._clients import ChatClientProtocol
-
-from agent_framework_ag_ui import AgentFrameworkAgent
+from agent_framework import AIFunction, ChatAgent, ChatClientProtocol
+from agent_framework.ag_ui import AgentFrameworkAgent
 
 # Declaration-only tools (func=None) - actual rendering happens on the client side
 generate_haiku = AIFunction[Any, str](

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py

@@ -4,8 +4,7 @@
 
 from typing import Any
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework._clients import ChatClientProtocol
+from agent_framework import ChatAgent, ChatClientProtocol, ai_function
 
 
 @ai_function

python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py

@@ -2,11 +2,10 @@
 
 """Backend tool rendering endpoint."""
 
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework.azure import AzureOpenAIChatClient
 from fastapi import FastAPI
 
-from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
-
 from ...agents.weather_agent import weather_agent
 
 

python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py

@@ -6,12 +6,11 @@ import logging
 import os
 
 import uvicorn
+from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework.azure import AzureOpenAIChatClient
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
-from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
-
 from ..agents.document_writer_agent import document_writer_agent
 from ..agents.human_in_the_loop_agent import human_in_the_loop_agent
 from ..agents.recipe_agent import recipe_agent

python/packages/ag-ui/getting_started/client.py

@@ -10,7 +10,7 @@ standard chat interface.
 import asyncio
 import os
 
-from agent_framework_ag_ui import AGUIChatClient
+from agent_framework.ag_ui import AGUIChatClient
 
 
 async def main():

python/packages/ag-ui/getting_started/client_advanced.py

@@ -13,8 +13,7 @@ import asyncio
 import os
 
 from agent_framework import ai_function
-
-from agent_framework_ag_ui import AGUIChatClient
+from agent_framework.ag_ui import AGUIChatClient
 
 
 @ai_function

python/packages/ag-ui/getting_started/client_with_agent.py

@@ -23,8 +23,7 @@ import logging
 import os
 
 from agent_framework import ChatAgent, FunctionCallContent, FunctionResultContent, TextContent, ai_function
-
-from agent_framework_ag_ui import AGUIChatClient
+from agent_framework.ag_ui import AGUIChatClient
 
 # Enable debug logging
 logging.basicConfig(

python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py

@@ -11,7 +11,7 @@ from agent_framework._types import ChatResponseUpdate
 
 async def test_agent_initialization_basic():
     """Test basic agent initialization without state schema."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -28,7 +28,7 @@ async def test_agent_initialization_basic():
 
 async def test_agent_initialization_with_state_schema():
     """Test agent initialization with state_schema."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -43,7 +43,7 @@ async def test_agent_initialization_with_state_schema():
 
 async def test_agent_initialization_with_predict_state_config():
     """Test agent initialization with predict_state_config."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -58,7 +58,7 @@ async def test_agent_initialization_with_predict_state_config():
 
 async def test_run_started_event_emission():
     """Test RunStartedEvent is emitted at start of run."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -81,7 +81,7 @@ async def test_run_started_event_emission():
 
 async def test_predict_state_custom_event_emission():
     """Test PredictState CustomEvent is emitted when predict_state_config is present."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -112,7 +112,7 @@ async def test_predict_state_custom_event_emission():
 
 async def test_initial_state_snapshot_with_schema():
     """Test initial StateSnapshotEvent emission when state_schema present."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -141,7 +141,7 @@ async def test_initial_state_snapshot_with_schema():
 
 async def test_state_initialization_object_type():
     """Test state initialization with object type in schema."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -167,7 +167,7 @@ async def test_state_initialization_object_type():
 
 async def test_state_initialization_array_type():
     """Test state initialization with array type in schema."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -193,7 +193,7 @@ async def test_state_initialization_array_type():
 
 async def test_run_finished_event_emission():
     """Test RunFinishedEvent is emitted at end of run."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -214,7 +214,7 @@ async def test_run_finished_event_emission():
 
 async def test_tool_result_confirm_changes_accepted():
     """Test confirm_changes tool result handling when accepted."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -260,7 +260,7 @@ async def test_tool_result_confirm_changes_accepted():
 
 async def test_tool_result_confirm_changes_rejected():
     """Test confirm_changes tool result handling when rejected."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -293,7 +293,7 @@ async def test_tool_result_confirm_changes_rejected():
 
 async def test_tool_result_function_approval_accepted():
     """Test function approval tool result when steps are accepted."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -338,7 +338,7 @@ async def test_tool_result_function_approval_accepted():
 
 async def test_tool_result_function_approval_rejected():
     """Test function approval tool result when rejected."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -374,7 +374,7 @@ async def test_tool_result_function_approval_rejected():
 
 async def test_thread_metadata_tracking():
     """Test that thread metadata includes ag_ui_thread_id and ag_ui_run_id."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     thread_metadata = {}
 
@@ -405,7 +405,7 @@ async def test_thread_metadata_tracking():
 
 async def test_state_context_injection():
     """Test that current state is injected into thread metadata."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     thread_metadata = {}
 
@@ -436,7 +436,7 @@ async def test_state_context_injection():
 
 async def test_no_messages_provided():
     """Test handling when no messages are provided."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -459,7 +459,7 @@ async def test_no_messages_provided():
 
 async def test_message_end_event_emission():
     """Test TextMessageEndEvent is emitted for assistant messages."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -486,7 +486,7 @@ async def test_message_end_event_emission():
 
 async def test_error_handling_with_exception():
     """Test that exceptions during agent execution are re-raised."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class FailingChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -506,7 +506,7 @@ async def test_error_handling_with_exception():
 
 async def test_json_decode_error_in_tool_result():
     """Test handling of orphaned tool result - should be sanitized out."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -543,7 +543,7 @@ async def test_json_decode_error_in_tool_result():
 
 async def test_suppressed_summary_with_document_state():
     """Test suppressed summary uses document state for confirmation message."""
-    from agent_framework_ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
+    from agent_framework.ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):

python/packages/ag-ui/tests/test_structured_output.py

@@ -32,7 +32,7 @@ class GenericOutput(BaseModel):
 
 async def test_structured_output_with_recipe():
     """Test structured output processing with recipe state."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -70,7 +70,7 @@ async def test_structured_output_with_recipe():
 
 async def test_structured_output_with_steps():
     """Test structured output processing with steps state."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -109,7 +109,7 @@ async def test_structured_output_with_steps():
 
 async def test_structured_output_with_no_schema_match():
     """Test structured output when response fields don't match state_schema keys."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -138,7 +138,7 @@ async def test_structured_output_with_no_schema_match():
 
 async def test_structured_output_without_schema():
     """Test structured output without state_schema treats all fields as state."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class DataOutput(BaseModel):
         """Output with data and info fields."""
@@ -175,7 +175,7 @@ async def test_structured_output_without_schema():
 
 async def test_no_structured_output_when_no_response_format():
     """Test that structured output path is skipped when no response_format."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -200,7 +200,7 @@ async def test_no_structured_output_when_no_response_format():
 
 async def test_structured_output_with_message_field():
     """Test structured output that includes a message field."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
@@ -234,7 +234,7 @@ async def test_structured_output_with_message_field():
 
 async def test_empty_updates_no_structured_processing():
     """Test that empty updates don't trigger structured output processing."""
-    from agent_framework_ag_ui import AgentFrameworkAgent
+    from agent_framework.ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):

python/packages/azurefunctions/agent_framework_azurefunctions/_callbacks.py

@@ -6,8 +6,6 @@ This module enables callers of AgentFunctionApp to supply streaming and final-re
 invoked during durable entity execution.
 """
 
-from __future__ import annotations
-
 from dataclasses import dataclass
 from typing import Protocol
 

python/packages/azurefunctions/agent_framework_azurefunctions/_errors.py

@@ -2,8 +2,6 @@
 
 """Custom exception types for the durable agent framework."""
 
-from __future__ import annotations
-
 
 class IncomingRequestError(ValueError):
     """Raised when an incoming HTTP request cannot be parsed or validated."""

python/packages/core/agent_framework/observability.py

@@ -1021,7 +1021,7 @@ def use_observability(
         .. code-block:: python
 
             from agent_framework import use_observability, setup_observability
-            from agent_framework._clients import ChatClientProtocol
+            from agent_framework import ChatClientProtocol
 
 
             # Decorate a custom chat client class

python/packages/core/tests/core/test_middleware.py

@@ -1693,7 +1693,7 @@ def mock_function() -> AIFunction[Any, Any]:
 @pytest.fixture
 def mock_chat_client() -> Any:
     """Mock chat client for testing."""
-    from agent_framework._clients import ChatClientProtocol
+    from agent_framework import ChatClientProtocol
 
     client = MagicMock(spec=ChatClientProtocol)
     client.service_url = MagicMock(return_value="mock://test")

python/packages/devui/agent_framework_devui/_discovery.py

@@ -2,8 +2,6 @@
 
 """Agent Framework entity discovery implementation."""
 
-from __future__ import annotations
-
 import ast
 import importlib
 import importlib.util

python/packages/devui/agent_framework_devui/models/_discovery_models.py

@@ -2,8 +2,6 @@
 
 """Discovery API models for entity information."""
 
-from __future__ import annotations
-
 import re
 from typing import Any
 

python/packages/purview/agent_framework_purview/_cache.py

@@ -1,8 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 """Cache provider for Purview data."""
 
-from __future__ import annotations
-
 import hashlib
 import heapq
 import json

python/packages/purview/agent_framework_purview/_client.py

@@ -1,5 +1,4 @@
 # Copyright (c) Microsoft. All rights reserved.
-from __future__ import annotations
 
 import base64
 import inspect

python/packages/purview/agent_framework_purview/_exceptions.py

@@ -1,8 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 """Purview specific exceptions (minimal error shaping)."""
 
-from __future__ import annotations
-
 from agent_framework.exceptions import ServiceResponseException
 
 __all__ = [

python/packages/purview/agent_framework_purview/_middleware.py

@@ -1,5 +1,4 @@
 # Copyright (c) Microsoft. All rights reserved.
-from __future__ import annotations
 
 from collections.abc import Awaitable, Callable
 

python/packages/purview/agent_framework_purview/_processor.py

@@ -1,5 +1,4 @@
 # Copyright (c) Microsoft. All rights reserved.
-from __future__ import annotations
 
 import asyncio
 import uuid

python/packages/purview/agent_framework_purview/_settings.py

@@ -1,7 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from __future__ import annotations
-
 from enum import Enum
 
 from agent_framework._pydantic import AFBaseSettings