import fix

Author: eavanvalkenburg

python/semantic_kernel/connectors/memory/chroma/chroma.py

@@ -5,14 +5,12 @@
 from collections.abc import Sequence
 from typing import Any, ClassVar, Generic, TypeVar
 
-from semantic_kernel.utils.experimental_decorator import experimental_class
-
 if sys.version_info >= (3, 12):
     from typing import override  # pragma: no cover
 else:
     from typing_extensions import override  # pragma: no cover
 
-import chromadb
+from chromadb import Client, Collection, QueryResult
 from chromadb.api import ClientAPI
 from chromadb.config import Settings
 
@@ -32,6 +30,7 @@
     VectorStoreInitializationException,
     VectorStoreModelValidationError,
 )
+from semantic_kernel.utils.experimental_decorator import experimental_class
 
 logger = logging.getLogger(__name__)
 
@@ -72,7 +71,7 @@ def __init__(
             if persist_directory is not None:
                 settings.is_persistent = True
                 settings.persist_directory = persist_directory
-            client = chromadb.Client(settings)
+            client = Client(settings)
         super().__init__(
             collection_name=collection_name,
             data_model_type=data_model_type,
@@ -82,7 +81,7 @@ def __init__(
             **kwargs,
         )
 
-    def _get_collection(self) -> chromadb.Collection:
+    def _get_collection(self) -> Collection:
         try:
             return self.client.get_collection(name=self.collection_name)
         except Exception as e:
@@ -200,7 +199,7 @@ async def _inner_get(self, keys: Sequence[str], **kwargs: Any) -> Sequence[Any]:
         return self._unpack_results(results, include_vectors)
 
     def _unpack_results(
-        self, results: chromadb.QueryResult, include_vectors: bool, include_distances: bool = False
+        self, results: QueryResult, include_vectors: bool, include_distances: bool = False
     ) -> Sequence[dict[str, Any]]:
         try:
             if isinstance(results["ids"][0], str):
@@ -330,7 +329,7 @@ def __init__(
             settings.is_persistent = True
             settings.persist_directory = persist_directory
         if client is None:
-            client = chromadb.Client(settings)
+            client = Client(settings)
         super().__init__(client=client, managed_client=managed_client, **kwargs)
 
     @override