diff --git a/libs/partners/chroma/tests/integration_tests/test_vectorstores.py b/libs/partners/chroma/tests/integration_tests/test_vectorstores.py index 99943c2f5a3..8b039231fc9 100644 --- a/libs/partners/chroma/tests/integration_tests/test_vectorstores.py +++ b/libs/partners/chroma/tests/integration_tests/test_vectorstores.py @@ -124,6 +124,230 @@ async def test_chroma_async_with_ids() -> None: assert output[0].id == "id_0" +async def test_chroma_async_http_client() -> None: + """Test async operations with AsyncHttpClient.""" + try: + import chromadb + from chromadb import AsyncHttpClient + except ImportError: + pytest.skip("AsyncHttpClient not available in this ChromaDB version") + + # Try to connect to a local Chroma server + # This test requires a running Chroma server + try: + async_client = await AsyncHttpClient( + host="localhost", + port=8000, + ) + # Test connection + await async_client.heartbeat() + except Exception: + pytest.skip("Chroma server not running on localhost:8000") + + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + + # Create Chroma instance with async client + collection_name = f"test_async_{uuid.uuid4().hex}" + docsearch = Chroma( + collection_name=collection_name, + embedding_function=FakeEmbeddings(), + async_client=async_client, + ) + + # Test async add_texts + ids = await docsearch.aadd_texts(texts=texts, metadatas=metadatas) + assert len(ids) == 3 + + # Test async similarity search + output = await docsearch.asimilarity_search("foo", k=2) + assert len(output) <= 2 + assert output[0].page_content in texts + + # Test async similarity search with score + output_with_scores = await docsearch.asimilarity_search_with_score("bar", k=2) + assert len(output_with_scores) <= 2 + for doc, score in output_with_scores: + assert doc.page_content in texts + assert isinstance(score, float) + + # Test async delete + await docsearch.adelete(ids=[ids[0]]) + + # Clean up + await docsearch.adelete_collection() + + +async def test_chroma_async_add_documents() -> None: + """Test async add_documents functionality.""" + documents = [ + Document(page_content="foo", metadata={"page": "0"}), + Document(page_content="bar", metadata={"page": "1"}), + Document(page_content="baz", metadata={"page": "2"}), + ] + + docsearch = Chroma.from_documents( + collection_name="test_async_docs", + documents=documents, + embedding=FakeEmbeddings(), + ) + + # Add more documents asynchronously + new_documents = [ + Document(page_content="qux", metadata={"page": "3"}), + Document(page_content="quux", metadata={"page": "4"}), + ] + ids = await docsearch.aadd_documents(documents=new_documents) + assert len(ids) == 2 + + # Verify documents were added + output = await docsearch.asimilarity_search("qux", k=1) + assert len(output) == 1 + assert output[0].page_content == "qux" + assert output[0].metadata == {"page": "3"} + + docsearch.delete_collection() + + +async def test_chroma_both_sync_and_async_clients() -> None: + """Test that both sync and async clients can coexist.""" + # Create sync client + sync_client = chromadb.Client(chromadb.config.Settings()) + + # Try to create async client if available + try: + from chromadb import AsyncHttpClient + async_client = await AsyncHttpClient( + host="localhost", + port=8000, + ) + await async_client.heartbeat() + except (ImportError, Exception): + # If AsyncHttpClient is not available or server not running, + # create a mock async client for testing + async_client = None + pytest.skip("AsyncHttpClient not available or server not running") + + collection_name = f"test_both_{uuid.uuid4().hex}" + + # Initialize with both clients + docsearch = Chroma( + collection_name=collection_name, + embedding_function=FakeEmbeddings(), + client=sync_client, + async_client=async_client, + ) + + texts = ["foo", "bar", "baz"] + + # Test sync operations + sync_ids = docsearch.add_texts(texts=texts) + assert len(sync_ids) == 3 + + sync_results = docsearch.similarity_search("foo", k=1) + assert len(sync_results) == 1 + + # Test async operations + async_texts = ["async_foo", "async_bar"] + async_ids = await docsearch.aadd_texts(texts=async_texts) + assert len(async_ids) == 2 + + async_results = await docsearch.asimilarity_search("async_foo", k=1) + assert len(async_results) == 1 + assert async_results[0].page_content == "async_foo" + + # Clean up + docsearch.delete_collection() + + +async def test_chroma_async_error_handling() -> None: + """Test error handling in async operations.""" + # Test without async client + docsearch = Chroma( + collection_name="test_error", + embedding_function=FakeEmbeddings(), + ) + + # Should raise error when trying to use async methods without async_client + with pytest.raises(ValueError, match="async_client"): + await docsearch.aadd_texts(["test"]) + + with pytest.raises(ValueError, match="async_client"): + await docsearch.asimilarity_search("test") + + with pytest.raises(ValueError, match="async_client"): + await docsearch.asimilarity_search_with_score("test") + + with pytest.raises(ValueError, match="async_client"): + await docsearch.adelete(["test_id"]) + + docsearch.delete_collection() + + +async def test_chroma_async_collection_operations() -> None: + """Test async collection management operations.""" + try: + from chromadb import AsyncHttpClient + async_client = await AsyncHttpClient( + host="localhost", + port=8000, + ) + await async_client.heartbeat() + except (ImportError, Exception): + pytest.skip("AsyncHttpClient not available or server not running") + + collection_name = f"test_collection_ops_{uuid.uuid4().hex}" + + docsearch = Chroma( + collection_name=collection_name, + embedding_function=FakeEmbeddings(), + async_client=async_client, + ) + + # Add initial data + texts = ["foo", "bar", "baz"] + await docsearch.aadd_texts(texts=texts) + + # Test reset collection + await docsearch.areset_collection() + + # Collection should be empty after reset + # Add new data to verify collection was reset + new_texts = ["new_foo"] + ids = await docsearch.aadd_texts(texts=new_texts) + assert len(ids) == 1 + + # Clean up + await docsearch.adelete_collection() + + +async def test_chroma_async_with_metadata_filter() -> None: + """Test async operations with metadata filtering.""" + texts = ["foo", "bar", "baz", "qux"] + metadatas = [ + {"type": "a", "page": "0"}, + {"type": "b", "page": "1"}, + {"type": "a", "page": "2"}, + {"type": "b", "page": "3"}, + ] + + docsearch = Chroma.from_texts( + collection_name="test_async_filter", + texts=texts, + embedding=FakeEmbeddings(), + metadatas=metadatas, + ) + + # Test async search with filter + output = await docsearch.asimilarity_search("foo", k=10, filter={"type": "a"}) + + # Should only return documents with type="a" + for doc in output: + assert doc.metadata.get("type") == "a" + + docsearch.delete_collection() + + def test_chroma_with_metadatas() -> None: """Test end to end construction and search.""" texts = ["foo", "bar", "baz"] @@ -827,3 +1051,4 @@ def test_delete_where_clause(client: chromadb.ClientAPI) -> None: assert vectorstore._collection.count() == 1 # Clean up vectorstore.delete_collection() +