perf(lib): stop pulling openai.types.beta into every import openai

src/openai/__init__.py

@@ -108,10 +108,27 @@
 from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI
 from .lib.bedrock import BedrockOpenAI as BedrockOpenAI, AsyncBedrockOpenAI as AsyncBedrockOpenAI
 from .lib._old_api import *
-from .lib.streaming import (
-    AssistantEventHandler as AssistantEventHandler,
-    AsyncAssistantEventHandler as AsyncAssistantEventHandler,
-)
+
+if _t.TYPE_CHECKING:
+    from .lib.streaming import (
+        AssistantEventHandler as AssistantEventHandler,
+        AsyncAssistantEventHandler as AsyncAssistantEventHandler,
+    )
+else:
+    # `openai.lib.streaming` reaches `openai.types.beta`, which is 318 modules
+    # and about a third of the cost of `import openai`, for the Assistants API.
+    # Deferring the module keeps the names available on `openai` while leaving
+    # them -- and their annotations -- fully resolvable once anything asks.
+    _STREAMING_EXPORTS = ("AssistantEventHandler", "AsyncAssistantEventHandler")
+
+    def __getattr__(__name: str) -> _t.Any:
+        if __name in _STREAMING_EXPORTS:
+            import importlib
+
+            value = getattr(importlib.import_module("openai.lib.streaming"), __name)
+            globals()[__name] = value
+            return value
+        raise AttributeError(f"module {__name__!r} has no attribute {__name!r}")
 
 _setup_logging()
 

tests/lib/test_streaming_lazy_types.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import sys
+import typing
+import subprocess
+
+import pytest
+
+import openai
+
+
+def _modules_after_import_openai() -> set[str]:
+    """Return the module names loaded by a bare `import openai` in a fresh interpreter."""
+    output = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys\nimport openai\nprint('\\n'.join(sys.modules))\n",
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    ).stdout
+    return set(output.split())
+
+
+def test_import_openai_does_not_load_beta_types() -> None:
+    # `openai.lib.streaming` is only needed by callers using the Assistants API,
+    # and it reaches `openai.types.beta`, so importing the package must not pull
+    # the namespace in.
+    modules = _modules_after_import_openai()
+
+    assert "openai" in modules
+    assert not [module for module in modules if module.startswith("openai.types.beta")]
+
+
+def test_assistant_event_handlers_are_still_exported() -> None:
+    assert openai.AssistantEventHandler.__name__ == "AssistantEventHandler"
+    assert openai.AsyncAssistantEventHandler.__name__ == "AsyncAssistantEventHandler"
+
+
+def test_handler_annotations_stay_resolvable() -> None:
+    # Deferring the module must not make the handlers' postponed annotations
+    # unresolvable: `get_type_hints` has to keep working for annotation-aware
+    # integrations and documentation tooling.
+    for name in ("on_event", "on_run_step_delta", "on_tool_call_delta"):
+        hints = typing.get_type_hints(getattr(openai.AssistantEventHandler, name))
+        assert "return" in hints
+
+
+def test_unknown_attribute_still_raises_attribute_error() -> None:
+    with pytest.raises(AttributeError, match="definitely_not_an_export"):
+        openai.definitely_not_an_export  # type: ignore[attr-defined]  # noqa: B018