chore: lazy-load imports (#345)

* chore: lazy-load imports

* fix: lazy-load pdfplumber

src/dreamer/trees/graph.py

@@ -4,19 +4,22 @@ Graph-theoretic surprisal using k-NN graph and random walk.
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 from numpy.typing import NDArray
-from sklearn.neighbors import NearestNeighbors
 
 from .base import SurprisalTree
 
+if TYPE_CHECKING:
+    from sklearn.neighbors import NearestNeighbors
+
 
 def _knn_indices(
     points: NDArray[np.floating[Any]], n_neighbors: int
 ) -> NDArray[np.intp]:
     """Get k-nearest neighbor indices for each point."""
+    from sklearn.neighbors import NearestNeighbors
 
     knn: NearestNeighbors = NearestNeighbors(n_neighbors=n_neighbors, algorithm="auto")
     knn.fit(points)  # pyright: ignore[reportUnknownMemberType]
@@ -26,6 +29,7 @@ def _knn_indices(
 
 def _nearest_index(points: NDArray[np.floating[Any]], query: np.ndarray) -> int:
     """Find index of nearest point to query."""
+    from sklearn.neighbors import NearestNeighbors
 
     knn: NearestNeighbors = NearestNeighbors(n_neighbors=1)
     knn.fit(points)  # pyright: ignore[reportUnknownMemberType]

src/dreamer/trees/sklearn_wrapper.py

@@ -2,17 +2,19 @@
 Wrapper for sklearn's KDTree and BallTree.
 """
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 from numpy.typing import NDArray
-from sklearn.neighbors import (
-    BallTree,  # pyright: ignore[reportUnknownVariableType]
-    KDTree,  # pyright: ignore[reportUnknownVariableType]
-)
 
 from .base import SurprisalTree
 
+if TYPE_CHECKING:
+    from sklearn.neighbors import (
+        BallTree,  # pyright: ignore[reportUnknownVariableType]
+        KDTree,  # pyright: ignore[reportUnknownVariableType]
+    )
+
 
 class SklearnTreeWrapper(SurprisalTree):
     """
@@ -23,7 +25,7 @@ class SklearnTreeWrapper(SurprisalTree):
     tree_type: str
     k: int
     points: list[NDArray[np.floating[Any]]]
-    tree: KDTree | BallTree | None
+    tree: "KDTree | BallTree | None"
     total_points: int
 
     def __init__(
@@ -50,6 +52,11 @@ class SklearnTreeWrapper(SurprisalTree):
         if len(self.points) == 0:
             return
 
+        from sklearn.neighbors import (
+            BallTree,  # pyright: ignore[reportUnknownVariableType]
+            KDTree,  # pyright: ignore[reportUnknownVariableType]
+        )
+
         points_array: NDArray[np.floating[Any]] = np.array(self.points)
         if self.tree_type == "kd":
             self.tree = KDTree(points_array)

src/utils/files.py

@@ -3,7 +3,6 @@ import logging
 from io import BytesIO
 from typing import Any, Protocol
 
-import pdfplumber
 from fastapi import UploadFile
 from nanoid import generate as generate_nanoid
 from sqlalchemy import Integer, select
@@ -27,6 +26,8 @@ class PDFProcessor:
         return content_type == "application/pdf"
 
     async def extract_text(self, content: bytes) -> str:
+        import pdfplumber
+
         with pdfplumber.open(BytesIO(content)) as pdf_reader:
             text_parts: list[str] = []
             for page_num, page in enumerate(pdf_reader.pages):

src/vector_store/__init__.py

@@ -192,17 +192,18 @@ class VectorStore(ABC):
         ...
 
 
-# Import implementations after base classes are defined to avoid circular imports
-from src.vector_store.lancedb import LanceDBVectorStore  # noqa: E402
-from src.vector_store.turbopuffer import TurbopufferVectorStore  # noqa: E402
 from src.vector_store.utils import upsert_with_retry  # noqa: E402
 
 
 def _create_store_by_type(store_type: str) -> VectorStore:
     """Create a vector store instance by type name."""
     if store_type == "turbopuffer":
+        from src.vector_store.turbopuffer import TurbopufferVectorStore
+
         return TurbopufferVectorStore()
     elif store_type == "lancedb":
+        from src.vector_store.lancedb import LanceDBVectorStore
+
         return LanceDBVectorStore()
     else:
         raise ValueError(f"Unknown vector store type: {store_type}")