test_chroma_update_document_with_id() — langchain Function Reference

Source Code

libs/partners/chroma/tests/integration_tests/test_vectorstores.py lines 475–533

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[index]
        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[index]
        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)