citadeldb-google-adk 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,116 @@
1
+ Metadata-Version: 2.5
2
+ Name: citadeldb-google-adk
3
+ Version: 2.0.0
4
+ Summary: Google ADK memory 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: adk,agent,agents,encryption,google-adk,memory
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: google-adk<3,>=2.0
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-google-adk
24
+
25
+ A [Google ADK](https://github.com/google/adk-python) `BaseMemoryService` 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
+ ```
30
+ pip install citadeldb-google-adk
31
+ ```
32
+
33
+ ```python
34
+ from google.adk.runners import Runner
35
+ from citadeldb_google_adk import CitadelMemoryService
36
+
37
+ memory = CitadelMemoryService("adk_memory.cdl", key="your-passphrase")
38
+
39
+ runner = Runner(
40
+ app_name="my_app",
41
+ agent=agent, # your root agent
42
+ session_service=session_service, # your session service
43
+ memory_service=memory,
44
+ )
45
+ ```
46
+
47
+ ## Search is ranked recall, not word matching
48
+
49
+ ADK hands the service a query string, so Citadel embeds it and runs hybrid recall: vector
50
+ distance, keyword rank and recency, fused into one score. The reference
51
+ `InMemoryMemoryService` returns only turns sharing a word with the query; nothing here is
52
+ dropped for lacking one. With the default `MockEmbedder` that ranking is still lexical:
53
+
54
+ ```python
55
+ await memory.add_session_to_memory(session) # a Session your Runner already ran
56
+
57
+ await memory.search_memory(app_name="my_app", user_id="alice",
58
+ query="why did the release break?")
59
+ # SearchMemoryResponse(memories=[MemoryEntry(...disk was full...)])
60
+ ```
61
+
62
+ ## Deletes destroy the key
63
+
64
+ Every event is sealed under its own key. Erasing destroys those keys, so the bytes on disk
65
+ stay unreadable. A backup taken before the delete carries its own copy of the wrapped key
66
+ and is out of scope.
67
+
68
+ ```python
69
+ memory.forget_user("my_app", "alice") # returns the number erased
70
+ memory.forget_session("my_app", "alice", "s-42")
71
+ ```
72
+
73
+ ADK's own memory services expose no erasure method.
74
+
75
+ ## Direct writes are supported
76
+
77
+ `add_memory` writes memories without going through a session. The reference
78
+ `InMemoryMemoryService` raises `NotImplementedError` for it.
79
+
80
+ ```python
81
+ from google.adk.memory.memory_entry import MemoryEntry
82
+ from google.genai import types
83
+
84
+ entry = MemoryEntry(content=types.Content(parts=[types.Part(text="prefers dark mode")]))
85
+ await memory.add_memory(app_name="my_app", user_id="alice", memories=[entry])
86
+ ```
87
+
88
+ `add_events_to_memory` likewise persists the events you pass rather than a whole session.
89
+
90
+ ## Notes
91
+
92
+ `add_session_to_memory` sets the session's events, as `InMemoryMemoryService` does:
93
+ re-adding never duplicates rows, and an event dropped from the session is dropped from
94
+ memory. `add_events_to_memory` is the additive one, skipping ids already stored.
95
+
96
+ Citadel is embedded and one process owns the file. A path already open on this thread,
97
+ under the same passphrase, is shared, so this can sit on the same database as another
98
+ Citadel adapter; construct them on the same thread.
99
+
100
+ `MockEmbedder` is the default and needs no download, which is enough to run an agent and to
101
+ test. For semantic recall pass a real embedder. `CandleEmbedder` is not in the default
102
+ `citadeldb` wheel and needs a source build (`maturin build --features candle-embed`); any
103
+ object exposing `dim`, `metric`, `model_id`, `embed` and `embed_queries` works too:
104
+
105
+ ```python
106
+ import citadeldb
107
+ memory = CitadelMemoryService(
108
+ "adk_memory.cdl",
109
+ key="your-passphrase",
110
+ embedder=citadeldb.CandleEmbedder("/path/to/e5-large", preset="e5-large"),
111
+ )
112
+ ```
113
+
114
+ ## License
115
+
116
+ Apache-2.0
@@ -0,0 +1,94 @@
1
+ # citadeldb-google-adk
2
+
3
+ A [Google ADK](https://github.com/google/adk-python) `BaseMemoryService` 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
+ ```
8
+ pip install citadeldb-google-adk
9
+ ```
10
+
11
+ ```python
12
+ from google.adk.runners import Runner
13
+ from citadeldb_google_adk import CitadelMemoryService
14
+
15
+ memory = CitadelMemoryService("adk_memory.cdl", key="your-passphrase")
16
+
17
+ runner = Runner(
18
+ app_name="my_app",
19
+ agent=agent, # your root agent
20
+ session_service=session_service, # your session service
21
+ memory_service=memory,
22
+ )
23
+ ```
24
+
25
+ ## Search is ranked recall, not word matching
26
+
27
+ ADK hands the service a query string, so Citadel embeds it and runs hybrid recall: vector
28
+ distance, keyword rank and recency, fused into one score. The reference
29
+ `InMemoryMemoryService` returns only turns sharing a word with the query; nothing here is
30
+ dropped for lacking one. With the default `MockEmbedder` that ranking is still lexical:
31
+
32
+ ```python
33
+ await memory.add_session_to_memory(session) # a Session your Runner already ran
34
+
35
+ await memory.search_memory(app_name="my_app", user_id="alice",
36
+ query="why did the release break?")
37
+ # SearchMemoryResponse(memories=[MemoryEntry(...disk was full...)])
38
+ ```
39
+
40
+ ## Deletes destroy the key
41
+
42
+ Every event is sealed under its own key. Erasing destroys those keys, so the bytes on disk
43
+ stay unreadable. A backup taken before the delete carries its own copy of the wrapped key
44
+ and is out of scope.
45
+
46
+ ```python
47
+ memory.forget_user("my_app", "alice") # returns the number erased
48
+ memory.forget_session("my_app", "alice", "s-42")
49
+ ```
50
+
51
+ ADK's own memory services expose no erasure method.
52
+
53
+ ## Direct writes are supported
54
+
55
+ `add_memory` writes memories without going through a session. The reference
56
+ `InMemoryMemoryService` raises `NotImplementedError` for it.
57
+
58
+ ```python
59
+ from google.adk.memory.memory_entry import MemoryEntry
60
+ from google.genai import types
61
+
62
+ entry = MemoryEntry(content=types.Content(parts=[types.Part(text="prefers dark mode")]))
63
+ await memory.add_memory(app_name="my_app", user_id="alice", memories=[entry])
64
+ ```
65
+
66
+ `add_events_to_memory` likewise persists the events you pass rather than a whole session.
67
+
68
+ ## Notes
69
+
70
+ `add_session_to_memory` sets the session's events, as `InMemoryMemoryService` does:
71
+ re-adding never duplicates rows, and an event dropped from the session is dropped from
72
+ memory. `add_events_to_memory` is the additive one, skipping ids already stored.
73
+
74
+ Citadel is embedded and one process owns the file. A path already open on this thread,
75
+ under the same passphrase, is shared, so this can sit on the same database as another
76
+ Citadel adapter; construct them on the same thread.
77
+
78
+ `MockEmbedder` is the default and needs no download, which is enough to run an agent and to
79
+ test. For semantic recall pass a real embedder. `CandleEmbedder` is not in the default
80
+ `citadeldb` wheel and needs a source build (`maturin build --features candle-embed`); any
81
+ object exposing `dim`, `metric`, `model_id`, `embed` and `embed_queries` works too:
82
+
83
+ ```python
84
+ import citadeldb
85
+ memory = CitadelMemoryService(
86
+ "adk_memory.cdl",
87
+ key="your-passphrase",
88
+ embedder=citadeldb.CandleEmbedder("/path/to/e5-large", preset="e5-large"),
89
+ )
90
+ ```
91
+
92
+ ## License
93
+
94
+ 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-google-adk"
7
+ dynamic = ["version"]
8
+ description = "Google ADK memory 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 = ["google-adk", "adk", "agents", "memory", "agent", "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
+ # An event replaces under its own id, which needs the keyed write added in 2.0.
22
+ dependencies = ["citadeldb>=2.0,<3", "google-adk>=2.0,<3"]
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_google_adk"]
38
+
39
+ [tool.pytest.ini_options]
40
+ asyncio_mode = "auto"
@@ -0,0 +1,13 @@
1
+ """Google ADK memory backed by Citadel, encrypted at rest."""
2
+
3
+ from importlib.metadata import PackageNotFoundError, version
4
+
5
+ from .memory import CitadelMemoryService
6
+
7
+ __all__ = ["CitadelMemoryService", "__version__"]
8
+
9
+
10
+ try:
11
+ __version__ = version("citadeldb-google-adk")
12
+ except PackageNotFoundError: # running from a source tree, never installed
13
+ __version__ = "0+unknown"
@@ -0,0 +1,330 @@
1
+ """Google ADK memory service over an encrypted Citadel region."""
2
+ from __future__ import annotations
3
+
4
+ import asyncio
5
+ import functools
6
+ from collections.abc import Mapping, Sequence
7
+ from datetime import datetime
8
+ from typing import TYPE_CHECKING, Any
9
+
10
+ import citadeldb
11
+ from google.adk.memory.base_memory_service import BaseMemoryService, SearchMemoryResponse
12
+ from google.adk.memory.memory_entry import MemoryEntry
13
+ from google.genai import types
14
+
15
+ if TYPE_CHECKING:
16
+ from google.adk.events.event import Event
17
+ from google.adk.sessions.session import Session
18
+
19
+ KIND = "event"
20
+ DEFAULT_PATH = "adk_memory.cdl"
21
+ DEFAULT_REGION = "adk_memory"
22
+ # A scope is read whole and filtered in Python; recall over-fetches too.
23
+ PAGE = 10_000
24
+ # `search_memory` takes no limit, so the service owns one; raise it on the constructor.
25
+ DEFAULT_SEARCH_LIMIT = 64
26
+
27
+
28
+ def _page(mem: Any, region: str, criterion: dict[str, Any] | None = None) -> list[Any]:
29
+ """Page to the end: one fetch is bounded, and a partial erase must not look whole."""
30
+ out: list[Any] = []
31
+ after = None
32
+ while True:
33
+ got = mem.fetch(region, KIND, payload_filter=criterion, limit=PAGE, after_id=after)
34
+ out.extend(got)
35
+ if len(got) < PAGE:
36
+ return out
37
+ after = got[-1].id
38
+ # ADK partitions memory by (app_name, user_id); nothing crosses that boundary.
39
+ _UNKNOWN_SESSION = "__unknown_session_id__"
40
+
41
+
42
+ def _text_of(content: types.Content | None) -> str:
43
+ """The searchable text of an event, which is what Citadel embeds."""
44
+ if content is None or not content.parts:
45
+ return ""
46
+ return " ".join(p.text for p in content.parts if getattr(p, "text", None))
47
+
48
+
49
+ def _stamp(timestamp: float | None) -> str | None:
50
+ """ADK formats memory timestamps as local-time ISO 8601, so match it."""
51
+ return datetime.fromtimestamp(timestamp).isoformat() if timestamp else None
52
+
53
+
54
+ def _scope(app_name: str, user_id: str) -> dict[str, str]:
55
+ return {"app": app_name, "user": user_id}
56
+
57
+
58
+ # A Database is pinned to its opening thread, so workers take Memory, not self.
59
+
60
+
61
+ def _memory_key(*parts: str) -> str:
62
+ """One engine key from several names, length-prefixed so that no two name
63
+ tuples can spell the same key."""
64
+ return "".join(f"{len(p)}\x1f{p}" for p in parts)
65
+
66
+
67
+ def _write_events(
68
+ mem: Any,
69
+ region: str,
70
+ app_name: str,
71
+ user_id: str,
72
+ session_id: str,
73
+ events: Sequence[Event],
74
+ custom_metadata: Mapping[str, object] | None,
75
+ *,
76
+ replace_session: bool = False,
77
+ ) -> None:
78
+ scope = _scope(app_name, user_id)
79
+ # Read before the write, so the erase below cannot take a row this call just
80
+ # made. Rows superseded by the keyed write are already gone by then, and
81
+ # forgetting an id twice is a no-op.
82
+ prior = (
83
+ _page(mem, region, scope | {"sid": session_id}) if replace_session else []
84
+ )
85
+ keyed: dict[str, dict[str, Any]] = {}
86
+ anonymous: list[dict[str, Any]] = []
87
+ for event in events:
88
+ if not event.content or not event.content.parts:
89
+ continue # the reference skips contentless events
90
+ text = _text_of(event.content)
91
+ if not text:
92
+ continue # nothing to embed, so nothing recall could ever return
93
+ atom = {
94
+ "kind": KIND,
95
+ "text": text,
96
+ "payload": {
97
+ **scope,
98
+ "sid": session_id,
99
+ "eid": event.id,
100
+ "author": event.author,
101
+ "ts": _stamp(event.timestamp),
102
+ # Stored whole so a memory hands back the content ADK gave us.
103
+ "content": event.content.model_dump(mode="json", exclude_none=True),
104
+ "meta": dict(custom_metadata) if custom_metadata else {},
105
+ },
106
+ }
107
+ if event.id:
108
+ # The first copy of an id in one call stands, as it did when a read
109
+ # of the stored ids decided this.
110
+ keyed.setdefault(_memory_key(app_name, user_id, session_id, event.id), atom)
111
+ else:
112
+ anonymous.append(atom) # no stable name to key on
113
+ if keyed:
114
+ # Keyed, so re-ingesting a session converges on the stored events instead
115
+ # of duplicating them, and two concurrent ingests of one event cannot
116
+ # both observe it absent and both write it.
117
+ mem.remember_replacing_keyed_batch(
118
+ region, [(atom, k) for k, atom in keyed.items()]
119
+ )
120
+ if anonymous:
121
+ mem.remember_batch(region, anonymous)
122
+ if replace_session:
123
+ # add_session_to_memory SETS the session's events rather than merging
124
+ # them, so an event dropped from the session is dropped from memory too.
125
+ # add_events_to_memory is the additive one and leaves this alone.
126
+ written = {atom["payload"]["eid"] for atom in keyed.values()}
127
+ stale = [h for h in prior if h.payload.get("eid") not in written]
128
+ if stale:
129
+ mem.forget(region, [h.id for h in stale])
130
+
131
+
132
+ def _write_memories(
133
+ mem: Any,
134
+ region: str,
135
+ app_name: str,
136
+ user_id: str,
137
+ memories: Sequence[MemoryEntry],
138
+ custom_metadata: Mapping[str, object] | None,
139
+ ) -> None:
140
+ scope = _scope(app_name, user_id)
141
+ # An id replaces rather than skips: ignoring an update is a silent no-op.
142
+ keyed: dict[str, dict[str, Any]] = {}
143
+ anonymous: list[dict[str, Any]] = []
144
+ for m in memories:
145
+ text = _text_of(m.content)
146
+ if not text:
147
+ continue
148
+ meta = dict(m.custom_metadata)
149
+ if custom_metadata:
150
+ meta.update(custom_metadata)
151
+ atom = {
152
+ "kind": KIND,
153
+ "text": text,
154
+ "payload": {
155
+ **scope,
156
+ "sid": _UNKNOWN_SESSION,
157
+ "eid": m.id,
158
+ "author": m.author,
159
+ "ts": m.timestamp,
160
+ "content": m.content.model_dump(mode="json", exclude_none=True),
161
+ "meta": meta,
162
+ },
163
+ }
164
+ if m.id:
165
+ keyed[m.id] = atom
166
+ else:
167
+ anonymous.append(atom) # no identity to replace, so it is always new
168
+ if keyed:
169
+ # Keyed, so an id replaces in one transaction. Two concurrent adds of one
170
+ # id supersede rather than both landing, which the read-then-write-then-
171
+ # erase this replaces could not prevent.
172
+ mem.remember_replacing_keyed_batch(
173
+ region,
174
+ [(atom, _memory_key(app_name, user_id, eid)) for eid, atom in keyed.items()],
175
+ )
176
+ if anonymous:
177
+ mem.remember_batch(region, anonymous)
178
+
179
+
180
+ def _entry(hit) -> MemoryEntry:
181
+ p = hit.payload
182
+ return MemoryEntry(
183
+ content=types.Content.model_validate(p["content"]),
184
+ custom_metadata=p.get("meta") or {},
185
+ id=p.get("eid"),
186
+ author=p.get("author"),
187
+ timestamp=p.get("ts"),
188
+ )
189
+
190
+
191
+ def _search(
192
+ mem: Any, region: str, app_name: str, user_id: str, query: str, limit: int
193
+ ) -> SearchMemoryResponse:
194
+ hits = mem.recall(
195
+ region,
196
+ text=query,
197
+ k=limit,
198
+ kinds=[KIND],
199
+ options=citadeldb.RecallOptions(payload_filter=_scope(app_name, user_id)),
200
+ )
201
+ return SearchMemoryResponse(memories=[_entry(h) for h in hits])
202
+
203
+
204
+ class CitadelMemoryService(BaseMemoryService):
205
+ """An ADK `BaseMemoryService` backed by one encrypted Citadel region."""
206
+
207
+ def __init__(
208
+ self,
209
+ path: str = DEFAULT_PATH,
210
+ key: str = "",
211
+ *,
212
+ region: str = DEFAULT_REGION,
213
+ embedder: Any | None = None,
214
+ search_limit: int = DEFAULT_SEARCH_LIMIT,
215
+ ) -> None:
216
+ if not key:
217
+ raise ValueError("a passphrase is required: conversations are the payload")
218
+ self._search_limit = search_limit
219
+ try:
220
+ self._db = citadeldb.connect(path, key=key, region_keys=True)
221
+ except citadeldb.OperationalError as e:
222
+ if "locked" not in str(e):
223
+ raise
224
+ raise RuntimeError(
225
+ f"{path} is open in another process. Citadel is embedded, so one "
226
+ f"process owns the file."
227
+ ) from e
228
+ self._mem = self._db.memory()
229
+ self._region = region
230
+ # Idempotent for a region of the same width, so a dim clash raises here.
231
+ self._mem.create_encrypted_region(
232
+ region, embedder or citadeldb.MockEmbedder(dim=64)
233
+ )
234
+
235
+ # ---- storage helpers --------------------------------------------------
236
+
237
+ def _owned(self, app_name: str, user_id: str) -> list[Any]:
238
+ """Every atom for one (app, user). Served by the payload index."""
239
+ return _page(self._mem, self._region, _scope(app_name, user_id))
240
+
241
+ # ---- the abstract surface --------------------------------------------
242
+
243
+ async def add_session_to_memory(self, session: Session) -> None:
244
+ await asyncio.to_thread(
245
+ functools.partial(
246
+ _write_events,
247
+ self._mem,
248
+ self._region,
249
+ session.app_name,
250
+ session.user_id,
251
+ session.id,
252
+ session.events,
253
+ None,
254
+ replace_session=True,
255
+ )
256
+ )
257
+
258
+ async def search_memory(
259
+ self, *, app_name: str, user_id: str, query: str
260
+ ) -> SearchMemoryResponse:
261
+ return await asyncio.to_thread(
262
+ _search,
263
+ self._mem,
264
+ self._region,
265
+ app_name,
266
+ user_id,
267
+ query,
268
+ self._search_limit,
269
+ )
270
+
271
+ # ---- optional surface the reference refuses ---------------------------
272
+
273
+ async def add_events_to_memory(
274
+ self,
275
+ *,
276
+ app_name: str,
277
+ user_id: str,
278
+ events: Sequence[Event],
279
+ session_id: str | None = None,
280
+ custom_metadata: Mapping[str, object] | None = None,
281
+ ) -> None:
282
+ await asyncio.to_thread(
283
+ _write_events,
284
+ self._mem,
285
+ self._region,
286
+ app_name,
287
+ user_id,
288
+ session_id or _UNKNOWN_SESSION,
289
+ events,
290
+ custom_metadata,
291
+ )
292
+
293
+ async def add_memory(
294
+ self,
295
+ *,
296
+ app_name: str,
297
+ user_id: str,
298
+ memories: Sequence[MemoryEntry],
299
+ custom_metadata: Mapping[str, object] | None = None,
300
+ ) -> None:
301
+ """Write memories directly, without a session."""
302
+ await asyncio.to_thread(
303
+ _write_memories,
304
+ self._mem,
305
+ self._region,
306
+ app_name,
307
+ user_id,
308
+ memories,
309
+ custom_metadata,
310
+ )
311
+
312
+ # ---- beyond the protocol ----------------------------------------------
313
+
314
+ def forget_user(self, app_name: str, user_id: str) -> int:
315
+ """Destroy everything one user owns, returning the number erased."""
316
+ doomed = self._owned(app_name, user_id)
317
+ if not doomed:
318
+ return 0
319
+ return self._mem.forget(self._region, [h.id for h in doomed]).erased_count
320
+
321
+ def forget_session(self, app_name: str, user_id: str, session_id: str) -> int:
322
+ """The same for one session of one user."""
323
+ scope = _scope(app_name, user_id) | {"sid": session_id}
324
+ doomed = _page(self._mem, self._region, scope)
325
+ if not doomed:
326
+ return 0
327
+ return self._mem.forget(self._region, [h.id for h in doomed]).erased_count
328
+
329
+ def count(self, app_name: str, user_id: str) -> int:
330
+ return len(self._owned(app_name, user_id))
@@ -0,0 +1,380 @@
1
+ import pytest
2
+ from google.adk.events.event import Event
3
+ from google.adk.memory.base_memory_service import BaseMemoryService
4
+ from google.adk.memory.memory_entry import MemoryEntry
5
+ from google.adk.sessions.session import Session
6
+ from google.genai import types
7
+
8
+ from citadeldb_google_adk import CitadelMemoryService
9
+
10
+ APP = "app"
11
+ USER = "user-1"
12
+
13
+
14
+ @pytest.fixture(scope="module")
15
+ def svc(tmp_path_factory):
16
+ # Citadel takes an exclusive lock, so the whole module shares one handle.
17
+ path = tmp_path_factory.mktemp("adk") / "m.cdl"
18
+ return CitadelMemoryService(str(path), key="test-passphrase")
19
+
20
+
21
+ def ev(text, *, author="user", eid=None, ts=1.0):
22
+ return Event(
23
+ id=eid or f"e-{text[:8]}-{ts}",
24
+ author=author,
25
+ timestamp=ts,
26
+ content=types.Content(role=author, parts=[types.Part(text=text)]),
27
+ )
28
+
29
+
30
+ def sess(sid, events, *, app=APP, user=USER):
31
+ return Session(id=sid, app_name=app, user_id=user, events=events)
32
+
33
+
34
+ # ---- protocol conformance ------------------------------------------------
35
+
36
+
37
+ def test_is_a_memory_service(svc):
38
+ assert isinstance(svc, BaseMemoryService)
39
+
40
+
41
+ # ---- ingest --------------------------------------------------------------
42
+
43
+
44
+ async def test_a_session_becomes_searchable(svc):
45
+ await svc.add_session_to_memory(
46
+ sess("s1", [ev("the deployment failed because the disk was full")])
47
+ )
48
+ found = await svc.search_memory(app_name=APP, user_id=USER, query="disk was full")
49
+ assert found.memories
50
+ assert "disk was full" in found.memories[0].content.parts[0].text
51
+
52
+
53
+ async def test_search_is_semantic_not_word_overlap(svc):
54
+ """The reference service matches shared words; recall shares none here."""
55
+ await svc.add_session_to_memory(
56
+ sess("sem", [ev("the deployment failed because the disk was full")])
57
+ )
58
+ found = await svc.search_memory(
59
+ app_name=APP, user_id=USER, query="why did the release break?"
60
+ )
61
+ assert found.memories, "hybrid recall returned nothing"
62
+
63
+
64
+ async def test_contentless_events_are_skipped(svc):
65
+ empty = Event(id="empty", author="user", timestamp=1.0, content=None)
66
+ await svc.add_session_to_memory(sess("skip", [empty]))
67
+ assert svc.count(APP, USER) == svc.count(APP, USER) # no crash, nothing stored
68
+
69
+
70
+ async def test_re_adding_a_session_does_not_duplicate(svc):
71
+ """ADK says a session may be added many times, so ingest must be a delta."""
72
+ s = sess("dup", [ev("only once", eid="fixed-1")])
73
+ await svc.add_session_to_memory(s)
74
+ before = svc.count(APP, USER)
75
+ await svc.add_session_to_memory(s)
76
+ assert svc.count(APP, USER) == before
77
+
78
+
79
+ async def test_a_growing_session_adds_only_the_new_events(svc):
80
+ first = ev("turn one", eid="g-1")
81
+ s = sess("grow", [first])
82
+ await svc.add_session_to_memory(s)
83
+ before = svc.count(APP, USER)
84
+ s.events.append(ev("turn two", eid="g-2"))
85
+ await svc.add_session_to_memory(s)
86
+ assert svc.count(APP, USER) == before + 1
87
+
88
+
89
+ # ---- scope ---------------------------------------------------------------
90
+
91
+
92
+ async def test_users_do_not_see_each_other(svc):
93
+ await svc.add_session_to_memory(sess("u1", [ev("mine alone")], user="alice"))
94
+ await svc.add_session_to_memory(sess("u2", [ev("yours alone")], user="bob"))
95
+ alice = await svc.search_memory(app_name=APP, user_id="alice", query="alone")
96
+ texts = [m.content.parts[0].text for m in alice.memories]
97
+ assert "mine alone" in texts and "yours alone" not in texts
98
+
99
+
100
+ async def test_apps_do_not_see_each_other(svc):
101
+ await svc.add_session_to_memory(sess("a1", [ev("app one secret")], app="one"))
102
+ found = await svc.search_memory(app_name="two", user_id=USER, query="secret")
103
+ assert [m for m in found.memories if "app one" in m.content.parts[0].text] == []
104
+
105
+
106
+ # ---- projection ----------------------------------------------------------
107
+
108
+
109
+ async def test_memory_entry_carries_author_and_timestamp(svc):
110
+ await svc.add_session_to_memory(
111
+ sess("meta", [ev("stamped", author="assistant", eid="m-1", ts=1700000000.0)])
112
+ )
113
+ found = await svc.search_memory(app_name=APP, user_id=USER, query="stamped")
114
+ entry = next(m for m in found.memories if m.id == "m-1")
115
+ assert entry.author == "assistant"
116
+ assert entry.timestamp and entry.timestamp.startswith("20")
117
+
118
+
119
+ async def test_non_text_parts_round_trip_verbatim(svc):
120
+ """Content is a union owned by google.genai; nothing may normalise it."""
121
+ content = types.Content(
122
+ role="model",
123
+ parts=[
124
+ types.Part(text="looking that up"),
125
+ types.Part(
126
+ function_call=types.FunctionCall(name="lookup", args={"q": "berlin"})
127
+ ),
128
+ ],
129
+ )
130
+ e = Event(id="fc-1", author="model", timestamp=3.0, content=content)
131
+ await svc.add_session_to_memory(sess("fc", [e]))
132
+ found = await svc.search_memory(app_name=APP, user_id=USER, query="looking that up")
133
+ entry = next(m for m in found.memories if m.id == "fc-1")
134
+ call = entry.content.parts[1].function_call
135
+ assert call.name == "lookup"
136
+ assert call.args == {"q": "berlin"}
137
+
138
+
139
+ async def test_an_event_with_no_text_at_all_is_skipped(svc):
140
+ """Nothing to embed means recall could never return it."""
141
+ only_call = types.Content(
142
+ role="model",
143
+ parts=[types.Part(function_call=types.FunctionCall(name="noop", args={}))],
144
+ )
145
+ e = Event(id="silent", author="model", timestamp=4.0, content=only_call)
146
+ before = svc.count(APP, USER)
147
+ await svc.add_session_to_memory(sess("silent", [e]))
148
+ assert svc.count(APP, USER) == before
149
+
150
+
151
+ async def test_search_for_an_unknown_user_is_empty(svc):
152
+ found = await svc.search_memory(app_name=APP, user_id="nobody", query="anything")
153
+ assert found.memories == []
154
+
155
+
156
+ async def test_search_on_an_empty_service_is_empty(tmp_path):
157
+ fresh = CitadelMemoryService(str(tmp_path / "empty.cdl"), key="pw")
158
+ found = await fresh.search_memory(app_name=APP, user_id=USER, query="anything")
159
+ assert found.memories == []
160
+
161
+
162
+ async def test_the_event_loop_is_not_blocked(svc):
163
+ """ADK's Runner is async, so the sync bindings must run off the loop."""
164
+ import asyncio
165
+
166
+ ticks = 0
167
+
168
+ async def tick():
169
+ nonlocal ticks
170
+ while True:
171
+ ticks += 1
172
+ await asyncio.sleep(0)
173
+
174
+ ticker = asyncio.create_task(tick())
175
+ await asyncio.sleep(0)
176
+ await svc.add_session_to_memory(
177
+ sess("loop", [ev(f"loop turn {n}", eid=f"l-{n}") for n in range(40)])
178
+ )
179
+ await svc.search_memory(app_name=APP, user_id=USER, query="loop turn")
180
+ ticker.cancel()
181
+ assert ticks > 1, "the loop made no progress during a memory call"
182
+
183
+
184
+ async def test_multi_part_content_round_trips(svc):
185
+ """Content is stored whole, so a memory hands back what ADK gave us."""
186
+ content = types.Content(
187
+ role="user",
188
+ parts=[types.Part(text="first part"), types.Part(text="second part")],
189
+ )
190
+ e = Event(id="multi", author="user", timestamp=2.0, content=content)
191
+ await svc.add_session_to_memory(sess("multi", [e]))
192
+ found = await svc.search_memory(app_name=APP, user_id=USER, query="second part")
193
+ entry = next(m for m in found.memories if m.id == "multi")
194
+ assert [p.text for p in entry.content.parts] == ["first part", "second part"]
195
+
196
+
197
+ # ---- the surface the reference refuses -----------------------------------
198
+
199
+
200
+ async def test_events_can_be_added_without_a_session(svc):
201
+ await svc.add_events_to_memory(
202
+ app_name=APP, user_id=USER, events=[ev("delta only", eid="d-1")],
203
+ session_id="delta", custom_metadata={"source": "webhook"},
204
+ )
205
+ found = await svc.search_memory(app_name=APP, user_id=USER, query="delta only")
206
+ entry = next(m for m in found.memories if m.id == "d-1")
207
+ assert entry.custom_metadata["source"] == "webhook"
208
+
209
+
210
+ async def test_re_adding_a_memory_id_replaces_it(svc):
211
+ """An id makes a second write a replacement, not a second copy."""
212
+ def entry(text):
213
+ return MemoryEntry(
214
+ content=types.Content(role="user", parts=[types.Part(text=text)]),
215
+ id="dup-1", author="user",
216
+ )
217
+
218
+ await svc.add_memory(app_name=APP, user_id="dupuser", memories=[entry("first")])
219
+ await svc.add_memory(app_name=APP, user_id="dupuser", memories=[entry("second")])
220
+ assert svc.count(APP, "dupuser") == 1
221
+ found = await svc.search_memory(app_name=APP, user_id="dupuser", query="first second")
222
+ assert [m.content.parts[0].text for m in found.memories] == ["second"]
223
+
224
+
225
+ async def test_duplicate_ids_within_one_call_keep_the_last(svc):
226
+ def entry(text):
227
+ return MemoryEntry(
228
+ content=types.Content(role="user", parts=[types.Part(text=text)]),
229
+ id="batch-dup", author="user",
230
+ )
231
+
232
+ await svc.add_memory(
233
+ app_name=APP, user_id="batchuser", memories=[entry("one"), entry("two")]
234
+ )
235
+ assert svc.count(APP, "batchuser") == 1
236
+
237
+
238
+ async def test_memories_without_an_id_are_always_new(svc):
239
+ """Nothing identifies them, so they cannot replace anything."""
240
+ def entry():
241
+ return MemoryEntry(
242
+ content=types.Content(role="user", parts=[types.Part(text="anonymous")]),
243
+ author="user",
244
+ )
245
+
246
+ await svc.add_memory(app_name=APP, user_id="anon", memories=[entry()])
247
+ await svc.add_memory(app_name=APP, user_id="anon", memories=[entry()])
248
+ assert svc.count(APP, "anon") == 2
249
+
250
+
251
+ async def test_memories_can_be_written_directly(svc):
252
+ """The reference service raises NotImplementedError here."""
253
+ m = MemoryEntry(
254
+ content=types.Content(role="user", parts=[types.Part(text="written directly")]),
255
+ id="direct-1",
256
+ author="user",
257
+ )
258
+ await svc.add_memory(app_name=APP, user_id=USER, memories=[m])
259
+ found = await svc.search_memory(app_name=APP, user_id=USER, query="written directly")
260
+ assert any(x.id == "direct-1" for x in found.memories)
261
+
262
+
263
+ # ---- erasure -------------------------------------------------------------
264
+
265
+
266
+ async def test_forgetting_a_session_leaves_the_rest(svc):
267
+ await svc.add_session_to_memory(sess("keep", [ev("keep me", eid="k-1")], user="carol"))
268
+ await svc.add_session_to_memory(sess("drop", [ev("drop me", eid="k-2")], user="carol"))
269
+ assert svc.forget_session(APP, "carol", "drop") == 1
270
+ left = await svc.search_memory(app_name=APP, user_id="carol", query="me")
271
+ assert [m.id for m in left.memories] == ["k-1"]
272
+
273
+
274
+ async def test_forgetting_a_user_erases_everything_they_own(svc):
275
+ await svc.add_session_to_memory(sess("g1", [ev("gone one")], user="dave"))
276
+ await svc.add_session_to_memory(sess("g2", [ev("gone two")], user="dave"))
277
+ assert svc.forget_user(APP, "dave") == 2
278
+ assert svc.count(APP, "dave") == 0
279
+
280
+
281
+ def test_forgetting_nothing_is_zero_not_an_error(svc):
282
+ assert svc.forget_user(APP, "nobody") == 0
283
+
284
+
285
+ async def test_it_survives_a_reopen(tmp_path):
286
+ """A region's embedder must survive the reopen, not just the events."""
287
+ import gc
288
+
289
+ p = str(tmp_path / "reopen.cdl")
290
+ first = CitadelMemoryService(p, key="pw")
291
+ await first.add_session_to_memory(
292
+ sess("s", [ev("the deployment failed because the disk was full")], user="ru")
293
+ )
294
+ del first
295
+ gc.collect()
296
+
297
+ again = CitadelMemoryService(p, key="pw")
298
+ assert again.count(APP, "ru") == 1
299
+ found = await again.search_memory(
300
+ app_name=APP, user_id="ru", query="why did the release break?"
301
+ )
302
+ assert found.memories, "recall did not survive the reopen"
303
+
304
+
305
+ def test_a_passphrase_is_required(tmp_path):
306
+ with pytest.raises(ValueError, match="passphrase"):
307
+ CitadelMemoryService(str(tmp_path / "nokey.cdl"), key="")
308
+
309
+
310
+ async def test_concurrent_ingestion_all_lands(tmp_path):
311
+ """Many in-flight sessions means many threads against one engine."""
312
+ import asyncio
313
+
314
+ svc = CitadelMemoryService(str(tmp_path / "conc.cdl"), key="pw")
315
+ await asyncio.gather(*(
316
+ svc.add_session_to_memory(sess(f"s{i}", [ev(f"event {i}")], user="cu"))
317
+ for i in range(40)
318
+ ))
319
+ assert svc.count(APP, "cu") == 40
320
+
321
+
322
+ async def test_re_adding_a_session_drops_events_removed_from_it(tmp_path):
323
+ """add_session_to_memory SETS the session's events; add_events_to_memory is
324
+ the additive one. Merging on both leaves a retracted turn searchable."""
325
+ svc = CitadelMemoryService(str(tmp_path / "retract.cdl"), key="pw")
326
+ keep, retract = ev("keep this"), ev("retract this")
327
+ await svc.add_session_to_memory(sess("s", [keep, retract], user="ru"))
328
+ assert svc.count(APP, "ru") == 2
329
+
330
+ await svc.add_session_to_memory(sess("s", [keep], user="ru"))
331
+ assert svc.count(APP, "ru") == 1
332
+ found = await svc.search_memory(app_name=APP, user_id="ru", query="retract this")
333
+ texts = [p.text for m in found.memories for p in m.content.parts]
334
+ assert "retract this" not in texts
335
+
336
+ # The additive path still adds rather than replacing.
337
+ await svc.add_events_to_memory(
338
+ app_name=APP, user_id="ru", session_id="s", events=[ev("added later")]
339
+ )
340
+ assert svc.count(APP, "ru") == 2
341
+
342
+
343
+ async def test_concurrent_ingestion_of_one_session_stores_it_once(tmp_path):
344
+ """Re-ingesting a session converges; racing ingests of it must too, or the
345
+ same turn comes back from search several times."""
346
+ import asyncio
347
+
348
+ svc = CitadelMemoryService(str(tmp_path / "onesess.cdl"), key="pw")
349
+ session = sess("s", [ev("the only turn")], user="ru")
350
+ await asyncio.gather(*(svc.add_session_to_memory(session) for _ in range(8)))
351
+ assert svc.count(APP, "ru") == 1
352
+
353
+ found = await svc.search_memory(app_name=APP, user_id="ru", query="the only turn")
354
+ assert len(found.memories) == 1
355
+
356
+
357
+ async def test_a_wrong_passphrase_cannot_reopen(tmp_path):
358
+ """Pins the encryption claim rather than inferring it from a reopen."""
359
+ import gc
360
+
361
+ import citadeldb
362
+
363
+ p = str(tmp_path / "enc.cdl")
364
+ first = CitadelMemoryService(p, key="right")
365
+ await first.add_session_to_memory(sess("s", [ev("secret")], user="ru"))
366
+ del first
367
+ gc.collect()
368
+
369
+ with pytest.raises(citadeldb.EncryptionError):
370
+ CitadelMemoryService(p, key="wrong")
371
+
372
+
373
+ async def test_shares_one_database_with_another_adapter(tmp_path):
374
+ """Two Citadel-backed services on one file must not fight over the lock."""
375
+ path = str(tmp_path / "shared.cdl")
376
+ a = CitadelMemoryService(path, key="pw", region="memory")
377
+ b = CitadelMemoryService(path, key="pw", region="other")
378
+ await a.add_session_to_memory(sess("s", [ev("in region a")]))
379
+ assert a.count(APP, USER) == 1
380
+ assert b.count(APP, USER) == 0, "regions must stay separate"