From 1170b10ee0c4429bd0c393d8f3e0306a5679b109 Mon Sep 17 00:00:00 2001
From: Ray <mailtangyu@gmail.com>
Date: Fri, 24 Jul 2026 01:28:18 +0800
Subject: [PATCH] perf: lazy-load the legacy indexing exports from
 pageindex/__init__
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`import pageindex` eagerly pulled the whole indexing stack — litellm
(which fetches a remote model cost map at import) and PyPDF2 — costing
~3.3s plus a network attempt for cloud-only 0.2.x users who never touch
local indexing. Legacy pre-SDK names now resolve via PEP 562
__getattr__; a TYPE_CHECKING block keeps real signatures for IDEs.
Import drops to ~0.3s with zero network and no litellm/PyPDF2 in
sys.modules until a legacy indexing name is actually used.
---
 pageindex/__init__.py | 58 ++++++++++++++++++++++++++++++++++---------
 1 file changed, 46 insertions(+), 12 deletions(-)

diff --git a/pageindex/__init__.py b/pageindex/__init__.py
index ace93bd..5c392e0 100644
--- a/pageindex/__init__.py
+++ b/pageindex/__init__.py
@@ -5,16 +5,20 @@ _load_dotenv()
 
 # Backward compatibility: honor CHATGPT_API_KEY as an alias for OPENAI_API_KEY.
 import os as _os
-if not _os.getenv("OPENAI_API_KEY") and _os.getenv("CHATGPT_API_KEY"):
-    _os.environ["OPENAI_API_KEY"] = _os.getenv("CHATGPT_API_KEY")
+_chatgpt_key = _os.getenv("CHATGPT_API_KEY")
+if not _os.getenv("OPENAI_API_KEY") and _chatgpt_key:
+    _os.environ["OPENAI_API_KEY"] = _chatgpt_key
 
-# Upstream exports (backward compatibility); import from the canonical index.*
-# modules so plain `import pageindex` doesn't trip the deprecation shims.
-from .index.page_index import *  # noqa: E402
-from .index.page_index_md import md_to_tree
-from .retrieve import get_document, get_document_structure, get_page_content
+from typing import TYPE_CHECKING as _TYPE_CHECKING
+if _TYPE_CHECKING:
+    # Static-only bindings for the lazily-loaded legacy names below, so type
+    # checkers and IDEs see real signatures without the runtime import cost.
+    from .index.page_index import page_index, page_index_main, tree_parser
+    from .index.page_index_md import md_to_tree
+    from .index.utils import ConfigLoader, llm_completion, llm_acompletion
+    from .retrieve import get_document, get_document_structure, get_page_content
 
-# SDK exports
+# SDK exports — lightweight, no LLM/indexing libraries.
 from .client import PageIndexClient, LocalClient, CloudClient
 from .config import IndexConfig, set_llm_params
 from .collection import Collection
@@ -68,11 +72,41 @@ __all__ = [
     "get_page_content",
 ]
 
+# Legacy (pre-SDK) exports resolve lazily via PEP 562: importing the indexing
+# stack costs seconds (litellm fetches a remote model cost map at import), so
+# plain `import pageindex` for cloud-only use must not pay for it.
+_LAZY_LEGACY = {
+    "md_to_tree": ".index.page_index_md",
+    "get_document": ".retrieve",
+    "get_document_structure": ".retrieve",
+    "get_page_content": ".retrieve",
+}
+
 
 def __getattr__(name):
-    # Lazy so plain `import pageindex` never trips the shims' deprecation
-    # warnings; they fire only when the legacy attribute is actually used.
+    if name.startswith("_"):
+        # Tooling probes dunders (pickle, IPython, …) — never let those
+        # trigger the heavy legacy import.
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    import importlib
     if name in ("utils", "page_index_md"):
-        import importlib
         return importlib.import_module(f".{name}", __name__)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module = _LAZY_LEGACY.get(name)
+    if module is not None:
+        value = getattr(importlib.import_module(module, __name__), name)
+        globals()[name] = value
+        return value
+    # Everything else from the pre-SDK surface (page_index, page_index_main,
+    # tree_parser, ConfigLoader, llm_completion, …) lives in the canonical
+    # index.page_index namespace (which star-exports index.utils).
+    legacy = importlib.import_module(".index.page_index", __name__)
+    try:
+        value = getattr(legacy, name)
+    except AttributeError:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__) | {"utils", "page_index_md"})