Merge branch 'main' into main

Author: garybadwal

integrations/chroma/CHANGELOG.md

@@ -1,5 +1,12 @@
 # Changelog
 
+## [integrations/chroma-v3.3.0] - 2025-06-09
+
+### 🚀 Features
+
+- Chroma - fix typing + ship types by adding py.typed files (#1910)
+
+
 ## [integrations/chroma-v3.2.0] - 2025-06-06
 
 ### 🚀 Features

integrations/chroma/pyproject.toml

@@ -66,7 +66,13 @@ integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
 cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
 
-types = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = "mypy -p haystack_integrations.components.retrievers.chroma -p haystack_integrations.document_stores.chroma {args}"
+
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -160,7 +166,3 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 [tool.pytest.ini_options]
 minversion = "6.0"
 markers = ["integration: integration tests"]
-
-[[tool.mypy.overrides]]
-module = ["haystack_integrations.*"]
-ignore_missing_imports = true

integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py

@@ -71,7 +71,7 @@ def run(
         query: str,
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, Any]:
         """
         Run the retriever on the given input data.
 
@@ -96,7 +96,7 @@ async def run_async(
         query: str,
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, Any]:
         """
         Asynchronously run the retriever on the given input data.
 
@@ -115,7 +115,7 @@ async def run_async(
         """
         filters = apply_filter_policy(self.filter_policy, self.filters, filters)
         top_k = top_k or self.top_k
-        return {"documents": await self.document_store.search_async([query], top_k, filters)[0]}
+        return {"documents": (await self.document_store.search_async([query], top_k, filters))[0]}
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "ChromaQueryTextRetriever":
@@ -184,7 +184,7 @@ def run(
         query_embedding: List[float],
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, Any]:
         """
         Run the retriever on the given input data.
 
@@ -211,7 +211,7 @@ async def run_async(
         query_embedding: List[float],
         filters: Optional[Dict[str, Any]] = None,
         top_k: Optional[int] = None,
-    ):
+    ) -> Dict[str, Any]:
         """
         Asynchronously run the retriever on the given input data.
 
@@ -232,7 +232,7 @@ async def run_async(
         top_k = top_k or self.top_k
 
         query_embeddings = [query_embedding]
-        return {"documents": await self.document_store.search_embeddings_async(query_embeddings, top_k, filters)[0]}
+        return {"documents": (await self.document_store.search_embeddings_async(query_embeddings, top_k, filters))[0]}
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "ChromaEmbeddingRetriever":

integrations/chroma/src/haystack_integrations/components/retrievers/py.typed

@@ -2,9 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Dict, List, Literal, Optional, Sequence, cast
 
 import chromadb
+from chromadb.api.models.AsyncCollection import AsyncCollection
 from chromadb.api.types import GetResult, QueryResult
 from haystack import default_from_dict, default_to_dict, logging
 from haystack.dataclasses import Document
@@ -37,7 +38,7 @@ def __init__(
         port: Optional[int] = None,
         distance_function: Literal["l2", "cosine", "ip"] = "l2",
         metadata: Optional[dict] = None,
-        **embedding_function_params,
+        **embedding_function_params: Any,
     ):
         """
         Creates a new ChromaDocumentStore instance.
@@ -86,8 +87,8 @@ def __init__(
         self._host = host
         self._port = port
 
-        self._collection = None
-        self._async_collection = None
+        self._collection: Optional[chromadb.Collection] = None
+        self._async_collection: Optional[AsyncCollection] = None
 
     def _ensure_initialized(self):
         if not self._collection:
@@ -482,7 +483,7 @@ def search_embeddings(
 
         kwargs = self._prepare_query_kwargs(filters)
         results = self._collection.query(
-            query_embeddings=query_embeddings,
+            query_embeddings=cast(List[Sequence[float]], query_embeddings),
             n_results=top_k,
             **kwargs,
         )
@@ -513,7 +514,7 @@ async def search_embeddings_async(
 
         kwargs = self._prepare_query_kwargs(filters)
         results = await self._async_collection.query(
-            query_embeddings=query_embeddings,
+            query_embeddings=cast(List[Sequence[float]], query_embeddings),
             n_results=top_k,
             **kwargs,
         )

integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py

@@ -2,6 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Any
+
 from chromadb.api.types import EmbeddingFunction
 from chromadb.utils.embedding_functions import (
     CohereEmbeddingFunction,
@@ -34,7 +36,7 @@
 }
 
 
-def get_embedding_function(function_name: str, **kwargs) -> EmbeddingFunction:
+def get_embedding_function(function_name: str, **kwargs: Any) -> EmbeddingFunction:
     """Load an embedding function by name.
 
     :param function_name: the name of the embedding function.