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.
- citadeldb_haystack-2.0.0/.gitignore +18 -0
- citadeldb_haystack-2.0.0/PKG-INFO +109 -0
- citadeldb_haystack-2.0.0/README.md +87 -0
- citadeldb_haystack-2.0.0/pyproject.toml +40 -0
- citadeldb_haystack-2.0.0/src/citadeldb_haystack/__init__.py +14 -0
- citadeldb_haystack-2.0.0/src/citadeldb_haystack/document_store.py +346 -0
- citadeldb_haystack-2.0.0/tests/test_conformance.py +23 -0
- citadeldb_haystack-2.0.0/tests/test_document_store.py +528 -0
|
@@ -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"
|