Python: improved loading

python/.vscode/settings.json

@@ -8,8 +8,8 @@
     ],
     "[python]": {
         "editor.codeActionsOnSave": {
-            "source.organizeImports": "explicit",
-            "source.fixAll": "explicit"
+            "source.organizeImports.ruff": "always",
+            "source.fixAll.ruff": "always"
         },
         "editor.formatOnSave": true,
         "editor.formatOnPaste": true,
@@ -36,4 +36,4 @@
             "depth": 2
         }
     ]
-}
+}

python/packages/azure/tests/unit/test_azure_chat_client.py

@@ -14,7 +14,7 @@
     TextContent,
 )
 from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
-from agent_framework.openai.exceptions import (
+from agent_framework.openai import (
     ContentFilterResultSeverity,
     OpenAIContentFilterException,
 )

python/packages/main/agent_framework/__init__.py

@@ -2,68 +2,15 @@
 
 import importlib
 import importlib.metadata
-from typing import Any
 
 try:
     __version__ = importlib.metadata.version(__name__)
 except importlib.metadata.PackageNotFoundError:
     __version__ = "0.0.0"  # Fallback for development mode
 
-_IMPORTS = {
-    "AFBaseModel": "._pydantic",
-    "AFBaseSettings": "._pydantic",
-    "Agent": "._agents",
-    "AgentRunResponse": "._types",
-    "AgentRunResponseUpdate": "._types",
-    "AgentThread": "._agents",
-    "AITool": "._tools",
-    "AIFunction": "._tools",
-    "AIContent": "._types",
-    "AIContents": "._types",
-    "ChatClient": "._clients",
-    "ChatClientAgent": "._agents",
-    "ChatClientAgentThread": "._agents",
-    "ChatClientAgentThreadType": "._agents",
-    "ChatClientBase": "._clients",
-    "ChatFinishReason": "._types",
-    "ChatMessage": "._types",
-    "ChatOptions": "._types",
-    "ChatResponse": "._types",
-    "ChatResponseUpdate": "._types",
-    "ChatRole": "._types",
-    "ChatToolMode": "._types",
-    "DataContent": "._types",
-    "EmbeddingGenerator": "._clients",
-    "ErrorContent": "._types",
-    "FunctionCallContent": "._types",
-    "FunctionResultContent": "._types",
-    "GeneratedEmbeddings": "._types",
-    "HttpsUrl": "._pydantic",
-    "InputGuardrail": ".guard_rails",
-    "OutputGuardrail": ".guard_rails",
-    "SpeechToTextOptions": "._types",
-    "StructuredResponse": "._types",
-    "TextContent": "._types",
-    "TextReasoningContent": "._types",
-    "TextToSpeechOptions": "._types",
-    "UriContent": "._types",
-    "UsageContent": "._types",
-    "UsageDetails": "._types",
-    "ai_function": "._tools",
-    "get_logger": "._logging",
-    "use_tool_calling": "._clients",
-}
-
-
-def __getattr__(name: str) -> Any:
-    if name == "__version__":
-        return __version__
-    if name in _IMPORTS:
-        submod_name = _IMPORTS[name]
-        module = importlib.import_module(submod_name, package=__name__)
-        return getattr(module, name)
-    raise AttributeError(f"module {__name__} has no attribute {name}")
-
-
-def __dir__() -> list[str]:
-    return [*list(_IMPORTS.keys()), "__version__"]
+from ._agents import *  # noqa: F403
+from ._clients import *  # noqa: F403
+from ._logging import *  # noqa: F403
+from ._pydantic import *  # noqa: F403
+from ._tools import *  # noqa: F403
+from ._types import *  # noqa: F403

python/packages/main/agent_framework/_agents.py

@@ -26,6 +26,16 @@
 
 # region AgentThread
 
+__all__ = [
+    "Agent",
+    "AgentBase",
+    "AgentThread",
+    "ChatClientAgent",
+    "ChatClientAgentThread",
+    "ChatClientAgentThreadType",
+    "MessagesRetrievableThread",
+]
+
 
 class AgentThread(AFBaseModel):
     """Base class for agent threads."""

python/packages/main/agent_framework/_clients.py

@@ -34,6 +34,13 @@
 
 logger = get_logger()
 
+__all__ = [
+    "ChatClient",
+    "ChatClientBase",
+    "EmbeddingGenerator",
+    "use_tool_calling",
+]
+
 # region: Tool Calling Functions and Decorators
 
 
@@ -88,7 +95,8 @@ def _prepare_tools_and_tool_choice(chat_options: ChatOptions) -> None:
         chat_options.tool_choice = ChatToolMode.NONE.mode
         return
     chat_options.tools = [
-        (_tool_to_json_schema_spec(t) if isinstance(t, AITool) else t) for t in chat_options._ai_tools or []
+        (_tool_to_json_schema_spec(t) if isinstance(t, AITool) else t)
+        for t in chat_options._ai_tools or []  # type: ignore[reportPrivateUsage]
     ]
     if not chat_options.tools:
         chat_options.tool_choice = ChatToolMode.NONE.mode
@@ -126,7 +134,7 @@ async def wrapper(
                     _auto_invoke_function(
                         function_call,
                         custom_args=kwargs,
-                        tool_map={t.name: t for t in chat_options._ai_tools or [] if isinstance(t, AIFunction)},
+                        tool_map={t.name: t for t in chat_options._ai_tools or [] if isinstance(t, AIFunction)},  # type: ignore[reportPrivateUsage]
                         sequence_index=seq_idx,
                         request_index=attempt_idx,
                     )
@@ -203,7 +211,7 @@ async def wrapper(
                     _auto_invoke_function(
                         function_call,
                         custom_args=kwargs,
-                        tool_map={t.name: t for t in chat_options._ai_tools or [] if isinstance(t, AIFunction)},
+                        tool_map={t.name: t for t in chat_options._ai_tools or [] if isinstance(t, AIFunction)},  # type: ignore[reportPrivateUsage]
                         sequence_index=seq_idx,
                         request_index=attempt_idx,
                     )

python/packages/main/agent_framework/_logging.py

@@ -9,6 +9,8 @@
     datefmt="%Y-%m-%d %H:%M:%S",
 )
 
+__all__ = ["get_logger"]
+
 
 def get_logger(name: str = "agent_framework") -> logging.Logger:
     """Get a logger with the specified name, defaulting to 'agent_framework'.

python/packages/main/agent_framework/_pydantic.py

@@ -9,6 +9,8 @@
 
 HttpsUrl = Annotated[AnyUrl, UrlConstraints(max_length=2083, allowed_schemes=["https"])]
 
+__all__ = ["AFBaseModel", "AFBaseSettings", "HttpsUrl"]
+
 
 class AFBaseModel(BaseModel):
     """Base class for all pydantic models in the Agent Framework."""

python/packages/main/agent_framework/_tools.py

@@ -7,6 +7,8 @@
 
 from pydantic import BaseModel, create_model
 
+__all__ = ["AIFunction", "AITool", "ai_function"]
+
 
 @runtime_checkable
 class AITool(Protocol):

python/packages/main/agent_framework/_types.py

@@ -75,6 +75,34 @@
 ]
 
 
+__all__ = [
+    "AIContent",
+    "AIContents",
+    "AgentRunResponse",
+    "AgentRunResponseUpdate",
+    "ChatFinishReason",
+    "ChatMessage",
+    "ChatOptions",
+    "ChatResponse",
+    "ChatResponseUpdate",
+    "ChatRole",
+    "ChatToolMode",
+    "DataContent",
+    "ErrorContent",
+    "FunctionCallContent",
+    "FunctionResultContent",
+    "GeneratedEmbeddings",
+    "SpeechToTextOptions",
+    "StructuredResponse",
+    "TextContent",
+    "TextReasoningContent",
+    "TextToSpeechOptions",
+    "UriContent",
+    "UsageContent",
+    "UsageDetails",
+]
+
+
 class UsageDetails(AFBaseModel):
     """Provides usage details about a request/response.
 

python/packages/main/agent_framework/azure/__init__.py

@@ -1,34 +1,28 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import importlib
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
+PACKAGE_NAME = "agent_framework_azure"
 PACKAGE_EXTRA = "azure"
-
-_IMPORTS = {
-    "__version__": "agent_framework_azure",
-    "AzureChatClient": "agent_framework_azure",
-    "get_entra_auth_token": "agent_framework_azure",
-}
+_IMPORTS = [
+    "AzureChatClient",
+    "get_entra_auth_token",
+    "__version__",
+]
 
 
 def __getattr__(name: str) -> Any:
     if name in _IMPORTS:
-        submod_name = _IMPORTS[name]
         try:
-            module = importlib.import_module(submod_name, package=__name__)
-            return getattr(module, name)
+            return getattr(importlib.import_module(PACKAGE_NAME), name)
         except ModuleNotFoundError as exc:
             raise ModuleNotFoundError(
                 f"The '{PACKAGE_EXTRA}' extra is not installed, "
                 f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
             ) from exc
-    raise AttributeError(f"module {__name__} has no attribute {name}")
+    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
 
 
 def __dir__() -> list[str]:
-    return list(_IMPORTS.keys())
-
-
-if TYPE_CHECKING:
-    from agent_framework_azure import __version__  # noqa: F401
+    return _IMPORTS

python/packages/main/agent_framework/azure/__init__.pyi

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_azure import (
+    AzureChatClient,
+    __version__,
+    get_entra_auth_token,
+)
+
+__all__ = [
+    "AzureChatClient",
+    "__version__",
+    "get_entra_auth_token",
+]

python/packages/main/agent_framework/guard_rails.py

@@ -6,6 +6,8 @@
 TInput = TypeVar("TInput")
 TResponse = TypeVar("TResponse")
 
+__all__ = ["InputGuardrail", "OutputGuardrail"]
+
 
 @runtime_checkable
 class InputGuardrail(Protocol, Generic[TInput]):

python/packages/main/agent_framework/openai/__init__.py

@@ -1,12 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 
-from ._chat_client import OpenAIChatClient
-from ._shared import OpenAIHandler, OpenAIModelTypes, OpenAISettings
-
-__all__ = [
-    "OpenAIChatClient",
-    "OpenAIHandler",
-    "OpenAIModelTypes",
-    "OpenAISettings",
-]
+from ._chat_client import *  # noqa: F403
+from ._exceptions import *  # noqa: F403
+from ._shared import *  # noqa: F403

python/packages/main/agent_framework/openai/_chat_client.py

@@ -31,6 +31,8 @@
 from ..exceptions import ServiceInitializationError, ServiceInvalidResponseError
 from ._shared import OpenAIConfigBase, OpenAIHandler, OpenAIModelTypes, OpenAISettings
 
+__all__ = ["OpenAIChatClient"]
+
 
 # Implements agent_framework.ChatClient protocol, through ChatClientBase
 @use_tool_calling