Chroma fix mmr (#3897)

Fixes #3628, thanks @derekmoeller for the issue!

langchain/vectorstores/chroma.py

@@ -104,8 +104,17 @@ class Chroma(VectorStore):
         query_embeddings: Optional[List[List[float]]] = None,
         n_results: int = 4,
         where: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
     ) -> List[Document]:
         """Query the chroma collection."""
+        try:
+            import chromadb
+        except ImportError:
+            raise ValueError(
+                "Could not import chromadb python package. "
+                "Please install it with `pip install chromadb`."
+            )
+
         for i in range(n_results, 0, -1):
             try:
                 return self._collection.query(
@@ -113,6 +122,7 @@ class Chroma(VectorStore):
                     query_embeddings=query_embeddings,
                     n_results=i,
                     where=where,
+                    **kwargs,
                 )
             except chromadb.errors.NotEnoughElementsException:
                 logger.error(

tests/integration_tests/vectorstores/test_chroma.py

@@ -126,3 +126,25 @@ def test_chroma_with_persistence() -> None:
     # Persist doesn't need to be called again
     # Data will be automatically persisted on object deletion
     # Or on program exit
+
+
+def test_chroma_mmr() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    docsearch = Chroma.from_texts(
+        collection_name="test_collection", texts=texts, embedding=FakeEmbeddings()
+    )
+    output = docsearch.max_marginal_relevance_search("foo", k=1)
+    assert output == [Document(page_content="foo")]
+
+
+def test_chroma_mmr_by_vector() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    embeddings = FakeEmbeddings()
+    docsearch = Chroma.from_texts(
+        collection_name="test_collection", texts=texts, embedding=embeddings
+    )
+    embedded_query = embeddings.embed_query("foo")
+    output = docsearch.max_marginal_relevance_search_by_vector(embedded_query, k=1)
+    assert output == [Document(page_content="foo")]