Separate BaseProvider for faster import (#1338)

* Separate `BaseProvider` for faster import

* Use `base_provider` imports across ai_magics package

* Vendor the `import_attr` code

* Support typing and add test to ensure integrity

* Add an import for typing block integrity

* Import `import_attr` as private

* Rename variables as requested in review

Co-authored-by: David L. Qiu <[email protected]>

---------

Co-authored-by: David L. Qiu <[email protected]>

Author: krassowski

packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py

@@ -1,34 +1,96 @@
+from typing import TYPE_CHECKING
+
+from ._import_utils import import_attr as _import_attr
 from ._version import __version__
 
-# expose embedding model providers on the package root
-from .embedding_providers import (
-    BaseEmbeddingsProvider,
-    GPT4AllEmbeddingsProvider,
-    HfHubEmbeddingsProvider,
-    QianfanEmbeddingsEndpointProvider,
-)
-from .exception import store_exception
-from .magics import AiMagics
-
-# expose JupyternautPersona on the package root
-# required by `jupyter-ai`.
-from .models.persona import JupyternautPersona, Persona
-
-# expose model providers on the package root
-from .providers import (
-    AI21Provider,
-    BaseProvider,
-    GPT4AllProvider,
-    HfHubProvider,
-    QianfanProvider,
-    TogetherAIProvider,
-)
+if TYPE_CHECKING:
+    # same as dynamic imports but understood by mypy
+    from .embedding_providers import (
+        BaseEmbeddingsProvider,
+        GPT4AllEmbeddingsProvider,
+        HfHubEmbeddingsProvider,
+        QianfanEmbeddingsEndpointProvider,
+    )
+    from .exception import store_exception
+    from .magics import AiMagics
+    from .models.persona import JupyternautPersona, Persona
+    from .providers import (
+        AI21Provider,
+        BaseProvider,
+        GPT4AllProvider,
+        HfHubProvider,
+        QianfanProvider,
+        TogetherAIProvider,
+    )
+else:
+    _exports_by_module = {
+        # expose embedding model providers on the package root
+        "embedding_providers": [
+            "BaseEmbeddingsProvider",
+            "GPT4AllEmbeddingsProvider",
+            "HfHubEmbeddingsProvider",
+            "QianfanEmbeddingsEndpointProvider",
+        ],
+        "exception": ["store_exception"],
+        "magics": ["AiMagics"],
+        # expose JupyternautPersona on the package root
+        # required by `jupyter-ai`.
+        "models.persona": ["JupyternautPersona", "Persona"],
+        # expose model providers on the package root
+        "providers": [
+            "AI21Provider",
+            "BaseProvider",
+            "GPT4AllProvider",
+            "HfHubProvider",
+            "QianfanProvider",
+            "TogetherAIProvider",
+        ],
+    }
+
+    _modules_by_export = {
+        import_name: module
+        for module, imports in _exports_by_module.items()
+        for import_name in imports
+    }
+
+    def __getattr__(export_name: str) -> object:
+        module_name = _modules_by_export.get(export_name)
+        result = _import_attr(export_name, module_name, __spec__.parent)
+        globals()[export_name] = result
+        return result
 
 
 def load_ipython_extension(ipython):
-    ipython.register_magics(AiMagics)
-    ipython.set_custom_exc((BaseException,), store_exception)
+    ipython.register_magics(__getattr__("AiMagics"))
+    ipython.set_custom_exc((BaseException,), __getattr__("store_exception"))
 
 
 def unload_ipython_extension(ipython):
     ipython.set_custom_exc((BaseException,), ipython.CustomTB)
+
+
+# required to preserve backward compatibility with `from jupyter_ai_magics import *`
+__all__ = [
+    "__version__",
+    "load_ipython_extension",
+    "unload_ipython_extension",
+    "BaseEmbeddingsProvider",
+    "GPT4AllEmbeddingsProvider",
+    "HfHubEmbeddingsProvider",
+    "QianfanEmbeddingsEndpointProvider",
+    "store_exception",
+    "AiMagics",
+    "JupyternautPersona",
+    "Persona",
+    "AI21Provider",
+    "BaseProvider",
+    "GPT4AllProvider",
+    "HfHubProvider",
+    "QianfanProvider",
+    "TogetherAIProvider",
+]
+
+
+def __dir__():
+    # Allows more editors (e.g. IPython) to complete on `jupyter_ai_magics.<tab>`
+    return list(__all__)

packages/jupyter-ai-magics/jupyter_ai_magics/_import_utils.py

@@ -0,0 +1,58 @@
+"""
+MIT License
+
+Copyright (c) LangChain, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+from importlib import import_module
+from typing import Union
+
+
+def import_attr(
+    attr_name: str,
+    module_name: Union[str, None],
+    package: Union[str, None],
+) -> object:
+    """Import an attribute from a module located in a package.
+
+    This utility function is used in custom __getattr__ methods within __init__.py
+    files to dynamically import attributes.
+
+    Args:
+        attr_name: The name of the attribute to import.
+        module_name: The name of the module to import from. If None, the attribute
+            is imported from the package itself.
+        package: The name of the package where the module is located.
+    """
+    if module_name == "__module__" or module_name is None:
+        try:
+            result = import_module(f".{attr_name}", package=package)
+        except ModuleNotFoundError:
+            msg = f"module '{package!r}' has no attribute {attr_name!r}"
+            raise AttributeError(msg) from None
+    else:
+        try:
+            module = import_module(f".{module_name}", package=package)
+        except ModuleNotFoundError:
+            msg = f"module '{package!r}.{module_name!r}' not found"
+            raise ImportError(msg) from None
+        result = getattr(module, attr_name)
+    return result