Lazy load chroma (#529)

mentat/feature_filters/embedding_similarity_filter.py

@@ -1,5 +1,4 @@
 from mentat.code_feature import CodeFeature
-from mentat.embeddings import get_feature_similarity_scores
 from mentat.feature_filters.feature_filter import FeatureFilter
 
 
@@ -11,6 +10,10 @@ async def score(
         self,
         features: list[CodeFeature],
     ) -> list[tuple[CodeFeature, float]]:
+        from mentat.embeddings import (  # dynamic import to improve startup time.
+            get_feature_similarity_scores,
+        )
+
         if self.query == "":
             return [(f, 0.0) for f in features]
 

mentat/llm_api_handler.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 from timeit import default_timer
 from typing import (
+    TYPE_CHECKING,
     Any,
     AsyncIterator,
     Callable,
@@ -22,7 +23,6 @@
 import attr
 import sentry_sdk
 import tiktoken
-from chromadb.api.types import Embeddings
 from dotenv import load_dotenv
 from openai import (
     APIConnectionError,
@@ -48,6 +48,10 @@
 
 TOKEN_COUNT_WARNING = 32000
 
+if TYPE_CHECKING:
+    # This import is slow
+    from chromadb.api.types import Embeddings
+
 
 def is_test_environment():
     """Returns True if in pytest and not benchmarks"""