diff --git a/src/dreamer/trees/graph.py b/src/dreamer/trees/graph.py index b4778f13..1899ee27 100644 --- a/src/dreamer/trees/graph.py +++ b/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] diff --git a/src/dreamer/trees/sklearn_wrapper.py b/src/dreamer/trees/sklearn_wrapper.py index 42aca2d5..86b9522e 100644 --- a/src/dreamer/trees/sklearn_wrapper.py +++ b/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) diff --git a/src/utils/files.py b/src/utils/files.py index 1cc40ba1..4dff50f4 100644 --- a/src/utils/files.py +++ b/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): diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 7a04eaa2..c77064a7 100644 --- a/src/vector_store/__init__.py +++ b/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}")