From fcccd0ad6a61ced9ab8ffaa759f4855656efde4a Mon Sep 17 00:00:00 2001 From: "open-swe[bot]" Date: Wed, 10 Sep 2025 00:09:45 +0000 Subject: [PATCH] Apply patch [skip ci] --- .../chroma/langchain_chroma/vectorstores.py | 59 ++++----- .../integration_tests/test_vectorstores.py | 75 +++++------ .../unit_tests/test_async_vectorstores.py | 122 +++++++++--------- 3 files changed, 125 insertions(+), 131 deletions(-) diff --git a/libs/partners/chroma/langchain_chroma/vectorstores.py b/libs/partners/chroma/langchain_chroma/vectorstores.py index 9cc0d2f3dee..67ccf062883 100644 --- a/libs/partners/chroma/langchain_chroma/vectorstores.py +++ b/libs/partners/chroma/langchain_chroma/vectorstores.py @@ -419,7 +419,7 @@ class Chroma(VectorStore): self._collection_name = collection_name self._collection_metadata = collection_metadata self._collection_configuration = collection_configuration - + # Only initialize collection synchronously if we have a sync client and not using async if create_collection_if_not_exists: if self._async_client is None and self._client is not None: @@ -429,8 +429,10 @@ class Chroma(VectorStore): self.__ensure_collection() else: if self._client is not None: - self._chroma_collection = self._client.get_collection(name=collection_name) - + self._chroma_collection = self._client.get_collection( + name=collection_name + ) + self.override_relevance_score_fn = relevance_score_fn def __ensure_collection(self) -> None: @@ -456,13 +458,15 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + # Get or create the collection asynchronously - self._async_chroma_collection = await self._async_client.get_or_create_collection( - name=self._collection_name, - embedding_function=None, - metadata=self._collection_metadata, - configuration=self._collection_configuration, + self._async_chroma_collection = ( + await self._async_client.get_or_create_collection( + name=self._collection_name, + embedding_function=None, + metadata=self._collection_metadata, + configuration=self._collection_configuration, + ) ) self._async_initialized = True @@ -474,17 +478,20 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + if not self._async_initialized: await self._aensure_collection() - - if not hasattr(self, '_async_chroma_collection') or self._async_chroma_collection is None: + + if ( + not hasattr(self, "_async_chroma_collection") + or self._async_chroma_collection is None + ): msg = ( "Async Chroma collection not initialized. " "Call _aensure_collection() first." ) raise ValueError(msg) - + return self._async_chroma_collection @property @@ -805,9 +812,9 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + collection = await self._aget_collection() - + if ids is None: ids = [str(uuid.uuid4()) for _ in texts] else: @@ -1056,7 +1063,7 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + if self._embedding_function is None: results = await self.__aquery_collection( query_texts=[query], @@ -1362,7 +1369,7 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + if self._async_initialized: collection = await self._aget_collection() await self._async_client.delete_collection(collection.name) @@ -1388,7 +1395,7 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + await self.adelete_collection() await self._aensure_collection() @@ -1728,7 +1735,7 @@ class Chroma(VectorStore): "Provide an async_client when initializing the Chroma instance." ) raise ValueError(msg) - + collection = await self._aget_collection() await collection.delete(ids=ids, **kwargs) @@ -1749,17 +1756,3 @@ class Chroma(VectorStore): texts = [doc.page_content for doc in documents] metadatas = [doc.metadata for doc in documents] return await self.aadd_texts(texts, metadatas=metadatas, **kwargs) - - - - - - - - - - - - - - diff --git a/libs/partners/chroma/tests/integration_tests/test_vectorstores.py b/libs/partners/chroma/tests/integration_tests/test_vectorstores.py index 8b039231fc9..d510a4ec9a9 100644 --- a/libs/partners/chroma/tests/integration_tests/test_vectorstores.py +++ b/libs/partners/chroma/tests/integration_tests/test_vectorstores.py @@ -131,7 +131,7 @@ async def test_chroma_async_http_client() -> None: 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: @@ -143,10 +143,10 @@ async def test_chroma_async_http_client() -> None: 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( @@ -154,26 +154,26 @@ async def test_chroma_async_http_client() -> None: 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() @@ -185,13 +185,13 @@ async def test_chroma_async_add_documents() -> None: 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"}), @@ -199,13 +199,13 @@ async def test_chroma_async_add_documents() -> None: ] 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() @@ -213,10 +213,11 @@ 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, @@ -227,9 +228,9 @@ async def test_chroma_both_sync_and_async_clients() -> None: # 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, @@ -237,25 +238,25 @@ async def test_chroma_both_sync_and_async_clients() -> None: 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() @@ -267,20 +268,20 @@ async def test_chroma_async_error_handling() -> None: 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() @@ -288,6 +289,7 @@ 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, @@ -295,28 +297,28 @@ async def test_chroma_async_collection_operations() -> None: 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() @@ -330,21 +332,21 @@ async def test_chroma_async_with_metadata_filter() -> None: {"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() @@ -1051,4 +1053,3 @@ def test_delete_where_clause(client: chromadb.ClientAPI) -> None: assert vectorstore._collection.count() == 1 # Clean up vectorstore.delete_collection() - diff --git a/libs/partners/chroma/tests/unit_tests/test_async_vectorstores.py b/libs/partners/chroma/tests/unit_tests/test_async_vectorstores.py index 7b7dbeec109..b2b7f245d26 100644 --- a/libs/partners/chroma/tests/unit_tests/test_async_vectorstores.py +++ b/libs/partners/chroma/tests/unit_tests/test_async_vectorstores.py @@ -2,7 +2,7 @@ import asyncio from typing import Any, Optional -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest from langchain_core.documents import Document @@ -61,7 +61,7 @@ class MockAsyncCollection: documents = [self.data[id_]["document"] for id_ in ids] metadatas = [self.data[id_]["metadata"] for id_ in ids] distances = [[0.1 * i for i in range(len(ids))]] - + return { "ids": [ids], "documents": [documents], @@ -104,14 +104,14 @@ class MockAsyncClient: async def test_async_client_initialization(): """Test that Chroma can be initialized with an async client.""" async_client = MockAsyncClient() - + # Initialize with async client only chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + assert chroma._async_client is not None assert chroma._client is None assert chroma._async_initialized is False @@ -121,16 +121,16 @@ async def test_async_client_initialization(): async def test_async_collection_initialization(): """Test async collection initialization.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Collection should be initialized on first async operation await chroma._aensure_collection() - + assert chroma._async_initialized is True assert hasattr(chroma, "_async_chroma_collection") @@ -139,19 +139,19 @@ async def test_async_collection_initialization(): async def test_aadd_texts(): """Test async add_texts functionality.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + texts = ["foo", "bar", "baz"] metadatas = [{"page": str(i)} for i in range(len(texts))] - + # Add texts asynchronously ids = await chroma.aadd_texts(texts=texts, metadatas=metadatas) - + assert len(ids) == 3 assert chroma._async_initialized is True @@ -160,22 +160,22 @@ async def test_aadd_texts(): async def test_aadd_documents(): """Test async add_documents functionality.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + documents = [ Document(page_content="foo", metadata={"page": "0"}), Document(page_content="bar", metadata={"page": "1"}), Document(page_content="baz", metadata={"page": "2"}), ] - + # Add documents asynchronously ids = await chroma.aadd_documents(documents=documents) - + assert len(ids) == 3 assert chroma._async_initialized is True @@ -184,21 +184,21 @@ async def test_aadd_documents(): async def test_asimilarity_search(): """Test async similarity search.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add some texts first texts = ["foo", "bar", "baz"] metadatas = [{"page": str(i)} for i in range(len(texts))] await chroma.aadd_texts(texts=texts, metadatas=metadatas) - + # Perform similarity search results = await chroma.asimilarity_search("foo", k=2) - + assert len(results) <= 2 assert all(isinstance(doc, Document) for doc in results) @@ -207,21 +207,21 @@ async def test_asimilarity_search(): async def test_asimilarity_search_with_score(): """Test async similarity search with scores.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add some texts first texts = ["foo", "bar", "baz"] metadatas = [{"page": str(i)} for i in range(len(texts))] await chroma.aadd_texts(texts=texts, metadatas=metadatas) - + # Perform similarity search with scores results = await chroma.asimilarity_search_with_score("foo", k=2) - + assert len(results) <= 2 for doc, score in results: assert isinstance(doc, Document) @@ -232,20 +232,20 @@ async def test_asimilarity_search_with_score(): async def test_adelete(): """Test async delete functionality.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add some texts first texts = ["foo", "bar", "baz"] ids = await chroma.aadd_texts(texts=texts) - + # Delete specific IDs await chroma.adelete(ids=[ids[0]]) - + # Verify deletion worked (this would need actual verification in real tests) assert chroma._async_initialized is True @@ -254,20 +254,20 @@ async def test_adelete(): async def test_adelete_collection(): """Test async delete collection functionality.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Initialize collection await chroma._aensure_collection() assert chroma._async_initialized is True - + # Delete collection await chroma.adelete_collection() - + assert chroma._async_initialized is False assert chroma._async_chroma_collection is None @@ -276,20 +276,20 @@ async def test_adelete_collection(): async def test_areset_collection(): """Test async reset collection functionality.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add some texts texts = ["foo", "bar", "baz"] await chroma.aadd_texts(texts=texts) - + # Reset collection await chroma.areset_collection() - + # Collection should be re-initialized but empty assert chroma._async_initialized is True @@ -302,23 +302,23 @@ async def test_error_without_async_client(): collection_name="test_collection", embedding_function=FakeEmbeddings(size=10), ) - + # All async methods should raise ValueError with pytest.raises(ValueError, match="async_client"): await chroma.aadd_texts(["test"]) - + with pytest.raises(ValueError, match="async_client"): await chroma.asimilarity_search("test") - + with pytest.raises(ValueError, match="async_client"): await chroma.asimilarity_search_with_score("test") - + with pytest.raises(ValueError, match="async_client"): await chroma.adelete(["test_id"]) - + with pytest.raises(ValueError, match="async_client"): await chroma.adelete_collection() - + with pytest.raises(ValueError, match="async_client"): await chroma.areset_collection() @@ -327,17 +327,17 @@ async def test_error_without_async_client(): async def test_sync_methods_error_with_only_async_client(): """Test that sync methods raise errors when only async_client is provided.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=FakeEmbeddings(size=10), async_client=async_client, ) - + # Sync methods that require collection should raise ValueError with pytest.raises(ValueError, match="async_client"): chroma._collection # Accessing sync collection property - + with pytest.raises(ValueError, match="async_client"): chroma.add_texts(["test"]) @@ -346,9 +346,9 @@ def test_both_sync_and_async_clients(): """Test that both sync and async clients can be provided.""" sync_client = MagicMock() sync_client.get_or_create_collection.return_value = MagicMock() - + async_client = MockAsyncClient() - + # Initialize with both clients chroma = Chroma( collection_name="test_collection", @@ -356,10 +356,10 @@ def test_both_sync_and_async_clients(): client=sync_client, async_client=async_client, ) - + assert chroma._client is not None assert chroma._async_client is not None - + # Sync collection should be initialized assert chroma._chroma_collection is not None @@ -368,13 +368,13 @@ def test_both_sync_and_async_clients(): async def test_async_with_metadata_filtering(): """Test async operations with metadata filtering.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add texts with metadata texts = ["foo", "bar", "baz", "qux"] metadatas = [ @@ -384,10 +384,10 @@ async def test_async_with_metadata_filtering(): {"type": "b", "page": "3"}, ] await chroma.aadd_texts(texts=texts, metadatas=metadatas) - + # Search with filter results = await chroma.asimilarity_search("foo", k=2, filter={"type": "a"}) - + assert len(results) <= 2 @@ -395,19 +395,19 @@ async def test_async_with_metadata_filtering(): async def test_async_empty_metadata_handling(): """Test async operations with empty metadata.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add texts with mixed empty and non-empty metadata texts = ["foo", "bar", "baz"] metadatas = [{"page": "0"}, {}, {"page": "2"}] - + ids = await chroma.aadd_texts(texts=texts, metadatas=metadatas) - + assert len(ids) == 3 @@ -415,25 +415,25 @@ async def test_async_empty_metadata_handling(): async def test_concurrent_async_operations(): """Test that multiple async operations can run concurrently.""" async_client = MockAsyncClient() - + chroma = Chroma( collection_name="test_collection", embedding_function=AsyncFakeEmbeddings(size=10), async_client=async_client, ) - + # Add initial data await chroma.aadd_texts(["initial"]) - + # Run multiple operations concurrently tasks = [ chroma.aadd_texts(["text1"]), chroma.aadd_texts(["text2"]), chroma.asimilarity_search("initial"), ] - + results = await asyncio.gather(*tasks) - + assert len(results) == 3 assert all(results[0]) # First add_texts returned IDs assert all(results[1]) # Second add_texts returned IDs