Apply patch [skip ci]

This commit is contained in:
open-swe[bot]
2025-09-10 00:09:45 +00:00
parent 1f2e074138
commit fcccd0ad6a
3 changed files with 125 additions and 131 deletions

View File

@@ -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)

View File

@@ -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()

View File

@@ -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