ordeq_chromadb

ChromaDBCollection dataclass

Bases: Input[Collection], Output[dict[str, Any]]

IO to load from and save collection using ChromaDB. Creates a chromadb.PersistentClient by default during initialization and calls client.get_collection under the hood.

Example usage:

>>> index = ChromaDBCollection.from_path(
...     Path("/path/to/chromadb"),
...     name="test_collection"
... )  # doctest: +SKIP
>>> index.save({
...    "ids": [0, 1],
...    "embeddings": [
...       [0.1, 0.2, 0.3],
...       [0.4, 0.5, 0.6]],
...    "metadatas": [{"page_number": 1}, {"page_number": 2}],
...    "documents": ["Document 1", "Document 2"],
... })  # doctest: +SKIP
>>> # collection_name should not be None
>>> collection = index.load()  # doctest: +SKIP
>>> results = collection.query(
...    query_embeddings=[[0.1, 0.2, 0.3]],
...    n_results=2
... )  # doctest: +SKIP