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

src/openai/lib/streaming/_assistants.py

@@ -10,18 +10,20 @@
 from ..._httpx2 import timeout_exceptions
 from ..._models import construct_type
 from ..._streaming import Stream, AsyncStream
-from ...types.beta import AssistantStreamEvent
-from ...types.beta.threads import (
-    Run,
-    Text,
-    Message,
-    ImageFile,
-    TextDelta,
-    MessageDelta,
-    MessageContent,
-    MessageContentDelta,
-)
-from ...types.beta.threads.runs import RunStep, ToolCall, RunStepDelta, ToolCallDelta
+
+if TYPE_CHECKING:
+    from ...types.beta import AssistantStreamEvent
+    from ...types.beta.threads import (
+        Run,
+        Text,
+        Message,
+        ImageFile,
+        TextDelta,
+        MessageDelta,
+        MessageContent,
+        MessageContentDelta,
+    )
+    from ...types.beta.threads.runs import RunStep, ToolCall, RunStepDelta, ToolCallDelta
 
 
 def _timeout_exceptions() -> tuple[type[Exception], ...]:
@@ -903,6 +905,8 @@ def accumulate_run_step(
         return
 
     if event.event == "thread.run.step.delta":
+        from ...types.beta.threads.runs import RunStep
+
         data = event.data
         snapshot = run_step_snapshots[data.id]
 
@@ -928,6 +932,8 @@ def accumulate_event(
     current_message_snapshot: Message | None,
 ) -> tuple[Message | None, list[MessageContentDelta]]:
     """Returns a tuple of message snapshot and newly created text message deltas"""
+    from ...types.beta.threads import MessageContent
+
     if event.event == "thread.message.created":
         return event.data, []
 

tests/lib/test_streaming_lazy_types.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import sys
+import subprocess
+
+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` only needs `openai.types.beta` for annotations and for two
+    # narrow runtime paths, 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_eagerly_exported() -> None:
+    assert openai.AssistantEventHandler.__name__ == "AssistantEventHandler"
+    assert openai.AsyncAssistantEventHandler.__name__ == "AsyncAssistantEventHandler"