Python: improved loading (#202)

* improved loading

* updated openai folder imports

* fixed import

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",
+]