From 0f923483f7869e27915bec410e2e09d9b941cf4e Mon Sep 17 00:00:00 2001 From: Thor Whalen <1906276+thorwhalen@users.noreply.github.com> Date: Mon, 7 Sep 2026 01:27:20 +0200 Subject: [PATCH] fix: apply `appendable` outside the value codec so append/extend work `ChromaDocuments` and `ChromaUris` stacked `appendable` *inside* the `ValueCodecs.single_nested_value` wrapper. dol's class-wrapping then re-installs `append`/`extend` as delegated attributes bound to the un-codec'd leaf store, so appended values skipped the codec and reached `ChromaCollection.__setitem__` raw -- which unpacks them as `**v`: TypeError: Collection.upsert() argument after ** must be a mapping, not str Swapping the two decorators puts `appendable` outermost, so `append` and `extend` now write the same values `__setitem__` writes. A comment records the ordering constraint so it doesn't get flipped back. Adds `AppendableChromaCollection` as the escape hatch for auto-keyed appends of raw chromadb kwargs (what `ChromaDocuments.append` accidentally did for mapping arguments before, minus the codec inconsistency). Verified every other observable surface is byte-identical before/after (name/MRO/dir, get/set/del/len/contains/iter/items/values, `.collection`, `.store`). Top-level `chromadol` exports are unchanged. Closes #2 Claude-Session: https://claude.ai/code/session_01L1aQPB34n7PU7jmbztSjBe --- chromadol/base.py | 30 +++++++-- chromadol/tests/test_appendable.py | 102 +++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 4 deletions(-) create mode 100644 chromadol/tests/test_appendable.py diff --git a/chromadol/base.py b/chromadol/base.py index edf4d08..53495b2 100644 --- a/chromadol/base.py +++ b/chromadol/base.py @@ -102,20 +102,42 @@ def get_collection( return codec(c) -@ValueCodecs.single_nested_value("documents") @appendable(item2kv=uuid_key) +class AppendableChromaCollection(ChromaCollection): + """ChromaCollection with ``append`` and ``extend``, auto-generating uuid keys. + + Items are the raw ``chromadb`` kwargs mappings that + :meth:`ChromaCollection.__setitem__` accepts (e.g. + ``{"documents": ..., "metadatas": ...}``) -- use this when you need to append + more than the single field the codec-ed stores below expose. + """ + + +# NOTE: `appendable` must be applied OUTSIDE any value codec. Applied inside, +# dol's class-wrapping re-installs `append`/`extend` as delegated attributes bound +# to the un-codec'd leaf store, so appended values bypass the codec and reach +# `ChromaCollection.__setitem__` raw. See https://github.com/i2mint/chromadol/issues/2 + + +@appendable(item2kv=uuid_key) +@ValueCodecs.single_nested_value("documents") class ChromaDocuments(ChromaCollection): - """ChromaCollection but reading and writing only the 'documents' field.""" + """ChromaCollection but reading and writing only the 'documents' field. + + ``append`` and ``extend`` take the same values ``__setitem__`` takes (that is, + the 'documents' field's value), generating uuid keys for them. To write raw + ``chromadb`` kwargs instead, use ``AppendableChromaCollection``. + """ -@ValueCodecs.single_nested_value("uris") @appendable(item2kv=uuid_key) +@ValueCodecs.single_nested_value("uris") class ChromaUris(ChromaCollection): """ChromaCollection but reading and writing only the 'uris' field.""" -@ValueCodecs.single_nested_value("metadata") @appendable(item2kv=uuid_key) +@ValueCodecs.single_nested_value("metadata") class ChromaUris(ChromaCollection): """ChromaCollection but reading and writing only the 'uris' field.""" diff --git a/chromadol/tests/test_appendable.py b/chromadol/tests/test_appendable.py new file mode 100644 index 0000000..036ffa6 --- /dev/null +++ b/chromadol/tests/test_appendable.py @@ -0,0 +1,102 @@ +"""Test that ``append``/``extend`` go through the value codec. + +``ChromaDocuments`` & friends stack ``appendable`` on top of a ``dol`` value +codec. The stacking order matters: when ``appendable`` sits *inside* the codec, +``dol``'s class-wrapping re-installs ``append``/``extend`` as delegated +attributes bound to the (un-codec'd) leaf store, so appended values bypass the +codec entirely. See https://github.com/i2mint/chromadol/issues/2 +""" + +import chromadb +import pytest + +from chromadol.base import ( + AppendableChromaCollection, + ChromaCollection, + ChromaDocuments, + ChromaUris, +) + +# The appendable stores, i.e. those whose value codec keeps a single field. +appendable_stores = [ChromaDocuments, ChromaUris] + + +class _RecordingCollection: + """Minimal stand-in for a ``chromadb`` Collection, recording ``upsert`` calls.""" + + def __init__(self): + self.upserts = [] + + def upsert(self, ids, **kwargs): + self.upserts.append((ids, kwargs)) + + def get(self, ids=None): + return {"ids": [ids for ids, _ in self.upserts]} + + def count(self): + return len(self.upserts) + + def delete(self, ids): + pass + + +def _documents_store(tmp_path, name): + """A ``ChromaDocuments`` over a fresh, empty, on-disk collection.""" + client = chromadb.PersistentClient(str(tmp_path / name)) + return ChromaDocuments(client.create_collection(name, get_or_create=True)) + + +@pytest.mark.parametrize("store_cls", appendable_stores, ids=lambda c: c.__name__) +def test_append_writes_what_setitem_writes(store_cls): + """``append`` must speak the same value language as ``__setitem__``. + + Uses a recording stand-in rather than a real collection so the invariant is + checked for every appendable store, independently of which ``chromadb`` + field its codec happens to target. + """ + collection = _RecordingCollection() + store = store_cls(collection) + + store["a_key"] = "a value" + store.append("a value") + + (_, via_setitem), (_, via_append) = collection.upserts + assert via_append == via_setitem + + +def test_append_goes_through_value_codec(tmp_path): + docs = _documents_store(tmp_path, "appendtest") + docs["k1"] = "via setitem" + assert docs["k1"] == ["via setitem"] + + docs.append("via append") + + assert len(docs) == 2 + (appended_key,) = (k for k in docs if k != "k1") + assert docs[appended_key] == ["via append"] + + +def test_extend_goes_through_value_codec(tmp_path): + docs = _documents_store(tmp_path, "extendtest") + + docs.extend(["first", "second"]) + + assert len(docs) == 2 + assert sorted(docs[k][0] for k in docs) == ["first", "second"] + + +def test_appendable_chroma_collection_appends_raw_chromadb_kwargs(tmp_path): + """The escape hatch: auto-keyed appends of raw ``chromadb`` kwargs.""" + client = chromadb.PersistentClient(str(tmp_path / "raw")) + raw = AppendableChromaCollection( + client.create_collection("raw", get_or_create=True) + ) + + raw.append({"documents": "raw document", "metadatas": {"author": "me"}}) + + assert issubclass(AppendableChromaCollection, ChromaCollection) + assert len(raw) == 1 + (key,) = raw + record = raw[key] + assert record["documents"] == ["raw document"] + assert record["metadatas"] == [{"author": "me"}]