langchain_chroma: added document.id support (#27995)

Description:
* Added internal `Document.id` support to Chroma VectorStore

Dependencies:
* #27968 should be merged
first and this PR should be re-based on top of those changes.

Tests:
* Modified/Added tests for `Document.id` support. All tests are passing.


Note: I am not a member of the Chroma team.

---------

Co-authored-by: Eugene Yurtsev <[email protected]>
Co-authored-by: Erick Friis <[email protected]>

libs/partners/chroma/langchain_chroma/vectorstores.py

@@ -44,10 +44,14 @@ def _results_to_docs_and_scores(results: Any) -> List[Tuple[Document, float]]:
     return [
         # TODO: Chroma can do batch querying,
         # we shouldn't hard code to the 1st result
-        (Document(page_content=result[0], metadata=result[1] or {}), result[2])
+        (
+            Document(page_content=result[0], metadata=result[1] or {}, id=result[2]),
+            result[3],
+        )
         for result in zip(
             results["documents"][0],
             results["metadatas"][0],
+            results["ids"][0],
             results["distances"][0],
         )
     ]
@@ -1185,6 +1189,8 @@ def from_documents(
         """
         texts = [doc.page_content for doc in documents]
         metadatas = [doc.metadata for doc in documents]
+        if ids is None:
+            ids = [doc.id if doc.id else "" for doc in documents]
         return cls.from_texts(
             texts=texts,
             embedding=embedding,

libs/partners/chroma/tests/integration_tests/test_vectorstores.py

@@ -46,8 +46,27 @@ def test_chroma() -> None:
     output = docsearch.similarity_search("foo", k=1)
 
     docsearch.delete_collection()
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].id is not None
 
-    assert output == [Document(page_content="foo")]
+
+def test_chroma_with_ids() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    ids = [f"id_{i}" for i in range(len(texts))]
+    docsearch = Chroma.from_texts(
+        collection_name="test_collection",
+        texts=texts,
+        embedding=FakeEmbeddings(),
+        ids=ids,
+    )
+    output = docsearch.similarity_search("foo", k=1)
+
+    docsearch.delete_collection()
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].id == "id_0"
 
 
 async def test_chroma_async() -> None:
@@ -59,7 +78,27 @@ async def test_chroma_async() -> None:
     output = await docsearch.asimilarity_search("foo", k=1)
 
     docsearch.delete_collection()
-    assert output == [Document(page_content="foo")]
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].id is not None
+
+
+async def test_chroma_async_with_ids() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    ids = [f"id_{i}" for i in range(len(texts))]
+    docsearch = Chroma.from_texts(
+        collection_name="test_collection",
+        texts=texts,
+        embedding=FakeEmbeddings(),
+        ids=ids,
+    )
+    output = await docsearch.asimilarity_search("foo", k=1)
+
+    docsearch.delete_collection()
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].id == "id_0"
 
 
 def test_chroma_with_metadatas() -> None:
@@ -74,22 +113,49 @@ def test_chroma_with_metadatas() -> None:
     )
     output = docsearch.similarity_search("foo", k=1)
     docsearch.delete_collection()
-    assert output == [Document(page_content="foo", metadata={"page": "0"})]
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].metadata == {"page": "0"}
+    assert output[0].id is not None
 
 
-def test_chroma_with_metadatas_with_scores() -> None:
+def test_chroma_with_metadatas_and_ids() -> None:
+    """Test end to end construction and search."""
+    texts = ["foo", "bar", "baz"]
+    metadatas = [{"page": str(i)} for i in range(len(texts))]
+    ids = [f"id_{i}" for i in range(len(texts))]
+    docsearch = Chroma.from_texts(
+        collection_name="test_collection",
+        texts=texts,
+        embedding=FakeEmbeddings(),
+        metadatas=metadatas,
+        ids=ids,
+    )
+    output = docsearch.similarity_search("foo", k=1)
+    docsearch.delete_collection()
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].metadata == {"page": "0"}
+    assert output[0].id == "id_0"
+
+
+def test_chroma_with_metadatas_with_scores_and_ids() -> None:
     """Test end to end construction and scored search."""
     texts = ["foo", "bar", "baz"]
     metadatas = [{"page": str(i)} for i in range(len(texts))]
+    ids = [f"id_{i}" for i in range(len(texts))]
     docsearch = Chroma.from_texts(
         collection_name="test_collection",
         texts=texts,
         embedding=FakeEmbeddings(),
         metadatas=metadatas,
+        ids=ids,
     )
     output = docsearch.similarity_search_with_score("foo", k=1)
     docsearch.delete_collection()
-    assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)]
+    assert output == [
+        (Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0)
+    ]
 
 
 def test_chroma_with_metadatas_with_vectors() -> None:
@@ -114,48 +180,60 @@ def test_chroma_with_metadatas_with_scores_using_vector() -> None:
     """Test end to end construction and scored search, using embedding vector."""
     texts = ["foo", "bar", "baz"]
     metadatas = [{"page": str(i)} for i in range(len(texts))]
+    ids = [f"id_{i}" for i in range(len(texts))]
     embeddings = FakeEmbeddings()
 
     docsearch = Chroma.from_texts(
         collection_name="test_collection",
         texts=texts,
         embedding=embeddings,
         metadatas=metadatas,
+        ids=ids,
     )
     embedded_query = embeddings.embed_query("foo")
     output = docsearch.similarity_search_by_vector_with_relevance_scores(
         embedding=embedded_query, k=1
     )
     docsearch.delete_collection()
-    assert output == [(Document(page_content="foo", metadata={"page": "0"}), 0.0)]
+    assert output == [
+        (Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0)
+    ]
 
 
 def test_chroma_search_filter() -> None:
     """Test end to end construction and search with metadata filtering."""
     texts = ["far", "bar", "baz"]
     metadatas = [{"first_letter": "{}".format(text[0])} for text in texts]
+    ids = [f"id_{i}" for i in range(len(texts))]
     docsearch = Chroma.from_texts(
         collection_name="test_collection",
         texts=texts,
         embedding=FakeEmbeddings(),
         metadatas=metadatas,
+        ids=ids,
     )
     output1 = docsearch.similarity_search("far", k=1, filter={"first_letter": "f"})
     output2 = docsearch.similarity_search("far", k=1, filter={"first_letter": "b"})
     docsearch.delete_collection()
-    assert output1 == [Document(page_content="far", metadata={"first_letter": "f"})]
-    assert output2 == [Document(page_content="bar", metadata={"first_letter": "b"})]
+    assert output1 == [
+        Document(page_content="far", metadata={"first_letter": "f"}, id="id_0")
+    ]
+    assert output2 == [
+        Document(page_content="bar", metadata={"first_letter": "b"}, id="id_1")
+    ]
 
 
 def test_chroma_search_filter_with_scores() -> None:
     """Test end to end construction and scored search with metadata filtering."""
     texts = ["far", "bar", "baz"]
     metadatas = [{"first_letter": "{}".format(text[0])} for text in texts]
+    ids = [f"id_{i}" for i in range(len(texts))]
     docsearch = Chroma.from_texts(
         collection_name="test_collection",
         texts=texts,
         embedding=FakeEmbeddings(),
         metadatas=metadatas,
+        ids=ids,
     )
     output1 = docsearch.similarity_search_with_score(
         "far", k=1, filter={"first_letter": "f"}
@@ -165,10 +243,10 @@ def test_chroma_search_filter_with_scores() -> None:
     )
     docsearch.delete_collection()
     assert output1 == [
-        (Document(page_content="far", metadata={"first_letter": "f"}), 0.0)
+        (Document(page_content="far", metadata={"first_letter": "f"}, id="id_0"), 0.0)
     ]
     assert output2 == [
-        (Document(page_content="bar", metadata={"first_letter": "b"}), 1.0)
+        (Document(page_content="bar", metadata={"first_letter": "b"}, id="id_1"), 1.0)
     ]
 
 
@@ -177,15 +255,18 @@ def test_chroma_with_persistence() -> None:
     chroma_persist_dir = "./tests/persist_dir"
     collection_name = "test_collection"
     texts = ["foo", "bar", "baz"]
+    ids = [f"id_{i}" for i in range(len(texts))]
+
     docsearch = Chroma.from_texts(
         collection_name=collection_name,
         texts=texts,
         embedding=FakeEmbeddings(),
         persist_directory=chroma_persist_dir,
+        ids=ids,
     )
 
     output = docsearch.similarity_search("foo", k=1)
-    assert output == [Document(page_content="foo")]
+    assert output == [Document(page_content="foo", id="id_0")]
 
     # Get a new VectorStore from the persisted directory
     docsearch = Chroma(
@@ -194,6 +275,7 @@ def test_chroma_with_persistence() -> None:
         persist_directory=chroma_persist_dir,
     )
     output = docsearch.similarity_search("foo", k=1)
+    assert output == [Document(page_content="foo", id="id_0")]
 
     # Clean up
     docsearch.delete_collection()
@@ -211,7 +293,9 @@ def test_chroma_mmr() -> None:
     )
     output = docsearch.max_marginal_relevance_search("foo", k=1)
     docsearch.delete_collection()
-    assert output == [Document(page_content="foo")]
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].id is not None
 
 
 def test_chroma_mmr_by_vector() -> None:
@@ -224,7 +308,9 @@ def test_chroma_mmr_by_vector() -> None:
     embedded_query = embeddings.embed_query("foo")
     output = docsearch.max_marginal_relevance_search_by_vector(embedded_query, k=1)
     docsearch.delete_collection()
-    assert output == [Document(page_content="foo")]
+    assert len(output) == 1
+    assert output[0].page_content == "foo"
+    assert output[0].id is not None
 
 
 def test_chroma_with_include_parameter() -> None:
@@ -241,7 +327,10 @@ def test_chroma_with_include_parameter() -> None:
 
 
 def test_chroma_update_document() -> None:
-    """Test the update_document function in the Chroma class."""
+    """Test the update_document function in the Chroma class.
+
+    Uses an external document id.
+    """
     # Make a consistent embedding
     embedding = ConsistentFakeEmbeddings()
 
@@ -283,7 +372,66 @@ def test_chroma_update_document() -> None:
     docsearch.delete_collection()
 
     # Assert that the updated document is returned by the search
-    assert output == [Document(page_content=updated_content, metadata={"page": "0"})]
+    assert output == [
+        Document(page_content=updated_content, metadata={"page": "0"}, id=document_id)
+    ]
+
+    assert list(new_embedding) == list(embedding.embed_documents([updated_content])[0])
+    assert list(new_embedding) != list(old_embedding)
+
+
+def test_chroma_update_document_with_id() -> None:
+    """Test the update_document function in the Chroma class.
+
+    Uses an internal document id.
+    """
+    # Make a consistent embedding
+    embedding = ConsistentFakeEmbeddings()
+
+    # Initial document content and id
+    initial_content = "foo"
+    document_id = "doc1"
+
+    # Create an instance of Document with initial content and metadata
+    original_doc = Document(
+        page_content=initial_content, metadata={"page": "0"}, id=document_id
+    )
+
+    # Initialize a Chroma instance with the original document
+    docsearch = Chroma.from_documents(
+        collection_name="test_collection",
+        documents=[original_doc],
+        embedding=embedding,
+    )
+    old_embedding = docsearch._collection.peek()["embeddings"][  # type: ignore
+        docsearch._collection.peek()["ids"].index(document_id)
+    ]
+
+    # Define updated content for the document
+    updated_content = "updated foo"
+
+    # Create a new Document instance with the updated content and the same id
+    updated_doc = Document(
+        page_content=updated_content, metadata={"page": "0"}, id=document_id
+    )
+
+    # Update the document in the Chroma instance
+    docsearch.update_document(document_id=document_id, document=updated_doc)
+
+    # Perform a similarity search with the updated content
+    output = docsearch.similarity_search(updated_content, k=1)
+
+    # Assert that the new embedding is correct
+    new_embedding = docsearch._collection.peek()["embeddings"][  # type: ignore
+        docsearch._collection.peek()["ids"].index(document_id)
+    ]
+
+    docsearch.delete_collection()
+
+    # Assert that the updated document is returned by the search
+    assert output == [
+        Document(page_content=updated_content, metadata={"page": "0"}, id=document_id)
+    ]
 
     assert list(new_embedding) == list(embedding.embed_documents([updated_content])[0])
     assert list(new_embedding) != list(old_embedding)
@@ -294,20 +442,22 @@ def test_chroma_with_relevance_score_custom_normalization_fn() -> None:
     """Test searching with relevance score and custom normalization function."""
     texts = ["foo", "bar", "baz"]
     metadatas = [{"page": str(i)} for i in range(len(texts))]
+    ids = [f"id_{i}" for i in range(len(texts))]
     docsearch = Chroma.from_texts(
         collection_name="test1_collection",
         texts=texts,
         embedding=FakeEmbeddings(),
         metadatas=metadatas,
+        ids=ids,
         relevance_score_fn=lambda d: d * 0,
         collection_metadata={"hnsw:space": "l2"},
     )
     output = docsearch.similarity_search_with_relevance_scores("foo", k=3)
     docsearch.delete_collection()
     assert output == [
-        (Document(page_content="foo", metadata={"page": "0"}), 0.0),
-        (Document(page_content="bar", metadata={"page": "1"}), 0.0),
-        (Document(page_content="baz", metadata={"page": "2"}), 0.0),
+        (Document(page_content="foo", metadata={"page": "0"}, id="id_0"), 0.0),
+        (Document(page_content="bar", metadata={"page": "1"}, id="id_1"), 0.0),
+        (Document(page_content="baz", metadata={"page": "2"}, id="id_2"), 0.0),
     ]
 
 
@@ -332,11 +482,11 @@ def test_chroma_add_documents_no_metadata() -> None:
 def test_chroma_add_documents_mixed_metadata() -> None:
     db = Chroma(embedding_function=FakeEmbeddings())
     docs = [
-        Document(page_content="foo"),
-        Document(page_content="bar", metadata={"baz": 1}),
+        Document(page_content="foo", id="0"),
+        Document(page_content="bar", metadata={"baz": 1}, id="1"),
     ]
     ids = ["0", "1"]
-    actual_ids = db.add_documents(docs, ids=ids)
+    actual_ids = db.add_documents(docs)
     search = db.similarity_search("foo bar")
     db.delete_collection()