chore: 添加虚拟环境到仓库
- 添加 backend_service/venv 虚拟环境 - 包含所有Python依赖包 - 注意：虚拟环境约393MB，包含12655个文件
This commit is contained in:
@@ -0,0 +1,620 @@
|
||||
import gc
|
||||
import math
|
||||
import os.path
|
||||
from uuid import UUID
|
||||
from contextlib import contextmanager
|
||||
|
||||
from chromadb.api.segment import SegmentAPI
|
||||
from chromadb.db.system import SysDB
|
||||
from chromadb.ingest.impl.utils import create_topic_name
|
||||
|
||||
from chromadb.config import System
|
||||
from chromadb.db.base import get_sql
|
||||
from chromadb.db.impl.sqlite import SqliteDB
|
||||
from time import sleep
|
||||
import psutil
|
||||
|
||||
from chromadb.segment import SegmentType
|
||||
from chromadb.test.property.strategies import NormalizedRecordSet, RecordSet
|
||||
from typing import Callable, Optional, Tuple, Union, List, TypeVar, cast, Any, Dict
|
||||
from typing_extensions import Literal
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
from chromadb.api import types, ClientAPI
|
||||
from chromadb.api.models.Collection import Collection
|
||||
from hypothesis import note
|
||||
from hypothesis.errors import InvalidArgument
|
||||
from pypika import Table, functions
|
||||
|
||||
from chromadb.utils import distance_functions
|
||||
from chromadb.execution.expression.plan import Search
|
||||
from chromadb.execution.expression.operator import Knn, Select, Limit, Key
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def wrap(value: Union[T, List[T]]) -> List[T]:
|
||||
"""Wrap a value in a list if it is not a list"""
|
||||
if value is None:
|
||||
raise InvalidArgument("value cannot be None")
|
||||
elif isinstance(value, List):
|
||||
return value
|
||||
else:
|
||||
return [value]
|
||||
|
||||
|
||||
def wrap_all(record_set: RecordSet) -> NormalizedRecordSet:
|
||||
"""Ensure that an embedding set has lists for all its values"""
|
||||
|
||||
embedding_list: Optional[types.Embeddings]
|
||||
if record_set["embeddings"] is None:
|
||||
embedding_list = None
|
||||
elif isinstance(record_set["embeddings"], list):
|
||||
assert record_set["embeddings"] is not None
|
||||
if len(record_set["embeddings"]) > 0:
|
||||
if all(
|
||||
isinstance(embedding, list) for embedding in record_set["embeddings"]
|
||||
):
|
||||
embedding_list = cast(types.Embeddings, record_set["embeddings"])
|
||||
elif all(
|
||||
isinstance(embedding, np.ndarray)
|
||||
for embedding in record_set["embeddings"]
|
||||
):
|
||||
embedding_list = cast(types.Embeddings, record_set["embeddings"])
|
||||
else:
|
||||
if all(
|
||||
isinstance(e, (int, float, np.integer, np.floating))
|
||||
for e in record_set["embeddings"]
|
||||
):
|
||||
embedding_list = cast(types.Embeddings, [record_set["embeddings"]])
|
||||
else:
|
||||
raise InvalidArgument(
|
||||
"an embedding must be a list of floats or ints"
|
||||
)
|
||||
else:
|
||||
embedding_list = cast(types.Embeddings, record_set["embeddings"])
|
||||
else:
|
||||
raise InvalidArgument(
|
||||
"embeddings must be a list of lists, a list of numpy arrays, a list of numbers, or None"
|
||||
)
|
||||
|
||||
return {
|
||||
"ids": wrap(record_set["ids"]),
|
||||
"documents": wrap(record_set["documents"])
|
||||
if record_set["documents"] is not None
|
||||
else None,
|
||||
"metadatas": wrap(record_set["metadatas"])
|
||||
if record_set["metadatas"] is not None
|
||||
else None,
|
||||
"embeddings": embedding_list,
|
||||
}
|
||||
|
||||
|
||||
def check_metadata(
|
||||
expected: Optional[types.Metadata], got: Optional[types.Metadata]
|
||||
) -> None:
|
||||
assert (expected is None and got is None) or (
|
||||
expected is not None and got is not None
|
||||
)
|
||||
if expected is not None and got is not None:
|
||||
assert len(expected) == len(got)
|
||||
for key, val in expected.items():
|
||||
assert key in got
|
||||
if isinstance(expected[key], float) and isinstance(got[key], float):
|
||||
assert abs(cast(float, expected[key]) - cast(float, got[key])) < 1e-6
|
||||
else:
|
||||
assert expected[key] == got[key]
|
||||
|
||||
|
||||
def count(collection: Collection, record_set: RecordSet) -> None:
|
||||
"""The given collection count is equal to the number of embeddings"""
|
||||
count = collection.count()
|
||||
normalized_record_set = wrap_all(record_set)
|
||||
if count != len(normalized_record_set["ids"]):
|
||||
print("count mismatch:", count, "=!", len(normalized_record_set["ids"]))
|
||||
assert count == len(normalized_record_set["ids"])
|
||||
|
||||
|
||||
def _field_matches(
|
||||
collection: Collection,
|
||||
normalized_record_set: NormalizedRecordSet,
|
||||
field_name: Union[
|
||||
Literal["documents"], Literal["metadatas"], Literal["embeddings"]
|
||||
],
|
||||
) -> None:
|
||||
"""
|
||||
The actual embedding field is equal to the expected field
|
||||
field_name: one of [documents, metadatas]
|
||||
"""
|
||||
result = collection.get(ids=normalized_record_set["ids"], include=[field_name]) # type: ignore[list-item]
|
||||
# The test_out_of_order_ids test fails because of this in test_add.py
|
||||
# Here we sort by the ids to match the input order
|
||||
embedding_id_to_index = {id: i for i, id in enumerate(normalized_record_set["ids"])}
|
||||
actual_field = result[field_name]
|
||||
|
||||
if len(normalized_record_set["ids"]) == 0:
|
||||
if field_name == "embeddings":
|
||||
assert cast(npt.NDArray[Any], actual_field).size == 0
|
||||
else:
|
||||
assert actual_field == []
|
||||
return
|
||||
|
||||
# This assert should never happen, if we include metadatas/documents it will be
|
||||
# [None, None..] if there is no metadata. It will not be just None.
|
||||
assert actual_field is not None
|
||||
sorted_field = sorted(
|
||||
enumerate(actual_field),
|
||||
key=lambda index_and_field_value: embedding_id_to_index[
|
||||
result["ids"][index_and_field_value[0]]
|
||||
],
|
||||
)
|
||||
field_values = [field_value for _, field_value in sorted_field]
|
||||
|
||||
expected_field = normalized_record_set[field_name]
|
||||
if expected_field is None:
|
||||
# Since an RecordSet is the user input, we need to convert the documents to
|
||||
# a List since thats what the API returns -> none per entry
|
||||
expected_field = [None] * len(normalized_record_set["ids"]) # type: ignore
|
||||
if field_name == "embeddings":
|
||||
assert np.allclose(np.array(field_values), np.array(expected_field))
|
||||
else:
|
||||
assert len(field_values) == len(expected_field)
|
||||
|
||||
for field_value, expected_field in zip(field_values, expected_field):
|
||||
if isinstance(expected_field, dict):
|
||||
check_metadata(
|
||||
cast(types.Metadata, field_value),
|
||||
cast(types.Metadata, expected_field),
|
||||
)
|
||||
else:
|
||||
assert field_value == expected_field
|
||||
|
||||
|
||||
def ids_match(collection: Collection, record_set: RecordSet) -> None:
|
||||
"""The actual embedding ids is equal to the expected ids"""
|
||||
normalized_record_set = wrap_all(record_set)
|
||||
actual_ids = collection.get(ids=normalized_record_set["ids"], include=[])["ids"]
|
||||
# The test_out_of_order_ids test fails because of this in test_add.py
|
||||
# Here we sort the ids to match the input order
|
||||
embedding_id_to_index = {id: i for i, id in enumerate(normalized_record_set["ids"])}
|
||||
actual_ids = sorted(actual_ids, key=lambda id: embedding_id_to_index[id])
|
||||
assert actual_ids == normalized_record_set["ids"]
|
||||
|
||||
|
||||
def metadatas_match(collection: Collection, record_set: RecordSet) -> None:
|
||||
"""The actual embedding metadata is equal to the expected metadata"""
|
||||
normalized_record_set = wrap_all(record_set)
|
||||
_field_matches(collection, normalized_record_set, "metadatas")
|
||||
|
||||
|
||||
def documents_match(collection: Collection, record_set: RecordSet) -> None:
|
||||
"""The actual embedding documents is equal to the expected documents"""
|
||||
normalized_record_set = wrap_all(record_set)
|
||||
_field_matches(collection, normalized_record_set, "documents")
|
||||
|
||||
|
||||
def embeddings_match(collection: Collection, record_set: RecordSet) -> None:
|
||||
"""The actual embedding documents is equal to the expected documents"""
|
||||
normalized_record_set = wrap_all(record_set)
|
||||
_field_matches(collection, normalized_record_set, "embeddings")
|
||||
|
||||
|
||||
def no_duplicates(collection: Collection) -> None:
|
||||
ids = collection.get()["ids"]
|
||||
assert len(ids) == len(set(ids))
|
||||
|
||||
|
||||
def _exact_distances(
|
||||
query: types.Embeddings,
|
||||
targets: types.Embeddings,
|
||||
distance_fn: Callable[
|
||||
[npt.ArrayLike, npt.ArrayLike], float
|
||||
] = distance_functions.l2,
|
||||
) -> Tuple[List[List[int]], List[List[float]]]:
|
||||
"""Return the ordered indices and distances from each query to each target"""
|
||||
np_query = np.array(query, dtype=np.float32)
|
||||
np_targets = np.array(targets, dtype=np.float32)
|
||||
|
||||
# Compute the distance between each query and each target, using the distance function
|
||||
distances = np.apply_along_axis(
|
||||
lambda query: np.apply_along_axis(distance_fn, 1, np_targets, query),
|
||||
1,
|
||||
np_query,
|
||||
)
|
||||
# Sort the distances and return the indices
|
||||
return np.argsort(distances).tolist(), distances.tolist()
|
||||
|
||||
|
||||
def fd_not_exceeding_threadpool_size(threadpool_size: int) -> None:
|
||||
"""
|
||||
Checks that the open file descriptors are not exceeding the threadpool size
|
||||
works only for SegmentAPI
|
||||
"""
|
||||
current_process = psutil.Process()
|
||||
open_files = current_process.open_files()
|
||||
max_retries = 5
|
||||
retry_count = 0
|
||||
# we probably don't need the below but we keep it to avoid flaky tests.
|
||||
while (
|
||||
len([p.path for p in open_files if "sqlite3" in p.path]) - 1 > threadpool_size
|
||||
and retry_count < max_retries
|
||||
):
|
||||
gc.collect() # GC to collect the orphaned TLS objects
|
||||
open_files = current_process.open_files()
|
||||
retry_count += 1
|
||||
sleep(1)
|
||||
assert (
|
||||
len([p.path for p in open_files if "sqlite3" in p.path]) - 1 <= threadpool_size
|
||||
)
|
||||
|
||||
|
||||
def get_space(collection: Collection):
|
||||
# TODO: this is a hack to get the space
|
||||
# We should update the tests to not pass space via metadata instead use collection
|
||||
# configuration_json
|
||||
space = None
|
||||
if "hnsw:space" in collection.metadata:
|
||||
space = collection.metadata["hnsw:space"]
|
||||
if collection._model.configuration_json is None:
|
||||
return space
|
||||
if (
|
||||
"spann" in collection._model.configuration_json
|
||||
and collection._model.configuration_json.get("spann") is not None
|
||||
and "space" in collection._model.configuration_json.get("spann")
|
||||
):
|
||||
space = collection._model.configuration_json.get("spann").get("space")
|
||||
elif (
|
||||
"hnsw" in collection._model.configuration_json
|
||||
and collection._model.configuration_json.get("hnsw") is not None
|
||||
and "space" in collection._model.configuration_json.get("hnsw")
|
||||
):
|
||||
if space is None:
|
||||
space = collection._model.configuration_json.get("hnsw").get("space")
|
||||
return space
|
||||
|
||||
|
||||
def ann_accuracy(
|
||||
collection: Collection,
|
||||
record_set: RecordSet,
|
||||
n_results: int = 1,
|
||||
min_recall: float = 0.95,
|
||||
embedding_function: Optional[types.EmbeddingFunction] = None, # type: ignore[type-arg]
|
||||
query_indices: Optional[List[int]] = None,
|
||||
query_embeddings: Optional[types.Embeddings] = None,
|
||||
use_search: bool = False,
|
||||
) -> None:
|
||||
"""Validate that the API performs nearest_neighbor searches correctly"""
|
||||
normalized_record_set = wrap_all(record_set)
|
||||
|
||||
if len(normalized_record_set["ids"]) == 0:
|
||||
return # nothing to test here
|
||||
|
||||
embeddings: Optional[types.Embeddings] = normalized_record_set["embeddings"]
|
||||
have_embeddings = embeddings is not None and len(embeddings) > 0
|
||||
if not have_embeddings:
|
||||
assert embedding_function is not None
|
||||
assert normalized_record_set["documents"] is not None
|
||||
assert isinstance(normalized_record_set["documents"], list)
|
||||
# Compute the embeddings for the documents
|
||||
embeddings = embedding_function(normalized_record_set["documents"])
|
||||
|
||||
space = get_space(collection)
|
||||
if space is None:
|
||||
distance_function = distance_functions.l2
|
||||
elif space == "cosine":
|
||||
distance_function = distance_functions.cosine
|
||||
elif space == "ip":
|
||||
distance_function = distance_functions.ip
|
||||
elif space == "l2":
|
||||
distance_function = distance_functions.l2
|
||||
|
||||
accuracy_threshold = 1e-6
|
||||
assert collection.metadata is not None
|
||||
assert embeddings is not None
|
||||
# TODO: ip and cosine are numerically unstable in HNSW.
|
||||
# The higher the dimensionality, the more noise is introduced, since each float element
|
||||
# of the vector has noise added, which is then subsequently included in all normalization calculations.
|
||||
# This means that higher dimensions will have more noise, and thus more error.
|
||||
assert all(isinstance(e, (list, np.ndarray)) for e in embeddings)
|
||||
dim = len(embeddings[0])
|
||||
accuracy_threshold = accuracy_threshold * math.pow(10, int(math.log10(dim)))
|
||||
|
||||
# Perform exact distance computation
|
||||
if query_embeddings is None:
|
||||
query_embeddings = (
|
||||
embeddings
|
||||
if query_indices is None
|
||||
else [embeddings[i] for i in query_indices]
|
||||
)
|
||||
query_documents = normalized_record_set["documents"]
|
||||
if query_indices is not None and query_documents is not None:
|
||||
query_documents = [query_documents[i] for i in query_indices]
|
||||
|
||||
indices, distances = _exact_distances(
|
||||
query_embeddings, embeddings, distance_fn=distance_function
|
||||
)
|
||||
|
||||
if use_search:
|
||||
# Use search API instead of query
|
||||
search_requests = []
|
||||
for query_embedding in query_embeddings:
|
||||
# Convert numpy array to list if needed
|
||||
if isinstance(query_embedding, np.ndarray):
|
||||
query_embedding_list = query_embedding.tolist()
|
||||
else:
|
||||
query_embedding_list = query_embedding
|
||||
search = Search(
|
||||
rank=Knn(query=query_embedding_list),
|
||||
limit=Limit(limit=n_results),
|
||||
).select_all()
|
||||
search_requests.append(search)
|
||||
|
||||
# Call _search API
|
||||
api = collection._client # type: ignore
|
||||
search_results = api._search(
|
||||
collection_id=collection.id,
|
||||
searches=search_requests,
|
||||
tenant='default_tenant',
|
||||
database='default_database',
|
||||
)
|
||||
|
||||
# Convert search results to query-like format
|
||||
query_results = cast(types.QueryResult, {
|
||||
"ids": search_results["ids"],
|
||||
"distances": search_results["scores"], # scores is distances in search API
|
||||
"embeddings": search_results["embeddings"],
|
||||
"documents": search_results["documents"],
|
||||
"metadatas": search_results["metadatas"],
|
||||
})
|
||||
else:
|
||||
query_results = collection.query(
|
||||
query_embeddings=query_embeddings if have_embeddings else None,
|
||||
query_texts=query_documents if not have_embeddings else None,
|
||||
n_results=n_results,
|
||||
include=["embeddings", "documents", "metadatas", "distances"], # type: ignore[list-item]
|
||||
)
|
||||
|
||||
_query_results_are_correct_shape(query_results, n_results)
|
||||
|
||||
# Assert fields are not None for type checking
|
||||
assert query_results["ids"] is not None
|
||||
assert query_results["distances"] is not None
|
||||
assert query_results["embeddings"] is not None
|
||||
assert query_results["documents"] is not None
|
||||
assert query_results["metadatas"] is not None
|
||||
|
||||
# Dict of ids to indices
|
||||
id_to_index = {id: i for i, id in enumerate(normalized_record_set["ids"])}
|
||||
missing = 0
|
||||
for i, (indices_i, distances_i) in enumerate(zip(indices, distances)):
|
||||
expected_ids = np.array(normalized_record_set["ids"])[indices_i[:n_results]]
|
||||
missing += len(set(expected_ids) - set(query_results["ids"][i]))
|
||||
|
||||
# For each id in the query results, find the index in the embeddings set
|
||||
# and assert that the embeddings are the same
|
||||
for j, id in enumerate(query_results["ids"][i]):
|
||||
# This may be because the true nth nearest neighbor didn't get returned by the ANN query
|
||||
unexpected_id = id not in expected_ids
|
||||
index = id_to_index[id]
|
||||
|
||||
correct_distance = np.allclose(
|
||||
distances_i[index],
|
||||
query_results["distances"][i][j],
|
||||
atol=accuracy_threshold,
|
||||
)
|
||||
if unexpected_id:
|
||||
# If the ID is unexpcted, but the distance is correct, then we
|
||||
# have a duplicate in the data. In this case, we should not reduce recall.
|
||||
if correct_distance:
|
||||
missing -= 1
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
assert correct_distance
|
||||
|
||||
assert np.allclose(embeddings[index], query_results["embeddings"][i][j])
|
||||
if normalized_record_set["documents"] is not None:
|
||||
assert (
|
||||
normalized_record_set["documents"][index]
|
||||
== query_results["documents"][i][j]
|
||||
)
|
||||
if normalized_record_set["metadatas"] is not None:
|
||||
check_metadata(
|
||||
normalized_record_set["metadatas"][index],
|
||||
query_results["metadatas"][i][j],
|
||||
)
|
||||
|
||||
size = len(normalized_record_set["ids"])
|
||||
recall = (size - missing) / size
|
||||
|
||||
try:
|
||||
note(
|
||||
f"# recall: {recall}, missing {missing} out of {size}, accuracy threshold {accuracy_threshold}"
|
||||
)
|
||||
except InvalidArgument:
|
||||
pass # it's ok if we're running outside hypothesis
|
||||
|
||||
assert recall >= min_recall
|
||||
|
||||
# Ensure that the query results are sorted by distance
|
||||
for distance_result in query_results["distances"]:
|
||||
assert np.allclose(np.sort(distance_result), distance_result)
|
||||
|
||||
|
||||
def _query_results_are_correct_shape(
|
||||
query_results: types.QueryResult, n_results: int
|
||||
) -> None:
|
||||
for result_type in ["distances", "embeddings", "documents", "metadatas"]:
|
||||
assert query_results[result_type] is not None # type: ignore[literal-required]
|
||||
assert all(
|
||||
len(result) == n_results for result in query_results[result_type] # type: ignore[literal-required]
|
||||
)
|
||||
|
||||
|
||||
def _total_embedding_queue_log_size(sqlite: SqliteDB) -> int:
|
||||
t = Table("embeddings_queue")
|
||||
q = sqlite.querybuilder().from_(t)
|
||||
|
||||
with sqlite.tx() as cur:
|
||||
sql, params = get_sql(
|
||||
q.select(functions.Count(t.seq_id)), sqlite.parameter_format()
|
||||
)
|
||||
result = cur.execute(sql, params)
|
||||
return cast(int, result.fetchone()[0])
|
||||
|
||||
|
||||
def log_size_below_max(
|
||||
system: System, collections: List[Collection], has_collection_mutated: bool
|
||||
) -> None:
|
||||
sqlite = system.instance(SqliteDB)
|
||||
|
||||
# Ephemeral Rust client is using its own sqlite impl, which cannot be accessed from Python
|
||||
if (
|
||||
not system.settings.is_persistent
|
||||
and system.settings.chroma_api_impl == "chromadb.api.rust.RustBindingsAPI"
|
||||
):
|
||||
return
|
||||
|
||||
if has_collection_mutated:
|
||||
# Must always keep one entry to avoid reusing seq_ids
|
||||
assert _total_embedding_queue_log_size(sqlite) >= 1
|
||||
|
||||
# We purge per-collection as the sync_threshold is a per-collection setting
|
||||
sync_threshold_sum = sum(
|
||||
collection.metadata.get("hnsw:sync_threshold", 1000)
|
||||
if collection.metadata is not None
|
||||
else 1000
|
||||
for collection in collections
|
||||
)
|
||||
batch_size_sum = sum(
|
||||
collection.metadata.get("hnsw:batch_size", 100)
|
||||
if collection.metadata is not None
|
||||
else 100
|
||||
for collection in collections
|
||||
)
|
||||
|
||||
limit = (
|
||||
sync_threshold_sum
|
||||
if system.settings.chroma_api_impl == "chromadb.api.rust.RustBindingsAPI"
|
||||
else sync_threshold_sum + batch_size_sum
|
||||
)
|
||||
|
||||
# -1 is used because the queue is always at least 1 entry long, so deletion stops before the max ack'ed sequence ID.
|
||||
# And for python impl if the batch_size != sync_threshold, the queue can have up to batch_size more entries.
|
||||
assert _total_embedding_queue_log_size(sqlite) - 1 <= limit
|
||||
else:
|
||||
assert _total_embedding_queue_log_size(sqlite) == 0
|
||||
|
||||
|
||||
def _total_embedding_queue_log_size_per_collection(
|
||||
system: System,
|
||||
collections: List[Collection],
|
||||
) -> Dict[UUID, int]:
|
||||
sqlite = system.instance(SqliteDB)
|
||||
t = Table("embeddings_queue")
|
||||
q = sqlite.querybuilder().from_(t)
|
||||
_tenant = system.settings.require("tenant_id")
|
||||
_topic_namespace = system.settings.require("topic_namespace")
|
||||
topic_mappings = {
|
||||
create_topic_name(_tenant, _topic_namespace, collection.id): collection
|
||||
for collection in collections
|
||||
}
|
||||
with sqlite.tx() as cur:
|
||||
sql, params = get_sql(
|
||||
q.select(t.topic, functions.Count(t.seq_id)).groupby("topic"),
|
||||
sqlite.parameter_format(),
|
||||
)
|
||||
result = cur.execute(sql, params)
|
||||
out = {}
|
||||
for res in result.fetchall():
|
||||
out[topic_mappings[res[0]].id] = res[1]
|
||||
return out
|
||||
|
||||
|
||||
def log_size_for_collections_match_expected(
|
||||
system: System, collections: List[Collection], has_collection_mutated: bool
|
||||
) -> None:
|
||||
if system.settings.chroma_api_impl == "chromadb.api.rust.RustBindingsAPI":
|
||||
# The rust impl does not use batch size
|
||||
return
|
||||
|
||||
sqlite = system.instance(SqliteDB)
|
||||
|
||||
if has_collection_mutated:
|
||||
# Must always keep one entry to avoid reusing seq_ids
|
||||
assert _total_embedding_queue_log_size(sqlite) >= 1
|
||||
|
||||
batch_size_sum = {
|
||||
collection.id: collection.metadata.get("hnsw:batch_size", 100)
|
||||
if collection.metadata is not None
|
||||
else 100
|
||||
for collection in collections
|
||||
}
|
||||
expected_sizes = {
|
||||
collection.id: collection.count() % batch_size_sum[collection.id] + 1
|
||||
for collection in collections
|
||||
}
|
||||
|
||||
actual_sizes = _total_embedding_queue_log_size_per_collection(
|
||||
system, collections
|
||||
)
|
||||
assert set(actual_sizes.keys()) == set(expected_sizes.keys())
|
||||
assert all(
|
||||
actual_sizes[collection.id] == expected_sizes[collection.id]
|
||||
for collection in collections
|
||||
)
|
||||
|
||||
else:
|
||||
assert _total_embedding_queue_log_size(sqlite) == 0
|
||||
|
||||
|
||||
@contextmanager
|
||||
def collection_deleted(client: ClientAPI, collection_name: str):
|
||||
# Invariant checks before deletion
|
||||
collection_names = [c.name for c in client.list_collections()]
|
||||
assert collection_name in collection_names
|
||||
collection = client.get_collection(collection_name)
|
||||
segments = []
|
||||
if isinstance(client._server, SegmentAPI): # type: ignore
|
||||
sysdb: SysDB = client._server._sysdb # type: ignore
|
||||
segments = sysdb.get_segments(collection=collection.id)
|
||||
segment_types = {}
|
||||
should_have_hnsw = False
|
||||
for segment in segments:
|
||||
segment_types[segment["type"]] = True
|
||||
if segment["type"] == SegmentType.HNSW_LOCAL_PERSISTED.value:
|
||||
sync_threshold = (
|
||||
collection.metadata["hnsw:sync_threshold"]
|
||||
if collection.metadata is not None
|
||||
and "hnsw:sync_threshold" in collection.metadata
|
||||
else 1000
|
||||
)
|
||||
if (
|
||||
collection.count() > sync_threshold
|
||||
): # we only check if vector segment dir exists if we've synced at least once
|
||||
should_have_hnsw = True
|
||||
assert os.path.exists(
|
||||
os.path.join(
|
||||
client.get_settings().persist_directory, str(segment["id"])
|
||||
)
|
||||
)
|
||||
if should_have_hnsw:
|
||||
assert segment_types[SegmentType.HNSW_LOCAL_PERSISTED.value]
|
||||
assert segment_types[SegmentType.SQLITE.value]
|
||||
|
||||
yield
|
||||
|
||||
# Invariant checks after deletion
|
||||
collection_names = [c.name for c in client.list_collections()]
|
||||
assert collection_name not in collection_names
|
||||
if len(segments) > 0:
|
||||
sysdb: SysDB = client._server._sysdb # type: ignore
|
||||
segments_after = sysdb.get_segments(collection=collection.id)
|
||||
assert len(segments_after) == 0
|
||||
for segment in segments:
|
||||
if segment["type"] == SegmentType.HNSW_LOCAL_PERSISTED.value:
|
||||
assert not os.path.exists(
|
||||
os.path.join(
|
||||
client.get_settings().persist_directory, str(segment["id"])
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,720 @@
|
||||
import hashlib
|
||||
import hypothesis
|
||||
import hypothesis.strategies as st
|
||||
from typing import Any, Optional, List, Dict, Union, cast
|
||||
from typing_extensions import TypedDict
|
||||
import uuid
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import chromadb.api.types as types
|
||||
import re
|
||||
from hypothesis.strategies._internal.strategies import SearchStrategy
|
||||
from chromadb.test.conftest import NOT_CLUSTER_ONLY
|
||||
from dataclasses import dataclass
|
||||
from chromadb.api.types import (
|
||||
Documents,
|
||||
Embeddable,
|
||||
EmbeddingFunction,
|
||||
Embeddings,
|
||||
Metadata,
|
||||
)
|
||||
from chromadb.types import LiteralValue, WhereOperator, LogicalOperator
|
||||
from chromadb.test.conftest import is_spann_disabled_mode, skip_reason_spann_disabled
|
||||
from chromadb.api.collection_configuration import (
|
||||
CreateCollectionConfiguration,
|
||||
CreateSpannConfiguration,
|
||||
)
|
||||
|
||||
# Set the random seed for reproducibility
|
||||
np.random.seed(0) # unnecessary, hypothesis does this for us
|
||||
|
||||
# See Hypothesis documentation for creating strategies at
|
||||
# https://hypothesis.readthedocs.io/en/latest/data.html
|
||||
|
||||
# NOTE: Because these strategies are used in state machines, we need to
|
||||
# work around an issue with state machines, in which strategies that frequently
|
||||
# are marked as invalid (i.e. through the use of `assume` or `.filter`) can cause the
|
||||
# state machine tests to fail with an hypothesis.errors.Unsatisfiable.
|
||||
|
||||
# Ultimately this is because the entire state machine is run as a single Hypothesis
|
||||
# example, which ends up drawing from the same strategies an enormous number of times.
|
||||
# Whenever a strategy marks itself as invalid, Hypothesis tries to start the entire
|
||||
# state machine run over. See https://github.com/HypothesisWorks/hypothesis/issues/3618
|
||||
|
||||
# Because strategy generation is all interrelated, seemingly small changes (especially
|
||||
# ones called early in a test) can have an outside effect. Generating lists with
|
||||
# unique=True, or dictionaries with a min size seems especially bad.
|
||||
|
||||
# Please make changes to these strategies incrementally, testing to make sure they don't
|
||||
# start generating unsatisfiable examples.
|
||||
|
||||
test_hnsw_config = {
|
||||
"hnsw:construction_ef": 128,
|
||||
"hnsw:search_ef": 128,
|
||||
"hnsw:M": 128,
|
||||
}
|
||||
|
||||
|
||||
class RecordSet(TypedDict):
|
||||
"""
|
||||
A generated set of embeddings, ids, metadatas, and documents that
|
||||
represent what a user would pass to the API.
|
||||
"""
|
||||
|
||||
ids: Union[types.ID, List[types.ID]]
|
||||
embeddings: Optional[Union[types.Embeddings, types.Embedding]]
|
||||
metadatas: Optional[Union[List[Optional[types.Metadata]], types.Metadata]]
|
||||
documents: Optional[Union[List[types.Document], types.Document]]
|
||||
|
||||
|
||||
class NormalizedRecordSet(TypedDict):
|
||||
"""
|
||||
A RecordSet, with all fields normalized to lists.
|
||||
"""
|
||||
|
||||
ids: List[types.ID]
|
||||
embeddings: Optional[types.Embeddings]
|
||||
metadatas: Optional[List[Optional[types.Metadata]]]
|
||||
documents: Optional[List[types.Document]]
|
||||
|
||||
|
||||
class StateMachineRecordSet(TypedDict):
|
||||
"""
|
||||
Represents the internal state of a state machine in hypothesis tests.
|
||||
"""
|
||||
|
||||
ids: List[types.ID]
|
||||
embeddings: types.Embeddings
|
||||
metadatas: List[Optional[types.Metadata]]
|
||||
documents: List[Optional[types.Document]]
|
||||
|
||||
|
||||
class Record(TypedDict):
|
||||
"""
|
||||
A single generated record.
|
||||
"""
|
||||
|
||||
id: types.ID
|
||||
embedding: Optional[types.Embedding]
|
||||
metadata: Optional[types.Metadata]
|
||||
document: Optional[types.Document]
|
||||
|
||||
|
||||
# TODO: support arbitrary text everywhere so we don't SQL-inject ourselves.
|
||||
# TODO: support empty strings everywhere
|
||||
sql_alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_"
|
||||
safe_text = st.text(alphabet=sql_alphabet, min_size=1)
|
||||
sql_alphabet_minus_underscore = (
|
||||
"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-"
|
||||
)
|
||||
safe_text_min_size_3 = st.text(alphabet=sql_alphabet_minus_underscore, min_size=3)
|
||||
tenant_database_name = st.text(alphabet=sql_alphabet, min_size=3)
|
||||
|
||||
# Workaround for FastAPI json encoding peculiarities
|
||||
# https://github.com/tiangolo/fastapi/blob/8ac8d70d52bb0dd9eb55ba4e22d3e383943da05c/fastapi/encoders.py#L104
|
||||
safe_text = safe_text.filter(lambda s: not s.startswith("_sa"))
|
||||
safe_text_min_size_3 = safe_text_min_size_3.filter(lambda s: not s.startswith("_sa"))
|
||||
tenant_database_name = tenant_database_name.filter(lambda s: not s.startswith("_sa"))
|
||||
|
||||
safe_integers = st.integers(
|
||||
min_value=-(2**31), max_value=2**31 - 1
|
||||
) # TODO: handle longs
|
||||
# In distributed chroma, floats are 32 bit hence we need to
|
||||
# restrict the generation to generate only 32 bit floats.
|
||||
safe_floats = st.floats(
|
||||
allow_infinity=False,
|
||||
allow_nan=False,
|
||||
allow_subnormal=False,
|
||||
width=32,
|
||||
min_value=-1e6,
|
||||
max_value=1e6,
|
||||
) # TODO: handle infinity and NAN
|
||||
|
||||
safe_values: List[SearchStrategy[Union[int, float, str, bool]]] = [
|
||||
safe_text,
|
||||
safe_integers,
|
||||
safe_floats,
|
||||
st.booleans(),
|
||||
]
|
||||
|
||||
|
||||
def one_or_both(
|
||||
strategy_a: st.SearchStrategy[Any], strategy_b: st.SearchStrategy[Any]
|
||||
) -> st.SearchStrategy[Any]:
|
||||
return st.one_of(
|
||||
st.tuples(strategy_a, strategy_b),
|
||||
st.tuples(strategy_a, st.none()),
|
||||
st.tuples(st.none(), strategy_b),
|
||||
)
|
||||
|
||||
|
||||
# Temporarily generate only these to avoid SQL formatting issues.
|
||||
legal_id_characters = (
|
||||
"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_./+"
|
||||
)
|
||||
|
||||
float_types = [np.float16, np.float32, np.float64]
|
||||
int_types = [np.int16, np.int32, np.int64] # TODO: handle int types
|
||||
|
||||
|
||||
@st.composite
|
||||
def collection_name(draw: st.DrawFn) -> str:
|
||||
_collection_name_re = re.compile(r"^[a-zA-Z][a-zA-Z0-9-]{1,60}[a-zA-Z0-9]$")
|
||||
_ipv4_address_re = re.compile(r"^([0-9]{1,3}\.){3}[0-9]{1,3}$")
|
||||
_two_periods_re = re.compile(r"\.\.")
|
||||
|
||||
name: str = draw(st.from_regex(_collection_name_re)).strip()
|
||||
hypothesis.assume(not _ipv4_address_re.match(name))
|
||||
hypothesis.assume(not _two_periods_re.search(name))
|
||||
|
||||
return name
|
||||
|
||||
|
||||
collection_metadata = st.one_of(
|
||||
st.none(), st.dictionaries(safe_text, st.one_of(*safe_values))
|
||||
)
|
||||
|
||||
|
||||
# TODO: Use a hypothesis strategy while maintaining embedding uniqueness
|
||||
# Or handle duplicate embeddings within a known epsilon
|
||||
def create_embeddings(
|
||||
dim: int,
|
||||
count: int,
|
||||
dtype: npt.DTypeLike,
|
||||
) -> types.Embeddings:
|
||||
embeddings: types.Embeddings = cast(
|
||||
types.Embeddings,
|
||||
(
|
||||
np.random.uniform(
|
||||
low=-1.0,
|
||||
high=1.0,
|
||||
size=(count, dim),
|
||||
)
|
||||
.astype(dtype)
|
||||
.tolist()
|
||||
),
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
def create_embeddings_ndarray(
|
||||
dim: int,
|
||||
count: int,
|
||||
dtype: npt.DTypeLike,
|
||||
) -> np.typing.NDArray[Any]:
|
||||
return np.random.uniform(
|
||||
low=-1.0,
|
||||
high=1.0,
|
||||
size=(count, dim),
|
||||
).astype(dtype)
|
||||
|
||||
|
||||
class hashing_embedding_function(types.EmbeddingFunction[Documents]):
|
||||
def __init__(self, dim: int, dtype: npt.DTypeLike) -> None:
|
||||
self.dim = dim
|
||||
self.dtype = dtype
|
||||
|
||||
def __call__(self, input: types.Documents) -> types.Embeddings:
|
||||
# Hash the texts and convert to hex strings
|
||||
hashed_texts = [
|
||||
list(hashlib.sha256(text.encode("utf-8")).hexdigest()) for text in input
|
||||
]
|
||||
# Pad with repetition, or truncate the hex strings to the desired dimension
|
||||
padded_texts = [
|
||||
text * (self.dim // len(text)) + text[: self.dim % len(text)]
|
||||
for text in hashed_texts
|
||||
]
|
||||
|
||||
# Convert the hex strings to dtype
|
||||
embeddings: types.Embeddings = [
|
||||
np.array([int(char, 16) / 15.0 for char in text], dtype=self.dtype)
|
||||
for text in padded_texts
|
||||
]
|
||||
|
||||
return embeddings
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"hashing_embedding_function(dim={self.dim}, dtype={self.dtype})"
|
||||
|
||||
|
||||
class not_implemented_embedding_function(types.EmbeddingFunction[Documents]):
|
||||
def __call__(self, input: Documents) -> Embeddings:
|
||||
assert False, "This embedding function is not implemented"
|
||||
|
||||
|
||||
def embedding_function_strategy(
|
||||
dim: int, dtype: npt.DTypeLike
|
||||
) -> st.SearchStrategy[types.EmbeddingFunction[Embeddable]]:
|
||||
return st.just(
|
||||
cast(EmbeddingFunction[Embeddable], hashing_embedding_function(dim, dtype))
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExternalCollection:
|
||||
"""
|
||||
An external view of a collection.
|
||||
|
||||
This strategy only contains information about a collection that a client of Chroma
|
||||
sees -- that is, it contains none of Chroma's internal bookkeeping. It should
|
||||
be used to test the API and client code.
|
||||
"""
|
||||
|
||||
name: str
|
||||
metadata: Optional[types.Metadata]
|
||||
embedding_function: Optional[types.EmbeddingFunction[Embeddable]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Collection(ExternalCollection):
|
||||
"""
|
||||
An internal view of a collection.
|
||||
|
||||
This strategy contains all the information Chroma uses internally to manage a
|
||||
collection. It is a superset of ExternalCollection and should be used to test
|
||||
internal Chroma logic.
|
||||
"""
|
||||
|
||||
id: uuid.UUID
|
||||
dimension: int
|
||||
dtype: npt.DTypeLike
|
||||
known_metadata_keys: types.Metadata
|
||||
known_document_keywords: List[str]
|
||||
has_documents: bool = False
|
||||
has_embeddings: bool = False
|
||||
collection_config: Optional[CreateCollectionConfiguration] = None
|
||||
|
||||
|
||||
@st.composite
|
||||
def collections(
|
||||
draw: st.DrawFn,
|
||||
add_filterable_data: bool = False,
|
||||
with_hnsw_params: bool = False,
|
||||
has_embeddings: Optional[bool] = None,
|
||||
has_documents: Optional[bool] = None,
|
||||
with_persistent_hnsw_params: st.SearchStrategy[bool] = st.just(False),
|
||||
max_hnsw_batch_size: int = 2000,
|
||||
max_hnsw_sync_threshold: int = 2000,
|
||||
) -> Collection:
|
||||
"""Strategy to generate a Collection object. If add_filterable_data is True, then known_metadata_keys and known_document_keywords will be populated with consistent data."""
|
||||
|
||||
assert not ((has_embeddings is False) and (has_documents is False))
|
||||
|
||||
name = draw(collection_name())
|
||||
metadata = draw(collection_metadata)
|
||||
dimension = draw(st.integers(min_value=2, max_value=2048))
|
||||
dtype = draw(st.sampled_from(float_types))
|
||||
|
||||
use_persistent_hnsw_params = draw(with_persistent_hnsw_params)
|
||||
|
||||
if use_persistent_hnsw_params and not with_hnsw_params:
|
||||
raise ValueError(
|
||||
"with_persistent_hnsw_params requires with_hnsw_params to be true"
|
||||
)
|
||||
|
||||
if with_hnsw_params:
|
||||
if metadata is None:
|
||||
metadata = {}
|
||||
metadata.update(test_hnsw_config)
|
||||
if use_persistent_hnsw_params:
|
||||
metadata["hnsw:sync_threshold"] = draw(
|
||||
st.integers(min_value=3, max_value=max_hnsw_sync_threshold)
|
||||
)
|
||||
metadata["hnsw:batch_size"] = draw(
|
||||
st.integers(
|
||||
min_value=3,
|
||||
max_value=min(
|
||||
[metadata["hnsw:sync_threshold"], max_hnsw_batch_size]
|
||||
),
|
||||
)
|
||||
)
|
||||
# Sometimes, select a space at random
|
||||
if draw(st.booleans()):
|
||||
# TODO: pull the distance functions from a source of truth that lives not
|
||||
# in tests once https://github.com/chroma-core/issues/issues/61 lands
|
||||
metadata["hnsw:space"] = draw(st.sampled_from(["cosine", "l2", "ip"]))
|
||||
|
||||
collection_config: Optional[CreateCollectionConfiguration] = None
|
||||
# Generate a spann config if in spann mode
|
||||
if not is_spann_disabled_mode:
|
||||
# Use metadata["hnsw:space"] if it exists, otherwise default to "l2"
|
||||
spann_space = metadata.get("hnsw:space", "l2") if metadata else "l2"
|
||||
|
||||
spann_config: CreateSpannConfiguration = {
|
||||
"space": spann_space,
|
||||
"write_nprobe": 4,
|
||||
"reassign_neighbor_count": 4
|
||||
}
|
||||
collection_config = {
|
||||
"spann": spann_config,
|
||||
}
|
||||
|
||||
known_metadata_keys: Dict[str, Union[int, str, float]] = {}
|
||||
if add_filterable_data:
|
||||
while len(known_metadata_keys) < 5:
|
||||
key = draw(safe_text)
|
||||
known_metadata_keys[key] = draw(st.one_of(*safe_values))
|
||||
|
||||
if has_documents is None:
|
||||
has_documents = draw(st.booleans())
|
||||
assert has_documents is not None
|
||||
# For cluster tests, we want to avoid generating documents and where_document
|
||||
# clauses of length < 3. We also don't want them to contain certan special
|
||||
# characters like _ and % that implicitly involve searching for a regex in sqlite.
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
if has_documents and add_filterable_data:
|
||||
known_document_keywords = draw(
|
||||
st.lists(safe_text_min_size_3, min_size=5, max_size=5)
|
||||
)
|
||||
else:
|
||||
known_document_keywords = []
|
||||
else:
|
||||
if has_documents and add_filterable_data:
|
||||
known_document_keywords = draw(st.lists(safe_text, min_size=5, max_size=5))
|
||||
else:
|
||||
known_document_keywords = []
|
||||
|
||||
if not has_documents:
|
||||
has_embeddings = True
|
||||
else:
|
||||
if has_embeddings is None:
|
||||
has_embeddings = draw(st.booleans())
|
||||
assert has_embeddings is not None
|
||||
|
||||
embedding_function = draw(embedding_function_strategy(dimension, dtype))
|
||||
|
||||
return Collection(
|
||||
id=uuid.uuid4(),
|
||||
name=name,
|
||||
metadata=metadata,
|
||||
dimension=dimension,
|
||||
dtype=dtype,
|
||||
known_metadata_keys=known_metadata_keys,
|
||||
has_documents=has_documents,
|
||||
known_document_keywords=known_document_keywords,
|
||||
has_embeddings=has_embeddings,
|
||||
embedding_function=embedding_function,
|
||||
collection_config=collection_config
|
||||
)
|
||||
|
||||
|
||||
@st.composite
|
||||
def metadata(
|
||||
draw: st.DrawFn,
|
||||
collection: Collection,
|
||||
min_size: int = 0,
|
||||
max_size: Optional[int] = None,
|
||||
) -> Optional[types.Metadata]:
|
||||
"""Strategy for generating metadata that could be a part of the given collection"""
|
||||
# First draw a random dictionary.
|
||||
metadata: types.Metadata = draw(
|
||||
st.dictionaries(
|
||||
safe_text, st.one_of(*safe_values), min_size=min_size, max_size=max_size
|
||||
)
|
||||
)
|
||||
# Then, remove keys that overlap with the known keys for the coll
|
||||
# to avoid type errors when comparing.
|
||||
if collection.known_metadata_keys:
|
||||
for key in collection.known_metadata_keys.keys():
|
||||
if key in metadata:
|
||||
del metadata[key] # type: ignore
|
||||
# Finally, add in some of the known keys for the collection
|
||||
sampling_dict: Dict[str, st.SearchStrategy[Union[str, int, float]]] = {
|
||||
k: st.just(v) for k, v in collection.known_metadata_keys.items()
|
||||
}
|
||||
metadata.update(draw(st.fixed_dictionaries({}, optional=sampling_dict))) # type: ignore
|
||||
# We don't allow submitting empty metadata
|
||||
if metadata == {}:
|
||||
return None
|
||||
return metadata
|
||||
|
||||
|
||||
@st.composite
|
||||
def document(draw: st.DrawFn, collection: Collection) -> types.Document:
|
||||
"""Strategy for generating documents that could be a part of the given collection"""
|
||||
# For cluster tests, we want to avoid generating documents of length < 3.
|
||||
# We also don't want them to contain certan special
|
||||
# characters like _ and % that implicitly involve searching for a regex in sqlite.
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
# Blacklist certain unicode characters that affect sqlite processing.
|
||||
# For example, the null (/x00) character makes sqlite stop processing a string.
|
||||
# Also, blacklist _ and % for cluster tests.
|
||||
blacklist_categories = ("Cc", "Cs", "Pc", "Po")
|
||||
if collection.known_document_keywords:
|
||||
known_words_st = st.sampled_from(collection.known_document_keywords)
|
||||
else:
|
||||
known_words_st = st.text(
|
||||
min_size=3,
|
||||
alphabet=st.characters(blacklist_categories=blacklist_categories), # type: ignore
|
||||
)
|
||||
|
||||
random_words_st = st.text(
|
||||
min_size=3, alphabet=st.characters(blacklist_categories=blacklist_categories) # type: ignore
|
||||
)
|
||||
words = draw(st.lists(st.one_of(known_words_st, random_words_st), min_size=1))
|
||||
return " ".join(words)
|
||||
|
||||
# Blacklist certain unicode characters that affect sqlite processing.
|
||||
# For example, the null (/x00) character makes sqlite stop processing a string.
|
||||
blacklist_categories = ("Cc", "Cs") # type: ignore
|
||||
if collection.known_document_keywords:
|
||||
known_words_st = st.sampled_from(collection.known_document_keywords)
|
||||
else:
|
||||
known_words_st = st.text(
|
||||
min_size=1,
|
||||
alphabet=st.characters(blacklist_categories=blacklist_categories), # type: ignore
|
||||
)
|
||||
|
||||
random_words_st = st.text(
|
||||
min_size=1, alphabet=st.characters(blacklist_categories=blacklist_categories) # type: ignore
|
||||
)
|
||||
words = draw(st.lists(st.one_of(known_words_st, random_words_st), min_size=1))
|
||||
return " ".join(words)
|
||||
|
||||
|
||||
@st.composite
|
||||
def recordsets(
|
||||
draw: st.DrawFn,
|
||||
collection_strategy: SearchStrategy[Collection] = collections(),
|
||||
id_strategy: SearchStrategy[str] = safe_text,
|
||||
min_size: int = 1,
|
||||
max_size: int = 50,
|
||||
# If num_unique_metadata is not None, then the number of metadata generations
|
||||
# will be the size of the record set. If set, the number of metadata
|
||||
# generations will be the value of num_unique_metadata.
|
||||
num_unique_metadata: Optional[int] = None,
|
||||
min_metadata_size: int = 0,
|
||||
max_metadata_size: Optional[int] = None,
|
||||
) -> RecordSet:
|
||||
collection = draw(collection_strategy)
|
||||
|
||||
ids = list(
|
||||
draw(st.lists(id_strategy, min_size=min_size, max_size=max_size, unique=True))
|
||||
)
|
||||
|
||||
embeddings: Optional[Embeddings] = None
|
||||
if collection.has_embeddings:
|
||||
embeddings = create_embeddings(collection.dimension, len(ids), collection.dtype)
|
||||
num_metadata = num_unique_metadata if num_unique_metadata is not None else len(ids)
|
||||
generated_metadatas = draw(
|
||||
st.lists(
|
||||
metadata(
|
||||
collection, min_size=min_metadata_size, max_size=max_metadata_size
|
||||
),
|
||||
min_size=num_metadata,
|
||||
max_size=num_metadata,
|
||||
)
|
||||
)
|
||||
metadatas = []
|
||||
for i in range(len(ids)):
|
||||
metadatas.append(generated_metadatas[i % len(generated_metadatas)])
|
||||
|
||||
documents: Optional[Documents] = None
|
||||
if collection.has_documents:
|
||||
documents = draw(
|
||||
st.lists(document(collection), min_size=len(ids), max_size=len(ids))
|
||||
)
|
||||
|
||||
# in the case where we have a single record, sometimes exercise
|
||||
# the code that handles individual values rather than lists.
|
||||
# In this case, any field may be a list or a single value.
|
||||
if len(ids) == 1:
|
||||
single_id: Union[str, List[str]] = ids[0] if draw(st.booleans()) else ids
|
||||
single_embedding = (
|
||||
embeddings[0]
|
||||
if embeddings is not None and draw(st.booleans())
|
||||
else embeddings
|
||||
)
|
||||
single_metadata: Union[Optional[Metadata], List[Optional[Metadata]]] = (
|
||||
metadatas[0] if draw(st.booleans()) else metadatas
|
||||
)
|
||||
single_document = (
|
||||
documents[0] if documents is not None and draw(st.booleans()) else documents
|
||||
)
|
||||
return {
|
||||
"ids": single_id,
|
||||
"embeddings": single_embedding,
|
||||
"metadatas": single_metadata,
|
||||
"documents": single_document,
|
||||
}
|
||||
return {
|
||||
"ids": ids,
|
||||
"embeddings": embeddings,
|
||||
"metadatas": metadatas,
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
|
||||
def opposite_value(value: LiteralValue) -> SearchStrategy[Any]:
|
||||
"""
|
||||
Returns a strategy that will generate all valid values except the input value - testing of $nin
|
||||
"""
|
||||
if isinstance(value, float):
|
||||
return safe_floats.filter(lambda x: x != value)
|
||||
elif isinstance(value, str):
|
||||
return safe_text.filter(lambda x: x != value)
|
||||
elif isinstance(value, bool):
|
||||
return st.booleans().filter(lambda x: x != value)
|
||||
elif isinstance(value, int):
|
||||
return st.integers(min_value=-(2**31), max_value=2**31 - 1).filter(
|
||||
lambda x: x != value
|
||||
)
|
||||
else:
|
||||
return st.from_type(type(value)).filter(lambda x: x != value)
|
||||
|
||||
|
||||
@st.composite
|
||||
def where_clause(draw: st.DrawFn, collection: Collection) -> types.Where:
|
||||
"""Generate a filter that could be used in a query against the given collection"""
|
||||
|
||||
known_keys = sorted(collection.known_metadata_keys.keys())
|
||||
|
||||
key = draw(st.sampled_from(known_keys))
|
||||
value = collection.known_metadata_keys[key]
|
||||
|
||||
legal_ops: List[Optional[str]] = [None]
|
||||
|
||||
if isinstance(value, bool):
|
||||
legal_ops.extend(["$eq", "$ne", "$in", "$nin"])
|
||||
elif isinstance(value, float):
|
||||
legal_ops.extend(["$gt", "$lt", "$lte", "$gte"])
|
||||
elif isinstance(value, int):
|
||||
legal_ops.extend(["$gt", "$lt", "$lte", "$gte", "$eq", "$ne", "$in", "$nin"])
|
||||
elif isinstance(value, str):
|
||||
legal_ops.extend(["$eq", "$ne", "$in", "$nin"])
|
||||
else:
|
||||
assert False, f"Unsupported type: {type(value)}"
|
||||
|
||||
if isinstance(value, float):
|
||||
# Add or subtract a small number to avoid floating point rounding errors
|
||||
value = value + draw(st.sampled_from([1e-6, -1e-6]))
|
||||
# Truncate to 32 bit
|
||||
value = float(np.float32(value))
|
||||
|
||||
op: WhereOperator = draw(st.sampled_from(legal_ops))
|
||||
|
||||
if op is None:
|
||||
return {key: value}
|
||||
elif op == "$in": # type: ignore
|
||||
if isinstance(value, str) and not value:
|
||||
return {}
|
||||
return {key: {op: [value, *[draw(opposite_value(value)) for _ in range(3)]]}}
|
||||
elif op == "$nin": # type: ignore
|
||||
if isinstance(value, str) and not value:
|
||||
return {}
|
||||
return {key: {op: [draw(opposite_value(value)) for _ in range(3)]}}
|
||||
else:
|
||||
return {key: {op: value}} # type: ignore
|
||||
|
||||
|
||||
@st.composite
|
||||
def where_doc_clause(draw: st.DrawFn, collection: Collection) -> types.WhereDocument:
|
||||
"""Generate a where_document filter that could be used against the given collection"""
|
||||
# For cluster tests, we want to avoid generating where_document
|
||||
# clauses of length < 3. We also don't want them to contain certan special
|
||||
# characters like _ and % that implicitly involve searching for a regex in sqlite.
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
if collection.known_document_keywords:
|
||||
word = draw(st.sampled_from(collection.known_document_keywords))
|
||||
else:
|
||||
word = draw(safe_text_min_size_3)
|
||||
else:
|
||||
if collection.known_document_keywords:
|
||||
word = draw(st.sampled_from(collection.known_document_keywords))
|
||||
else:
|
||||
word = draw(safe_text)
|
||||
|
||||
# This is hacky, but the distributed system does not support $not_contains
|
||||
# so we need to avoid generating these operators for now in that case.
|
||||
# TODO: Remove this once the distributed system supports $not_contains
|
||||
op = draw(st.sampled_from(["$contains", "$not_contains"]))
|
||||
|
||||
if op == "$contains":
|
||||
return {"$contains": word}
|
||||
else:
|
||||
assert op == "$not_contains"
|
||||
return {"$not_contains": word}
|
||||
|
||||
|
||||
def binary_operator_clause(
|
||||
base_st: SearchStrategy[types.Where],
|
||||
) -> SearchStrategy[types.Where]:
|
||||
op: SearchStrategy[LogicalOperator] = st.sampled_from(["$and", "$or"])
|
||||
return st.dictionaries(
|
||||
keys=op,
|
||||
values=st.lists(base_st, max_size=2, min_size=2),
|
||||
min_size=1,
|
||||
max_size=1,
|
||||
)
|
||||
|
||||
|
||||
def binary_document_operator_clause(
|
||||
base_st: SearchStrategy[types.WhereDocument],
|
||||
) -> SearchStrategy[types.WhereDocument]:
|
||||
op: SearchStrategy[LogicalOperator] = st.sampled_from(["$and", "$or"])
|
||||
return st.dictionaries(
|
||||
keys=op,
|
||||
values=st.lists(base_st, max_size=2, min_size=2),
|
||||
min_size=1,
|
||||
max_size=1,
|
||||
)
|
||||
|
||||
|
||||
@st.composite
|
||||
def recursive_where_clause(draw: st.DrawFn, collection: Collection) -> types.Where:
|
||||
base_st = where_clause(collection)
|
||||
where: types.Where = draw(st.recursive(base_st, binary_operator_clause))
|
||||
return where
|
||||
|
||||
|
||||
@st.composite
|
||||
def recursive_where_doc_clause(
|
||||
draw: st.DrawFn, collection: Collection
|
||||
) -> types.WhereDocument:
|
||||
base_st = where_doc_clause(collection)
|
||||
where: types.WhereDocument = draw(
|
||||
st.recursive(base_st, binary_document_operator_clause)
|
||||
)
|
||||
return where
|
||||
|
||||
|
||||
class Filter(TypedDict):
|
||||
where: Optional[types.Where]
|
||||
ids: Optional[Union[str, List[str]]]
|
||||
where_document: Optional[types.WhereDocument]
|
||||
|
||||
|
||||
@st.composite
|
||||
def filters(
|
||||
draw: st.DrawFn,
|
||||
collection_st: st.SearchStrategy[Collection],
|
||||
recordset_st: st.SearchStrategy[RecordSet],
|
||||
include_all_ids: bool = False,
|
||||
) -> Filter:
|
||||
collection = draw(collection_st)
|
||||
recordset = draw(recordset_st)
|
||||
|
||||
where_clause = draw(st.one_of(st.none(), recursive_where_clause(collection)))
|
||||
where_document_clause = draw(
|
||||
st.one_of(st.none(), recursive_where_doc_clause(collection))
|
||||
)
|
||||
|
||||
ids: Optional[Union[List[types.ID], types.ID]]
|
||||
# Record sets can be a value instead of a list of values if there is only one record
|
||||
if isinstance(recordset["ids"], str):
|
||||
ids = [recordset["ids"]]
|
||||
else:
|
||||
ids = recordset["ids"]
|
||||
|
||||
if not include_all_ids:
|
||||
ids = draw(st.one_of(st.none(), st.lists(st.sampled_from(ids), min_size=1)))
|
||||
if ids is not None:
|
||||
# Remove duplicates since hypothesis samples with replacement
|
||||
ids = list(set(ids))
|
||||
|
||||
# Test both the single value list and the unwrapped single value case
|
||||
if ids is not None and len(ids) == 1 and draw(st.booleans()):
|
||||
ids = ids[0]
|
||||
|
||||
return {"where": where_clause, "where_document": where_document_clause, "ids": ids}
|
||||
@@ -0,0 +1,371 @@
|
||||
import uuid
|
||||
from random import randint
|
||||
from typing import cast, List, Any, Dict
|
||||
import hypothesis
|
||||
import numpy as np
|
||||
import pytest
|
||||
import hypothesis.strategies as st
|
||||
from hypothesis import given, settings
|
||||
from chromadb.api import ClientAPI
|
||||
from chromadb.api.types import Embeddings, Metadatas
|
||||
from chromadb.test.conftest import (
|
||||
NOT_CLUSTER_ONLY,
|
||||
override_hypothesis_profile,
|
||||
create_isolated_database,
|
||||
)
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import chromadb.test.property.invariants as invariants
|
||||
from chromadb.test.utils.wait_for_version_increase import wait_for_version_increase
|
||||
from chromadb.utils.batch_utils import create_batches
|
||||
|
||||
|
||||
collection_st = st.shared(strategies.collections(with_hnsw_params=True), key="coll")
|
||||
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=strategies.recordsets(collection_st, min_size=1, max_size=5),
|
||||
)
|
||||
@settings(
|
||||
deadline=None,
|
||||
parent=override_hypothesis_profile(
|
||||
normal=hypothesis.settings(max_examples=500),
|
||||
fast=hypothesis.settings(max_examples=200),
|
||||
),
|
||||
max_examples=2
|
||||
)
|
||||
def test_add_miniscule(
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
) -> None:
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
_test_add(client, collection, record_set, True, always_compact=True)
|
||||
|
||||
|
||||
# Hypothesis tends to generate smaller values so we explicitly segregate the
|
||||
# the tests into tiers, Small, Medium. Hypothesis struggles to generate large
|
||||
# record sets so we explicitly create a large record set without using Hypothesis
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=strategies.recordsets(collection_st, min_size=1, max_size=500),
|
||||
should_compact=st.booleans(),
|
||||
)
|
||||
@settings(
|
||||
deadline=None,
|
||||
parent=override_hypothesis_profile(
|
||||
normal=hypothesis.settings(max_examples=500),
|
||||
fast=hypothesis.settings(max_examples=200),
|
||||
),
|
||||
)
|
||||
def test_add_small(
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
should_compact: bool,
|
||||
) -> None:
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
_test_add(client, collection, record_set, should_compact)
|
||||
|
||||
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=strategies.recordsets(
|
||||
collection_st,
|
||||
min_size=250,
|
||||
max_size=500,
|
||||
num_unique_metadata=5,
|
||||
min_metadata_size=1,
|
||||
max_metadata_size=5,
|
||||
),
|
||||
should_compact=st.booleans(),
|
||||
)
|
||||
@settings(
|
||||
deadline=None,
|
||||
parent=override_hypothesis_profile(
|
||||
normal=hypothesis.settings(max_examples=10),
|
||||
fast=hypothesis.settings(max_examples=5),
|
||||
),
|
||||
suppress_health_check=[
|
||||
hypothesis.HealthCheck.too_slow,
|
||||
hypothesis.HealthCheck.data_too_large,
|
||||
hypothesis.HealthCheck.large_base_example,
|
||||
hypothesis.HealthCheck.function_scoped_fixture,
|
||||
],
|
||||
)
|
||||
def test_add_medium(
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
should_compact: bool,
|
||||
) -> None:
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
# Cluster tests transmit their results over grpc, which has a payload limit
|
||||
# This breaks the ann_accuracy invariant by default, since
|
||||
# the vector reader returns a payload of dataset size. So we need to batch
|
||||
# the queries in the ann_accuracy invariant
|
||||
_test_add(client, collection, record_set, should_compact, batch_ann_accuracy=True)
|
||||
|
||||
|
||||
def _test_add(
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
should_compact: bool,
|
||||
batch_ann_accuracy: bool = False,
|
||||
always_compact: bool = False,
|
||||
) -> None:
|
||||
create_isolated_database(client)
|
||||
|
||||
# TODO: Generative embedding functions
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
configuration=collection.collection_config,
|
||||
)
|
||||
initial_version = cast(int, coll.get_model()["version"])
|
||||
|
||||
normalized_record_set = invariants.wrap_all(record_set)
|
||||
|
||||
# TODO: The type of add() is incorrect as it does not allow for metadatas
|
||||
# like [{"a": 1}, None, {"a": 3}]
|
||||
for batch in create_batches(
|
||||
api=client,
|
||||
ids=cast(List[str], record_set["ids"]),
|
||||
embeddings=cast(Embeddings, record_set["embeddings"]),
|
||||
metadatas=cast(Metadatas, record_set["metadatas"]),
|
||||
documents=cast(List[str], record_set["documents"]),
|
||||
):
|
||||
coll.add(*batch)
|
||||
# Only wait for compaction if the size of the collection is
|
||||
# some minimal size
|
||||
if (
|
||||
not NOT_CLUSTER_ONLY
|
||||
and should_compact
|
||||
and (len(normalized_record_set["ids"]) > 10 or always_compact)
|
||||
):
|
||||
# Wait for the model to be updated
|
||||
wait_for_version_increase(client, collection.name, initial_version)
|
||||
|
||||
invariants.count(coll, cast(strategies.RecordSet, normalized_record_set))
|
||||
n_results = max(1, (len(normalized_record_set["ids"]) // 10))
|
||||
|
||||
if batch_ann_accuracy:
|
||||
batch_size = 10
|
||||
for i in range(0, len(normalized_record_set["ids"]), batch_size):
|
||||
invariants.ann_accuracy(
|
||||
coll,
|
||||
cast(strategies.RecordSet, normalized_record_set),
|
||||
n_results=n_results,
|
||||
embedding_function=collection.embedding_function,
|
||||
query_indices=list(
|
||||
range(i, min(i + batch_size, len(normalized_record_set["ids"])))
|
||||
),
|
||||
)
|
||||
else:
|
||||
invariants.ann_accuracy(
|
||||
coll,
|
||||
cast(strategies.RecordSet, normalized_record_set),
|
||||
n_results=n_results,
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
|
||||
|
||||
# Hypothesis struggles to generate large record sets so we explicitly create
|
||||
# a large record set
|
||||
def create_large_recordset(
|
||||
min_size: int = 45000,
|
||||
max_size: int = 50000,
|
||||
) -> strategies.RecordSet:
|
||||
size = randint(min_size, max_size)
|
||||
|
||||
ids = [str(uuid.uuid4()) for _ in range(size)]
|
||||
metadatas = [{"some_key": f"{i}"} for i in range(size)]
|
||||
documents = [f"Document {i}" for i in range(size)]
|
||||
embeddings = [[1, 2, 3] for _ in range(size)]
|
||||
record_set: Dict[str, List[Any]] = {
|
||||
"ids": ids,
|
||||
"embeddings": cast(Embeddings, embeddings),
|
||||
"metadatas": metadatas,
|
||||
"documents": documents,
|
||||
}
|
||||
return cast(strategies.RecordSet, record_set)
|
||||
|
||||
|
||||
@given(collection=collection_st, should_compact=st.booleans())
|
||||
@settings(deadline=None, max_examples=5)
|
||||
def test_add_large(
|
||||
client: ClientAPI, collection: strategies.Collection, should_compact: bool
|
||||
) -> None:
|
||||
create_isolated_database(client)
|
||||
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
|
||||
record_set = create_large_recordset(
|
||||
min_size=10000,
|
||||
max_size=50000,
|
||||
)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
normalized_record_set = invariants.wrap_all(record_set)
|
||||
initial_version = cast(int, coll.get_model()["version"])
|
||||
|
||||
for batch in create_batches(
|
||||
api=client,
|
||||
ids=cast(List[str], record_set["ids"]),
|
||||
embeddings=cast(Embeddings, record_set["embeddings"]),
|
||||
metadatas=cast(Metadatas, record_set["metadatas"]),
|
||||
documents=cast(List[str], record_set["documents"]),
|
||||
):
|
||||
coll.add(*batch)
|
||||
|
||||
if (
|
||||
not NOT_CLUSTER_ONLY
|
||||
and should_compact
|
||||
and len(normalized_record_set["ids"]) > 10
|
||||
):
|
||||
# Wait for the model to be updated, since the record set is larger, add some additional time
|
||||
wait_for_version_increase(
|
||||
client, collection.name, initial_version, additional_time=240
|
||||
)
|
||||
|
||||
invariants.count(coll, cast(strategies.RecordSet, normalized_record_set))
|
||||
|
||||
|
||||
@given(collection=collection_st)
|
||||
@settings(deadline=None, max_examples=1)
|
||||
def test_add_large_exceeding(
|
||||
client: ClientAPI, collection: strategies.Collection
|
||||
) -> None:
|
||||
create_isolated_database(client)
|
||||
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
|
||||
record_set = create_large_recordset(
|
||||
min_size=client.get_max_batch_size(),
|
||||
max_size=client.get_max_batch_size()
|
||||
+ 100, # Exceed the max batch size by 100 records
|
||||
)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
|
||||
with pytest.raises(Exception) as e:
|
||||
coll.add(**record_set) # type: ignore[arg-type]
|
||||
assert "batch size" in str(e.value)
|
||||
|
||||
|
||||
# TODO: This test fails right now because the ids are not sorted by the input order
|
||||
@pytest.mark.xfail(
|
||||
reason="This is expected to fail right now. We should change the API to sort the \
|
||||
ids by input order."
|
||||
)
|
||||
def test_out_of_order_ids(client: ClientAPI) -> None:
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
ooo_ids = [
|
||||
"40",
|
||||
"05",
|
||||
"8",
|
||||
"6",
|
||||
"10",
|
||||
"01",
|
||||
"00",
|
||||
"3",
|
||||
"04",
|
||||
"20",
|
||||
"02",
|
||||
"9",
|
||||
"30",
|
||||
"11",
|
||||
"13",
|
||||
"2",
|
||||
"0",
|
||||
"7",
|
||||
"06",
|
||||
"5",
|
||||
"50",
|
||||
"12",
|
||||
"03",
|
||||
"4",
|
||||
"1",
|
||||
]
|
||||
|
||||
coll = client.create_collection(
|
||||
"test", embedding_function=lambda input: [[1, 2, 3] for _ in input] # type: ignore
|
||||
)
|
||||
embeddings: Embeddings = [np.array([1, 2, 3]) for _ in ooo_ids]
|
||||
coll.add(ids=ooo_ids, embeddings=embeddings)
|
||||
get_ids = coll.get(ids=ooo_ids)["ids"]
|
||||
assert get_ids == ooo_ids
|
||||
|
||||
|
||||
def test_add_partial(client: ClientAPI) -> None:
|
||||
"""Tests adding a record set with some of the fields set to None."""
|
||||
|
||||
create_isolated_database(client)
|
||||
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"TODO @jai, come back and debug why CI runners fail with async + sync"
|
||||
)
|
||||
|
||||
coll = client.create_collection("test")
|
||||
# TODO: We need to clean up the api types to support this typing
|
||||
coll.add(
|
||||
ids=["1", "2", "3"],
|
||||
# All embeddings must be provided, or else None - no partial lists allowed
|
||||
embeddings=[[1, 2, 3], [1, 2, 3], [1, 2, 3]], # type: ignore
|
||||
# Metadatas can always be partial
|
||||
metadatas=[{"a": 1}, None, {"a": 3}], # type: ignore
|
||||
# Documents are optional if embeddings are provided
|
||||
documents=["a", "b", None], # type: ignore
|
||||
)
|
||||
|
||||
results = coll.get()
|
||||
assert results["ids"] == ["1", "2", "3"]
|
||||
assert results["metadatas"] == [{"a": 1}, None, {"a": 3}]
|
||||
assert results["documents"] == ["a", "b", None]
|
||||
@@ -0,0 +1,91 @@
|
||||
from hypothesis import given, strategies as st
|
||||
from chromadb.api.types import (
|
||||
optional_embeddings_to_base64_strings,
|
||||
optional_base64_strings_to_embeddings,
|
||||
)
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
|
||||
@given(st.lists(st.lists(st.integers(min_value=-128, max_value=127))))
|
||||
def test_base64_conversion_is_identity_i8(embeddings) -> None: # type: ignore
|
||||
b64_strings = optional_embeddings_to_base64_strings(embeddings)
|
||||
assert b64_strings is not None
|
||||
assert len(b64_strings) == len(embeddings)
|
||||
decoded_embeddings = optional_base64_strings_to_embeddings(b64_strings)
|
||||
for orig, decoded in zip(embeddings, decoded_embeddings): # type: ignore
|
||||
np.testing.assert_allclose(orig, decoded, rtol=1e-6)
|
||||
|
||||
|
||||
@given(st.lists(st.lists(st.floats(width=16))))
|
||||
def test_base64_conversion_is_identity_f16(embeddings) -> None: # type: ignore
|
||||
b64_strings = optional_embeddings_to_base64_strings(embeddings)
|
||||
assert b64_strings is not None
|
||||
assert len(b64_strings) == len(embeddings)
|
||||
decoded_embeddings = optional_base64_strings_to_embeddings(b64_strings)
|
||||
for orig, decoded in zip(embeddings, decoded_embeddings): # type: ignore
|
||||
np.testing.assert_allclose(orig, decoded, rtol=1e-6)
|
||||
|
||||
|
||||
@given(st.lists(st.lists(st.floats(width=32))))
|
||||
def test_base64_conversion_is_identity_f32(embeddings) -> None: # type: ignore
|
||||
b64_strings = optional_embeddings_to_base64_strings(embeddings)
|
||||
assert b64_strings is not None
|
||||
assert len(b64_strings) == len(embeddings)
|
||||
decoded_embeddings = optional_base64_strings_to_embeddings(b64_strings)
|
||||
for orig, decoded in zip(embeddings, decoded_embeddings): # type: ignore
|
||||
np.testing.assert_allclose(orig, decoded, rtol=1e-6)
|
||||
|
||||
|
||||
@given(st.lists(st.lists(st.floats(width=64))))
|
||||
def test_base64_conversion_is_identity_f64(embeddings) -> None: # type: ignore
|
||||
b64_strings = optional_embeddings_to_base64_strings(embeddings)
|
||||
assert b64_strings is not None
|
||||
assert len(b64_strings) == len(embeddings)
|
||||
decoded_embeddings = optional_base64_strings_to_embeddings(b64_strings)
|
||||
|
||||
expected_embeddings = []
|
||||
for embedding in embeddings:
|
||||
expected_embedding = []
|
||||
for value in embedding:
|
||||
if math.isnan(value):
|
||||
expected_embedding.append(float("nan"))
|
||||
elif value > np.finfo(np.float32).max:
|
||||
expected_embedding.append(float("inf"))
|
||||
elif value < np.finfo(np.float32).min:
|
||||
expected_embedding.append(float("-inf"))
|
||||
else:
|
||||
f32_value = np.float32(value)
|
||||
expected_embedding.append(float(f32_value))
|
||||
expected_embeddings.append(expected_embedding)
|
||||
|
||||
for orig, decoded in zip(expected_embeddings, decoded_embeddings): # type: ignore
|
||||
np.testing.assert_allclose(orig, decoded, rtol=1e-6)
|
||||
|
||||
|
||||
@given(st.lists(st.lists(st.floats(width=32))))
|
||||
def test_base64_conversion_numpy_is_identity_f32(embeddings) -> None: # type: ignore
|
||||
b64_strings = optional_embeddings_to_base64_strings(
|
||||
[np.array(embedding, dtype=np.float32) for embedding in embeddings]
|
||||
)
|
||||
assert b64_strings is not None
|
||||
assert len(b64_strings) == len(embeddings)
|
||||
decoded_embeddings = optional_base64_strings_to_embeddings(b64_strings)
|
||||
|
||||
expected_embeddings = []
|
||||
for embedding in embeddings:
|
||||
expected_embedding = []
|
||||
for value in embedding:
|
||||
if math.isnan(value):
|
||||
expected_embedding.append(float("nan"))
|
||||
elif value > np.finfo(np.float32).max:
|
||||
expected_embedding.append(float("inf"))
|
||||
elif value < np.finfo(np.float32).min:
|
||||
expected_embedding.append(float("-inf"))
|
||||
else:
|
||||
f32_value = np.float32(value)
|
||||
expected_embedding.append(float(f32_value))
|
||||
expected_embeddings.append(expected_embedding)
|
||||
|
||||
for orig, decoded in zip(expected_embeddings, decoded_embeddings): # type: ignore
|
||||
np.testing.assert_allclose(orig, decoded, rtol=1e-6)
|
||||
@@ -0,0 +1,134 @@
|
||||
from typing import Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pytest
|
||||
from hypothesis import given, strategies as st
|
||||
|
||||
from chromadb.api.fastapi import FastAPI
|
||||
|
||||
|
||||
def hostname_strategy() -> st.SearchStrategy[str]:
|
||||
label = st.text(
|
||||
alphabet=st.characters(min_codepoint=97, max_codepoint=122),
|
||||
min_size=1,
|
||||
max_size=63,
|
||||
)
|
||||
return st.lists(label, min_size=1, max_size=3).map("-".join)
|
||||
|
||||
|
||||
tld_list = ["com", "org", "net", "edu"]
|
||||
|
||||
|
||||
def domain_strategy() -> st.SearchStrategy[str]:
|
||||
label = st.text(
|
||||
alphabet=st.characters(min_codepoint=97, max_codepoint=122),
|
||||
min_size=1,
|
||||
max_size=63,
|
||||
)
|
||||
tld = st.sampled_from(tld_list)
|
||||
return st.tuples(label, tld).map(".".join)
|
||||
|
||||
|
||||
port_strategy = st.one_of(st.integers(min_value=1, max_value=65535), st.none())
|
||||
|
||||
ssl_enabled_strategy = st.booleans()
|
||||
|
||||
|
||||
def url_path_strategy() -> st.SearchStrategy[str]:
|
||||
path_segment = st.text(
|
||||
alphabet=st.sampled_from("abcdefghijklmnopqrstuvwxyz/-_"),
|
||||
min_size=1,
|
||||
max_size=10,
|
||||
)
|
||||
return (
|
||||
st.lists(path_segment, min_size=1, max_size=5)
|
||||
.map("/".join)
|
||||
.map(lambda x: "/" + x)
|
||||
)
|
||||
|
||||
|
||||
def is_valid_url(url: str) -> bool:
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
return all([parsed.scheme, parsed.netloc])
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def generate_valid_domain_url() -> st.SearchStrategy[str]:
|
||||
return st.builds(
|
||||
lambda url_scheme, hostname, url_path: f"{url_scheme}{hostname}{url_path}",
|
||||
url_scheme=st.sampled_from(["http://", "https://"]),
|
||||
hostname=domain_strategy(),
|
||||
url_path=url_path_strategy(),
|
||||
)
|
||||
|
||||
|
||||
def generate_invalid_domain_url() -> st.SearchStrategy[str]:
|
||||
return st.builds(
|
||||
lambda url_scheme, hostname, url_path: f"{url_scheme}{hostname}{url_path}",
|
||||
url_scheme=st.builds(
|
||||
lambda scheme, suffix: f"{scheme}{suffix}",
|
||||
scheme=st.text(max_size=10),
|
||||
suffix=st.sampled_from(["://", ":///", ":////", ""]),
|
||||
),
|
||||
hostname=domain_strategy(),
|
||||
url_path=url_path_strategy(),
|
||||
)
|
||||
|
||||
|
||||
host_or_domain_strategy = st.one_of(
|
||||
generate_valid_domain_url(), domain_strategy(), st.sampled_from(["localhost"])
|
||||
)
|
||||
|
||||
|
||||
@given(
|
||||
hostname=host_or_domain_strategy,
|
||||
port=port_strategy,
|
||||
ssl_enabled=ssl_enabled_strategy,
|
||||
default_api_path=st.sampled_from(["/api/v1", "/api/v2", None]),
|
||||
)
|
||||
def test_url_resolve(
|
||||
hostname: str,
|
||||
port: Optional[int],
|
||||
ssl_enabled: bool,
|
||||
default_api_path: Optional[str],
|
||||
) -> None:
|
||||
_url = FastAPI.resolve_url(
|
||||
chroma_server_host=hostname,
|
||||
chroma_server_http_port=port,
|
||||
chroma_server_ssl_enabled=ssl_enabled,
|
||||
default_api_path=default_api_path,
|
||||
)
|
||||
assert is_valid_url(_url), f"Invalid URL: {_url}"
|
||||
assert (
|
||||
_url.startswith("https") if ssl_enabled else _url.startswith("http")
|
||||
), f"Invalid URL: {_url} - SSL Enabled: {ssl_enabled}"
|
||||
if hostname.startswith("http"):
|
||||
assert ":" + str(port) not in _url, f"Port in URL not expected: {_url}"
|
||||
else:
|
||||
assert ":" + str(port) in _url, f"Port in URL expected: {_url}"
|
||||
if default_api_path:
|
||||
assert _url.endswith(default_api_path), f"Invalid URL: {_url}"
|
||||
|
||||
|
||||
@given(
|
||||
hostname=generate_invalid_domain_url(),
|
||||
port=port_strategy,
|
||||
ssl_enabled=ssl_enabled_strategy,
|
||||
default_api_path=st.sampled_from(["/api/v1", "/api/v2", None]),
|
||||
)
|
||||
def test_resolve_invalid(
|
||||
hostname: str,
|
||||
port: Optional[int],
|
||||
ssl_enabled: bool,
|
||||
default_api_path: Optional[str],
|
||||
) -> None:
|
||||
with pytest.raises(ValueError) as e:
|
||||
FastAPI.resolve_url(
|
||||
chroma_server_host=hostname,
|
||||
chroma_server_http_port=port,
|
||||
chroma_server_ssl_enabled=ssl_enabled,
|
||||
default_api_path=default_api_path,
|
||||
)
|
||||
assert "Invalid URL" in str(e.value)
|
||||
@@ -0,0 +1,338 @@
|
||||
import pytest
|
||||
import logging
|
||||
import hypothesis.strategies as st
|
||||
from chromadb.test.property.invariants import check_metadata
|
||||
import chromadb.test.property.strategies as strategies
|
||||
from chromadb.api import ClientAPI
|
||||
import chromadb.api.types as types
|
||||
from hypothesis.stateful import (
|
||||
Bundle,
|
||||
RuleBasedStateMachine,
|
||||
rule,
|
||||
initialize,
|
||||
multiple,
|
||||
consumes,
|
||||
run_state_machine_as_test,
|
||||
MultipleResults,
|
||||
)
|
||||
import chromadb.test.property.invariants as invariants
|
||||
from typing import Any, Dict, Mapping, Optional
|
||||
import numpy
|
||||
from chromadb.test.property.strategies import hashing_embedding_function
|
||||
|
||||
|
||||
class CollectionStateMachine(RuleBasedStateMachine):
|
||||
collections: Bundle[strategies.ExternalCollection]
|
||||
_model: Dict[str, Optional[types.CollectionMetadata]]
|
||||
|
||||
collections = Bundle("collections")
|
||||
|
||||
def __init__(self, client: ClientAPI):
|
||||
super().__init__()
|
||||
self._model = {}
|
||||
self.client = client
|
||||
|
||||
@initialize()
|
||||
def initialize(self) -> None:
|
||||
self.client.reset()
|
||||
self._model = {}
|
||||
|
||||
@rule(target=collections, coll=strategies.collections())
|
||||
def create_coll(
|
||||
self, coll: strategies.ExternalCollection
|
||||
) -> MultipleResults[strategies.ExternalCollection]:
|
||||
# Metadata can either be None or a non-empty dict
|
||||
if coll.name in self.model or (
|
||||
coll.metadata is not None and len(coll.metadata) == 0
|
||||
):
|
||||
with pytest.raises(Exception):
|
||||
c = self.client.create_collection(
|
||||
name=coll.name,
|
||||
metadata=coll.metadata, # type: ignore[arg-type]
|
||||
embedding_function=coll.embedding_function,
|
||||
)
|
||||
return multiple()
|
||||
|
||||
c = self.client.create_collection(
|
||||
name=coll.name,
|
||||
metadata=coll.metadata, # type: ignore[arg-type]
|
||||
embedding_function=coll.embedding_function,
|
||||
)
|
||||
self.set_model(coll.name, coll.metadata) # type: ignore[arg-type]
|
||||
|
||||
assert c.name == coll.name
|
||||
check_metadata(self.model[coll.name], c.metadata)
|
||||
return multiple(coll)
|
||||
|
||||
@rule(coll=collections)
|
||||
def get_coll(self, coll: strategies.ExternalCollection) -> None:
|
||||
if coll.name in self.model:
|
||||
c = self.client.get_collection(name=coll.name)
|
||||
assert c.name == coll.name
|
||||
check_metadata(self.model[coll.name], c.metadata)
|
||||
else:
|
||||
with pytest.raises(Exception):
|
||||
self.client.get_collection(name=coll.name)
|
||||
|
||||
@rule(coll=consumes(collections))
|
||||
def delete_coll(self, coll: strategies.ExternalCollection) -> None:
|
||||
if coll.name in self.model:
|
||||
with invariants.collection_deleted(self.client, coll.name):
|
||||
self.client.delete_collection(name=coll.name)
|
||||
self.delete_from_model(coll.name)
|
||||
else:
|
||||
with pytest.raises(Exception):
|
||||
self.client.delete_collection(name=coll.name)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
self.client.get_collection(name=coll.name)
|
||||
|
||||
@rule()
|
||||
def list_collections(self) -> None:
|
||||
colls = self.client.list_collections()
|
||||
assert len(colls) == len(self.model)
|
||||
for c in colls:
|
||||
assert c.name in self.model
|
||||
|
||||
# @rule for list_collections with limit and offset
|
||||
@rule(
|
||||
limit=st.integers(min_value=1, max_value=5),
|
||||
offset=st.integers(min_value=0, max_value=5),
|
||||
)
|
||||
def list_collections_with_limit_offset(self, limit: int, offset: int) -> None:
|
||||
colls = self.client.list_collections(limit=limit, offset=offset)
|
||||
total_collections = self.client.count_collections()
|
||||
|
||||
# get all collections
|
||||
all_colls = self.client.list_collections()
|
||||
# manually slice the collections based on the given limit and offset
|
||||
man_colls = all_colls[offset : offset + limit]
|
||||
|
||||
# given limit and offset, make various assertions regarding the total number of collections
|
||||
if limit + offset > total_collections:
|
||||
assert len(colls) == max(total_collections - offset, 0)
|
||||
# assert that our manually sliced collections are the same as the ones returned by the API
|
||||
assert colls == man_colls
|
||||
|
||||
else:
|
||||
assert len(colls) == limit
|
||||
|
||||
@rule(
|
||||
target=collections,
|
||||
new_metadata=st.one_of(st.none(), strategies.collection_metadata),
|
||||
coll=st.one_of(consumes(collections), strategies.collections()),
|
||||
)
|
||||
def get_or_create_coll(
|
||||
self,
|
||||
coll: strategies.ExternalCollection,
|
||||
new_metadata: Optional[types.Metadata],
|
||||
) -> MultipleResults[strategies.ExternalCollection]:
|
||||
# Cases for get_or_create
|
||||
|
||||
# Case 0
|
||||
# new_metadata is none, coll is an existing collection
|
||||
# get_or_create should return the existing collection with existing metadata
|
||||
|
||||
# Case 1
|
||||
# new_metadata is none, coll is a new collection
|
||||
# get_or_create should create a new collection with the metadata of None
|
||||
|
||||
# Case 2
|
||||
# new_metadata is not none, coll is an existing collection
|
||||
# get_or_create should return the existing collection with the original metadata
|
||||
|
||||
# Case 3
|
||||
# new_metadata is not none, coll is a new collection
|
||||
# get_or_create should create a new collection with the new metadata
|
||||
|
||||
if new_metadata is not None and len(new_metadata) == 0:
|
||||
with pytest.raises(Exception):
|
||||
c = self.client.get_or_create_collection(
|
||||
name=coll.name,
|
||||
metadata=new_metadata, # type: ignore[arg-type]
|
||||
embedding_function=coll.embedding_function,
|
||||
)
|
||||
return multiple()
|
||||
|
||||
# Update model
|
||||
if coll.name not in self.model:
|
||||
# Handles case 1 and 3
|
||||
coll.metadata = new_metadata
|
||||
self.set_model(coll.name, coll.metadata) # type: ignore[arg-type]
|
||||
|
||||
# Update API
|
||||
c = self.client.get_or_create_collection(
|
||||
name=coll.name,
|
||||
metadata=new_metadata, # type: ignore[arg-type]
|
||||
embedding_function=coll.embedding_function,
|
||||
)
|
||||
|
||||
# Check that model and API are in sync
|
||||
assert c.name == coll.name
|
||||
check_metadata(self.model[coll.name], c.metadata)
|
||||
return multiple(coll)
|
||||
|
||||
@rule(
|
||||
target=collections,
|
||||
coll=consumes(collections),
|
||||
new_metadata=strategies.collection_metadata,
|
||||
new_name=st.one_of(st.none(), strategies.collection_name()),
|
||||
)
|
||||
def modify_coll(
|
||||
self,
|
||||
coll: strategies.ExternalCollection,
|
||||
new_metadata: types.Metadata,
|
||||
new_name: Optional[str],
|
||||
) -> MultipleResults[strategies.ExternalCollection]:
|
||||
if coll.name not in self.model:
|
||||
with pytest.raises(Exception):
|
||||
c = self.client.get_collection(name=coll.name)
|
||||
return multiple()
|
||||
|
||||
c = self.client.get_collection(name=coll.name)
|
||||
_metadata: Optional[Mapping[str, Any]] = self.model[coll.name]
|
||||
_name: str = coll.name
|
||||
if new_metadata is not None:
|
||||
# Can't set metadata to an empty dict
|
||||
if len(new_metadata) == 0:
|
||||
with pytest.raises(Exception):
|
||||
c = self.client.get_or_create_collection(
|
||||
name=coll.name,
|
||||
metadata=new_metadata, # type: ignore[arg-type]
|
||||
embedding_function=coll.embedding_function,
|
||||
)
|
||||
return multiple()
|
||||
|
||||
coll.metadata = new_metadata
|
||||
_metadata = new_metadata
|
||||
|
||||
if new_name is not None:
|
||||
if new_name in self.model and new_name != coll.name:
|
||||
with pytest.raises(Exception):
|
||||
c.modify(metadata=new_metadata, name=new_name) # type: ignore[arg-type]
|
||||
return multiple()
|
||||
|
||||
self.delete_from_model(coll.name)
|
||||
coll.name = new_name
|
||||
_name = new_name
|
||||
|
||||
self.set_model(_name, _metadata) # type: ignore[arg-type]
|
||||
c.modify(metadata=_metadata, name=_name) # type: ignore[arg-type]
|
||||
c = self.client.get_collection(name=coll.name)
|
||||
|
||||
assert c.name == coll.name
|
||||
check_metadata(self.model[coll.name], c.metadata)
|
||||
return multiple(coll)
|
||||
|
||||
def set_model(
|
||||
self,
|
||||
name: str,
|
||||
metadata: Optional[types.CollectionMetadata],
|
||||
) -> None:
|
||||
model = self.model
|
||||
model[name] = metadata
|
||||
|
||||
def delete_from_model(self, name: str) -> None:
|
||||
model = self.model
|
||||
del model[name]
|
||||
|
||||
@property
|
||||
def model(self) -> Dict[str, Optional[types.CollectionMetadata]]:
|
||||
return self._model
|
||||
|
||||
|
||||
def test_collections(caplog: pytest.LogCaptureFixture, client: ClientAPI) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
run_state_machine_as_test(lambda: CollectionStateMachine(client)) # type: ignore
|
||||
|
||||
|
||||
# Below are tests that have failed in the past. If your test fails, please add
|
||||
# it to protect against regressions in the test harness itself. If you need
|
||||
# help doing so, talk to anton.
|
||||
|
||||
|
||||
def test_previously_failing_one(client: ClientAPI) -> None:
|
||||
state = CollectionStateMachine(client)
|
||||
state.initialize()
|
||||
# I don't know why the typechecker is red here. This code is correct and is
|
||||
# pulled from the logs.
|
||||
(v1,) = state.get_or_create_coll( # type: ignore[misc]
|
||||
coll=strategies.ExternalCollection(
|
||||
name="jjn2yjLW1zp2T",
|
||||
metadata=None,
|
||||
embedding_function=hashing_embedding_function(dtype=numpy.float32, dim=863), # type: ignore[arg-type]
|
||||
),
|
||||
new_metadata=None,
|
||||
)
|
||||
(v6,) = state.get_or_create_coll( # type: ignore[misc]
|
||||
coll=strategies.ExternalCollection(
|
||||
name="jjn2yjLW1zp2T",
|
||||
metadata=None,
|
||||
embedding_function=hashing_embedding_function(dtype=numpy.float32, dim=863), # type: ignore[arg-type]
|
||||
),
|
||||
new_metadata=None,
|
||||
)
|
||||
state.modify_coll(
|
||||
coll=v1, new_metadata={"7": -1281, "fGe": -0.0, "K5j": "im"}, new_name=None
|
||||
)
|
||||
state.modify_coll(coll=v6, new_metadata=None, new_name=None)
|
||||
|
||||
|
||||
# https://github.com/chroma-core/chroma/commit/cf476d70f0cebb7c87cb30c7172ba74d6ea175cd#diff-e81868b665d149bb315d86890dea6fc6a9fc9fc9ea3089aa7728142b54f622c5R210
|
||||
def test_previously_failing_two(client: ClientAPI) -> None:
|
||||
state = CollectionStateMachine(client)
|
||||
state.initialize()
|
||||
(v13,) = state.get_or_create_coll( # type: ignore[misc]
|
||||
coll=strategies.ExternalCollection(
|
||||
name="C1030",
|
||||
metadata={},
|
||||
embedding_function=hashing_embedding_function(dim=2, dtype=numpy.float32), # type: ignore[arg-type]
|
||||
),
|
||||
new_metadata=None,
|
||||
)
|
||||
(v15,) = state.modify_coll( # type: ignore[misc]
|
||||
coll=v13,
|
||||
new_metadata={
|
||||
"0": "10",
|
||||
"40": "0",
|
||||
"p1nviWeL7fO": "qN",
|
||||
"7b": "YS",
|
||||
"VYWq4LEMWjCo": True,
|
||||
},
|
||||
new_name="OF5F0MzbQg",
|
||||
)
|
||||
state.get_or_create_coll(
|
||||
coll=strategies.ExternalCollection(
|
||||
name="VS0QGh",
|
||||
metadata={
|
||||
"h": 5.681951615025145e-227,
|
||||
"A1": 61126,
|
||||
"uhUhLEEMfeC_kN": 2147483647,
|
||||
"weF": "pSP",
|
||||
"B3DSaP": False,
|
||||
"6H533K": 1.192092896e-07,
|
||||
},
|
||||
embedding_function=hashing_embedding_function( # type: ignore[arg-type]
|
||||
dim=1915, dtype=numpy.float32
|
||||
),
|
||||
),
|
||||
new_metadata={
|
||||
"xVW09xUpDZA": 31734,
|
||||
"g": 1.1,
|
||||
"n1dUTalF-MY": -1000000.0,
|
||||
"y": "G3EtXTZ",
|
||||
"ugXZ_hK": 5494,
|
||||
},
|
||||
)
|
||||
v17 = state.modify_coll( # noqa: F841
|
||||
coll=v15, new_metadata={"L35J2S": "K0l026"}, new_name="Ai1"
|
||||
)
|
||||
v18 = state.get_or_create_coll(coll=v13, new_metadata=None) # noqa: F841
|
||||
state.get_or_create_coll(
|
||||
coll=strategies.ExternalCollection(
|
||||
name="VS0QGh",
|
||||
metadata=None,
|
||||
embedding_function=hashing_embedding_function(dim=326, dtype=numpy.float16), # type: ignore[arg-type]
|
||||
),
|
||||
new_metadata=None,
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
import logging
|
||||
from typing import Dict, Optional, Tuple
|
||||
import pytest
|
||||
from chromadb.api import AdminAPI
|
||||
import chromadb.api.types as types
|
||||
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT
|
||||
from chromadb.test.conftest import ClientFactories
|
||||
from chromadb.test.property.test_collections import CollectionStateMachine
|
||||
from hypothesis.stateful import (
|
||||
Bundle,
|
||||
rule,
|
||||
initialize,
|
||||
multiple,
|
||||
run_state_machine_as_test,
|
||||
MultipleResults,
|
||||
)
|
||||
import chromadb.test.property.strategies as strategies
|
||||
|
||||
|
||||
class TenantDatabaseCollectionStateMachine(CollectionStateMachine):
|
||||
"""A collection state machine test that includes tenant and database information,
|
||||
and switches between them."""
|
||||
|
||||
tenants: Bundle # [str]
|
||||
databases: Bundle # [Tuple[str, str]] # database to tenant it belongs to
|
||||
tenant_to_database_to_model: Dict[
|
||||
str, Dict[str, Dict[str, Optional[types.CollectionMetadata]]]
|
||||
]
|
||||
admin_client: AdminAPI
|
||||
curr_tenant: str
|
||||
curr_database: str
|
||||
|
||||
tenants = Bundle("tenants")
|
||||
databases = Bundle("databases")
|
||||
|
||||
def __init__(self, client_factories: ClientFactories):
|
||||
client = client_factories.create_client()
|
||||
super().__init__(client)
|
||||
self.client = client
|
||||
self.admin_client = client_factories.create_admin_client_from_system()
|
||||
|
||||
@initialize()
|
||||
def initialize(self) -> None:
|
||||
self.client.reset()
|
||||
self.tenant_to_database_to_model = {}
|
||||
self.curr_tenant = DEFAULT_TENANT
|
||||
self.curr_database = DEFAULT_DATABASE
|
||||
self.client.set_tenant(DEFAULT_TENANT, DEFAULT_DATABASE)
|
||||
self.set_tenant_model(self.curr_tenant, {})
|
||||
self.set_database_model_for_tenant(self.curr_tenant, self.curr_database, {})
|
||||
|
||||
@rule(target=tenants, name=strategies.tenant_database_name)
|
||||
def create_tenant(self, name: str) -> MultipleResults: # [str]:
|
||||
tenant = self.overwrite_tenant(name)
|
||||
# Check if tenant already exists
|
||||
if self.has_tenant(tenant):
|
||||
with pytest.raises(Exception):
|
||||
self.admin_client.create_tenant(tenant)
|
||||
return multiple()
|
||||
|
||||
self.admin_client.create_tenant(tenant)
|
||||
# When we create a tenant, create a default database for it just for testing
|
||||
# since the state machine could call collection operations before creating a
|
||||
# database
|
||||
self.admin_client.create_database(DEFAULT_DATABASE, tenant=tenant)
|
||||
self.set_tenant_model(tenant, {})
|
||||
self.set_database_model_for_tenant(tenant, DEFAULT_DATABASE, {})
|
||||
return multiple(tenant)
|
||||
|
||||
@rule(target=databases, name=strategies.tenant_database_name)
|
||||
def create_database(self, name: str) -> MultipleResults: # [Tuple[str, str]]:
|
||||
database = self.overwrite_database(name)
|
||||
tenant = self.overwrite_tenant(self.curr_tenant)
|
||||
# If database already exists in current tenant, raise an error
|
||||
if self.has_database_for_tenant(tenant, database):
|
||||
with pytest.raises(Exception):
|
||||
self.admin_client.create_database(name=database, tenant=tenant)
|
||||
return multiple()
|
||||
|
||||
self.admin_client.create_database(name=database, tenant=tenant)
|
||||
self.set_database_model_for_tenant(
|
||||
tenant=tenant, database=database, database_model={}
|
||||
)
|
||||
return multiple((database, tenant))
|
||||
|
||||
@rule(database=databases)
|
||||
def set_database_and_tenant(self, database: Tuple[str, str]) -> None:
|
||||
# Get a database and switch to the database and the tenant it belongs to
|
||||
database_name = database[0]
|
||||
tenant_name = database[1]
|
||||
self.set_api_tenant_database(tenant_name, database_name)
|
||||
self.curr_database = database_name
|
||||
self.curr_tenant = tenant_name
|
||||
|
||||
@rule(tenant=tenants)
|
||||
def set_tenant(self, tenant: str) -> None:
|
||||
self.set_api_tenant_database(tenant, DEFAULT_DATABASE)
|
||||
self.curr_tenant = tenant
|
||||
self.curr_database = DEFAULT_DATABASE
|
||||
|
||||
# These methods allow other tests, namely
|
||||
# test_collections_with_database_tenant_override.py, to swap out the model
|
||||
# without needing to do a bunch of pythonic cleverness to fake a dict which
|
||||
# preteds to have every key.
|
||||
def set_api_tenant_database(self, tenant: str, database: str) -> None:
|
||||
self.client.set_tenant(tenant, database)
|
||||
|
||||
# For calls to create_database, and create_tenant we may want to override the tenant and database
|
||||
# This is a leaky abstraction that exists soley for the purpose of
|
||||
# test_collections_with_database_tenant_override.py
|
||||
def overwrite_tenant(self, tenant: str) -> str:
|
||||
return tenant
|
||||
|
||||
def overwrite_database(self, database: str) -> str:
|
||||
return database
|
||||
|
||||
def has_tenant(self, tenant: str) -> bool:
|
||||
return tenant in self.tenant_to_database_to_model
|
||||
|
||||
def get_tenant_model(
|
||||
self, tenant: str
|
||||
) -> Dict[str, Dict[str, Optional[types.CollectionMetadata]]]:
|
||||
return self.tenant_to_database_to_model[tenant]
|
||||
|
||||
def set_tenant_model(
|
||||
self,
|
||||
tenant: str,
|
||||
model: Dict[str, Dict[str, Optional[types.CollectionMetadata]]],
|
||||
) -> None:
|
||||
self.tenant_to_database_to_model[tenant] = model
|
||||
|
||||
def has_database_for_tenant(self, tenant: str, database: str) -> bool:
|
||||
return database in self.tenant_to_database_to_model[tenant]
|
||||
|
||||
def set_database_model_for_tenant(
|
||||
self,
|
||||
tenant: str,
|
||||
database: str,
|
||||
database_model: Dict[str, Optional[types.CollectionMetadata]],
|
||||
) -> None:
|
||||
self.tenant_to_database_to_model[tenant][database] = database_model
|
||||
|
||||
@property
|
||||
def model(self) -> Dict[str, Optional[types.CollectionMetadata]]:
|
||||
return self.tenant_to_database_to_model[self.curr_tenant][self.curr_database]
|
||||
|
||||
|
||||
def test_collections(
|
||||
caplog: pytest.LogCaptureFixture, client_factories: ClientFactories
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
run_state_machine_as_test(lambda: TenantDatabaseCollectionStateMachine(client_factories)) # type: ignore
|
||||
@@ -0,0 +1,214 @@
|
||||
from typing import Dict, Optional, Tuple
|
||||
from overrides import overrides
|
||||
from hypothesis.stateful import (
|
||||
initialize,
|
||||
invariant,
|
||||
rule,
|
||||
run_state_machine_as_test,
|
||||
)
|
||||
|
||||
import uuid
|
||||
import logging
|
||||
import pytest
|
||||
from chromadb.api import AdminAPI
|
||||
from chromadb.api.client import AdminClient, Client
|
||||
from chromadb.config import Settings, System
|
||||
from chromadb.test.conftest import (
|
||||
ClientFactories,
|
||||
fastapi_fixture_admin_and_singleton_tenant_db_user,
|
||||
)
|
||||
from chromadb.test.property.test_collections_with_database_tenant import (
|
||||
TenantDatabaseCollectionStateMachine,
|
||||
)
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import numpy
|
||||
import chromadb.api.types as types
|
||||
|
||||
# See conftest.py
|
||||
SINGLETON_TENANT = "singleton_tenant"
|
||||
SINGLETON_DATABASE = "singleton_database"
|
||||
|
||||
|
||||
class SingletonTenantDatabaseCollectionStateMachine(
|
||||
TenantDatabaseCollectionStateMachine
|
||||
):
|
||||
singleton_client: Client
|
||||
singleton_admin_client: AdminAPI
|
||||
root_client: Client
|
||||
root_admin_client: AdminAPI
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
singleton_client: Client,
|
||||
root_client: Client,
|
||||
client_factories: ClientFactories,
|
||||
) -> None:
|
||||
super().__init__(client_factories)
|
||||
self.root_client = root_client
|
||||
self.root_admin_client = self.admin_client
|
||||
|
||||
self.singleton_client = singleton_client
|
||||
self.singleton_admin_client = AdminClient.from_system(singleton_client._system)
|
||||
|
||||
@initialize()
|
||||
def initialize(self) -> None:
|
||||
# Make sure we're back to the root client and admin client before
|
||||
# doing reset/initialize things.
|
||||
self.client = self.root_client
|
||||
self.admin_client = self.root_admin_client
|
||||
|
||||
super().initialize()
|
||||
|
||||
self.root_admin_client.create_tenant(SINGLETON_TENANT)
|
||||
self.root_admin_client.create_database(SINGLETON_DATABASE, SINGLETON_TENANT)
|
||||
|
||||
self.set_tenant_model(SINGLETON_TENANT, {})
|
||||
self.set_database_model_for_tenant(SINGLETON_TENANT, SINGLETON_DATABASE, {})
|
||||
|
||||
@invariant()
|
||||
def check_api_and_admin_client_are_in_sync(self) -> None:
|
||||
if self.client == self.singleton_client:
|
||||
assert self.admin_client == self.singleton_admin_client
|
||||
else:
|
||||
assert self.admin_client == self.root_admin_client
|
||||
|
||||
@rule()
|
||||
def change_clients(self) -> None:
|
||||
if self.client == self.singleton_client:
|
||||
self.client = self.root_client
|
||||
self.admin_client = self.root_admin_client
|
||||
else:
|
||||
self.client = self.singleton_client
|
||||
self.admin_client = self.singleton_admin_client
|
||||
|
||||
@overrides
|
||||
def set_api_tenant_database(self, tenant: str, database: str) -> None:
|
||||
self.root_client.set_tenant(tenant, database)
|
||||
|
||||
@overrides
|
||||
def get_tenant_model(
|
||||
self, tenant: str
|
||||
) -> Dict[str, Dict[str, Optional[types.CollectionMetadata]]]:
|
||||
if self.client == self.singleton_client:
|
||||
tenant = SINGLETON_TENANT
|
||||
return self.tenant_to_database_to_model[tenant]
|
||||
|
||||
@overrides
|
||||
def set_tenant_model(
|
||||
self,
|
||||
tenant: str,
|
||||
model: Dict[str, Dict[str, Optional[types.CollectionMetadata]]],
|
||||
) -> None:
|
||||
if self.client == self.singleton_client:
|
||||
# This never happens because we never actually issue a
|
||||
# create_tenant call on singleton_tenant:
|
||||
# thanks to the above overriding of get_tenant_model(),
|
||||
# the underlying state machine test should always expect an error
|
||||
# when it sends the request, so shouldn't try to update the model.
|
||||
raise ValueError("trying to overwrite the model for singleton??")
|
||||
self.tenant_to_database_to_model[tenant] = model
|
||||
|
||||
@overrides
|
||||
def set_database_model_for_tenant(
|
||||
self,
|
||||
tenant: str,
|
||||
database: str,
|
||||
database_model: Dict[str, Optional[types.CollectionMetadata]],
|
||||
) -> None:
|
||||
if self.client == self.singleton_client:
|
||||
# This never happens because we never actually issue a
|
||||
# create_database call on (singleton_tenant, singleton_database):
|
||||
# thanks to the above overriding of has_database_for_tenant(),
|
||||
# the underlying state machine test should always expect an error
|
||||
# when it sends the request, so shouldn't try to update the model.
|
||||
raise ValueError("trying to overwrite the model for singleton??")
|
||||
self.tenant_to_database_to_model[tenant][database] = database_model
|
||||
|
||||
@overrides
|
||||
def overwrite_database(self, database: str) -> str:
|
||||
if self.client == self.singleton_client:
|
||||
return SINGLETON_DATABASE
|
||||
return database
|
||||
|
||||
@overrides
|
||||
def overwrite_tenant(self, tenant: str) -> str:
|
||||
if self.client == self.singleton_client:
|
||||
return SINGLETON_TENANT
|
||||
return tenant
|
||||
|
||||
@property
|
||||
def model(self) -> Dict[str, Optional[types.CollectionMetadata]]:
|
||||
if self.client == self.singleton_client:
|
||||
return self.tenant_to_database_to_model[SINGLETON_TENANT][
|
||||
SINGLETON_DATABASE
|
||||
]
|
||||
return self.tenant_to_database_to_model[self.curr_tenant][self.curr_database]
|
||||
|
||||
|
||||
def _singleton_and_root_clients() -> Tuple[Client, Client, ClientFactories]:
|
||||
api_fixture = fastapi_fixture_admin_and_singleton_tenant_db_user()
|
||||
sys: System = next(api_fixture)
|
||||
sys.reset_state()
|
||||
client_factories = ClientFactories(sys)
|
||||
root_client = client_factories.create_client()
|
||||
_root_admin_client = client_factories.create_admin_client_from_system()
|
||||
|
||||
# This is a little awkward but we have to create the tenant and DB
|
||||
# before we can instantiate a Client which connects to them. This also
|
||||
# means we need to manually populate state in the state machine.
|
||||
_root_admin_client.create_tenant(SINGLETON_TENANT)
|
||||
_root_admin_client.create_database(SINGLETON_DATABASE, SINGLETON_TENANT)
|
||||
|
||||
singleton_settings = Settings(**dict(sys.settings))
|
||||
singleton_settings.chroma_client_auth_credentials = "singleton-token"
|
||||
singleton_system = System(singleton_settings)
|
||||
singleton_system.start()
|
||||
singleton_client = Client.from_system(singleton_system)
|
||||
|
||||
return singleton_client, root_client, client_factories
|
||||
|
||||
|
||||
def test_collections_with_tenant_database_overwrite(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
singleton_client, root_client, client_factories = _singleton_and_root_clients()
|
||||
run_state_machine_as_test(
|
||||
lambda: SingletonTenantDatabaseCollectionStateMachine(
|
||||
singleton_client, root_client, client_factories
|
||||
)
|
||||
) # type: ignore
|
||||
|
||||
|
||||
def test_repeat_failure(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
singleton_client, root_client, client_factories = _singleton_and_root_clients()
|
||||
|
||||
state = SingletonTenantDatabaseCollectionStateMachine(
|
||||
singleton_client, root_client, client_factories
|
||||
)
|
||||
state.initialize()
|
||||
state.check_api_and_admin_client_are_in_sync()
|
||||
state.change_clients()
|
||||
state.check_api_and_admin_client_are_in_sync()
|
||||
state.create_coll(
|
||||
coll=strategies.Collection(
|
||||
name="A00",
|
||||
metadata=None,
|
||||
embedding_function=strategies.hashing_embedding_function(
|
||||
dim=2, dtype=numpy.float16 # type: ignore
|
||||
),
|
||||
id=uuid.UUID("c9bcb72f-92b1-4604-a8cb-084162dfe98b"),
|
||||
dimension=2,
|
||||
dtype=numpy.float16,
|
||||
known_metadata_keys={},
|
||||
known_document_keywords=[],
|
||||
has_documents=False,
|
||||
has_embeddings=True,
|
||||
)
|
||||
)
|
||||
state.teardown() # type: ignore
|
||||
@@ -0,0 +1,350 @@
|
||||
from multiprocessing.connection import Connection
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from typing import Generator, List, Tuple, Dict, Any, Callable, Type
|
||||
from hypothesis import given, settings
|
||||
import hypothesis.strategies as st
|
||||
import pytest
|
||||
import json
|
||||
from urllib import request
|
||||
from chromadb import config
|
||||
from chromadb.api.configuration import (
|
||||
ConfigurationParameter,
|
||||
EmbeddingsQueueConfigurationInternal,
|
||||
)
|
||||
from chromadb.api.types import Documents, EmbeddingFunction, Embeddings
|
||||
from chromadb.db.impl.sqlite import SqliteDB
|
||||
from chromadb.ingest.impl.utils import trigger_vector_segments_max_seq_id_migration
|
||||
from chromadb.segment import SegmentManager
|
||||
from chromadb.segment.impl.manager.local import LocalSegmentManager
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import chromadb.test.property.invariants as invariants
|
||||
from packaging import version as packaging_version
|
||||
import re
|
||||
import multiprocessing
|
||||
from chromadb.config import Settings
|
||||
from chromadb.api.client import Client as ClientCreator
|
||||
from chromadb.test.utils.cross_version import (
|
||||
switch_to_version,
|
||||
install_version,
|
||||
get_path_to_version_install,
|
||||
)
|
||||
|
||||
# Minimum persisted version we support, and other substantial change versions
|
||||
# 0.4.1 is the first version with persistence
|
||||
# 0.5.3 is the first version with the new API where the serverapi and client api return types and arguments differ
|
||||
BASELINE_VERSIONS = ["0.4.1", "0.5.3"]
|
||||
version_re = re.compile(r"^[0-9]+\.[0-9]+\.[0-9]+$")
|
||||
|
||||
# Some modules do not work across versions, since we upgrade our support for them, and should be explicitly reimported in the subprocess
|
||||
VERSIONED_MODULES = ["pydantic", "numpy", "tokenizers"]
|
||||
|
||||
|
||||
def versions() -> List[str]:
|
||||
"""Returns the pinned minimum version and the latest version of chromadb."""
|
||||
url = "https://pypi.org/pypi/chromadb/json"
|
||||
data = json.load(request.urlopen(request.Request(url)))
|
||||
versions = list(data["releases"].keys())
|
||||
# Older versions on pypi contain "devXYZ" suffixes
|
||||
versions = [v for v in versions if version_re.match(v)]
|
||||
versions.sort(key=packaging_version.Version)
|
||||
return BASELINE_VERSIONS + [versions[-1]]
|
||||
|
||||
|
||||
def _bool_to_int(metadata: Dict[str, Any]) -> Dict[str, Any]:
|
||||
metadata.update((k, 1) for k, v in metadata.items() if v is True)
|
||||
metadata.update((k, 0) for k, v in metadata.items() if v is False)
|
||||
return metadata
|
||||
|
||||
|
||||
def _patch_boolean_metadata(
|
||||
collection: strategies.Collection,
|
||||
embeddings: strategies.RecordSet,
|
||||
settings: Settings,
|
||||
) -> None:
|
||||
# Since the old version does not support boolean value metadata, we will convert
|
||||
# boolean value metadata to int
|
||||
collection_metadata = collection.metadata
|
||||
if collection_metadata is not None:
|
||||
_bool_to_int(collection_metadata) # type: ignore
|
||||
|
||||
if embeddings["metadatas"] is not None:
|
||||
if isinstance(embeddings["metadatas"], list):
|
||||
for metadata in embeddings["metadatas"]:
|
||||
if metadata is not None and isinstance(metadata, dict):
|
||||
_bool_to_int(metadata)
|
||||
elif isinstance(embeddings["metadatas"], dict):
|
||||
metadata = embeddings["metadatas"]
|
||||
_bool_to_int(metadata)
|
||||
|
||||
|
||||
def _patch_telemetry_client(
|
||||
collection: strategies.Collection,
|
||||
embeddings: strategies.RecordSet,
|
||||
settings: Settings,
|
||||
) -> None:
|
||||
# chroma 0.4.14 added OpenTelemetry, distinct from ProductTelemetry. Before 0.4.14
|
||||
# ProductTelemetry was simply called Telemetry.
|
||||
settings.chroma_telemetry_impl = "chromadb.telemetry.posthog.Posthog"
|
||||
|
||||
|
||||
version_patches: List[
|
||||
Tuple[str, Callable[[strategies.Collection, strategies.RecordSet, Settings], None]]
|
||||
] = [
|
||||
("0.4.3", _patch_boolean_metadata),
|
||||
("0.4.14", _patch_telemetry_client),
|
||||
]
|
||||
|
||||
|
||||
def patch_for_version(
|
||||
version: str,
|
||||
collection: strategies.Collection,
|
||||
embeddings: strategies.RecordSet,
|
||||
settings: Settings,
|
||||
) -> None:
|
||||
"""Override aspects of the collection and embeddings, before testing, to account for
|
||||
breaking changes in old versions."""
|
||||
|
||||
for patch_version, patch in version_patches:
|
||||
if packaging_version.Version(version) <= packaging_version.Version(
|
||||
patch_version
|
||||
):
|
||||
patch(collection, embeddings, settings)
|
||||
|
||||
|
||||
def api_import_for_version(module: Any, version: str) -> Type: # type: ignore
|
||||
if packaging_version.Version(version) <= packaging_version.Version("0.4.14"):
|
||||
return module.api.API # type: ignore
|
||||
return module.api.ServerAPI # type: ignore
|
||||
|
||||
|
||||
def configurations(versions: List[str]) -> List[Tuple[str, Settings]]:
|
||||
return [
|
||||
(
|
||||
version,
|
||||
Settings(
|
||||
chroma_api_impl="chromadb.api.rust.RustBindingsAPI"
|
||||
if "CHROMA_RUST_BINDINGS_TEST_ONLY" in os.environ
|
||||
else "chromadb.api.segment.SegmentAPI",
|
||||
chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager",
|
||||
allow_reset=True,
|
||||
is_persistent=True,
|
||||
persist_directory=tempfile.mkdtemp(),
|
||||
),
|
||||
)
|
||||
for version in versions
|
||||
]
|
||||
|
||||
|
||||
test_old_versions = versions()
|
||||
base_install_dir = tempfile.mkdtemp()
|
||||
|
||||
|
||||
# This fixture is not shared with the rest of the tests because it is unique in how it
|
||||
# installs the versions of chromadb
|
||||
@pytest.fixture(scope="module", params=configurations(test_old_versions)) # type: ignore
|
||||
def version_settings(request) -> Generator[Tuple[str, Settings], None, None]:
|
||||
configuration = request.param
|
||||
version = configuration[0]
|
||||
|
||||
install_version(version, {})
|
||||
yield configuration
|
||||
# Cleanup the installed version
|
||||
path = get_path_to_version_install(version)
|
||||
shutil.rmtree(path)
|
||||
# Cleanup the persisted data
|
||||
data_path = configuration[1].persist_directory
|
||||
if os.path.exists(data_path):
|
||||
shutil.rmtree(data_path, ignore_errors=True)
|
||||
|
||||
|
||||
class not_implemented_ef(EmbeddingFunction[Documents]):
|
||||
def __call__(self, input: Documents) -> Embeddings:
|
||||
assert False, "Embedding function should not be called"
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def persist_generated_data_with_old_version(
|
||||
version: str,
|
||||
settings: Settings,
|
||||
collection_strategy: strategies.Collection,
|
||||
embeddings_strategy: strategies.RecordSet,
|
||||
conn: Connection,
|
||||
) -> None:
|
||||
try:
|
||||
old_module = switch_to_version(version, VERSIONED_MODULES)
|
||||
# In 0.7.0 we switch to Rust client. The old versions are using the the python SegmentAPI client
|
||||
if "CHROMA_RUST_BINDINGS_TEST_ONLY" in os.environ and packaging_version.Version(
|
||||
version
|
||||
) < packaging_version.Version("0.7.0"):
|
||||
settings.chroma_api_impl = "chromadb.api.segment.SegmentAPI"
|
||||
system = old_module.config.System(settings)
|
||||
api = system.instance(api_import_for_version(old_module, version))
|
||||
system.start()
|
||||
|
||||
api.reset()
|
||||
# In 0.5.4 we changed the API of the server api level to
|
||||
# deal with collection models instead of collections
|
||||
# in order to work with this we need to wrap the api in a client
|
||||
# for versions greater than or equal to 0.5.4
|
||||
if packaging_version.Version(version) >= packaging_version.Version("0.5.4"):
|
||||
api = old_module.api.client.Client.from_system(system)
|
||||
coll = api.create_collection(
|
||||
name=collection_strategy.name,
|
||||
metadata=collection_strategy.metadata,
|
||||
# In order to test old versions, we can't rely on the not_implemented function
|
||||
embedding_function=not_implemented_ef(),
|
||||
)
|
||||
coll.add(**embeddings_strategy)
|
||||
|
||||
# Just use some basic checks for sanity and manual testing where you break the new
|
||||
# version
|
||||
|
||||
check_embeddings = invariants.wrap_all(embeddings_strategy)
|
||||
# Check count
|
||||
assert coll.count() == len(check_embeddings["embeddings"] or [])
|
||||
# Check ids
|
||||
result = coll.get()
|
||||
actual_ids = result["ids"]
|
||||
embedding_id_to_index = {id: i for i, id in enumerate(check_embeddings["ids"])}
|
||||
actual_ids = sorted(actual_ids, key=lambda id: embedding_id_to_index[id])
|
||||
assert actual_ids == check_embeddings["ids"]
|
||||
|
||||
# Leave writes on the queue to be processed by the next version's
|
||||
# segment manager so we can test cross version serialization
|
||||
# compatibility.
|
||||
system.instance(LocalSegmentManager).stop()
|
||||
coll.upsert(**embeddings_strategy)
|
||||
|
||||
# Shutdown system
|
||||
system.stop()
|
||||
except Exception as e:
|
||||
conn.send(e)
|
||||
raise e
|
||||
|
||||
|
||||
# Since we can't pickle the embedding function, we always generate record sets with embeddings
|
||||
collection_st: st.SearchStrategy[strategies.Collection] = st.shared(
|
||||
strategies.collections(
|
||||
with_hnsw_params=True,
|
||||
has_embeddings=True,
|
||||
# By default, these are set to 2000, which makes it unlikely that index mutations will ever be fully flushed
|
||||
max_hnsw_sync_threshold=10,
|
||||
max_hnsw_batch_size=10,
|
||||
with_persistent_hnsw_params=st.booleans(),
|
||||
),
|
||||
key="coll",
|
||||
)
|
||||
|
||||
|
||||
@given(
|
||||
collection_strategy=collection_st,
|
||||
embeddings_strategy=strategies.recordsets(collection_st, max_size=200),
|
||||
)
|
||||
@settings(deadline=None)
|
||||
def test_cycle_versions(
|
||||
version_settings: Tuple[str, Settings],
|
||||
collection_strategy: strategies.Collection,
|
||||
embeddings_strategy: strategies.RecordSet,
|
||||
) -> None:
|
||||
# Test backwards compatibility
|
||||
# For the current version, ensure that we can load a collection from
|
||||
# the previous versions
|
||||
version, settings = version_settings
|
||||
# The strategies can generate metadatas of malformed inputs. Other tests
|
||||
# will error check and cover these cases to make sure they error. Here we
|
||||
# just convert them to valid values since the error cases are already tested
|
||||
if embeddings_strategy["metadatas"] == {}:
|
||||
embeddings_strategy["metadatas"] = None
|
||||
if embeddings_strategy["metadatas"] is not None and isinstance(
|
||||
embeddings_strategy["metadatas"], list
|
||||
):
|
||||
embeddings_strategy["metadatas"] = [
|
||||
m if m is None or len(m) > 0 else None
|
||||
for m in embeddings_strategy["metadatas"]
|
||||
]
|
||||
|
||||
patch_for_version(version, collection_strategy, embeddings_strategy, settings)
|
||||
|
||||
# Can't pickle a function, and we won't need them
|
||||
collection_strategy.embedding_function = None
|
||||
collection_strategy.known_metadata_keys = {}
|
||||
|
||||
# Run the task in a separate process to avoid polluting the current process
|
||||
# with the old version. Using spawn instead of fork to avoid sharing the
|
||||
# current process memory which would cause the old version to be loaded
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
conn1, conn2 = multiprocessing.Pipe()
|
||||
p = ctx.Process(
|
||||
target=persist_generated_data_with_old_version,
|
||||
args=(version, settings, collection_strategy, embeddings_strategy, conn2),
|
||||
)
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
if conn1.poll():
|
||||
e = conn1.recv()
|
||||
raise e
|
||||
|
||||
p.close()
|
||||
|
||||
# Switch to the current version (local working directory) and check the invariants
|
||||
# are preserved for the collection
|
||||
system = config.System(settings)
|
||||
system.start()
|
||||
client = ClientCreator.from_system(system)
|
||||
coll = client.get_collection(
|
||||
name=collection_strategy.name,
|
||||
embedding_function=not_implemented_ef(), # type: ignore
|
||||
)
|
||||
|
||||
embeddings_queue = system.instance(SqliteDB)
|
||||
|
||||
# Automatic pruning should be disabled since embeddings_queue is non-empty
|
||||
if packaging_version.Version(version) < packaging_version.Version(
|
||||
"0.5.7"
|
||||
): # (automatic pruning is enabled by default in 0.5.7 and later)
|
||||
assert (
|
||||
embeddings_queue.config.get_parameter("automatically_purge").value is False
|
||||
)
|
||||
|
||||
# Update to True so log_size_below_max() invariant will pass
|
||||
embeddings_queue.set_config(
|
||||
EmbeddingsQueueConfigurationInternal(
|
||||
[ConfigurationParameter("automatically_purge", True)]
|
||||
)
|
||||
)
|
||||
|
||||
# Should be able to clean log immediately after updating
|
||||
|
||||
# 07/29/24: the max_seq_id for vector segments was moved from the pickled metadata file to SQLite.
|
||||
# Cleaning the log is dependent on vector segments migrating their max_seq_id from the pickled metadata file to SQLite.
|
||||
# Vector segments migrate this field automatically on init, but at this point the segment has not been loaded yet.
|
||||
if "CHROMA_RUST_BINDINGS_TEST_ONLY" in os.environ:
|
||||
# Trigger log purge in Rust impl
|
||||
invariants.count(coll, embeddings_strategy)
|
||||
else:
|
||||
trigger_vector_segments_max_seq_id_migration(
|
||||
embeddings_queue, system.instance(SegmentManager)
|
||||
)
|
||||
embeddings_queue.purge_log(coll.id)
|
||||
invariants.log_size_below_max(system, [coll], True)
|
||||
|
||||
# Should be able to add embeddings
|
||||
coll.add(**embeddings_strategy) # type: ignore
|
||||
|
||||
invariants.count(coll, embeddings_strategy)
|
||||
invariants.metadatas_match(coll, embeddings_strategy)
|
||||
invariants.documents_match(coll, embeddings_strategy)
|
||||
invariants.ids_match(coll, embeddings_strategy)
|
||||
invariants.ann_accuracy(coll, embeddings_strategy)
|
||||
invariants.log_size_below_max(system, [coll], True)
|
||||
|
||||
# Shutdown system
|
||||
system.stop()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,839 @@
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
import uuid
|
||||
from hypothesis import example, given, settings, HealthCheck
|
||||
import pytest
|
||||
from typing import cast
|
||||
from chromadb.api import ClientAPI
|
||||
from chromadb.test.property import invariants
|
||||
from chromadb.api.types import (
|
||||
Document,
|
||||
Documents,
|
||||
Embedding,
|
||||
Embeddings,
|
||||
GetResult,
|
||||
IDs,
|
||||
Metadata,
|
||||
Metadatas,
|
||||
Where,
|
||||
WhereDocument,
|
||||
)
|
||||
from chromadb.test.conftest import reset, NOT_CLUSTER_ONLY
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import hypothesis.strategies as st
|
||||
from chromadb.execution.expression.plan import Search
|
||||
from chromadb.execution.expression.operator import Knn, In, Key, Eq, And, Or, Contains, NotContains
|
||||
import logging
|
||||
from chromadb.test.utils.wait_for_version_increase import wait_for_version_increase
|
||||
import numpy as np
|
||||
from chromadb.api.models.Collection import Collection
|
||||
from chromadb.execution.expression.operator import Where as WhereExpr
|
||||
|
||||
|
||||
def _filter_where_clause(clause: Where, metadata: Optional[Metadata]) -> bool:
|
||||
"""Return true if the where clause is true for the given metadata map"""
|
||||
metadata = metadata or dict()
|
||||
key, expr = list(clause.items())[0]
|
||||
|
||||
# Handle the shorthand for equal: {key: val} where val is a simple value
|
||||
if (
|
||||
isinstance(expr, str)
|
||||
or isinstance(expr, bool)
|
||||
or isinstance(expr, int)
|
||||
or isinstance(expr, float)
|
||||
):
|
||||
return _filter_where_clause({key: {"$eq": expr}}, metadata) # type: ignore[dict-item]
|
||||
|
||||
# expr is a list of clauses
|
||||
if key == "$and":
|
||||
assert isinstance(expr, list)
|
||||
return all(_filter_where_clause(clause, metadata) for clause in expr)
|
||||
|
||||
if key == "$or":
|
||||
assert isinstance(expr, list)
|
||||
return any(_filter_where_clause(clause, metadata) for clause in expr)
|
||||
|
||||
# expr is an operator expression
|
||||
assert isinstance(expr, dict)
|
||||
op, val = list(expr.items())[0]
|
||||
assert isinstance(metadata, dict)
|
||||
if op == "$eq":
|
||||
return key in metadata and metadata[key] == val
|
||||
elif op == "$ne":
|
||||
return key not in metadata or metadata[key] != val
|
||||
elif op == "$in":
|
||||
return key in metadata and metadata[key] in val # type: ignore[operator]
|
||||
elif op == "$nin":
|
||||
return key not in metadata or metadata[key] not in val # type: ignore[operator]
|
||||
|
||||
# The following conditions only make sense for numeric values
|
||||
assert (
|
||||
key not in metadata
|
||||
or isinstance(metadata[key], int)
|
||||
or isinstance(metadata[key], float)
|
||||
)
|
||||
assert isinstance(val, int) or isinstance(val, float)
|
||||
if op == "$gt":
|
||||
return key in metadata and metadata[key] > val
|
||||
elif op == "$gte":
|
||||
return key in metadata and metadata[key] >= val
|
||||
elif op == "$lt":
|
||||
return key in metadata and metadata[key] < val
|
||||
elif op == "$lte":
|
||||
return key in metadata and metadata[key] <= val
|
||||
else:
|
||||
raise ValueError("Unknown operator: {}".format(key))
|
||||
|
||||
|
||||
def _filter_where_doc_clause(clause: WhereDocument, doc: Document) -> bool:
|
||||
key, expr = list(clause.items())[0]
|
||||
|
||||
if key == "$and":
|
||||
assert isinstance(expr, list)
|
||||
return all(_filter_where_doc_clause(clause, doc) for clause in expr)
|
||||
if key == "$or":
|
||||
assert isinstance(expr, list)
|
||||
return any(_filter_where_doc_clause(clause, doc) for clause in expr)
|
||||
|
||||
# Simple $contains clause
|
||||
assert isinstance(expr, str)
|
||||
if key == "$contains":
|
||||
if not doc:
|
||||
return False
|
||||
return expr in doc
|
||||
elif key == "$not_contains":
|
||||
if not doc:
|
||||
return True
|
||||
return expr not in doc
|
||||
else:
|
||||
raise ValueError("Unknown operator: {}".format(key))
|
||||
|
||||
|
||||
EMPTY_DICT: Dict[Any, Any] = {}
|
||||
EMPTY_STRING: str = ""
|
||||
|
||||
|
||||
def _filter_embedding_set(
|
||||
record_set: strategies.RecordSet, filter: strategies.Filter
|
||||
) -> IDs:
|
||||
"""Return IDs from the embedding set that match the given filter object
|
||||
If none match, return an empty list
|
||||
"""
|
||||
|
||||
normalized_record_set = invariants.wrap_all(record_set)
|
||||
ids = set(normalized_record_set["ids"])
|
||||
|
||||
filter_ids = filter["ids"]
|
||||
|
||||
if filter_ids is not None:
|
||||
filter_ids = invariants.wrap(filter_ids)
|
||||
assert filter_ids is not None
|
||||
# If the filter ids is an empty list then we treat that as get all
|
||||
if len(filter_ids) != 0:
|
||||
ids = ids.intersection(filter_ids)
|
||||
|
||||
for i in range(len(normalized_record_set["ids"])):
|
||||
if filter["where"]:
|
||||
metadatas: Metadatas
|
||||
if isinstance(normalized_record_set["metadatas"], list):
|
||||
metadatas = normalized_record_set["metadatas"] # type: ignore[assignment]
|
||||
else:
|
||||
metadatas = [EMPTY_DICT] * len(normalized_record_set["ids"])
|
||||
filter_where: Where = filter["where"]
|
||||
if not _filter_where_clause(filter_where, metadatas[i]):
|
||||
ids.discard(normalized_record_set["ids"][i])
|
||||
|
||||
if filter["where_document"]:
|
||||
documents = normalized_record_set["documents"] or [EMPTY_STRING] * len(
|
||||
normalized_record_set["ids"]
|
||||
)
|
||||
if not _filter_where_doc_clause(filter["where_document"], documents[i]):
|
||||
ids.discard(normalized_record_set["ids"][i])
|
||||
|
||||
return list(ids)
|
||||
|
||||
|
||||
class LegacyWhereWrapper(WhereExpr):
|
||||
"""
|
||||
Wraps old-style where/where_document dicts for testing.
|
||||
Converts where_document to use #document field and combines with where using $and.
|
||||
"""
|
||||
def __init__(self, where: Optional[Where] = None, where_document: Optional[WhereDocument] = None):
|
||||
self.where = where
|
||||
self.where_document = where_document
|
||||
|
||||
def _convert_where_document(self, where_doc: WhereDocument) -> Dict[str, Any]:
|
||||
"""Convert where_document filters to use #document field."""
|
||||
if not where_doc:
|
||||
return {}
|
||||
|
||||
# Handle logical operators recursively
|
||||
if "$and" in where_doc:
|
||||
and_clauses = where_doc["$and"]
|
||||
if isinstance(and_clauses, list):
|
||||
return {"$and": [self._convert_where_document(clause) for clause in and_clauses]}
|
||||
elif "$or" in where_doc:
|
||||
or_clauses = where_doc["$or"]
|
||||
if isinstance(or_clauses, list):
|
||||
return {"$or": [self._convert_where_document(clause) for clause in or_clauses]}
|
||||
|
||||
# Handle document operators - convert to #document field
|
||||
if "$contains" in where_doc:
|
||||
return {"#document": {"$contains": where_doc["$contains"]}}
|
||||
elif "$not_contains" in where_doc:
|
||||
return {"#document": {"$not_contains": where_doc["$not_contains"]}}
|
||||
|
||||
if "$regex" in where_doc:
|
||||
return {"#document": {"$regex": where_doc["$regex"]}}
|
||||
elif "$not_regex" in where_doc:
|
||||
return {"#document": {"$not_regex": where_doc["$not_regex"]}}
|
||||
|
||||
# Cast to dict for return
|
||||
return cast(Dict[str, Any], where_doc)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
# Combine where and where_document into a single where clause
|
||||
combined_where = None
|
||||
|
||||
# Build list of conditions to AND together
|
||||
conditions = []
|
||||
|
||||
if self.where:
|
||||
conditions.append(self.where)
|
||||
|
||||
if self.where_document:
|
||||
# Convert where_document to use #document field
|
||||
converted_doc_filter = self._convert_where_document(self.where_document)
|
||||
if converted_doc_filter:
|
||||
conditions.append(converted_doc_filter)
|
||||
|
||||
# Combine conditions with $and if needed
|
||||
if len(conditions) == 1:
|
||||
combined_where = conditions[0]
|
||||
elif len(conditions) > 1:
|
||||
combined_where = {"$and": conditions}
|
||||
|
||||
# Return the combined where clause directly
|
||||
if combined_where:
|
||||
return combined_where
|
||||
return {}
|
||||
|
||||
|
||||
def _search_with_filter(
|
||||
collection: Collection,
|
||||
filter: strategies.Filter,
|
||||
query_embedding: Optional[Embedding] = None,
|
||||
n_results: int = 10
|
||||
) -> List[str]:
|
||||
"""Use the search API to retrieve results with filters - test helper function."""
|
||||
# Build Search object
|
||||
search = Search()
|
||||
|
||||
# Add KNN if embedding provided
|
||||
if query_embedding is not None:
|
||||
search = search.rank(Knn(query=query_embedding)) # type: ignore[arg-type]
|
||||
|
||||
# Add filters using the LegacyWhereWrapper
|
||||
if filter.get("where") or filter.get("where_document") or filter.get("ids"):
|
||||
# Convert ids to list if it's a string
|
||||
ids_val = filter.get("ids")
|
||||
if isinstance(ids_val, str):
|
||||
ids_val = [ids_val]
|
||||
|
||||
# Build the where clause
|
||||
where_expr = None
|
||||
|
||||
# Add legacy where/where_document if present
|
||||
if filter.get("where") or filter.get("where_document"):
|
||||
wrapper = LegacyWhereWrapper(
|
||||
where=filter.get("where"),
|
||||
where_document=filter.get("where_document"),
|
||||
)
|
||||
if wrapper.to_dict(): # Only use if it has content
|
||||
where_expr = wrapper
|
||||
|
||||
# Add ID filter if present
|
||||
if ids_val:
|
||||
id_expr = Key.ID.is_in(ids_val)
|
||||
if where_expr:
|
||||
where_expr = where_expr & id_expr # type: ignore[assignment]
|
||||
else:
|
||||
where_expr = id_expr
|
||||
|
||||
# Apply the where clause if we have one
|
||||
if where_expr:
|
||||
search = search.where(where_expr)
|
||||
|
||||
# Set limit and select only IDs
|
||||
search = search.limit(n_results).select("id")
|
||||
|
||||
# Execute search and return IDs
|
||||
result = collection.search(search)
|
||||
return result["ids"][0] if result["ids"] else []
|
||||
|
||||
|
||||
collection_st = st.shared(
|
||||
strategies.collections(add_filterable_data=True, with_hnsw_params=True),
|
||||
key="coll",
|
||||
)
|
||||
recordset_st = st.shared(
|
||||
strategies.recordsets(collection_st, max_size=1000), key="recordset"
|
||||
)
|
||||
|
||||
|
||||
@settings(
|
||||
deadline=90000,
|
||||
suppress_health_check=[
|
||||
HealthCheck.function_scoped_fixture,
|
||||
HealthCheck.large_base_example,
|
||||
HealthCheck.filter_too_much,
|
||||
],
|
||||
) # type: ignore
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=recordset_st,
|
||||
filters=st.lists(strategies.filters(collection_st, recordset_st), min_size=1),
|
||||
should_compact=st.booleans(),
|
||||
)
|
||||
def test_filterable_metadata_get(
|
||||
caplog,
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set,
|
||||
filters,
|
||||
should_compact: bool,
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
|
||||
initial_version = coll.get_model()["version"]
|
||||
|
||||
coll.add(**record_set)
|
||||
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
# Only wait for compaction if the size of the collection is
|
||||
# some minimal size
|
||||
if should_compact and len(invariants.wrap(record_set["ids"])) > 10:
|
||||
# Wait for the model to be updated
|
||||
wait_for_version_increase(client, collection.name, initial_version) # type: ignore
|
||||
|
||||
for filter in filters:
|
||||
result_ids = coll.get(**filter)["ids"]
|
||||
expected_ids = _filter_embedding_set(record_set, filter)
|
||||
assert sorted(result_ids) == sorted(expected_ids)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
NOT_CLUSTER_ONLY,
|
||||
reason="Search API only available in distributed mode"
|
||||
)
|
||||
@settings(
|
||||
deadline=90000,
|
||||
suppress_health_check=[
|
||||
HealthCheck.function_scoped_fixture,
|
||||
HealthCheck.large_base_example,
|
||||
HealthCheck.filter_too_much,
|
||||
],
|
||||
) # type: ignore
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=recordset_st,
|
||||
filters=st.lists(strategies.filters(collection_st, recordset_st), min_size=1),
|
||||
should_compact=st.booleans(),
|
||||
)
|
||||
def test_filterable_metadata_search(
|
||||
caplog,
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set,
|
||||
filters,
|
||||
should_compact: bool,
|
||||
) -> None:
|
||||
"""Test metadata filtering using search API endpoint."""
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
|
||||
initial_version = coll.get_model()["version"]
|
||||
coll.add(**record_set)
|
||||
|
||||
if should_compact and len(invariants.wrap(record_set["ids"])) > 10:
|
||||
wait_for_version_increase(client, collection.name, initial_version) # type: ignore
|
||||
|
||||
for filter in filters:
|
||||
# Use search API instead of get
|
||||
result_ids = _search_with_filter(coll, filter, n_results=1000)
|
||||
expected_ids = _filter_embedding_set(record_set, filter)
|
||||
assert sorted(result_ids) == sorted(expected_ids)
|
||||
|
||||
|
||||
@settings(
|
||||
deadline=90000,
|
||||
suppress_health_check=[
|
||||
HealthCheck.function_scoped_fixture,
|
||||
HealthCheck.large_base_example,
|
||||
HealthCheck.filter_too_much,
|
||||
],
|
||||
) # type: ignore
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=recordset_st,
|
||||
filters=st.lists(strategies.filters(collection_st, recordset_st), min_size=1),
|
||||
limit=st.integers(min_value=1, max_value=10),
|
||||
offset=st.integers(min_value=0, max_value=10),
|
||||
should_compact=st.booleans(),
|
||||
)
|
||||
# Repro of a former off-by-one error in distributed Chroma. Fixed in https://github.com/chroma-core/chroma/pull/3489.
|
||||
@example(
|
||||
collection=strategies.Collection(
|
||||
name="test",
|
||||
metadata={"test": "test"},
|
||||
embedding_function=None,
|
||||
id=uuid.uuid4(),
|
||||
dimension=2,
|
||||
dtype="float32",
|
||||
known_metadata_keys={},
|
||||
known_document_keywords=[],
|
||||
),
|
||||
record_set=strategies.RecordSet(
|
||||
ids=[str(i) for i in range(11)],
|
||||
embeddings=[np.random.rand(2).tolist() for _ in range(11)],
|
||||
metadatas=[{"test": "test"} for _ in range(11)],
|
||||
documents=None,
|
||||
),
|
||||
filters=[
|
||||
strategies.Filter(
|
||||
{
|
||||
"where_document": {"$not_contains": "foo"},
|
||||
"ids": None,
|
||||
"where": None,
|
||||
}
|
||||
)
|
||||
],
|
||||
limit=10,
|
||||
offset=10,
|
||||
should_compact=True,
|
||||
)
|
||||
def test_filterable_metadata_get_limit_offset(
|
||||
caplog,
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set,
|
||||
filters,
|
||||
limit,
|
||||
offset,
|
||||
should_compact: bool,
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
|
||||
initial_version = coll.get_model()["version"]
|
||||
|
||||
coll.add(**record_set)
|
||||
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
# Only wait for compaction if the size of the collection is
|
||||
# some minimal size
|
||||
if should_compact and len(invariants.wrap(record_set["ids"])) > 10:
|
||||
# Wait for the model to be updated
|
||||
wait_for_version_increase(client, collection.name, initial_version) # type: ignore
|
||||
|
||||
for filter in filters:
|
||||
# add limit and offset to filter
|
||||
filter["limit"] = limit
|
||||
filter["offset"] = offset
|
||||
result_ids = coll.get(**filter)["ids"]
|
||||
expected_ids = _filter_embedding_set(record_set, filter)
|
||||
if len(expected_ids) > 0:
|
||||
collection_ids = coll.get(ids=expected_ids)["ids"]
|
||||
offset_id_order = {id: index for index, id in enumerate(collection_ids)}
|
||||
assert (
|
||||
result_ids
|
||||
== sorted(expected_ids, key=lambda id: offset_id_order[id])[
|
||||
offset : offset + limit
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@settings(
|
||||
deadline=90000,
|
||||
suppress_health_check=[
|
||||
HealthCheck.function_scoped_fixture,
|
||||
HealthCheck.large_base_example,
|
||||
HealthCheck.filter_too_much,
|
||||
],
|
||||
)
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=recordset_st,
|
||||
filters=st.lists(
|
||||
strategies.filters(collection_st, recordset_st, include_all_ids=True),
|
||||
min_size=1,
|
||||
),
|
||||
should_compact=st.booleans(),
|
||||
data=st.data(),
|
||||
)
|
||||
def test_filterable_metadata_query(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
filters: List[strategies.Filter],
|
||||
should_compact: bool,
|
||||
data: st.DataObject,
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
initial_version = coll.get_model()["version"]
|
||||
normalized_record_set = invariants.wrap_all(record_set)
|
||||
|
||||
coll.add(**record_set) # type: ignore[arg-type]
|
||||
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
# Only wait for compaction if the size of the collection is
|
||||
# some minimal size
|
||||
if should_compact and len(invariants.wrap(record_set["ids"])) > 10:
|
||||
# Wait for the model to be updated
|
||||
wait_for_version_increase(client, collection.name, initial_version) # type: ignore
|
||||
|
||||
total_count = len(normalized_record_set["ids"])
|
||||
# Pick a random vector using Hypothesis data
|
||||
random_query: Embedding
|
||||
|
||||
query_index = data.draw(st.integers(min_value=0, max_value=total_count - 1))
|
||||
if collection.has_embeddings:
|
||||
assert normalized_record_set["embeddings"] is not None
|
||||
assert all(isinstance(e, list) for e in normalized_record_set["embeddings"])
|
||||
# Use data.draw to select index
|
||||
random_query = normalized_record_set["embeddings"][query_index]
|
||||
else:
|
||||
assert isinstance(normalized_record_set["documents"], list)
|
||||
assert collection.embedding_function is not None
|
||||
# Use data.draw to select index
|
||||
random_query = collection.embedding_function(
|
||||
[normalized_record_set["documents"][query_index]]
|
||||
)[0]
|
||||
for filter in filters:
|
||||
result_ids = set(
|
||||
coll.query(
|
||||
query_embeddings=random_query,
|
||||
n_results=total_count,
|
||||
where=filter["where"],
|
||||
where_document=filter["where_document"],
|
||||
)["ids"][0]
|
||||
)
|
||||
expected_ids = set(
|
||||
_filter_embedding_set(
|
||||
cast(strategies.RecordSet, normalized_record_set), filter
|
||||
)
|
||||
)
|
||||
assert len(result_ids.intersection(expected_ids)) == len(result_ids)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
NOT_CLUSTER_ONLY,
|
||||
reason="Search API only available in distributed mode"
|
||||
)
|
||||
@settings(
|
||||
deadline=90000,
|
||||
suppress_health_check=[
|
||||
HealthCheck.function_scoped_fixture,
|
||||
HealthCheck.large_base_example,
|
||||
HealthCheck.filter_too_much,
|
||||
],
|
||||
)
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=recordset_st,
|
||||
filters=st.lists(
|
||||
strategies.filters(collection_st, recordset_st, include_all_ids=True),
|
||||
min_size=1,
|
||||
),
|
||||
should_compact=st.booleans(),
|
||||
data=st.data(),
|
||||
)
|
||||
def test_filterable_metadata_query_via_search(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
filters: List[strategies.Filter],
|
||||
should_compact: bool,
|
||||
data: st.DataObject,
|
||||
) -> None:
|
||||
"""Test query-like filtering using search API endpoint."""
|
||||
caplog.set_level(logging.ERROR)
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
|
||||
initial_version = coll.get_model()["version"]
|
||||
normalized_record_set = invariants.wrap_all(record_set)
|
||||
coll.add(**record_set) # type: ignore[arg-type]
|
||||
|
||||
if should_compact and len(invariants.wrap(record_set["ids"])) > 10:
|
||||
wait_for_version_increase(client, collection.name, initial_version) # type: ignore
|
||||
|
||||
total_count = len(normalized_record_set["ids"])
|
||||
|
||||
# Pick a random query embedding
|
||||
query_index = data.draw(st.integers(min_value=0, max_value=total_count - 1))
|
||||
if collection.has_embeddings:
|
||||
assert normalized_record_set["embeddings"] is not None
|
||||
random_query = normalized_record_set["embeddings"][query_index]
|
||||
else:
|
||||
assert isinstance(normalized_record_set["documents"], list)
|
||||
assert collection.embedding_function is not None
|
||||
random_query = collection.embedding_function(
|
||||
[normalized_record_set["documents"][query_index]]
|
||||
)[0]
|
||||
|
||||
for filter in filters:
|
||||
# Use search API with query embedding
|
||||
result_ids = set(_search_with_filter(
|
||||
coll,
|
||||
filter,
|
||||
query_embedding=random_query,
|
||||
n_results=total_count
|
||||
))
|
||||
expected_ids = set(
|
||||
_filter_embedding_set(
|
||||
cast(strategies.RecordSet, normalized_record_set), filter
|
||||
)
|
||||
)
|
||||
assert len(result_ids.intersection(expected_ids)) == len(result_ids)
|
||||
|
||||
|
||||
def test_empty_filter(client: ClientAPI) -> None:
|
||||
"""Test that a filter where no document matches returns an empty result"""
|
||||
reset(client)
|
||||
coll = client.create_collection(name="test")
|
||||
|
||||
test_ids: IDs = ["1", "2", "3"]
|
||||
test_embeddings: Embeddings = [np.array([1, 1]), np.array([2, 2]), np.array([3, 3])]
|
||||
test_query_embedding: Embedding = np.array([1, 2])
|
||||
test_query_embeddings: Embeddings = [test_query_embedding, test_query_embedding]
|
||||
|
||||
coll.add(ids=test_ids, embeddings=test_embeddings)
|
||||
|
||||
res = coll.query(
|
||||
query_embeddings=test_query_embedding,
|
||||
where={"q": {"$eq": 4}}, # type: ignore[dict-item]
|
||||
n_results=3,
|
||||
include=["embeddings", "distances", "metadatas"],
|
||||
)
|
||||
assert res["ids"] == [[]]
|
||||
if res["embeddings"] is not None:
|
||||
assert cast(np.ndarray, res["embeddings"][0]).size == 0 # type: ignore
|
||||
assert res["distances"] == [[]]
|
||||
assert res["metadatas"] == [[]]
|
||||
assert set(res["included"]) == set(["embeddings", "distances", "metadatas"])
|
||||
|
||||
res = coll.query(
|
||||
query_embeddings=test_query_embeddings,
|
||||
where={"test": "yes"},
|
||||
n_results=3,
|
||||
)
|
||||
assert res["ids"] == [[], []]
|
||||
assert res["embeddings"] is None
|
||||
assert res["distances"] == [[], []]
|
||||
assert res["metadatas"] == [[], []]
|
||||
assert set(res["included"]) == set(["metadatas", "documents", "distances"])
|
||||
|
||||
|
||||
def test_boolean_metadata(client: ClientAPI) -> None:
|
||||
"""Test that metadata with boolean values is correctly filtered"""
|
||||
reset(client)
|
||||
coll = client.create_collection(name="test")
|
||||
|
||||
test_ids: IDs = ["1", "2", "3"]
|
||||
test_embeddings: Embeddings = [np.array([1, 1]), np.array([2, 2]), np.array([3, 3])]
|
||||
test_metadatas: Metadatas = [{"test": True}, {"test": False}, {"test": True}]
|
||||
|
||||
coll.add(ids=test_ids, embeddings=test_embeddings, metadatas=test_metadatas)
|
||||
|
||||
res = coll.get(where={"test": True})
|
||||
|
||||
assert res["ids"] == ["1", "3"]
|
||||
|
||||
|
||||
def test_get_empty(client: ClientAPI) -> None:
|
||||
"""Tests that calling get() with empty filters returns nothing"""
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(name="test")
|
||||
|
||||
test_ids: IDs = ["1", "2", "3"]
|
||||
test_embeddings: Embeddings = [np.array([1, 1]), np.array([2, 2]), np.array([3, 3])]
|
||||
test_metadatas: Metadatas = [{"test": 10}, {"test": 20}, {"test": 30}]
|
||||
|
||||
def check_empty_res(res: GetResult) -> None:
|
||||
assert len(res["ids"]) == 0
|
||||
assert res["embeddings"] is not None
|
||||
assert len(res["embeddings"]) == 0
|
||||
assert res["documents"] is not None
|
||||
assert len(res["documents"]) == 0
|
||||
assert res["metadatas"] is not None
|
||||
|
||||
coll.add(ids=test_ids, embeddings=test_embeddings, metadatas=test_metadatas)
|
||||
|
||||
res = coll.get(ids=["nope"], include=["embeddings", "metadatas", "documents"])
|
||||
check_empty_res(res)
|
||||
res = coll.get(
|
||||
include=["embeddings", "metadatas", "documents"], where={"test": 100}
|
||||
)
|
||||
check_empty_res(res)
|
||||
|
||||
|
||||
@settings(
|
||||
deadline=90000,
|
||||
suppress_health_check=[
|
||||
HealthCheck.function_scoped_fixture,
|
||||
HealthCheck.large_base_example,
|
||||
],
|
||||
)
|
||||
@given(
|
||||
collection=collection_st,
|
||||
record_set=recordset_st,
|
||||
n_results_st=st.integers(min_value=1, max_value=100),
|
||||
should_compact=st.booleans(),
|
||||
data=st.data(),
|
||||
)
|
||||
def test_query_ids_filter_property(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
client: ClientAPI,
|
||||
collection: strategies.Collection,
|
||||
record_set: strategies.RecordSet,
|
||||
n_results_st: int,
|
||||
should_compact: bool,
|
||||
data: st.DataObject,
|
||||
) -> None:
|
||||
"""Property test for querying with only the ids filter."""
|
||||
if (
|
||||
client.get_settings().chroma_api_impl
|
||||
== "chromadb.api.async_fastapi.AsyncFastAPI"
|
||||
):
|
||||
pytest.skip(
|
||||
"Skipping test for async client due to potential resource/timeout issues"
|
||||
)
|
||||
caplog.set_level(logging.ERROR)
|
||||
reset(client)
|
||||
coll = client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
initial_version = coll.get_model()["version"]
|
||||
normalized_record_set = invariants.wrap_all(record_set)
|
||||
|
||||
if len(normalized_record_set["ids"]) == 0:
|
||||
# Cannot add empty record set
|
||||
return
|
||||
|
||||
coll.add(**record_set) # type: ignore[arg-type]
|
||||
|
||||
if not NOT_CLUSTER_ONLY:
|
||||
if should_compact and len(normalized_record_set["ids"]) > 10:
|
||||
wait_for_version_increase(client, collection.name, initial_version) # type: ignore
|
||||
|
||||
total_count = len(normalized_record_set["ids"])
|
||||
n_results = min(n_results_st, total_count)
|
||||
|
||||
# Generate a random subset of ids to filter on using Hypothesis data
|
||||
ids_to_query = data.draw(
|
||||
st.lists(
|
||||
st.sampled_from(normalized_record_set["ids"]),
|
||||
min_size=0,
|
||||
max_size=total_count,
|
||||
unique=True,
|
||||
)
|
||||
)
|
||||
|
||||
# Pick a random query vector using Hypothesis data
|
||||
random_query: Embedding
|
||||
query_index = data.draw(st.integers(min_value=0, max_value=total_count - 1))
|
||||
if collection.has_embeddings:
|
||||
assert normalized_record_set["embeddings"] is not None
|
||||
assert all(isinstance(e, list) for e in normalized_record_set["embeddings"])
|
||||
# Use data.draw to select index
|
||||
random_query = normalized_record_set["embeddings"][query_index]
|
||||
else:
|
||||
assert isinstance(normalized_record_set["documents"], list)
|
||||
assert collection.embedding_function is not None
|
||||
# Use data.draw to select index
|
||||
random_query = collection.embedding_function(
|
||||
[normalized_record_set["documents"][query_index]]
|
||||
)[0]
|
||||
|
||||
# Perform the query with only the ids filter
|
||||
result = coll.query(
|
||||
query_embeddings=[random_query],
|
||||
ids=ids_to_query,
|
||||
n_results=n_results,
|
||||
)
|
||||
|
||||
result_ids = set(result["ids"][0])
|
||||
filter_ids_set = set(ids_to_query)
|
||||
|
||||
# The core assertion: all returned IDs must be within the filter set
|
||||
assert result_ids.issubset(filter_ids_set)
|
||||
|
||||
# Also check that the number of results is reasonable
|
||||
assert len(result_ids) <= n_results
|
||||
assert len(result_ids) <= len(filter_ids_set)
|
||||
|
||||
|
||||
def test_regex(client: ClientAPI) -> None:
|
||||
"""Tests that regex works"""
|
||||
|
||||
reset(client)
|
||||
coll = client.create_collection(name="test")
|
||||
|
||||
test_ids: IDs = ["1", "2", "3"]
|
||||
test_documents: Documents = ["cat", "Cat", "CAT"]
|
||||
test_embeddings: Embeddings = [np.array([1, 1]), np.array([2, 2]), np.array([3, 3])]
|
||||
test_metadatas: Metadatas = [{"test": 10}, {"test": 20}, {"test": 30}]
|
||||
|
||||
coll.add(
|
||||
ids=test_ids,
|
||||
documents=test_documents,
|
||||
embeddings=test_embeddings,
|
||||
metadatas=test_metadatas,
|
||||
)
|
||||
|
||||
res = coll.get(where_document={"$regex": "cat"})
|
||||
assert res["ids"] == ["1"]
|
||||
|
||||
res = coll.get(where_document={"$regex": "(?i)cat"})
|
||||
assert sorted(res["ids"]) == ["1", "2", "3"]
|
||||
|
||||
res = coll.get(
|
||||
where={"test": {"$ne": 10}}, where_document={"$regex": "(?i)c(?-i)at"} # type: ignore[dict-item]
|
||||
)
|
||||
assert res["ids"] == ["2"]
|
||||
@@ -0,0 +1,192 @@
|
||||
import chromadb
|
||||
import chromadb.test.property.invariants as invariants
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import copy
|
||||
import hypothesis.strategies as hyst
|
||||
import logging
|
||||
import pytest
|
||||
|
||||
from chromadb.api.models.Collection import Collection
|
||||
from chromadb.test.conftest import reset, skip_if_not_cluster
|
||||
from chromadb.test.utils.wait_for_version_increase import wait_for_version_increase
|
||||
from hypothesis.stateful import (
|
||||
Bundle,
|
||||
RuleBasedStateMachine,
|
||||
rule,
|
||||
initialize,
|
||||
multiple,
|
||||
consumes,
|
||||
run_state_machine_as_test,
|
||||
MultipleResults,
|
||||
)
|
||||
from overrides import overrides
|
||||
from typing import Dict, cast, Union, Tuple, Set
|
||||
|
||||
collection_st = hyst.shared(strategies.collections(with_hnsw_params=True), key="source")
|
||||
|
||||
|
||||
class ForkStateMachine(RuleBasedStateMachine):
|
||||
updated_collections: Bundle[
|
||||
Tuple[Collection, strategies.StateMachineRecordSet]
|
||||
] = Bundle("changing_collections")
|
||||
forked_collections: Bundle[
|
||||
Tuple[Collection, strategies.StateMachineRecordSet]
|
||||
] = Bundle("collections")
|
||||
collection_names: Set[str]
|
||||
|
||||
def __init__(self, client: chromadb.api.ClientAPI):
|
||||
super().__init__()
|
||||
self.client = client
|
||||
self.collection_names = set()
|
||||
|
||||
@initialize(collection=collection_st, target=updated_collections)
|
||||
def initialize(
|
||||
self, collection: strategies.Collection
|
||||
) -> Tuple[Collection, strategies.StateMachineRecordSet]:
|
||||
source = self.client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore[arg-type]
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
self.collection_names.add(source.name)
|
||||
return source, strategies.StateMachineRecordSet(
|
||||
ids=[], metadatas=[], documents=[], embeddings=[]
|
||||
)
|
||||
|
||||
@overrides
|
||||
def teardown(self) -> None:
|
||||
reset(self.client)
|
||||
|
||||
@rule(
|
||||
source=consumes(updated_collections),
|
||||
new_name=strategies.collection_name(),
|
||||
target=forked_collections,
|
||||
)
|
||||
def fork(
|
||||
self, source: Tuple[Collection, strategies.StateMachineRecordSet], new_name: str
|
||||
) -> MultipleResults[Tuple[Collection, strategies.StateMachineRecordSet]]:
|
||||
collection, record_set = source
|
||||
if new_name in self.collection_names:
|
||||
with pytest.raises(Exception):
|
||||
collection.fork(new_name)
|
||||
return multiple(source)
|
||||
|
||||
target = collection.fork(new_name)
|
||||
self.collection_names.add(target.name)
|
||||
return multiple(source, (target, copy.deepcopy(record_set)))
|
||||
|
||||
@rule(
|
||||
cursor=consumes(forked_collections),
|
||||
delta=strategies.recordsets(collection_st),
|
||||
target=updated_collections,
|
||||
)
|
||||
def upsert(
|
||||
self,
|
||||
cursor: Tuple[Collection, strategies.StateMachineRecordSet],
|
||||
delta: strategies.RecordSet,
|
||||
) -> Tuple[Collection, strategies.StateMachineRecordSet]:
|
||||
collection, record_set_state = cursor
|
||||
normalized_delta: strategies.NormalizedRecordSet = invariants.wrap_all(delta)
|
||||
collection.upsert(**normalized_delta) # type: ignore[arg-type]
|
||||
for idx, id in enumerate(normalized_delta["ids"]):
|
||||
if id in record_set_state["ids"]:
|
||||
target_idx = record_set_state["ids"].index(id)
|
||||
if normalized_delta["embeddings"] is not None:
|
||||
record_set_state["embeddings"][target_idx] = normalized_delta[
|
||||
"embeddings"
|
||||
][idx]
|
||||
else:
|
||||
assert normalized_delta["documents"] is not None
|
||||
assert collection._embedding_function is not None
|
||||
record_set_state["embeddings"][
|
||||
target_idx
|
||||
] = collection._embedding_function(
|
||||
[normalized_delta["documents"][idx]]
|
||||
)[
|
||||
0
|
||||
]
|
||||
if normalized_delta["metadatas"] is not None:
|
||||
record_set_state_metadata = cast(
|
||||
Dict[str, Union[str, int, float]],
|
||||
record_set_state["metadatas"][target_idx],
|
||||
)
|
||||
if record_set_state_metadata is not None:
|
||||
if normalized_delta["metadatas"][idx] is not None:
|
||||
record_set_state_metadata.update(
|
||||
normalized_delta["metadatas"][idx] # type: ignore[arg-type]
|
||||
)
|
||||
else:
|
||||
record_set_state["metadatas"][target_idx] = normalized_delta[
|
||||
"metadatas"
|
||||
][idx]
|
||||
if normalized_delta["documents"] is not None:
|
||||
record_set_state["documents"][target_idx] = normalized_delta[
|
||||
"documents"
|
||||
][idx]
|
||||
else:
|
||||
record_set_state["ids"].append(id)
|
||||
if normalized_delta["embeddings"] is not None:
|
||||
record_set_state["embeddings"].append(
|
||||
normalized_delta["embeddings"][idx]
|
||||
)
|
||||
else:
|
||||
assert collection._embedding_function is not None
|
||||
assert normalized_delta["documents"] is not None
|
||||
record_set_state["embeddings"].append(
|
||||
collection._embedding_function(
|
||||
[normalized_delta["documents"][idx]]
|
||||
)[0]
|
||||
)
|
||||
if normalized_delta["metadatas"] is not None:
|
||||
record_set_state["metadatas"].append(
|
||||
normalized_delta["metadatas"][idx]
|
||||
)
|
||||
else:
|
||||
record_set_state["metadatas"].append(None)
|
||||
if normalized_delta["documents"] is not None:
|
||||
record_set_state["documents"].append(
|
||||
normalized_delta["documents"][idx]
|
||||
)
|
||||
else:
|
||||
record_set_state["documents"].append(None)
|
||||
return collection, record_set_state
|
||||
|
||||
@rule(
|
||||
cursor=consumes(forked_collections),
|
||||
target=updated_collections,
|
||||
)
|
||||
def delete(
|
||||
self, cursor: Tuple[Collection, strategies.StateMachineRecordSet]
|
||||
) -> Tuple[Collection, strategies.StateMachineRecordSet]:
|
||||
collection, record_set_state = cursor
|
||||
boundary = len(record_set_state["ids"]) // 10
|
||||
if boundary == 0:
|
||||
return collection, record_set_state
|
||||
ids_to_delete = record_set_state["ids"][:boundary]
|
||||
collection.delete(ids_to_delete)
|
||||
record_set_state["ids"] = record_set_state["ids"][boundary:]
|
||||
record_set_state["embeddings"] = record_set_state["embeddings"][boundary:]
|
||||
record_set_state["metadatas"] = record_set_state["metadatas"][boundary:]
|
||||
record_set_state["documents"] = record_set_state["documents"][boundary:]
|
||||
return collection, record_set_state
|
||||
|
||||
@rule(
|
||||
cursor=forked_collections,
|
||||
)
|
||||
def verify(
|
||||
self, cursor: Tuple[Collection, strategies.StateMachineRecordSet]
|
||||
) -> None:
|
||||
collection, record_set_state = cursor
|
||||
if len(record_set_state["ids"]) == 0:
|
||||
assert collection.count() == 0
|
||||
else:
|
||||
record_set = cast(strategies.RecordSet, record_set_state)
|
||||
invariants.embeddings_match(collection, record_set)
|
||||
invariants.metadatas_match(collection, record_set)
|
||||
invariants.documents_match(collection, record_set)
|
||||
|
||||
|
||||
@skip_if_not_cluster()
|
||||
def test_fork(caplog: pytest.LogCaptureFixture, client: chromadb.api.ClientAPI) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
run_state_machine_as_test(lambda: ForkStateMachine(client)) # type: ignore
|
||||
@@ -0,0 +1,725 @@
|
||||
import logging
|
||||
import multiprocessing
|
||||
from multiprocessing.connection import Connection
|
||||
import multiprocessing.context
|
||||
import time
|
||||
from typing import Generator, Callable, List, Tuple, cast
|
||||
from uuid import UUID
|
||||
from hypothesis import given
|
||||
import hypothesis.strategies as st
|
||||
import pytest
|
||||
import chromadb
|
||||
from chromadb.api import ClientAPI, ServerAPI
|
||||
from chromadb.config import Settings, System
|
||||
from chromadb.segment import VectorReader
|
||||
from chromadb.segment.impl.manager.local import LocalSegmentManager
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import chromadb.test.property.invariants as invariants
|
||||
from chromadb.test.property.strategies import hashing_embedding_function
|
||||
from chromadb.test.property.test_embeddings import (
|
||||
EmbeddingStateMachineStates,
|
||||
trace,
|
||||
EmbeddingStateMachineBase,
|
||||
)
|
||||
from hypothesis.stateful import (
|
||||
run_state_machine_as_test,
|
||||
rule,
|
||||
precondition,
|
||||
initialize,
|
||||
MultipleResults,
|
||||
)
|
||||
import os
|
||||
from chromadb.api.client import Client as ClientCreator
|
||||
from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
|
||||
import numpy as np
|
||||
import tempfile
|
||||
|
||||
CreatePersistAPI = Callable[[], ServerAPI]
|
||||
|
||||
configurations = (
|
||||
[
|
||||
Settings(
|
||||
chroma_api_impl="chromadb.api.rust.RustBindingsAPI",
|
||||
chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager",
|
||||
allow_reset=True,
|
||||
is_persistent=True,
|
||||
persist_directory=tempfile.mkdtemp(),
|
||||
)
|
||||
]
|
||||
if "CHROMA_RUST_BINDINGS_TEST_ONLY" in os.environ
|
||||
else [
|
||||
Settings(
|
||||
chroma_api_impl="chromadb.api.segment.SegmentAPI",
|
||||
chroma_sysdb_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_producer_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_consumer_impl="chromadb.db.impl.sqlite.SqliteDB",
|
||||
chroma_segment_manager_impl="chromadb.segment.impl.manager.local.LocalSegmentManager",
|
||||
allow_reset=True,
|
||||
is_persistent=True,
|
||||
persist_directory=tempfile.mkdtemp(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=configurations)
|
||||
def settings(request: pytest.FixtureRequest) -> Generator[Settings, None, None]:
|
||||
yield request.param
|
||||
|
||||
|
||||
collection_st = st.shared(
|
||||
strategies.collections(
|
||||
with_hnsw_params=True,
|
||||
with_persistent_hnsw_params=st.just(True),
|
||||
# Makes it more likely to find persist-related bugs (by default these are set to 2000).
|
||||
# Lower values make it more likely that a test will trigger a persist to disk.
|
||||
max_hnsw_batch_size=10,
|
||||
max_hnsw_sync_threshold=10,
|
||||
),
|
||||
key="coll",
|
||||
)
|
||||
|
||||
|
||||
@st.composite
|
||||
def collection_and_recordset_strategy(
|
||||
draw: st.DrawFn,
|
||||
) -> Tuple[strategies.Collection, strategies.RecordSet]:
|
||||
collection = draw(
|
||||
strategies.collections(
|
||||
with_hnsw_params=True,
|
||||
with_persistent_hnsw_params=st.just(True),
|
||||
# Makes it more likely to find persist-related bugs (by default these are set to 2000).
|
||||
max_hnsw_batch_size=10,
|
||||
max_hnsw_sync_threshold=10,
|
||||
)
|
||||
)
|
||||
recordset = draw(strategies.recordsets(st.just(collection)))
|
||||
return collection, recordset
|
||||
|
||||
|
||||
@given(
|
||||
collection_and_recordset_strategies=st.lists(
|
||||
collection_and_recordset_strategy(),
|
||||
min_size=1,
|
||||
unique_by=(lambda x: x[0].name, lambda x: x[0].name),
|
||||
)
|
||||
)
|
||||
def test_persist(
|
||||
settings: Settings,
|
||||
collection_and_recordset_strategies: List[
|
||||
Tuple[strategies.Collection, strategies.RecordSet]
|
||||
],
|
||||
) -> None:
|
||||
system_1 = System(settings)
|
||||
system_1.start()
|
||||
client_1 = ClientCreator.from_system(system_1)
|
||||
|
||||
client_1.reset()
|
||||
for (
|
||||
collection_strategy,
|
||||
recordset_strategy,
|
||||
) in collection_and_recordset_strategies:
|
||||
coll = client_1.create_collection(
|
||||
name=collection_strategy.name,
|
||||
metadata=collection_strategy.metadata, # type: ignore[arg-type]
|
||||
embedding_function=collection_strategy.embedding_function,
|
||||
)
|
||||
|
||||
coll.add(**recordset_strategy) # type: ignore[arg-type]
|
||||
|
||||
invariants.count(coll, recordset_strategy)
|
||||
invariants.metadatas_match(coll, recordset_strategy)
|
||||
invariants.documents_match(coll, recordset_strategy)
|
||||
invariants.ids_match(coll, recordset_strategy)
|
||||
invariants.ann_accuracy(
|
||||
coll,
|
||||
recordset_strategy,
|
||||
embedding_function=collection_strategy.embedding_function,
|
||||
)
|
||||
|
||||
system_1.stop()
|
||||
del client_1
|
||||
del system_1
|
||||
|
||||
system_2 = System(settings)
|
||||
system_2.start()
|
||||
client_2 = ClientCreator.from_system(system_2)
|
||||
|
||||
for (
|
||||
collection_strategy,
|
||||
recordset_strategy,
|
||||
) in collection_and_recordset_strategies:
|
||||
coll = client_2.get_collection(
|
||||
name=collection_strategy.name,
|
||||
embedding_function=collection_strategy.embedding_function,
|
||||
)
|
||||
invariants.count(coll, recordset_strategy)
|
||||
invariants.metadatas_match(coll, recordset_strategy)
|
||||
invariants.documents_match(coll, recordset_strategy)
|
||||
invariants.ids_match(coll, recordset_strategy)
|
||||
invariants.ann_accuracy(
|
||||
coll,
|
||||
recordset_strategy,
|
||||
embedding_function=collection_strategy.embedding_function,
|
||||
)
|
||||
|
||||
system_2.stop()
|
||||
del client_2
|
||||
del system_2
|
||||
|
||||
|
||||
def test_sync_threshold(settings: Settings) -> None:
|
||||
system = System(settings)
|
||||
system.start()
|
||||
client = ClientCreator.from_system(system)
|
||||
|
||||
collection = client.create_collection(
|
||||
name="test", metadata={"hnsw:batch_size": 3, "hnsw:sync_threshold": 3}
|
||||
)
|
||||
|
||||
manager = system.instance(LocalSegmentManager)
|
||||
segment = manager.get_segment(collection.id, VectorReader)
|
||||
|
||||
def get_index_last_modified_at() -> float:
|
||||
# Time resolution on Windows can be up to 10ms
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
return os.path.getmtime(segment._get_metadata_file()) # type: ignore[attr-defined]
|
||||
except FileNotFoundError:
|
||||
return -1
|
||||
|
||||
last_modified_at = get_index_last_modified_at()
|
||||
|
||||
collection.add(ids=["1", "2"], embeddings=[[1.0], [2.0]]) # type: ignore[arg-type]
|
||||
|
||||
# Should not have yet persisted
|
||||
assert get_index_last_modified_at() == last_modified_at
|
||||
last_modified_at = get_index_last_modified_at()
|
||||
|
||||
# Now there's 3 additions, and the sync threshold is 3...
|
||||
collection.add(ids=["3"], embeddings=[[3.0]]) # type: ignore[arg-type]
|
||||
|
||||
# ...so it should have persisted
|
||||
assert get_index_last_modified_at() > last_modified_at
|
||||
last_modified_at = get_index_last_modified_at()
|
||||
|
||||
# The same thing should happen with upserts
|
||||
collection.upsert(ids=["1", "2", "3"], embeddings=[[1.0], [2.0], [3.0]]) # type: ignore[arg-type]
|
||||
|
||||
# Should have persisted
|
||||
assert get_index_last_modified_at() > last_modified_at
|
||||
last_modified_at = get_index_last_modified_at()
|
||||
|
||||
# Mixed usage should also trigger persistence
|
||||
collection.add(ids=["4"], embeddings=[[4.0]]) # type: ignore[arg-type]
|
||||
collection.upsert(ids=["1", "2"], embeddings=[[1.0], [2.0]]) # type: ignore[arg-type]
|
||||
|
||||
# Should have persisted
|
||||
assert get_index_last_modified_at() > last_modified_at
|
||||
last_modified_at = get_index_last_modified_at()
|
||||
|
||||
# Invalid updates should also trigger persistence
|
||||
collection.add(ids=["5"], embeddings=[[5.0]]) # type: ignore[arg-type]
|
||||
collection.add(ids=["1", "2"], embeddings=[[1.0], [2.0]]) # type: ignore[arg-type]
|
||||
|
||||
# Should have persisted
|
||||
assert get_index_last_modified_at() > last_modified_at
|
||||
last_modified_at = get_index_last_modified_at()
|
||||
|
||||
|
||||
def load_and_check(
|
||||
settings: Settings,
|
||||
collection_name: str,
|
||||
record_set: strategies.RecordSet,
|
||||
conn: Connection,
|
||||
) -> None:
|
||||
try:
|
||||
system = System(settings)
|
||||
system.start()
|
||||
client = ClientCreator.from_system(system)
|
||||
|
||||
coll = client.get_collection(
|
||||
name=collection_name,
|
||||
embedding_function=strategies.not_implemented_embedding_function(), # type: ignore[arg-type]
|
||||
)
|
||||
invariants.count(coll, record_set)
|
||||
invariants.metadatas_match(coll, record_set)
|
||||
invariants.documents_match(coll, record_set)
|
||||
invariants.ids_match(coll, record_set)
|
||||
invariants.ann_accuracy(coll, record_set)
|
||||
|
||||
system.stop()
|
||||
except Exception as e:
|
||||
conn.send(e)
|
||||
raise e
|
||||
|
||||
|
||||
def get_multiprocessing_context(): # type: ignore[no-untyped-def]
|
||||
try:
|
||||
# Run the invariants in a new process to bypass any shared state/caching (which would defeat the purpose of the test)
|
||||
# (forkserver is used because it's much faster than spawn—it will spawn a new, minimal singleton process and then fork that singleton)
|
||||
ctx = multiprocessing.get_context("forkserver")
|
||||
# This is like running `import chromadb` in the single process that is forked rather than importing it in each forked process.
|
||||
# Gives a ~3x speedup since importing chromadb is fairly expensive.
|
||||
ctx.set_forkserver_preload(["chromadb"])
|
||||
return ctx
|
||||
except Exception:
|
||||
# forkserver/fork is not available on Windows
|
||||
return multiprocessing.get_context("spawn")
|
||||
|
||||
|
||||
class PersistEmbeddingsStateMachineStates(EmbeddingStateMachineStates):
|
||||
persist = "persist"
|
||||
|
||||
|
||||
MIN_STATE_CHANGES_BEFORE_PERSIST = 5
|
||||
|
||||
|
||||
class PersistEmbeddingsStateMachine(EmbeddingStateMachineBase):
|
||||
def __init__(self, client: ClientAPI, settings: Settings):
|
||||
self.client = client
|
||||
self.settings = settings
|
||||
self.min_state_changes_left_before_persisting = MIN_STATE_CHANGES_BEFORE_PERSIST
|
||||
self.client.reset()
|
||||
super().__init__(self.client)
|
||||
|
||||
@initialize(collection=collection_st) # type: ignore
|
||||
def initialize(self, collection: strategies.Collection):
|
||||
self.client.reset()
|
||||
self.collection = self.client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore[arg-type]
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
self.embedding_function = collection.embedding_function
|
||||
trace("init")
|
||||
self.on_state_change(EmbeddingStateMachineStates.initialize)
|
||||
|
||||
self.record_set_state = strategies.StateMachineRecordSet(
|
||||
ids=[], metadatas=[], documents=[], embeddings=[]
|
||||
)
|
||||
|
||||
@precondition(
|
||||
lambda self: len(self.record_set_state["ids"]) >= 1
|
||||
and self.min_state_changes_left_before_persisting <= 0
|
||||
)
|
||||
@rule()
|
||||
def persist(self) -> None:
|
||||
self.on_state_change(PersistEmbeddingsStateMachineStates.persist)
|
||||
collection_name = self.collection.name
|
||||
conn1, conn2 = multiprocessing.Pipe()
|
||||
ctx = get_multiprocessing_context() # type: ignore[no-untyped-call]
|
||||
p = ctx.Process(
|
||||
target=load_and_check,
|
||||
args=(self.settings, collection_name, self.record_set_state, conn2),
|
||||
)
|
||||
p.start()
|
||||
p.join()
|
||||
|
||||
if conn1.poll():
|
||||
e = conn1.recv()
|
||||
raise e
|
||||
|
||||
p.close()
|
||||
|
||||
def on_state_change(self, new_state: str) -> None:
|
||||
super().on_state_change(new_state)
|
||||
if new_state == PersistEmbeddingsStateMachineStates.persist:
|
||||
self.min_state_changes_left_before_persisting = (
|
||||
MIN_STATE_CHANGES_BEFORE_PERSIST
|
||||
)
|
||||
else:
|
||||
self.min_state_changes_left_before_persisting -= 1
|
||||
|
||||
def teardown(self) -> None:
|
||||
self.client.reset()
|
||||
|
||||
|
||||
def test_persist_embeddings_state(
|
||||
caplog: pytest.LogCaptureFixture, settings: Settings
|
||||
) -> None:
|
||||
caplog.set_level(logging.ERROR)
|
||||
client = chromadb.Client(settings)
|
||||
run_state_machine_as_test(
|
||||
lambda: PersistEmbeddingsStateMachine(settings=settings, client=client),
|
||||
) # type: ignore
|
||||
|
||||
|
||||
def test_delete_less_than_k(
|
||||
caplog: pytest.LogCaptureFixture, settings: Settings
|
||||
) -> None:
|
||||
client = chromadb.Client(settings)
|
||||
state = PersistEmbeddingsStateMachine(settings=settings, client=client)
|
||||
state.initialize(
|
||||
collection=strategies.Collection(
|
||||
name="A00",
|
||||
metadata={
|
||||
"hnsw:construction_ef": 128,
|
||||
"hnsw:search_ef": 128,
|
||||
"hnsw:M": 128,
|
||||
"hnsw:sync_threshold": 3,
|
||||
"hnsw:batch_size": 3,
|
||||
},
|
||||
embedding_function=None,
|
||||
id=UUID("2d3eddc7-2314-45f4-a951-47a9a8e099d2"),
|
||||
dimension=2,
|
||||
dtype=np.float16,
|
||||
known_metadata_keys={},
|
||||
known_document_keywords=[],
|
||||
has_documents=False,
|
||||
has_embeddings=True,
|
||||
)
|
||||
)
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
(embedding_ids_0,) = state.add_embeddings(record_set={"ids": ["0"], "embeddings": [[0.09765625, 0.430419921875]], "metadatas": [None], "documents": None}) # type: ignore
|
||||
state.ann_accuracy()
|
||||
# recall: 1.0, missing 0 out of 1, accuracy threshold 1e-06
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
embedding_ids_1, embedding_ids_2 = state.add_embeddings(record_set={"ids": ["1", "2"], "embeddings": [[0.20556640625, 0.08978271484375], [-0.1527099609375, 0.291748046875]], "metadatas": [None, None], "documents": None}) # type: ignore
|
||||
state.ann_accuracy()
|
||||
# recall: 1.0, missing 0 out of 3, accuracy threshold 1e-06
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
state.delete_by_ids(ids=[embedding_ids_2])
|
||||
state.ann_accuracy()
|
||||
state.teardown()
|
||||
|
||||
|
||||
# Ideally this scenario would be exercised by Hypothesis, but most runs don't seem to trigger this particular state.
|
||||
def test_delete_add_after_persist(settings: Settings) -> None:
|
||||
client = chromadb.Client(settings)
|
||||
state = PersistEmbeddingsStateMachine(settings=settings, client=client)
|
||||
|
||||
state.initialize(
|
||||
collection=strategies.Collection(
|
||||
name="A00",
|
||||
metadata={
|
||||
"hnsw:construction_ef": 128,
|
||||
"hnsw:search_ef": 128,
|
||||
"hnsw:M": 128,
|
||||
# Important: both batch_size and sync_threshold are 3
|
||||
"hnsw:batch_size": 3,
|
||||
"hnsw:sync_threshold": 3,
|
||||
},
|
||||
embedding_function=DefaultEmbeddingFunction(), # type: ignore[arg-type]
|
||||
id=UUID("0851f751-2f11-4424-ab23-4ae97074887a"),
|
||||
dimension=2,
|
||||
dtype=None,
|
||||
known_metadata_keys={},
|
||||
known_document_keywords=[],
|
||||
has_documents=False,
|
||||
has_embeddings=True,
|
||||
)
|
||||
)
|
||||
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
# Add 3 records to hit the batch_size and sync_threshold
|
||||
"ids": ["0", "1", "2"],
|
||||
"embeddings": [[0, 0], [0, 0], [0, 0]],
|
||||
"metadatas": [None, None, None],
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
|
||||
# Delete and then re-add record
|
||||
state.delete_by_ids(ids=["0"])
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": ["0"],
|
||||
"embeddings": [[1, 1]],
|
||||
"metadatas": [None],
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
|
||||
# At this point, the changes above are not fully persisted
|
||||
state.fields_match()
|
||||
|
||||
|
||||
def test_batch_size_less_than_sync_with_duplicate_adds_results_in_skipped_seq_ids(
|
||||
caplog: pytest.LogCaptureFixture, settings: Settings
|
||||
) -> None:
|
||||
# NOTE(hammadb) this test was autogenerate by hypothesis and added here to ensure that the test is run
|
||||
# in the future. It tests a case where the max seq id was incorrect in response to the same
|
||||
# id being added multiple times in a bathc.
|
||||
client = chromadb.Client(settings)
|
||||
state = PersistEmbeddingsStateMachine(settings=settings, client=client)
|
||||
state.initialize(
|
||||
collection=strategies.Collection(
|
||||
name="JqzMs4pPm14c",
|
||||
metadata={
|
||||
"hnsw:construction_ef": 128,
|
||||
"hnsw:search_ef": 128,
|
||||
"hnsw:M": 128,
|
||||
"hnsw:sync_threshold": 9,
|
||||
"hnsw:batch_size": 7,
|
||||
},
|
||||
embedding_function=hashing_embedding_function(dim=92, dtype=np.float64), # type: ignore[arg-type]
|
||||
id=UUID("45c5c816-0a90-4293-8d01-4325ff860040"),
|
||||
dimension=92,
|
||||
dtype=np.float64,
|
||||
known_metadata_keys={},
|
||||
known_document_keywords=[],
|
||||
has_documents=False,
|
||||
has_embeddings=True,
|
||||
)
|
||||
)
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
(
|
||||
embedding_ids_0,
|
||||
embedding_ids_1,
|
||||
embedding_ids_2,
|
||||
embedding_ids_3,
|
||||
embedding_ids_4,
|
||||
embedding_ids_5,
|
||||
embedding_ids_6,
|
||||
) = cast(
|
||||
MultipleResults[str],
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": ["N", "e8r6", "4", "Yao", "qFjA2c", "jHCv", "2"],
|
||||
"embeddings": [
|
||||
[0.0, 0.0, 0.0],
|
||||
[1.0, 1.0, 1.0],
|
||||
[2.0, 2.0, 2.0],
|
||||
[3.0, 3.0, 3.0],
|
||||
[4.0, 4.0, 4.0],
|
||||
[5.0, 5.0, 5.0],
|
||||
[6.0, 6.0, 6.0],
|
||||
],
|
||||
"metadatas": None,
|
||||
"documents": None,
|
||||
}
|
||||
),
|
||||
)
|
||||
state.ann_accuracy()
|
||||
# recall: 1.0, missing 0 out of 7, accuracy threshold 1e-06
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
print("\n\n")
|
||||
(_) = state.add_embeddings(
|
||||
record_set={
|
||||
"ids": ["MVu393QTc"],
|
||||
"embeddings": [[7.0, 7.0, 7.0]],
|
||||
"metadatas": None,
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
state.ann_accuracy()
|
||||
# recall: 1.0, missing 0 out of 8, accuracy threshold 1e-06
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
(
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
embedding_ids_12,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
embedding_ids_17,
|
||||
embedding_ids_18,
|
||||
_,
|
||||
_,
|
||||
_,
|
||||
embedding_ids_22,
|
||||
_,
|
||||
_,
|
||||
) = cast(
|
||||
MultipleResults[str],
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": [
|
||||
"CyF0Mk-",
|
||||
"q_Fwu",
|
||||
"2D2sQSFogDgPLkcfT",
|
||||
"SrwuQHQ6w4f51qWr2enLPQw8uKYs1",
|
||||
"G",
|
||||
"wdzt",
|
||||
"5W",
|
||||
"8tpsn",
|
||||
"fJbV7z",
|
||||
"5",
|
||||
"V",
|
||||
"1iFkoJX",
|
||||
"Zw4u",
|
||||
"Fc",
|
||||
"7",
|
||||
"vEEwrP",
|
||||
"Yf",
|
||||
],
|
||||
"embeddings": [
|
||||
[8.0, 8.0, 8.0],
|
||||
[9.0, 9.0, 9.0],
|
||||
[10.0, 10.0, 10.0],
|
||||
[11.0, 11.0, 11.0],
|
||||
[12.0, 12.0, 12.0],
|
||||
[13.0, 13.0, 13.0],
|
||||
[14.0, 14.0, 14.0],
|
||||
[15.0, 15.0, 15.0],
|
||||
[16.0, 16.0, 16.0],
|
||||
[17.0, 17.0, 17.0],
|
||||
[18.0, 18.0, 18.0],
|
||||
[19.0, 19.0, 19.0],
|
||||
[20.0, 20.0, 20.0],
|
||||
[21.0, 21.0, 21.0],
|
||||
[22.0, 22.0, 22.0],
|
||||
[23.0, 23.0, 23.0],
|
||||
[24.0, 24.0, 24.0],
|
||||
],
|
||||
"metadatas": None,
|
||||
"documents": None,
|
||||
}
|
||||
),
|
||||
)
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": ["0", "df_RWhR0HelOcv"],
|
||||
"embeddings": [[25.0, 25.0, 25.0], [26.0, 26.0, 26.0]],
|
||||
"metadatas": [None, None],
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": ["3R", "9_", "44u", "3B", "MZCXZDS", "Uelx"],
|
||||
"embeddings": [
|
||||
[27.0, 27.0, 27.0],
|
||||
[28.0, 28.0, 28.0],
|
||||
[29.0, 29.0, 29.0],
|
||||
[30.0, 30.0, 30.0],
|
||||
[31.0, 31.0, 31.0],
|
||||
[32.0, 32.0, 32.0],
|
||||
],
|
||||
"metadatas": None,
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
state.persist()
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": "YlVm",
|
||||
"embeddings": [[33.0, 33.0, 33.0]],
|
||||
"metadatas": None,
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
state.ann_accuracy()
|
||||
# recall: 1.0, missing 0 out of 34, accuracy threshold 1e-06
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": ["Rk1", "TPL"],
|
||||
"embeddings": [[34.0, 34.0, 34.0], [35.0, 35.0, 35.0]],
|
||||
"metadatas": [None, None],
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
state.ann_accuracy()
|
||||
# recall: 1.0, missing 0 out of 36, accuracy threshold 1e-06
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.no_duplicates()
|
||||
|
||||
state.add_embeddings(
|
||||
record_set={
|
||||
"ids": [
|
||||
"CyF0Mk-",
|
||||
"q_Fwu",
|
||||
"2D2sQSFogDgPLkcfT",
|
||||
"SrwuQHQ6w4f51qWr2enLPQw8uKYs1",
|
||||
embedding_ids_12,
|
||||
"wdzt",
|
||||
"5W",
|
||||
"8tpsn",
|
||||
"fJbV7z",
|
||||
embedding_ids_17,
|
||||
embedding_ids_18,
|
||||
"1iFkoJX",
|
||||
"Zw4u",
|
||||
"Fc",
|
||||
embedding_ids_22,
|
||||
"vEEwrP",
|
||||
"Yf",
|
||||
],
|
||||
"embeddings": [
|
||||
[8.0, 8.0, 8.0],
|
||||
[9.0, 9.0, 9.0],
|
||||
[10.0, 10.0, 10.0],
|
||||
[11.0, 11.0, 11.0],
|
||||
[12.0, 12.0, 12.0],
|
||||
[13.0, 13.0, 13.0],
|
||||
[14.0, 14.0, 14.0],
|
||||
[15.0, 15.0, 15.0],
|
||||
[16.0, 16.0, 16.0],
|
||||
[17.0, 17.0, 17.0],
|
||||
[18.0, 18.0, 18.0],
|
||||
[19.0, 19.0, 19.0],
|
||||
[20.0, 20.0, 20.0],
|
||||
[21.0, 21.0, 21.0],
|
||||
[22.0, 22.0, 22.0],
|
||||
[23.0, 23.0, 23.0],
|
||||
[24.0, 24.0, 24.0],
|
||||
],
|
||||
"metadatas": None,
|
||||
"documents": None,
|
||||
}
|
||||
)
|
||||
state.ann_accuracy()
|
||||
state.count()
|
||||
state.fields_match()
|
||||
state.log_size_below_max()
|
||||
state.teardown()
|
||||
@@ -0,0 +1,86 @@
|
||||
from overrides import overrides
|
||||
from chromadb.api.client import Client
|
||||
from chromadb.config import System
|
||||
import hypothesis.strategies as st
|
||||
from hypothesis.stateful import (
|
||||
rule,
|
||||
run_state_machine_as_test,
|
||||
initialize,
|
||||
)
|
||||
|
||||
from chromadb.test.property.test_embeddings import (
|
||||
EmbeddingStateMachineBase,
|
||||
EmbeddingStateMachineStates,
|
||||
trace,
|
||||
)
|
||||
import chromadb.test.property.strategies as strategies
|
||||
import os
|
||||
|
||||
|
||||
collection_persistent_st = st.shared(
|
||||
strategies.collections(
|
||||
with_hnsw_params=True,
|
||||
with_persistent_hnsw_params=st.just(True),
|
||||
# Makes it more likely to find persist-related bugs (by default these are set to 2000).
|
||||
max_hnsw_batch_size=10,
|
||||
max_hnsw_sync_threshold=10,
|
||||
),
|
||||
key="coll_persistent",
|
||||
)
|
||||
|
||||
|
||||
# This machine shares a lot of similarity with the machine in chromadb/test/property/test_persist.py.
|
||||
# However, test_persist.py tests correctness under complete process isolation and therefore can only check invariants on a new system--whereas this machine does not have full process isolation between systems/clients but after a restart continues to exercise the state machine with the newly-created system.
|
||||
class RestartablePersistedEmbeddingStateMachine(EmbeddingStateMachineBase):
|
||||
system: System
|
||||
|
||||
def __init__(self, system: System) -> None:
|
||||
self.system = system
|
||||
client = Client.from_system(system)
|
||||
super().__init__(client)
|
||||
|
||||
@initialize(collection=collection_persistent_st) # type: ignore
|
||||
@overrides
|
||||
def initialize(self, collection: strategies.Collection):
|
||||
self.client.reset()
|
||||
|
||||
self.collection = self.client.create_collection(
|
||||
name=collection.name,
|
||||
metadata=collection.metadata, # type: ignore
|
||||
embedding_function=collection.embedding_function,
|
||||
)
|
||||
self.embedding_function = collection.embedding_function
|
||||
trace("init")
|
||||
self.on_state_change(EmbeddingStateMachineStates.initialize)
|
||||
|
||||
self.record_set_state = strategies.StateMachineRecordSet(
|
||||
ids=[], metadatas=[], documents=[], embeddings=[]
|
||||
)
|
||||
|
||||
@rule()
|
||||
def restart_system(self) -> None:
|
||||
# Simulates restarting chromadb
|
||||
self.system.stop()
|
||||
self.system = System(self.system.settings)
|
||||
self.system.start()
|
||||
self.client.clear_system_cache()
|
||||
self.client = Client.from_system(self.system)
|
||||
self.collection = self.client.get_collection(
|
||||
self.collection.name, embedding_function=self.embedding_function
|
||||
)
|
||||
|
||||
@overrides
|
||||
def teardown(self) -> None:
|
||||
super().teardown()
|
||||
# Need to manually stop the system to cleanup resources because we may have created a new system (above rule).
|
||||
# Normally, we wouldn't have to worry about this as the system from the fixture is shared between state machine runs.
|
||||
# (This helps avoid a "too many open files" error.)
|
||||
self.system.stop()
|
||||
|
||||
|
||||
def test_restart_persisted_client(sqlite_persistent: System) -> None:
|
||||
# TODO: This test is broken for rust bindings and should be fixed
|
||||
if sqlite_persistent.settings.chroma_api_impl != "chromadb.api.rust.RustBindingsAPI":
|
||||
run_state_machine_as_test(
|
||||
lambda: RestartablePersistedEmbeddingStateMachine(sqlite_persistent),
|
||||
) # type: ignore
|
||||
Reference in New Issue
Block a user