citadeldb-haystack 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,18 @@
1
+ /target
2
+ **/*.rs.bk
3
+ *.swp
4
+ *.swo
5
+ *~
6
+ .DS_Store
7
+ site/public/
8
+ site/static/wasm/*.wasm
9
+ site/static/wasm/*.js
10
+ /notes/
11
+ __pycache__/
12
+ *.py[cod]
13
+ # maturin build output; the .pyd is caught by the line above only by accident.
14
+ *.pdb
15
+ .pytest_cache/
16
+ .mypy_cache/
17
+ /dist/
18
+ packaging/*/dist/
@@ -0,0 +1,109 @@
1
+ Metadata-Version: 2.5
2
+ Name: citadeldb-haystack
3
+ Version: 2.0.0
4
+ Summary: Haystack document store backed by Citadel: encrypted at rest, with deletes that destroy the key
5
+ Project-URL: Homepage, https://citadeldb.dev
6
+ Project-URL: Repository, https://github.com/yp3y5akh0v/citadel
7
+ Author: Yuriy Peysakhov
8
+ License-Expression: Apache-2.0
9
+ Keywords: document-store,encryption,haystack,rag,vector-store
10
+ Classifier: Development Status :: 4 - Beta
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Topic :: Database
14
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
15
+ Requires-Python: >=3.10
16
+ Requires-Dist: citadeldb<3,>=2.0
17
+ Requires-Dist: haystack-ai<4,>=2.9
18
+ Provides-Extra: test
19
+ Requires-Dist: pytest-asyncio>=0.23; extra == 'test'
20
+ Requires-Dist: pytest>=8; extra == 'test'
21
+ Description-Content-Type: text/markdown
22
+
23
+ # citadeldb-haystack
24
+
25
+ A [Haystack](https://github.com/deepset-ai/haystack) `DocumentStore` backed by
26
+ [Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and deletes
27
+ that destroy the key, not just the row.
28
+
29
+ Passes deepset's own `DocumentStoreBaseTests` conformance suite.
30
+
31
+ ```
32
+ pip install citadeldb-haystack
33
+ ```
34
+
35
+ ```python
36
+ from haystack import Document
37
+ from haystack.utils import Secret
38
+ from citadeldb_haystack import CitadelDocumentStore
39
+
40
+ store = CitadelDocumentStore("corpus.cdl", Secret.from_env_var("CITADEL_KEY"))
41
+ # CITADEL_KEY must be set: an env-var secret is what lets a pipeline serialize.
42
+
43
+ store.write_documents([Document(id="d1", content="...", meta={"chapter": "intro"})])
44
+ store.filter_documents({"field": "meta.chapter", "operator": "==", "value": "intro"})
45
+ ```
46
+
47
+ `dim` defaults to 768 and must match your embedding model.
48
+
49
+ ## The passphrase never lands in a pipeline file
50
+
51
+ The passphrase is a Haystack `Secret`. Pipelines are serialized to disk, and a literal
52
+ token refuses to serialize, so a passphrase cannot be written into a pipeline by accident:
53
+
54
+ ```python
55
+ CitadelDocumentStore("literal.cdl", "literal-passphrase").to_dict()
56
+ # ValueError: Cannot serialize token-based secret.
57
+
58
+ CitadelDocumentStore("corpus.cdl", Secret.from_env_var("CITADEL_KEY")).to_dict()
59
+ # {... "key": {"type": "env_var", "env_vars": ["CITADEL_KEY"], ...}}
60
+ ```
61
+
62
+ Use `Secret.from_env_var` for any store that goes into a saved pipeline.
63
+
64
+ ## Deletes destroy the key
65
+
66
+ Every document is sealed under its own key. Deleting destroys that key and then removes the
67
+ row, so any ciphertext surviving elsewhere stays unreadable.
68
+
69
+ ```python
70
+ store.delete_documents(["d1"])
71
+ store.delete_all() # returns the number erased
72
+ ```
73
+
74
+ `DuplicatePolicy.NONE` falls back to `FAIL`, as `InMemoryDocumentStore` does, so an
75
+ accidental re-write is reported rather than silently replacing a document whose key would
76
+ then be destroyed.
77
+
78
+ ## Retrieval
79
+
80
+ ```python
81
+ query_embedding = [0.0] * 768 # from your Haystack text embedder, `dim` wide
82
+
83
+ store.embedding_retrieval(query_embedding, top_k=5,
84
+ filters={"field": "meta.chapter", "operator": "==", "value": "intro"})
85
+ ```
86
+
87
+ Filtering uses Haystack's own evaluator, so the whole filter language, date comparisons
88
+ included, matches `InMemoryDocumentStore` operator for operator. `top_k` is `top_k`: a
89
+ filter matching only distant documents still returns them, however many others outrank
90
+ them.
91
+
92
+ A top-level `AND` of string equality conditions is pushed into the scan, including nested
93
+ paths like `meta.person.name`. Everything else is evaluated afterwards, so the two agree:
94
+ nothing is pushed under `OR` or `NOT`, and numbers are not pushed either, because `==` here
95
+ is Python's (`1 == 1.0`) where the stored comparison is JSON-type exact.
96
+
97
+ ## Notes
98
+
99
+ Documents Haystack did not embed are stored, filterable and countable, but take no part in
100
+ `embedding_retrieval`, as in `InMemoryDocumentStore`, which retrieves only documents that
101
+ have embeddings.
102
+
103
+ Citadel is embedded and one process owns the file. A path already open on this thread,
104
+ under the same passphrase, is shared, so this can sit on the same database as another
105
+ Citadel adapter; construct them on the same thread.
106
+
107
+ ## License
108
+
109
+ Apache-2.0
@@ -0,0 +1,87 @@
1
+ # citadeldb-haystack
2
+
3
+ A [Haystack](https://github.com/deepset-ai/haystack) `DocumentStore` backed by
4
+ [Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and deletes
5
+ that destroy the key, not just the row.
6
+
7
+ Passes deepset's own `DocumentStoreBaseTests` conformance suite.
8
+
9
+ ```
10
+ pip install citadeldb-haystack
11
+ ```
12
+
13
+ ```python
14
+ from haystack import Document
15
+ from haystack.utils import Secret
16
+ from citadeldb_haystack import CitadelDocumentStore
17
+
18
+ store = CitadelDocumentStore("corpus.cdl", Secret.from_env_var("CITADEL_KEY"))
19
+ # CITADEL_KEY must be set: an env-var secret is what lets a pipeline serialize.
20
+
21
+ store.write_documents([Document(id="d1", content="...", meta={"chapter": "intro"})])
22
+ store.filter_documents({"field": "meta.chapter", "operator": "==", "value": "intro"})
23
+ ```
24
+
25
+ `dim` defaults to 768 and must match your embedding model.
26
+
27
+ ## The passphrase never lands in a pipeline file
28
+
29
+ The passphrase is a Haystack `Secret`. Pipelines are serialized to disk, and a literal
30
+ token refuses to serialize, so a passphrase cannot be written into a pipeline by accident:
31
+
32
+ ```python
33
+ CitadelDocumentStore("literal.cdl", "literal-passphrase").to_dict()
34
+ # ValueError: Cannot serialize token-based secret.
35
+
36
+ CitadelDocumentStore("corpus.cdl", Secret.from_env_var("CITADEL_KEY")).to_dict()
37
+ # {... "key": {"type": "env_var", "env_vars": ["CITADEL_KEY"], ...}}
38
+ ```
39
+
40
+ Use `Secret.from_env_var` for any store that goes into a saved pipeline.
41
+
42
+ ## Deletes destroy the key
43
+
44
+ Every document is sealed under its own key. Deleting destroys that key and then removes the
45
+ row, so any ciphertext surviving elsewhere stays unreadable.
46
+
47
+ ```python
48
+ store.delete_documents(["d1"])
49
+ store.delete_all() # returns the number erased
50
+ ```
51
+
52
+ `DuplicatePolicy.NONE` falls back to `FAIL`, as `InMemoryDocumentStore` does, so an
53
+ accidental re-write is reported rather than silently replacing a document whose key would
54
+ then be destroyed.
55
+
56
+ ## Retrieval
57
+
58
+ ```python
59
+ query_embedding = [0.0] * 768 # from your Haystack text embedder, `dim` wide
60
+
61
+ store.embedding_retrieval(query_embedding, top_k=5,
62
+ filters={"field": "meta.chapter", "operator": "==", "value": "intro"})
63
+ ```
64
+
65
+ Filtering uses Haystack's own evaluator, so the whole filter language, date comparisons
66
+ included, matches `InMemoryDocumentStore` operator for operator. `top_k` is `top_k`: a
67
+ filter matching only distant documents still returns them, however many others outrank
68
+ them.
69
+
70
+ A top-level `AND` of string equality conditions is pushed into the scan, including nested
71
+ paths like `meta.person.name`. Everything else is evaluated afterwards, so the two agree:
72
+ nothing is pushed under `OR` or `NOT`, and numbers are not pushed either, because `==` here
73
+ is Python's (`1 == 1.0`) where the stored comparison is JSON-type exact.
74
+
75
+ ## Notes
76
+
77
+ Documents Haystack did not embed are stored, filterable and countable, but take no part in
78
+ `embedding_retrieval`, as in `InMemoryDocumentStore`, which retrieves only documents that
79
+ have embeddings.
80
+
81
+ Citadel is embedded and one process owns the file. A path already open on this thread,
82
+ under the same passphrase, is shared, so this can sit on the same database as another
83
+ Citadel adapter; construct them on the same thread.
84
+
85
+ ## License
86
+
87
+ Apache-2.0
@@ -0,0 +1,40 @@
1
+ [build-system]
2
+ requires = ["hatchling", "hatch-vcs"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "citadeldb-haystack"
7
+ dynamic = ["version"]
8
+ description = "Haystack document store backed by Citadel: encrypted at rest, with deletes that destroy the key"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "Apache-2.0"
12
+ authors = [{ name = "Yuriy Peysakhov" }]
13
+ keywords = ["haystack", "rag", "document-store", "vector-store", "encryption"]
14
+ classifiers = [
15
+ "Development Status :: 4 - Beta",
16
+ "Intended Audience :: Developers",
17
+ "Programming Language :: Python :: 3",
18
+ "Topic :: Database",
19
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
20
+ ]
21
+ # The precomputed vector needs the `embedding` field added in citadeldb 2.0.
22
+ dependencies = ["citadeldb>=2.0,<3", "haystack-ai>=2.9,<4"]
23
+
24
+ [project.optional-dependencies]
25
+ test = ["pytest>=8", "pytest-asyncio>=0.23"]
26
+
27
+ [project.urls]
28
+ Homepage = "https://citadeldb.dev"
29
+ Repository = "https://github.com/yp3y5akh0v/citadel"
30
+
31
+ # The version comes from the release tag, so there is nothing to bump.
32
+ [tool.hatch.version]
33
+ source = "vcs"
34
+ raw-options = { root = "../..", tag_regex = '^v(?P<version>\d+\.\d+\.\d+)$' }
35
+
36
+ [tool.hatch.build.targets.wheel]
37
+ packages = ["src/citadeldb_haystack"]
38
+
39
+ [tool.pytest.ini_options]
40
+ asyncio_mode = "auto"
@@ -0,0 +1,14 @@
1
+ """Haystack document storage backed by Citadel, encrypted at rest."""
2
+
3
+ from importlib.metadata import PackageNotFoundError, version
4
+
5
+ from .document_store import CitadelDocumentStore
6
+
7
+ __all__ = ["CitadelDocumentStore", "__version__"]
8
+
9
+ # Haystack wires stores by construction, so there is nothing to register.
10
+
11
+ try:
12
+ __version__ = version("citadeldb-haystack")
13
+ except PackageNotFoundError: # running from a source tree, never installed
14
+ __version__ = "0+unknown"
@@ -0,0 +1,346 @@
1
+ """Haystack DocumentStore over an encrypted Citadel region."""
2
+ from __future__ import annotations
3
+
4
+ import asyncio
5
+ import hashlib
6
+ from dataclasses import replace
7
+ from typing import Any
8
+
9
+ import citadeldb
10
+ from haystack import default_from_dict, default_to_dict
11
+ from haystack.dataclasses import Document
12
+ from haystack.document_stores.errors import DuplicateDocumentError
13
+ from haystack.document_stores.types import DuplicatePolicy
14
+ from haystack.utils import Secret, deserialize_secrets_inplace
15
+ from haystack.utils.filters import document_matches_filter
16
+
17
+ KIND = "doc"
18
+ DEFAULT_PATH = "haystack.cdl"
19
+ DEFAULT_REGION = "documents"
20
+ # The default width of Haystack's own embedders and of its conformance fixtures.
21
+ DEFAULT_DIM = 768
22
+ PAGE = 10_000
23
+ # Haystack prefixes metadata fields; bare names address the document.
24
+ _META_PREFIX = "meta."
25
+
26
+
27
+ # A Database is pinned to its opening thread, so workers take Memory, not self.
28
+
29
+
30
+ def _placeholder(text: str, dim: int) -> list[float]:
31
+ """A deterministic vector for a document Haystack did not embed."""
32
+ digest = hashlib.sha256(text.encode()).digest()
33
+ return [digest[i % len(digest)] / 255.0 for i in range(dim)]
34
+
35
+
36
+ def _fetch(mem: Any, region: str, criterion: dict[str, Any] | None) -> list[Any]:
37
+ """Page to the end: one fetch is bounded, and a partial erase must not look whole."""
38
+ out: list[Any] = []
39
+ after = None
40
+ while True:
41
+ page = mem.fetch(region, KIND, payload_filter=criterion, limit=PAGE, after_id=after)
42
+ out.extend(page)
43
+ if len(page) < PAGE:
44
+ return out
45
+ after = page[-1].id
46
+
47
+
48
+ def _erase(mem: Any, region: str, hits: list[Any]) -> int:
49
+ if not hits:
50
+ return 0
51
+ return mem.forget(region, [h.id for h in hits]).erased_count
52
+
53
+
54
+ def _document(hit: Any) -> Document:
55
+ # Unflattened so the document round-trips whole.
56
+ return Document.from_dict(hit.payload["doc"])
57
+
58
+
59
+ def _pushdown(filters: dict[str, Any] | None) -> dict[str, Any] | None:
60
+ """Equality leaves of a top-level AND; this narrows and never decides."""
61
+ if not filters:
62
+ return None
63
+ if filters.get("operator") == "AND":
64
+ conditions = filters.get("conditions", [])
65
+ elif "field" in filters:
66
+ conditions = [filters] # a bare condition is an AND of one
67
+ else:
68
+ return None
69
+ leaves: list[tuple[tuple[str, ...], Any]] = []
70
+ for c in conditions:
71
+ field = c.get("field", "")
72
+ if (
73
+ c.get("operator") == "=="
74
+ and field.startswith(_META_PREFIX)
75
+ # Only strings: containment compares JSON types exactly where
76
+ # Haystack's `==` is Python's, which holds 1 == 1.0 == True, so
77
+ # pushing a number would decide rather than narrow. None means
78
+ # absent, which containment cannot express at all.
79
+ and isinstance(c.get("value"), str)
80
+ ):
81
+ path = tuple(p for p in field[len(_META_PREFIX):].split(".") if p)
82
+ if path:
83
+ leaves.append((path, c["value"]))
84
+ eq = _needle(leaves)
85
+ # The needle is shaped like the payload, not like the filter.
86
+ return {"doc": {"meta": eq}} if eq else None
87
+
88
+
89
+ def _needle(leaves: list[tuple[tuple[str, ...], Any]]) -> dict[str, Any]:
90
+ """Nested-object needle for dotted metadata paths.
91
+
92
+ `meta.person.name` addresses a nested value, so it has to become
93
+ `{"person": {"name": v}}`; one flat key spelled with a dot can never exist in
94
+ the payload and would match nothing. Two leaves that collide, or one whose
95
+ path runs through another's, are dropped rather than merged - containment
96
+ holds one value per position, and guessing which wins would decide.
97
+ """
98
+ out: dict[str, Any] = {}
99
+ for path, value in leaves:
100
+ if any(
101
+ other is not path and (other[: len(path)] == path or path[: len(other)] == other)
102
+ for other, _ in leaves
103
+ ):
104
+ continue
105
+ node = out
106
+ for part in path[:-1]:
107
+ node = node.setdefault(part, {})
108
+ node[path[-1]] = value
109
+ return out
110
+
111
+
112
+ def _filter(mem: Any, region: str, filters: dict[str, Any] | None) -> list[Document]:
113
+ docs = [_document(h) for h in _fetch(mem, region, _pushdown(filters))]
114
+ if not filters:
115
+ return docs
116
+ return [d for d in docs if document_matches_filter(filters, d)]
117
+
118
+
119
+ def _atom(doc: Document, dim: int) -> dict[str, Any]:
120
+ if doc.embedding is not None and len(doc.embedding) != dim:
121
+ raise ValueError(
122
+ f"document {doc.id} has a {len(doc.embedding)}-dimension embedding but "
123
+ f"this region is {dim}. Build the store with dim={len(doc.embedding)} to "
124
+ f"match your embedding model."
125
+ )
126
+ return {
127
+ "kind": KIND,
128
+ "text": doc.content or "",
129
+ "embedding": (
130
+ list(doc.embedding)
131
+ if doc.embedding is not None
132
+ else _placeholder(doc.content or doc.id, dim)
133
+ ),
134
+ # Unflattened so the document round-trips whole. `emb` records whether
135
+ # the vector above is Haystack's or the placeholder, which is what keeps
136
+ # a hash of the text out of embedding_retrieval's answers.
137
+ "payload": {
138
+ "did": doc.id,
139
+ "emb": doc.embedding is not None,
140
+ "doc": doc.to_dict(flatten=False),
141
+ },
142
+ }
143
+
144
+
145
+ def _write(
146
+ mem: Any,
147
+ region: str,
148
+ dim: int,
149
+ documents: list[Document],
150
+ policy: DuplicatePolicy,
151
+ ) -> int:
152
+ if policy == DuplicatePolicy.NONE:
153
+ policy = DuplicatePolicy.FAIL # as InMemoryDocumentStore defaults
154
+
155
+ # OVERWRITE decides nothing from a prior read, so it does not pay for one:
156
+ # the keyed write below supersedes whatever the id already named.
157
+ seen: set[str] = set()
158
+ if policy != DuplicatePolicy.OVERWRITE:
159
+ # A live id set: a second copy in one batch collides like a stored one.
160
+ seen = {
161
+ h.payload["did"]
162
+ for doc in documents
163
+ for h in _fetch(mem, region, {"did": doc.id})
164
+ }
165
+ written = len(documents)
166
+ atoms: dict[str, dict[str, Any]] = {}
167
+ for doc in documents:
168
+ if policy != DuplicatePolicy.OVERWRITE and doc.id in seen:
169
+ if policy == DuplicatePolicy.FAIL:
170
+ raise DuplicateDocumentError(f"ID '{doc.id}' already exists.")
171
+ written -= 1 # SKIP: the copy already accepted stands
172
+ continue
173
+ # Last copy of an id in one batch wins, as the reference's dict does.
174
+ atoms[doc.id] = _atom(doc, dim)
175
+ seen.add(doc.id)
176
+ if atoms:
177
+ # Keyed on the document id and committed as one transaction: two writers
178
+ # of one id supersede rather than each adding a row, which a read out
179
+ # here cannot prevent. Every policy binds the key, so a document written
180
+ # under SKIP is still the one a later OVERWRITE replaces.
181
+ mem.remember_replacing_keyed_batch(
182
+ region, [(atom, did) for did, atom in atoms.items()]
183
+ )
184
+ return written
185
+
186
+
187
+ def _delete(mem: Any, region: str, document_ids: list[str]) -> int:
188
+ # Contracted not to fail on an id that is not present.
189
+ return _erase(
190
+ mem,
191
+ region,
192
+ [h for did in dict.fromkeys(document_ids) for h in _fetch(mem, region, {"did": did})],
193
+ )
194
+
195
+
196
+ def _count(mem: Any, region: str) -> int:
197
+ return len(_fetch(mem, region, None))
198
+
199
+
200
+ class CitadelDocumentStore:
201
+ """A Haystack `DocumentStore` backed by one encrypted Citadel region."""
202
+
203
+ def __init__(
204
+ self,
205
+ path: str = DEFAULT_PATH,
206
+ key: Secret | str = Secret.from_env_var("CITADEL_KEY"),
207
+ *,
208
+ region: str = DEFAULT_REGION,
209
+ dim: int = DEFAULT_DIM,
210
+ ) -> None:
211
+ self._key = Secret.from_token(key) if isinstance(key, str) else key
212
+ passphrase = self._key.resolve_value()
213
+ if not passphrase:
214
+ raise ValueError("a passphrase is required: the corpus is the payload")
215
+ self._path = path
216
+ self._region = region
217
+ self._dim = dim
218
+ try:
219
+ self._db = citadeldb.connect(path, key=passphrase, region_keys=True)
220
+ except citadeldb.OperationalError as e:
221
+ if "locked" not in str(e):
222
+ raise
223
+ raise RuntimeError(
224
+ f"{path} is open in another process. Citadel is embedded, so one "
225
+ f"process owns the file."
226
+ ) from e
227
+ self._mem = self._db.memory()
228
+ # Idempotent for a region of the same width, so a dim clash raises here.
229
+ self._mem.create_encrypted_region(region, citadeldb.MockEmbedder(dim=dim))
230
+
231
+ # ---- serialization ----------------------------------------------------
232
+
233
+ def to_dict(self) -> dict[str, Any]:
234
+ """Serialize for a pipeline file; a literal passphrase refuses."""
235
+ return default_to_dict(
236
+ self,
237
+ path=self._path,
238
+ key=self._key.to_dict(),
239
+ region=self._region,
240
+ dim=self._dim,
241
+ )
242
+
243
+ @classmethod
244
+ def from_dict(cls, data: dict[str, Any]) -> CitadelDocumentStore:
245
+ deserialize_secrets_inplace(data["init_parameters"], keys=["key"])
246
+ return default_from_dict(cls, data)
247
+
248
+ # ---- the protocol -----------------------------------------------------
249
+
250
+ def count_documents(self) -> int:
251
+ return _count(self._mem, self._region)
252
+
253
+ def filter_documents(self, filters: dict[str, Any] | None = None) -> list[Document]:
254
+ return _filter(self._mem, self._region, filters)
255
+
256
+ def write_documents(
257
+ self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE
258
+ ) -> int:
259
+ if not isinstance(documents, list) or any(
260
+ not isinstance(d, Document) for d in documents
261
+ ):
262
+ raise ValueError("Please provide a list of Documents.")
263
+ return _write(self._mem, self._region, self._dim, documents, policy)
264
+
265
+ def delete_documents(self, document_ids: list[str]) -> None:
266
+ _delete(self._mem, self._region, document_ids)
267
+
268
+ # ---- async ------------------------------------------------------------
269
+ # The bindings are sync, so a worker thread keeps the event loop free.
270
+
271
+ async def count_documents_async(self) -> int:
272
+ return await asyncio.to_thread(_count, self._mem, self._region)
273
+
274
+ async def filter_documents_async(
275
+ self, filters: dict[str, Any] | None = None
276
+ ) -> list[Document]:
277
+ return await asyncio.to_thread(_filter, self._mem, self._region, filters)
278
+
279
+ async def write_documents_async(
280
+ self, documents: list[Document], policy: DuplicatePolicy = DuplicatePolicy.NONE
281
+ ) -> int:
282
+ if not isinstance(documents, list) or any(
283
+ not isinstance(d, Document) for d in documents
284
+ ):
285
+ raise ValueError("Please provide a list of Documents.")
286
+ return await asyncio.to_thread(
287
+ _write, self._mem, self._region, self._dim, documents, policy
288
+ )
289
+
290
+ async def delete_documents_async(self, document_ids: list[str]) -> None:
291
+ await asyncio.to_thread(_delete, self._mem, self._region, document_ids)
292
+
293
+ # ---- beyond the protocol ----------------------------------------------
294
+
295
+ def embedding_retrieval(
296
+ self,
297
+ query_embedding: list[float],
298
+ top_k: int = 10,
299
+ filters: dict[str, Any] | None = None,
300
+ ) -> list[Document]:
301
+ """Documents ranked by recall, best first, with their score set.
302
+
303
+ Only documents Haystack embedded take part, as its own store does: an
304
+ unembedded one carries a hash of its text, and ranking that against a
305
+ real query vector produces a plausible score with no meaning behind it.
306
+ """
307
+ criterion = _pushdown(filters) or {}
308
+ criterion = {**criterion, "emb": True}
309
+ options = citadeldb.RecallOptions(payload_filter=criterion)
310
+
311
+ def surviving(hits: list[Any]) -> list[Document]:
312
+ out: list[Document] = []
313
+ for h in hits:
314
+ doc = _document(h)
315
+ if filters and not document_matches_filter(filters, doc):
316
+ continue
317
+ # Mutating a shared Document would affect other pipeline steps.
318
+ out.append(
319
+ replace(
320
+ doc,
321
+ score=min(1.0, 1.0 - h.distance)
322
+ if h.distance is not None
323
+ else h.score,
324
+ )
325
+ )
326
+ return out
327
+
328
+ # A filter containment cannot express is settled above, so the window has
329
+ # to widen until top_k survive it rather than answer short.
330
+ k = max(top_k, 32)
331
+ while True:
332
+ hits = self._mem.recall(
333
+ self._region,
334
+ embedding=query_embedding,
335
+ k=k,
336
+ kinds=[KIND],
337
+ options=options,
338
+ )
339
+ out = surviving(hits)
340
+ if len(out) >= top_k or len(hits) < k:
341
+ return out[:top_k]
342
+ k *= 2
343
+
344
+ def delete_all(self) -> int:
345
+ """Destroy every document's key, returning the number erased."""
346
+ return _erase(self._mem, self._region, _fetch(self._mem, self._region, None))
@@ -0,0 +1,23 @@
1
+ """deepset's DocumentStoreBaseTests run against CitadelDocumentStore."""
2
+ import pytest
3
+ from haystack.dataclasses import Document
4
+ from haystack.document_stores.errors import DuplicateDocumentError
5
+ from haystack.testing.document_store import DocumentStoreBaseTests
6
+
7
+ from citadeldb_haystack import CitadelDocumentStore
8
+
9
+ # The conformance fixtures embed at 768.
10
+ DIM = 768
11
+
12
+
13
+ class TestCitadelDocumentStore(DocumentStoreBaseTests):
14
+ @pytest.fixture
15
+ def document_store(self, tmp_path) -> CitadelDocumentStore:
16
+ return CitadelDocumentStore(str(tmp_path / "conformance.cdl"), "pw", dim=DIM)
17
+
18
+ def test_write_documents(self, document_store: CitadelDocumentStore):
19
+ """NONE falls back to FAIL, so an accidental re-write is reported."""
20
+ doc = Document(content="test doc")
21
+ assert document_store.write_documents([doc]) == 1
22
+ with pytest.raises(DuplicateDocumentError):
23
+ document_store.write_documents([doc])
@@ -0,0 +1,528 @@
1
+ import pytest
2
+ from haystack import Pipeline
3
+ from haystack.dataclasses import Document
4
+ from haystack.document_stores.errors import DuplicateDocumentError
5
+ from haystack.document_stores.types import DuplicatePolicy
6
+ from haystack.utils import Secret
7
+
8
+ from citadeldb_haystack import CitadelDocumentStore
9
+
10
+ DIM = 8
11
+
12
+
13
+ def vec(axis: int) -> list[float]:
14
+ v = [0.0] * DIM
15
+ v[axis] = 1.0
16
+ return v
17
+
18
+
19
+ @pytest.fixture()
20
+ def store(tmp_path):
21
+ return CitadelDocumentStore(str(tmp_path / "d.cdl"), "pw", dim=DIM)
22
+
23
+
24
+ @pytest.fixture()
25
+ def reference():
26
+ """The store whose behaviour the protocol is defined by."""
27
+ from haystack.document_stores.in_memory import InMemoryDocumentStore
28
+
29
+ return InMemoryDocumentStore()
30
+
31
+
32
+ # ---- conformance ---------------------------------------------------------
33
+
34
+
35
+ def test_satisfies_the_protocol(store):
36
+ """DocumentStore is not runtime_checkable, so conformance is by shape."""
37
+ for name in (
38
+ "to_dict", "from_dict", "count_documents", "filter_documents",
39
+ "write_documents", "delete_documents",
40
+ ):
41
+ assert callable(getattr(store, name)), name
42
+
43
+
44
+ def test_an_equality_filter_returns_the_matching_documents(store):
45
+ """Exercises the pushed-down filter path, which conformance never hits."""
46
+ store.write_documents([
47
+ Document(id="a", content="1", meta={"cat": "x", "n": 1}),
48
+ Document(id="b", content="2", meta={"cat": "y", "n": 2}),
49
+ ])
50
+ got = store.filter_documents({"field": "meta.cat", "operator": "==", "value": "y"})
51
+ assert [d.id for d in got] == ["b"]
52
+
53
+
54
+ def test_an_and_of_equalities_returns_the_matching_documents(store):
55
+ store.write_documents([
56
+ Document(id="a", content="1", meta={"cat": "x", "n": 1}),
57
+ Document(id="b", content="2", meta={"cat": "x", "n": 2}),
58
+ ])
59
+ got = store.filter_documents({
60
+ "operator": "AND",
61
+ "conditions": [
62
+ {"field": "meta.cat", "operator": "==", "value": "x"},
63
+ {"field": "meta.n", "operator": "==", "value": 2},
64
+ ],
65
+ })
66
+ assert [d.id for d in got] == ["b"]
67
+
68
+
69
+ # ---- duplicate policy ----------------------------------------------------
70
+
71
+
72
+ def test_skip_leaves_the_original_and_reports_fewer_written(store):
73
+ store.write_documents([Document(id="a", content="first")])
74
+ written = store.write_documents(
75
+ [Document(id="a", content="second"), Document(id="b", content="new")],
76
+ policy=DuplicatePolicy.SKIP,
77
+ )
78
+ assert written == 1
79
+ assert store.filter_documents({"field": "id", "operator": "==", "value": "a"})[
80
+ 0
81
+ ].content == "first"
82
+
83
+
84
+ def test_overwrite_replaces_and_counts_every_input(store):
85
+ store.write_documents([Document(id="a", content="first")])
86
+ assert (
87
+ store.write_documents(
88
+ [Document(id="a", content="second")], policy=DuplicatePolicy.OVERWRITE
89
+ )
90
+ == 1
91
+ )
92
+ assert store.count_documents() == 1
93
+
94
+
95
+ def test_duplicate_ids_within_one_call_keep_the_last(store):
96
+ """The reference assigns storage[doc.id], so the last write wins."""
97
+ store.write_documents(
98
+ [Document(id="d", content="one"), Document(id="d", content="two")],
99
+ policy=DuplicatePolicy.OVERWRITE,
100
+ )
101
+ assert store.count_documents() == 1
102
+ assert store.filter_documents()[0].content == "two"
103
+
104
+
105
+ def test_fail_raises_before_anything_is_written(store):
106
+ store.write_documents([Document(id="a", content="first")])
107
+ with pytest.raises(DuplicateDocumentError):
108
+ store.write_documents(
109
+ [Document(id="b", content="new"), Document(id="a", content="dup")]
110
+ )
111
+ assert store.count_documents() == 1, "a failed batch must not half-apply"
112
+
113
+
114
+ def test_non_documents_are_refused(store):
115
+ with pytest.raises(ValueError, match="list of Documents"):
116
+ store.write_documents(["not a document"])
117
+
118
+
119
+ # These compare against the reference directly, not against its docstring.
120
+
121
+
122
+ def test_overwrite_with_duplicates_counts_every_input(store, reference):
123
+ """Two copies of one id is still two documents in input."""
124
+ docs = [Document(id="d", content="one"), Document(id="d", content="two")]
125
+ assert store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) == (
126
+ reference.write_documents(docs, policy=DuplicatePolicy.OVERWRITE)
127
+ )
128
+
129
+
130
+ def test_skip_with_duplicates_counts_like_the_reference(store, reference):
131
+ docs = [Document(id="d", content="one"), Document(id="d", content="two")]
132
+ assert store.write_documents(docs, policy=DuplicatePolicy.SKIP) == (
133
+ reference.write_documents(docs, policy=DuplicatePolicy.SKIP)
134
+ )
135
+
136
+
137
+ def test_fail_raises_on_a_duplicate_within_one_batch(store, reference):
138
+ """A second copy collides even though neither was stored before."""
139
+ docs = [Document(id="d", content="one"), Document(id="d", content="two")]
140
+ with pytest.raises(DuplicateDocumentError):
141
+ reference.write_documents(docs, policy=DuplicatePolicy.FAIL)
142
+ with pytest.raises(DuplicateDocumentError):
143
+ store.write_documents(docs, policy=DuplicatePolicy.FAIL)
144
+
145
+
146
+ def test_skip_of_a_fresh_duplicate_stores_the_first(store):
147
+ """SKIP drops the later copy, so the earlier one is what survives."""
148
+ store.write_documents(
149
+ [Document(id="d", content="first"), Document(id="d", content="second")],
150
+ policy=DuplicatePolicy.SKIP,
151
+ )
152
+ assert store.count_documents() == 1
153
+ assert store.filter_documents()[0].content == "first"
154
+
155
+
156
+ # ---- round-trip ----------------------------------------------------------
157
+
158
+
159
+ def test_embeddings_round_trip(store):
160
+ """The conformance suite compares whole Documents, embeddings included."""
161
+ store.write_documents([Document(id="e", content="x", embedding=vec(3))])
162
+ assert store.filter_documents()[0].embedding == vec(3)
163
+
164
+
165
+ def test_documents_without_embeddings_are_storable(store):
166
+ store.write_documents([Document(id="n", content="no vector")])
167
+ got = store.filter_documents()[0]
168
+ assert got.embedding is None and got.content == "no vector"
169
+
170
+
171
+ def test_nested_metadata_round_trips(store):
172
+ store.write_documents([Document(id="m", content="x", meta={"o": {"i": [1, 2]}})])
173
+ assert store.filter_documents()[0].meta["o"] == {"i": [1, 2]}
174
+
175
+
176
+ def test_a_wrong_width_embedding_names_the_fix(store):
177
+ with pytest.raises(ValueError, match="dim=3"):
178
+ store.write_documents([Document(id="w", content="x", embedding=[1.0, 2.0, 3.0])])
179
+
180
+
181
+ # ---- deletes -------------------------------------------------------------
182
+
183
+
184
+ def test_delete_destroys_keys(store):
185
+ store.write_documents([Document(id="s", content="secret")])
186
+ store.delete_documents(["s"])
187
+ assert store.count_documents() == 0
188
+
189
+
190
+ def test_delete_tolerates_unknown_and_duplicate_ids(store):
191
+ store.write_documents([Document(id="a", content="x")])
192
+ store.delete_documents(["missing", "a", "a"])
193
+ assert store.count_documents() == 0
194
+
195
+
196
+ def test_delete_all_empties_the_store(store):
197
+ store.write_documents([Document(content="a"), Document(content="b")])
198
+ assert store.delete_all() == 2
199
+ assert store.count_documents() == 0
200
+
201
+
202
+ # ---- serialization -------------------------------------------------------
203
+
204
+
205
+ def test_a_literal_passphrase_refuses_to_serialize(store):
206
+ """Haystack writes pipelines to disk; a passphrase must not travel."""
207
+ with pytest.raises(ValueError, match="Cannot serialize"):
208
+ store.to_dict()
209
+
210
+
211
+ def test_an_env_var_passphrase_serializes_by_reference(tmp_path, monkeypatch):
212
+ monkeypatch.setenv("CITADEL_TEST_KEY", "pw")
213
+ s = CitadelDocumentStore(
214
+ str(tmp_path / "s.cdl"), Secret.from_env_var("CITADEL_TEST_KEY"), dim=DIM
215
+ )
216
+ data = s.to_dict()
217
+ assert data["init_parameters"]["key"]["env_vars"] == ["CITADEL_TEST_KEY"]
218
+ assert "pw" not in str(data), "the passphrase itself must never be serialized"
219
+
220
+
221
+ def test_round_trip_through_from_dict(tmp_path, monkeypatch):
222
+ monkeypatch.setenv("CITADEL_TEST_KEY", "pw")
223
+ path = str(tmp_path / "r.cdl")
224
+ s = CitadelDocumentStore(path, Secret.from_env_var("CITADEL_TEST_KEY"), dim=DIM)
225
+ s.write_documents([Document(id="a", content="kept")])
226
+ revived = CitadelDocumentStore.from_dict(s.to_dict())
227
+ assert revived.count_documents() == 1
228
+
229
+
230
+ def test_a_missing_passphrase_is_refused(tmp_path, monkeypatch):
231
+ monkeypatch.setenv("CITADEL_EMPTY_KEY", "")
232
+ with pytest.raises(ValueError, match="passphrase"):
233
+ CitadelDocumentStore(
234
+ str(tmp_path / "k.cdl"), Secret.from_env_var("CITADEL_EMPTY_KEY"), dim=DIM
235
+ )
236
+
237
+
238
+ # ---- retrieval -----------------------------------------------------------
239
+
240
+
241
+ def test_embedding_retrieval_ranks_and_scores(store):
242
+ store.write_documents([
243
+ Document(id="near", content="a", embedding=vec(0)),
244
+ Document(id="far", content="b", embedding=vec(7)),
245
+ ])
246
+ got = store.embedding_retrieval(vec(0), top_k=2)
247
+ assert [d.id for d in got] == ["near", "far"]
248
+ assert got[0].score is not None and 0.0 <= got[0].score <= 1.0
249
+
250
+
251
+ def test_embedding_retrieval_respects_filters(store):
252
+ store.write_documents([
253
+ Document(id="x", content="a", meta={"cat": "x"}, embedding=vec(0)),
254
+ Document(id="y", content="b", meta={"cat": "y"}, embedding=vec(0)),
255
+ ])
256
+ got = store.embedding_retrieval(
257
+ vec(0), top_k=5, filters={"field": "meta.cat", "operator": "==", "value": "y"}
258
+ )
259
+ assert [d.id for d in got] == ["y"]
260
+
261
+
262
+ def test_embedding_retrieval_respects_top_k(store):
263
+ store.write_documents(
264
+ [Document(id=f"d{i}", content="t", embedding=vec(i % DIM)) for i in range(6)]
265
+ )
266
+ assert len(store.embedding_retrieval(vec(0), top_k=2)) == 2
267
+
268
+
269
+ # ---- edges ---------------------------------------------------------------
270
+
271
+
272
+ def test_writing_nothing_is_not_an_error(store):
273
+ assert store.write_documents([]) == 0
274
+ assert store.count_documents() == 0
275
+
276
+
277
+ def test_filtering_an_empty_store_is_empty(store):
278
+ assert store.filter_documents() == []
279
+ assert store.filter_documents({"field": "meta.x", "operator": "==", "value": 1}) == []
280
+
281
+
282
+ def test_an_or_filter_is_not_narrowed_by_pushdown(store):
283
+ """Pushing a leaf of an OR would drop rows the filter keeps."""
284
+ store.write_documents([
285
+ Document(id="a", content="1", meta={"cat": "x"}),
286
+ Document(id="b", content="2", meta={"cat": "y"}),
287
+ ])
288
+ got = store.filter_documents({
289
+ "operator": "OR",
290
+ "conditions": [
291
+ {"field": "meta.cat", "operator": "==", "value": "x"},
292
+ {"field": "meta.cat", "operator": "==", "value": "y"},
293
+ ],
294
+ })
295
+ assert sorted(d.id for d in got) == ["a", "b"]
296
+
297
+
298
+ def test_a_none_valued_equality_filter_is_not_pushed(store):
299
+ """None means absent, which containment cannot express."""
300
+ store.write_documents([
301
+ Document(id="has", content="1", meta={"n": 1}),
302
+ Document(id="hasnt", content="2", meta={}),
303
+ ])
304
+ got = store.filter_documents({"field": "meta.n", "operator": "==", "value": None})
305
+ assert [d.id for d in got] == ["hasnt"]
306
+
307
+
308
+ def test_a_document_with_no_content_is_storable(store):
309
+ store.write_documents([Document(id="empty", content=None, meta={"k": 1})])
310
+ got = store.filter_documents()[0]
311
+ assert got.content is None and got.meta["k"] == 1
312
+
313
+
314
+ def test_filtering_on_a_document_field_not_metadata(store):
315
+ store.write_documents([Document(id="wanted", content="x")])
316
+ got = store.filter_documents({"field": "id", "operator": "==", "value": "wanted"})
317
+ assert [d.id for d in got] == ["wanted"]
318
+
319
+
320
+ def test_a_not_filter_is_not_narrowed_by_pushdown(store):
321
+ store.write_documents([
322
+ Document(id="a", content="1", meta={"cat": "x"}),
323
+ Document(id="b", content="2", meta={"cat": "y"}),
324
+ ])
325
+ got = store.filter_documents({
326
+ "operator": "NOT",
327
+ "conditions": [{"field": "meta.cat", "operator": "==", "value": "x"}],
328
+ })
329
+ assert [d.id for d in got] == ["b"]
330
+
331
+
332
+ def test_retrieval_on_an_empty_store_is_empty(store):
333
+ assert store.embedding_retrieval(vec(0), top_k=5) == []
334
+
335
+
336
+ def test_scores_do_not_leak_onto_stored_documents(store):
337
+ """Mutating a shared Document affects other pipeline steps."""
338
+ store.write_documents([Document(id="a", content="x", embedding=vec(0))])
339
+ store.embedding_retrieval(vec(0), top_k=1)
340
+ assert store.filter_documents()[0].score is None
341
+
342
+
343
+ def test_it_survives_a_reopen(tmp_path):
344
+ """A region's embedder lives in memory, so a reopen must rebuild it."""
345
+ import gc
346
+
347
+ p = str(tmp_path / "reopen.cdl")
348
+ first = CitadelDocumentStore(p, "pw", dim=DIM)
349
+ first.write_documents([Document(id="d1", content="the disk was full",
350
+ meta={"chapter": "intro"}, embedding=vec(3))])
351
+ del first
352
+ gc.collect()
353
+
354
+ again = CitadelDocumentStore(p, "pw", dim=DIM)
355
+ assert again.count_documents() == 1
356
+ got = again.filter_documents(
357
+ {"field": "meta.chapter", "operator": "==", "value": "intro"}
358
+ )
359
+ assert [d.id for d in got] == ["d1"]
360
+ assert got[0].embedding == vec(3), "the embedding did not survive the reopen"
361
+ assert [d.id for d in again.embedding_retrieval(vec(3), top_k=1)] == ["d1"]
362
+
363
+
364
+ def test_a_wrong_passphrase_cannot_reopen(tmp_path):
365
+ """Pins the encryption claim rather than inferring it from a reopen."""
366
+ import gc
367
+
368
+ import citadeldb
369
+
370
+ p = str(tmp_path / "enc.cdl")
371
+ first = CitadelDocumentStore(p, "right", dim=DIM)
372
+ first.write_documents([Document(id="d1", content="the disk was full",
373
+ embedding=vec(3))])
374
+ del first
375
+ gc.collect()
376
+
377
+ with pytest.raises(citadeldb.EncryptionError):
378
+ CitadelDocumentStore(p, "wrong", dim=DIM)
379
+
380
+
381
+ def test_concurrent_writes_all_land(tmp_path):
382
+ """Indexing pipelines run in parallel; the engine is shared."""
383
+ import concurrent.futures as cf
384
+
385
+ s = CitadelDocumentStore(str(tmp_path / "conc.cdl"), "pw", dim=DIM)
386
+ with cf.ThreadPoolExecutor(max_workers=4) as ex:
387
+ list(ex.map(
388
+ lambda i: s.write_documents(
389
+ [Document(id=f"c{i}", content=f"body {i}", embedding=vec(i % DIM))],
390
+ policy=DuplicatePolicy.OVERWRITE,
391
+ ),
392
+ range(40),
393
+ ))
394
+ assert s.count_documents() == 40
395
+
396
+
397
+ def test_a_nested_metadata_filter_finds_its_document(tmp_path):
398
+ """`meta.person.name` addresses a nested value. Pushing it as one flat key
399
+ spelled with a dot matches nothing, and the post-filter never sees the row."""
400
+ s = CitadelDocumentStore(str(tmp_path / "nested.cdl"), "pw", dim=DIM)
401
+ s.write_documents([
402
+ Document(id="n1", content="nested", meta={"person": {"name": "ada"}},
403
+ embedding=vec(0)),
404
+ Document(id="n2", content="other", meta={"person": {"name": "bob"}},
405
+ embedding=vec(1)),
406
+ ])
407
+ f = {"field": "meta.person.name", "operator": "==", "value": "ada"}
408
+ assert [d.id for d in s.filter_documents(filters=f)] == ["n1"]
409
+ assert [d.id for d in s.embedding_retrieval(vec(0), top_k=5, filters=f)] == ["n1"]
410
+
411
+
412
+ def test_a_numeric_metadata_filter_is_not_decided_by_the_pushdown(tmp_path):
413
+ """Haystack's `==` is Python's, which holds 1 == 1.0; containment is type
414
+ strict, so pushing a number would drop a row the filter keeps."""
415
+ s = CitadelDocumentStore(str(tmp_path / "num.cdl"), "pw", dim=DIM)
416
+ s.write_documents([Document(id="i1", content="int", meta={"year": 2024},
417
+ embedding=vec(0))])
418
+ f = {"field": "meta.year", "operator": "==", "value": 2024.0}
419
+ assert [d.id for d in s.filter_documents(filters=f)] == ["i1"]
420
+
421
+
422
+ def test_unembedded_documents_stay_out_of_embedding_retrieval(tmp_path):
423
+ """Haystack's own store excludes them. A hash of the text scores plausibly
424
+ and would push the one real match out of top_k."""
425
+ s = CitadelDocumentStore(str(tmp_path / "unemb.cdl"), "pw", dim=DIM)
426
+ s.write_documents([Document(id=f"u{i}", content=f"unembedded {i}") for i in range(20)])
427
+ s.write_documents([Document(id="real", content="embedded", embedding=vec(0))])
428
+ assert s.count_documents() == 21
429
+ found = s.embedding_retrieval(vec(0), top_k=10)
430
+ assert [d.id for d in found] == ["real"]
431
+ # They are stored and readable, just not ranked against a query vector.
432
+ assert len(s.filter_documents()) == 21
433
+
434
+
435
+ def test_an_unpushable_filter_survives_a_window_of_other_documents(tmp_path):
436
+ """`in` cannot be pushed, so it is settled after ranking; the only matching
437
+ document must still be found under a corpus that outranks it."""
438
+ s = CitadelDocumentStore(str(tmp_path / "window.cdl"), "pw", dim=DIM)
439
+ s.write_documents([Document(id="gold", content="needle", meta={"tag": "keep"},
440
+ embedding=vec(1))])
441
+ s.write_documents([
442
+ Document(id=f"c{i}", content=f"nearer {i}", meta={"tag": "drop"},
443
+ embedding=vec(0))
444
+ for i in range(400)
445
+ ])
446
+ f = {"field": "meta.tag", "operator": "in", "value": ["keep"]}
447
+ assert [d.id for d in s.embedding_retrieval(vec(0), top_k=1, filters=f)] == ["gold"]
448
+
449
+
450
+ def test_concurrent_overwrites_of_one_id_leave_one_document(tmp_path):
451
+ """A document id is unique, so parallel pipelines overwriting the same one
452
+ must supersede rather than each add a row."""
453
+ import concurrent.futures as cf
454
+
455
+ s = CitadelDocumentStore(str(tmp_path / "oneid.cdl"), "pw", dim=DIM)
456
+ with cf.ThreadPoolExecutor(max_workers=8) as ex:
457
+ list(ex.map(
458
+ lambda i: s.write_documents(
459
+ [Document(id="same", content=f"v{i}", embedding=vec(i % DIM))],
460
+ policy=DuplicatePolicy.OVERWRITE,
461
+ ),
462
+ range(32),
463
+ ))
464
+ assert s.count_documents() == 1
465
+ # Retrieval must not fill top_k with copies of the one document either.
466
+ found = s.embedding_retrieval(query_embedding=vec(0), top_k=5)
467
+ assert [d.id for d in found] == ["same"]
468
+ s.delete_documents(["same"])
469
+ assert s.count_documents() == 0
470
+
471
+
472
+ def test_two_stores_share_one_database_file(tmp_path):
473
+ path = str(tmp_path / "shared.cdl")
474
+ a = CitadelDocumentStore(path, "pw", region="a", dim=DIM)
475
+ b = CitadelDocumentStore(path, "pw", region="b", dim=DIM)
476
+ a.write_documents([Document(id="1", content="in a")])
477
+ b.write_documents([Document(id="2", content="in b")])
478
+ assert a.count_documents() == 1 and b.count_documents() == 1
479
+
480
+
481
+ def test_a_large_batch_keeps_every_document(store):
482
+ store.write_documents([Document(id=f"b{i}", content=f"body {i}") for i in range(300)])
483
+ assert store.count_documents() == 300
484
+
485
+
486
+ def test_it_serializes_inside_a_pipeline(tmp_path, monkeypatch):
487
+ """A store is only useful if a pipeline carrying it can be written out."""
488
+ monkeypatch.setenv("CITADEL_TEST_KEY", "pw")
489
+ s = CitadelDocumentStore(
490
+ str(tmp_path / "p.cdl"), Secret.from_env_var("CITADEL_TEST_KEY"), dim=DIM
491
+ )
492
+ from haystack.components.writers import DocumentWriter
493
+
494
+ pipe = Pipeline()
495
+ pipe.add_component("writer", DocumentWriter(document_store=s))
496
+ assert "CitadelDocumentStore" in str(pipe.to_dict())
497
+
498
+
499
+ # ---- async ---------------------------------------------------------------
500
+
501
+
502
+ async def test_async_surface_round_trips(store):
503
+ assert await store.write_documents_async([Document(id="a", content="x")]) == 1
504
+ assert await store.count_documents_async() == 1
505
+ assert len(await store.filter_documents_async()) == 1
506
+ await store.delete_documents_async(["a"])
507
+ assert await store.count_documents_async() == 0
508
+
509
+
510
+ async def test_the_event_loop_is_not_blocked(store):
511
+ import asyncio
512
+
513
+ ticks = 0
514
+
515
+ async def tick():
516
+ nonlocal ticks
517
+ while True:
518
+ ticks += 1
519
+ await asyncio.sleep(0)
520
+
521
+ ticker = asyncio.create_task(tick())
522
+ await asyncio.sleep(0)
523
+ await store.write_documents_async(
524
+ [Document(id=f"L{i}", content=f"loop {i}") for i in range(40)]
525
+ )
526
+ await store.filter_documents_async()
527
+ ticker.cancel()
528
+ assert ticks > 1, "the loop made no progress during a store call"