citadeldb-strands-agents 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_strands_agents-2.0.0/.gitignore +18 -0
- citadeldb_strands_agents-2.0.0/PKG-INFO +87 -0
- citadeldb_strands_agents-2.0.0/README.md +66 -0
- citadeldb_strands_agents-2.0.0/pyproject.toml +38 -0
- citadeldb_strands_agents-2.0.0/src/citadeldb_strands_agents/__init__.py +13 -0
- citadeldb_strands_agents-2.0.0/src/citadeldb_strands_agents/session.py +332 -0
- citadeldb_strands_agents-2.0.0/tests/test_multi_agent.py +195 -0
- citadeldb_strands_agents-2.0.0/tests/test_session.py +362 -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,87 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: citadeldb-strands-agents
|
|
3
|
+
Version: 2.0.0
|
|
4
|
+
Summary: Strands Agents session storage 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: agents,encryption,memory,session,strands,strands-agents
|
|
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: strands-agents<2,>=1.0
|
|
18
|
+
Provides-Extra: test
|
|
19
|
+
Requires-Dist: pytest>=8; extra == 'test'
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# citadeldb-strands-agents
|
|
23
|
+
|
|
24
|
+
A [Strands Agents](https://strandsagents.com/) session manager backed by
|
|
25
|
+
[Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and
|
|
26
|
+
redaction that removes the redacted content rather than annotating it.
|
|
27
|
+
|
|
28
|
+
```
|
|
29
|
+
pip install citadeldb-strands-agents
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
from strands import Agent
|
|
34
|
+
from citadeldb_strands_agents import CitadelSessionManager
|
|
35
|
+
|
|
36
|
+
sessions = CitadelSessionManager("user-123", "sessions.cdl", key="your-passphrase")
|
|
37
|
+
agent = Agent(session_manager=sessions)
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Built the way the shipped managers are: it implements `SessionRepository` and inherits
|
|
41
|
+
`RepositorySessionManager`, exactly as `FileSessionManager` and `S3SessionManager` do.
|
|
42
|
+
|
|
43
|
+
## Redaction actually redacts
|
|
44
|
+
|
|
45
|
+
Strands redacts by setting a `redact_message` field beside the original and updating the
|
|
46
|
+
record:
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
# from strands RepositorySessionManager.redact_latest_message
|
|
50
|
+
latest_agent_message.redact_message = redact_message
|
|
51
|
+
return self.session_repository.update_message(...)
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
The original `message` stays in the record. Strands reads a redacted message through
|
|
55
|
+
`to_message()`, which returns the redaction, so the original is never read again. This
|
|
56
|
+
manager therefore does not keep it: the superseded record's key is destroyed and the
|
|
57
|
+
replacement carries the redaction in place of the original.
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
after = sessions.read_message("user-123", agent.agent_id, 0)
|
|
61
|
+
after.to_message() # {'content': [{'text': '[REDACTED]'}], 'role': 'user'}
|
|
62
|
+
after.message # the redaction, not the original
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
An ordinary update is unaffected: only a message carrying a redaction drops its original.
|
|
66
|
+
|
|
67
|
+
## Erasure
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
sessions.forget_session("user-123") # returns the number of records erased
|
|
71
|
+
```
|
|
72
|
+
|
|
73
|
+
Strands has no delete in its repository protocol. This destroys the session, its agents
|
|
74
|
+
and every message key.
|
|
75
|
+
|
|
76
|
+
## Notes
|
|
77
|
+
|
|
78
|
+
Messages are ordered by their conversation index rather than write time, so editing an
|
|
79
|
+
earlier turn does not move it to the end of the history.
|
|
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 many sessions and other Citadel adapters can share
|
|
83
|
+
one database; construct them on the same thread.
|
|
84
|
+
|
|
85
|
+
## License
|
|
86
|
+
|
|
87
|
+
Apache-2.0
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
# citadeldb-strands-agents
|
|
2
|
+
|
|
3
|
+
A [Strands Agents](https://strandsagents.com/) session manager backed by
|
|
4
|
+
[Citadel](https://citadeldb.dev). Encrypted at rest, embedded in your process, and
|
|
5
|
+
redaction that removes the redacted content rather than annotating it.
|
|
6
|
+
|
|
7
|
+
```
|
|
8
|
+
pip install citadeldb-strands-agents
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
```python
|
|
12
|
+
from strands import Agent
|
|
13
|
+
from citadeldb_strands_agents import CitadelSessionManager
|
|
14
|
+
|
|
15
|
+
sessions = CitadelSessionManager("user-123", "sessions.cdl", key="your-passphrase")
|
|
16
|
+
agent = Agent(session_manager=sessions)
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
Built the way the shipped managers are: it implements `SessionRepository` and inherits
|
|
20
|
+
`RepositorySessionManager`, exactly as `FileSessionManager` and `S3SessionManager` do.
|
|
21
|
+
|
|
22
|
+
## Redaction actually redacts
|
|
23
|
+
|
|
24
|
+
Strands redacts by setting a `redact_message` field beside the original and updating the
|
|
25
|
+
record:
|
|
26
|
+
|
|
27
|
+
```python
|
|
28
|
+
# from strands RepositorySessionManager.redact_latest_message
|
|
29
|
+
latest_agent_message.redact_message = redact_message
|
|
30
|
+
return self.session_repository.update_message(...)
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
The original `message` stays in the record. Strands reads a redacted message through
|
|
34
|
+
`to_message()`, which returns the redaction, so the original is never read again. This
|
|
35
|
+
manager therefore does not keep it: the superseded record's key is destroyed and the
|
|
36
|
+
replacement carries the redaction in place of the original.
|
|
37
|
+
|
|
38
|
+
```python
|
|
39
|
+
after = sessions.read_message("user-123", agent.agent_id, 0)
|
|
40
|
+
after.to_message() # {'content': [{'text': '[REDACTED]'}], 'role': 'user'}
|
|
41
|
+
after.message # the redaction, not the original
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
An ordinary update is unaffected: only a message carrying a redaction drops its original.
|
|
45
|
+
|
|
46
|
+
## Erasure
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
sessions.forget_session("user-123") # returns the number of records erased
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Strands has no delete in its repository protocol. This destroys the session, its agents
|
|
53
|
+
and every message key.
|
|
54
|
+
|
|
55
|
+
## Notes
|
|
56
|
+
|
|
57
|
+
Messages are ordered by their conversation index rather than write time, so editing an
|
|
58
|
+
earlier turn does not move it to the end of the history.
|
|
59
|
+
|
|
60
|
+
Citadel is embedded and one process owns the file. A path already open on this thread,
|
|
61
|
+
under the same passphrase, is shared, so many sessions and other Citadel adapters can share
|
|
62
|
+
one database; construct them on the same thread.
|
|
63
|
+
|
|
64
|
+
## License
|
|
65
|
+
|
|
66
|
+
Apache-2.0
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling", "hatch-vcs"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
# Bare "strands" is an unrelated package on PyPI.
|
|
7
|
+
name = "citadeldb-strands-agents"
|
|
8
|
+
dynamic = ["version"]
|
|
9
|
+
description = "Strands Agents session storage backed by Citadel: encrypted at rest, with deletes that destroy the key"
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
license = "Apache-2.0"
|
|
13
|
+
authors = [{ name = "Yuriy Peysakhov" }]
|
|
14
|
+
keywords = ["strands", "strands-agents", "agents", "session", "memory", "encryption"]
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Development Status :: 4 - Beta",
|
|
17
|
+
"Intended Audience :: Developers",
|
|
18
|
+
"Programming Language :: Python :: 3",
|
|
19
|
+
"Topic :: Database",
|
|
20
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
21
|
+
]
|
|
22
|
+
# A record replaces under its own id, which needs the keyed write added in 2.0.
|
|
23
|
+
dependencies = ["citadeldb>=2.0,<3", "strands-agents>=1.0,<2"]
|
|
24
|
+
|
|
25
|
+
[project.optional-dependencies]
|
|
26
|
+
test = ["pytest>=8"]
|
|
27
|
+
|
|
28
|
+
[project.urls]
|
|
29
|
+
Homepage = "https://citadeldb.dev"
|
|
30
|
+
Repository = "https://github.com/yp3y5akh0v/citadel"
|
|
31
|
+
|
|
32
|
+
# The version comes from the release tag, so there is nothing to bump.
|
|
33
|
+
[tool.hatch.version]
|
|
34
|
+
source = "vcs"
|
|
35
|
+
raw-options = { root = "../..", tag_regex = '^v(?P<version>\d+\.\d+\.\d+)$' }
|
|
36
|
+
|
|
37
|
+
[tool.hatch.build.targets.wheel]
|
|
38
|
+
packages = ["src/citadeldb_strands_agents"]
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
"""Strands Agents session storage backed by Citadel, encrypted at rest."""
|
|
2
|
+
|
|
3
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
4
|
+
|
|
5
|
+
from .session import CitadelSessionManager
|
|
6
|
+
|
|
7
|
+
__all__ = ["CitadelSessionManager", "__version__"]
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
try:
|
|
11
|
+
__version__ = version("citadeldb-strands-agents")
|
|
12
|
+
except PackageNotFoundError: # running from a source tree, never installed
|
|
13
|
+
__version__ = "0+unknown"
|
|
@@ -0,0 +1,332 @@
|
|
|
1
|
+
"""Strands Agents session store on an encrypted Citadel region."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
import citadeldb
|
|
7
|
+
from strands.session.repository_session_manager import RepositorySessionManager
|
|
8
|
+
from strands.session.session_repository import SessionRepository
|
|
9
|
+
from strands.types.exceptions import SessionException
|
|
10
|
+
from strands.types.session import (
|
|
11
|
+
Session,
|
|
12
|
+
SessionAgent,
|
|
13
|
+
SessionMessage,
|
|
14
|
+
encode_bytes_values,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
SESSION_KIND = "session"
|
|
18
|
+
AGENT_KIND = "agent"
|
|
19
|
+
MESSAGE_KIND = "message"
|
|
20
|
+
MULTI_AGENT_KIND = "multi_agent"
|
|
21
|
+
DEFAULT_PATH = "strands_sessions.cdl"
|
|
22
|
+
DEFAULT_REGION = "strands_sessions"
|
|
23
|
+
PAGE = 10_000
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _page(mem: Any, region: str, kind: str, criterion: dict[str, Any]) -> list[Any]:
|
|
27
|
+
"""Page to the end: one fetch is bounded, and a partial erase must not look whole."""
|
|
28
|
+
out: list[Any] = []
|
|
29
|
+
after = None
|
|
30
|
+
while True:
|
|
31
|
+
got = mem.fetch(region, kind, payload_filter=criterion, limit=PAGE, after_id=after)
|
|
32
|
+
out.extend(got)
|
|
33
|
+
if len(got) < PAGE:
|
|
34
|
+
return out
|
|
35
|
+
after = got[-1].id
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
# Memory holds no database reference, so no handle is pinned to a thread.
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def _one(mem: Any, region: str, kind: str, criterion: dict[str, Any]):
|
|
42
|
+
hits = mem.fetch(region, kind, payload_filter=criterion, limit=1)
|
|
43
|
+
return hits[0] if hits else None
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _key(*parts: Any) -> str:
|
|
47
|
+
"""One engine key from several names, length-prefixed so that no two name
|
|
48
|
+
tuples can spell the same key."""
|
|
49
|
+
return "".join(f"{len(s)}\x1f{s}" for s in (str(p) for p in parts))
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _put(mem: Any, region: str, atom: dict, key: str) -> None:
|
|
53
|
+
"""Store `atom` as the sole occupant of `key`. Every reference repository
|
|
54
|
+
keys its records - a file per message, an object per key - so a second write
|
|
55
|
+
replaces the first rather than standing beside it, and two request handlers
|
|
56
|
+
sharing a store cannot each add a copy."""
|
|
57
|
+
mem.remember_replacing_keyed(region, atom, key)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def _persisted(session_message: SessionMessage) -> dict:
|
|
61
|
+
"""The record to store; a redaction replaces the original message."""
|
|
62
|
+
data = session_message.to_dict()
|
|
63
|
+
if session_message.redact_message is not None:
|
|
64
|
+
# to_dict() runs the bytes encoder over the whole record, so the
|
|
65
|
+
# replacement has to go through it too. Assigning it raw sends live
|
|
66
|
+
# Python bytes to the engine whenever a redaction carries an image or
|
|
67
|
+
# a document block.
|
|
68
|
+
data["message"] = encode_bytes_values(session_message.redact_message)
|
|
69
|
+
return data
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def _text_of(session_message: SessionMessage) -> str:
|
|
73
|
+
"""An atom needs text to embed; the index identifies it otherwise."""
|
|
74
|
+
content = session_message.redact_message or session_message.message
|
|
75
|
+
parts = content.get("content") or [] if isinstance(content, dict) else []
|
|
76
|
+
text = " ".join(
|
|
77
|
+
p["text"] for p in parts if isinstance(p, dict) and isinstance(p.get("text"), str)
|
|
78
|
+
)
|
|
79
|
+
return text or f"message {session_message.message_id}"
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class CitadelSessionManager(RepositorySessionManager, SessionRepository):
|
|
83
|
+
"""A Strands SessionRepository backed by one encrypted Citadel region."""
|
|
84
|
+
|
|
85
|
+
def __init__(
|
|
86
|
+
self,
|
|
87
|
+
session_id: str,
|
|
88
|
+
path: str = DEFAULT_PATH,
|
|
89
|
+
key: str = "",
|
|
90
|
+
*,
|
|
91
|
+
region: str = DEFAULT_REGION,
|
|
92
|
+
embedder: Any | None = None,
|
|
93
|
+
**kwargs: Any,
|
|
94
|
+
) -> None:
|
|
95
|
+
if not key:
|
|
96
|
+
raise ValueError("a passphrase is required: transcripts are the payload")
|
|
97
|
+
try:
|
|
98
|
+
self._db = citadeldb.connect(path, key=key, region_keys=True)
|
|
99
|
+
except citadeldb.OperationalError as e:
|
|
100
|
+
if "locked" not in str(e):
|
|
101
|
+
raise
|
|
102
|
+
raise RuntimeError(
|
|
103
|
+
f"{path} is open in another process. Citadel is embedded, so one "
|
|
104
|
+
f"process owns the file."
|
|
105
|
+
) from e
|
|
106
|
+
self._mem = self._db.memory()
|
|
107
|
+
self._region = region
|
|
108
|
+
# Idempotent for a region of the same width, so a dim clash raises here.
|
|
109
|
+
self._mem.create_encrypted_region(
|
|
110
|
+
region, embedder or citadeldb.MockEmbedder(dim=64)
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
super().__init__(session_id=session_id, session_repository=self, **kwargs)
|
|
114
|
+
|
|
115
|
+
# ---- sessions ---------------------------------------------------------
|
|
116
|
+
|
|
117
|
+
def create_session(self, session: Session, **kwargs: Any) -> Session:
|
|
118
|
+
if _one(self._mem, self._region, SESSION_KIND, {"sid": session.session_id}):
|
|
119
|
+
raise SessionException(f"Session {session.session_id} already exists")
|
|
120
|
+
_put(
|
|
121
|
+
self._mem,
|
|
122
|
+
self._region,
|
|
123
|
+
{
|
|
124
|
+
"kind": SESSION_KIND,
|
|
125
|
+
"text": session.session_id,
|
|
126
|
+
"payload": {"sid": session.session_id, "session": session.to_dict()},
|
|
127
|
+
},
|
|
128
|
+
_key(session.session_id),
|
|
129
|
+
)
|
|
130
|
+
return session
|
|
131
|
+
|
|
132
|
+
def read_session(self, session_id: str, **kwargs: Any) -> Session | None:
|
|
133
|
+
hit = _one(self._mem, self._region, SESSION_KIND, {"sid": session_id})
|
|
134
|
+
return Session.from_dict(hit.payload["session"]) if hit else None
|
|
135
|
+
|
|
136
|
+
# ---- agents -----------------------------------------------------------
|
|
137
|
+
|
|
138
|
+
def create_agent(
|
|
139
|
+
self, session_id: str, session_agent: SessionAgent, **kwargs: Any
|
|
140
|
+
) -> None:
|
|
141
|
+
_put(
|
|
142
|
+
self._mem,
|
|
143
|
+
self._region,
|
|
144
|
+
{
|
|
145
|
+
"kind": AGENT_KIND,
|
|
146
|
+
"text": session_agent.agent_id,
|
|
147
|
+
"payload": {
|
|
148
|
+
"sid": session_id,
|
|
149
|
+
"aid": session_agent.agent_id,
|
|
150
|
+
"agent": session_agent.to_dict(),
|
|
151
|
+
},
|
|
152
|
+
},
|
|
153
|
+
_key(session_id, session_agent.agent_id),
|
|
154
|
+
)
|
|
155
|
+
|
|
156
|
+
def read_agent(
|
|
157
|
+
self, session_id: str, agent_id: str, **kwargs: Any
|
|
158
|
+
) -> SessionAgent | None:
|
|
159
|
+
hit = _one(
|
|
160
|
+
self._mem, self._region, AGENT_KIND, {"sid": session_id, "aid": agent_id}
|
|
161
|
+
)
|
|
162
|
+
return SessionAgent.from_dict(hit.payload["agent"]) if hit else None
|
|
163
|
+
|
|
164
|
+
def update_agent(
|
|
165
|
+
self, session_id: str, session_agent: SessionAgent, **kwargs: Any
|
|
166
|
+
) -> None:
|
|
167
|
+
agent_id = session_agent.agent_id
|
|
168
|
+
previous = self.read_agent(session_id, agent_id)
|
|
169
|
+
if previous is None:
|
|
170
|
+
raise SessionException(
|
|
171
|
+
f"Agent {agent_id} in session {session_id} does not exist"
|
|
172
|
+
)
|
|
173
|
+
# The reference carries the original creation time forward on update.
|
|
174
|
+
session_agent.created_at = previous.created_at
|
|
175
|
+
_put(
|
|
176
|
+
self._mem,
|
|
177
|
+
self._region,
|
|
178
|
+
{
|
|
179
|
+
"kind": AGENT_KIND,
|
|
180
|
+
"text": agent_id,
|
|
181
|
+
"payload": {
|
|
182
|
+
"sid": session_id,
|
|
183
|
+
"aid": agent_id,
|
|
184
|
+
"agent": session_agent.to_dict(),
|
|
185
|
+
},
|
|
186
|
+
},
|
|
187
|
+
_key(session_id, agent_id),
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
# ---- messages ---------------------------------------------------------
|
|
191
|
+
|
|
192
|
+
def create_message(
|
|
193
|
+
self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any
|
|
194
|
+
) -> None:
|
|
195
|
+
_put(
|
|
196
|
+
self._mem,
|
|
197
|
+
self._region,
|
|
198
|
+
{
|
|
199
|
+
"kind": MESSAGE_KIND,
|
|
200
|
+
"text": _text_of(session_message),
|
|
201
|
+
"payload": {
|
|
202
|
+
"sid": session_id,
|
|
203
|
+
"aid": agent_id,
|
|
204
|
+
"mid": session_message.message_id,
|
|
205
|
+
"message": _persisted(session_message),
|
|
206
|
+
},
|
|
207
|
+
},
|
|
208
|
+
_key(session_id, agent_id, session_message.message_id),
|
|
209
|
+
)
|
|
210
|
+
|
|
211
|
+
def read_message(
|
|
212
|
+
self, session_id: str, agent_id: str, message_id: int, **kwargs: Any
|
|
213
|
+
) -> SessionMessage | None:
|
|
214
|
+
hit = _one(
|
|
215
|
+
self._mem, self._region, MESSAGE_KIND,
|
|
216
|
+
{"sid": session_id, "aid": agent_id, "mid": message_id},
|
|
217
|
+
)
|
|
218
|
+
return SessionMessage.from_dict(hit.payload["message"]) if hit else None
|
|
219
|
+
|
|
220
|
+
def update_message(
|
|
221
|
+
self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any
|
|
222
|
+
) -> None:
|
|
223
|
+
message_id = session_message.message_id
|
|
224
|
+
previous = self.read_message(session_id, agent_id, message_id)
|
|
225
|
+
if previous is None:
|
|
226
|
+
raise SessionException(f"Message {message_id} does not exist")
|
|
227
|
+
session_message.created_at = previous.created_at
|
|
228
|
+
# A redaction lands here, so the superseded record's key is destroyed.
|
|
229
|
+
_put(
|
|
230
|
+
self._mem,
|
|
231
|
+
self._region,
|
|
232
|
+
{
|
|
233
|
+
"kind": MESSAGE_KIND,
|
|
234
|
+
"text": _text_of(session_message),
|
|
235
|
+
"payload": {
|
|
236
|
+
"sid": session_id,
|
|
237
|
+
"aid": agent_id,
|
|
238
|
+
"mid": message_id,
|
|
239
|
+
"message": _persisted(session_message),
|
|
240
|
+
},
|
|
241
|
+
},
|
|
242
|
+
_key(session_id, agent_id, message_id),
|
|
243
|
+
)
|
|
244
|
+
|
|
245
|
+
def list_messages(
|
|
246
|
+
self,
|
|
247
|
+
session_id: str,
|
|
248
|
+
agent_id: str,
|
|
249
|
+
limit: int | None = None,
|
|
250
|
+
offset: int = 0,
|
|
251
|
+
**kwargs: Any,
|
|
252
|
+
) -> list[SessionMessage]:
|
|
253
|
+
hits = _page(self._mem, self._region, MESSAGE_KIND,
|
|
254
|
+
{"sid": session_id, "aid": agent_id})
|
|
255
|
+
# The reference gates on the message container, which create_message
|
|
256
|
+
# makes: an agent with stored messages lists them whether or not an
|
|
257
|
+
# agent record was written first. Refusing on the record alone would
|
|
258
|
+
# hide writes this repository accepted.
|
|
259
|
+
if not hits and self.read_agent(session_id, agent_id) is None:
|
|
260
|
+
raise SessionException(
|
|
261
|
+
f"Messages missing from agent: {agent_id} in session {session_id}"
|
|
262
|
+
)
|
|
263
|
+
# Atom order follows write time and diverges once a message is updated.
|
|
264
|
+
ordered = sorted(hits, key=lambda h: h.payload["mid"])
|
|
265
|
+
window = ordered[offset : offset + limit] if limit is not None else ordered[offset:]
|
|
266
|
+
return [SessionMessage.from_dict(h.payload["message"]) for h in window]
|
|
267
|
+
|
|
268
|
+
# ---- multi-agent state ------------------------------------------------
|
|
269
|
+
# Base class raises NotImplementedError; a Swarm fails without these.
|
|
270
|
+
|
|
271
|
+
def create_multi_agent(
|
|
272
|
+
self, session_id: str, multi_agent: Any, **kwargs: Any
|
|
273
|
+
) -> None:
|
|
274
|
+
_put(
|
|
275
|
+
self._mem,
|
|
276
|
+
self._region,
|
|
277
|
+
{
|
|
278
|
+
"kind": MULTI_AGENT_KIND,
|
|
279
|
+
"text": multi_agent.id,
|
|
280
|
+
"payload": {
|
|
281
|
+
"sid": session_id,
|
|
282
|
+
"maid": multi_agent.id,
|
|
283
|
+
"state": multi_agent.serialize_state(),
|
|
284
|
+
},
|
|
285
|
+
},
|
|
286
|
+
_key(session_id, multi_agent.id),
|
|
287
|
+
)
|
|
288
|
+
|
|
289
|
+
def read_multi_agent(
|
|
290
|
+
self, session_id: str, multi_agent_id: str, **kwargs: Any
|
|
291
|
+
) -> dict[str, Any] | None:
|
|
292
|
+
hit = _one(
|
|
293
|
+
self._mem, self._region, MULTI_AGENT_KIND,
|
|
294
|
+
{"sid": session_id, "maid": multi_agent_id},
|
|
295
|
+
)
|
|
296
|
+
return hit.payload["state"] if hit else None
|
|
297
|
+
|
|
298
|
+
def update_multi_agent(
|
|
299
|
+
self, session_id: str, multi_agent: Any, **kwargs: Any
|
|
300
|
+
) -> None:
|
|
301
|
+
if self.read_multi_agent(session_id, multi_agent.id) is None:
|
|
302
|
+
raise SessionException(
|
|
303
|
+
f"MultiAgent state {multi_agent.id} in session {session_id} "
|
|
304
|
+
f"does not exist"
|
|
305
|
+
)
|
|
306
|
+
_put(
|
|
307
|
+
self._mem,
|
|
308
|
+
self._region,
|
|
309
|
+
{
|
|
310
|
+
"kind": MULTI_AGENT_KIND,
|
|
311
|
+
"text": multi_agent.id,
|
|
312
|
+
"payload": {
|
|
313
|
+
"sid": session_id,
|
|
314
|
+
"maid": multi_agent.id,
|
|
315
|
+
"state": multi_agent.serialize_state(),
|
|
316
|
+
},
|
|
317
|
+
},
|
|
318
|
+
_key(session_id, multi_agent.id),
|
|
319
|
+
)
|
|
320
|
+
|
|
321
|
+
# ---- beyond the repository --------------------------------------------
|
|
322
|
+
|
|
323
|
+
def forget_session(self, session_id: str) -> int:
|
|
324
|
+
"""Destroy a session's records, returning the number erased."""
|
|
325
|
+
doomed = [
|
|
326
|
+
h
|
|
327
|
+
for kind in (SESSION_KIND, AGENT_KIND, MESSAGE_KIND, MULTI_AGENT_KIND)
|
|
328
|
+
for h in _page(self._mem, self._region, kind, {"sid": session_id})
|
|
329
|
+
]
|
|
330
|
+
if not doomed:
|
|
331
|
+
return 0
|
|
332
|
+
return self._mem.forget(self._region, [h.id for h in doomed]).erased_count
|
|
@@ -0,0 +1,195 @@
|
|
|
1
|
+
"""Multi-agent state and the manager paths the framework drives."""
|
|
2
|
+
import os
|
|
3
|
+
import tempfile
|
|
4
|
+
|
|
5
|
+
import pytest
|
|
6
|
+
from strands.session.file_session_manager import FileSessionManager
|
|
7
|
+
from strands.types.exceptions import SessionException
|
|
8
|
+
from strands.types.session import SessionAgent, SessionMessage
|
|
9
|
+
|
|
10
|
+
from citadeldb_strands_agents import CitadelSessionManager
|
|
11
|
+
|
|
12
|
+
SID, AID = "s1", "a1"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def path(name="s.cdl"):
|
|
16
|
+
return os.path.join(tempfile.mkdtemp(), name)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def mgr(p=None, session_id=SID):
|
|
20
|
+
return CitadelSessionManager(session_id, p or path(), "pw")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def ref(session_id=SID):
|
|
24
|
+
return FileSessionManager(session_id, storage_dir=tempfile.mkdtemp())
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def msg(i, text="hello"):
|
|
28
|
+
return SessionMessage(
|
|
29
|
+
message={"role": "user", "content": [{"text": text}]}, message_id=i
|
|
30
|
+
)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def agent(aid=AID):
|
|
34
|
+
return SessionAgent(agent_id=aid, state={}, conversation_manager_state={})
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class FakeMultiAgent:
|
|
38
|
+
"""The shape RepositorySessionManager requires."""
|
|
39
|
+
|
|
40
|
+
def __init__(self, mid="swarm-1", state=None):
|
|
41
|
+
self.id = mid
|
|
42
|
+
self._state = state or {"nodes": ["a", "b"], "step": 1}
|
|
43
|
+
|
|
44
|
+
def serialize_state(self):
|
|
45
|
+
return dict(self._state)
|
|
46
|
+
|
|
47
|
+
def deserialize_state(self, state):
|
|
48
|
+
self._state = dict(state)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
# ---- multi-agent state: the base class raises without these --------------
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def test_multi_agent_state_round_trips():
|
|
55
|
+
m = mgr()
|
|
56
|
+
swarm = FakeMultiAgent()
|
|
57
|
+
m.create_multi_agent(SID, swarm)
|
|
58
|
+
assert m.read_multi_agent(SID, "swarm-1") == {"nodes": ["a", "b"], "step": 1}
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def test_reading_unknown_multi_agent_state_is_none_like_the_reference():
|
|
62
|
+
assert mgr().read_multi_agent(SID, "nope") is None
|
|
63
|
+
assert ref().read_multi_agent(SID, "nope") is None
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def test_updating_unknown_multi_agent_state_raises_like_the_reference():
|
|
67
|
+
m, r = mgr(), ref()
|
|
68
|
+
for store in (m, r):
|
|
69
|
+
with pytest.raises(SessionException, match="does not exist"):
|
|
70
|
+
store.update_multi_agent(SID, FakeMultiAgent("ghost"))
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def test_updating_multi_agent_state_does_not_duplicate():
|
|
74
|
+
m = mgr()
|
|
75
|
+
swarm = FakeMultiAgent()
|
|
76
|
+
m.create_multi_agent(SID, swarm)
|
|
77
|
+
for step in range(2, 6):
|
|
78
|
+
swarm._state["step"] = step
|
|
79
|
+
m.update_multi_agent(SID, swarm)
|
|
80
|
+
assert m.read_multi_agent(SID, "swarm-1")["step"] == 5
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def test_the_manager_hook_initializes_multi_agent_state():
|
|
84
|
+
"""initialize_multi_agent creates on first sight and restores afterwards."""
|
|
85
|
+
p = path()
|
|
86
|
+
first = mgr(p)
|
|
87
|
+
swarm = FakeMultiAgent(state={"step": 7})
|
|
88
|
+
first.initialize_multi_agent(swarm)
|
|
89
|
+
first.sync_multi_agent(swarm)
|
|
90
|
+
del first
|
|
91
|
+
|
|
92
|
+
import gc
|
|
93
|
+
|
|
94
|
+
gc.collect()
|
|
95
|
+
again = CitadelSessionManager(SID, p, "pw")
|
|
96
|
+
restored = FakeMultiAgent(state={"step": 0})
|
|
97
|
+
again.initialize_multi_agent(restored)
|
|
98
|
+
assert restored._state["step"] == 7, "multi-agent state did not restore"
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def test_forget_session_also_destroys_multi_agent_state():
|
|
102
|
+
m = mgr()
|
|
103
|
+
m.create_multi_agent(SID, FakeMultiAgent())
|
|
104
|
+
m.forget_session(SID)
|
|
105
|
+
assert m.read_multi_agent(SID, "swarm-1") is None
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
# ---- the path an Agent actually takes ------------------------------------
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def test_an_agent_restores_its_history_through_the_manager():
|
|
112
|
+
"""initialize() is what Agent construction calls."""
|
|
113
|
+
from strands import Agent
|
|
114
|
+
|
|
115
|
+
p = path()
|
|
116
|
+
first = mgr(p)
|
|
117
|
+
a = Agent(agent_id=AID, session_manager=first, messages=[
|
|
118
|
+
{"role": "user", "content": [{"text": "remembered turn"}]}
|
|
119
|
+
])
|
|
120
|
+
first.sync_agent(a)
|
|
121
|
+
del a, first
|
|
122
|
+
|
|
123
|
+
import gc
|
|
124
|
+
|
|
125
|
+
gc.collect()
|
|
126
|
+
again = CitadelSessionManager(SID, p, "pw")
|
|
127
|
+
restored = Agent(agent_id=AID, session_manager=again)
|
|
128
|
+
assert any(
|
|
129
|
+
"remembered turn" in str(m) for m in restored.messages
|
|
130
|
+
), f"history did not restore: {restored.messages}"
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def test_agent_state_survives_a_restore():
|
|
134
|
+
from strands import Agent
|
|
135
|
+
|
|
136
|
+
p = path()
|
|
137
|
+
first = mgr(p)
|
|
138
|
+
a = Agent(agent_id=AID, session_manager=first)
|
|
139
|
+
a.state.set("locale", "en-GB")
|
|
140
|
+
first.sync_agent(a)
|
|
141
|
+
del a, first
|
|
142
|
+
|
|
143
|
+
import gc
|
|
144
|
+
|
|
145
|
+
gc.collect()
|
|
146
|
+
again = CitadelSessionManager(SID, p, "pw")
|
|
147
|
+
restored = Agent(agent_id=AID, session_manager=again)
|
|
148
|
+
assert restored.state.get("locale") == "en-GB"
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
# ---- list_messages slice edges vs the reference --------------------------
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def seeded(store, n=4):
|
|
155
|
+
store.create_agent(SID, agent())
|
|
156
|
+
for i in range(n):
|
|
157
|
+
store.create_message(SID, AID, msg(i, f"turn {i}"))
|
|
158
|
+
return store
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def test_slice_edges_match_the_reference():
|
|
162
|
+
m, r = seeded(mgr()), seeded(ref())
|
|
163
|
+
for limit, offset in ((0, 0), (1, 0), (None, 0), (None, 10), (3, 2), (100, 1)):
|
|
164
|
+
ours = [x.message_id for x in m.list_messages(SID, AID, limit, offset)]
|
|
165
|
+
theirs = [x.message_id for x in r.list_messages(SID, AID, limit, offset)]
|
|
166
|
+
assert ours == theirs, (limit, offset, ours, theirs)
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def test_listing_an_agent_with_no_messages_is_empty_like_the_reference():
|
|
170
|
+
m, r = mgr(), ref()
|
|
171
|
+
for store in (m, r):
|
|
172
|
+
store.create_agent(SID, agent())
|
|
173
|
+
assert store.list_messages(SID, AID) == []
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
# ---- an update that fails must not lose the record ----------------------
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def test_a_failed_update_leaves_the_original_readable():
|
|
180
|
+
"""The write path erases before it writes; a failure must not orphan."""
|
|
181
|
+
m = seeded(mgr(), 2)
|
|
182
|
+
broken = msg(0, "replacement")
|
|
183
|
+
broken.message_id = 0
|
|
184
|
+
|
|
185
|
+
class Exploding(dict):
|
|
186
|
+
def __getitem__(self, k):
|
|
187
|
+
raise RuntimeError("boom")
|
|
188
|
+
|
|
189
|
+
# A message whose serialization explodes.
|
|
190
|
+
try:
|
|
191
|
+
m.update_message(SID, AID, msg(99)) # unknown id: raises before any erase
|
|
192
|
+
except SessionException:
|
|
193
|
+
pass
|
|
194
|
+
assert m.read_message(SID, AID, 0).message["content"][0]["text"] == "turn 0"
|
|
195
|
+
assert len(m.list_messages(SID, AID)) == 2
|
|
@@ -0,0 +1,362 @@
|
|
|
1
|
+
"""Tests for CitadelSessionManager, asserted against FileSessionManager."""
|
|
2
|
+
import concurrent.futures as cf
|
|
3
|
+
import gc
|
|
4
|
+
import os
|
|
5
|
+
import tempfile
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
from strands.session.file_session_manager import FileSessionManager
|
|
9
|
+
from strands.types.exceptions import SessionException
|
|
10
|
+
from strands.types.session import Session, SessionAgent, SessionMessage, SessionType
|
|
11
|
+
|
|
12
|
+
from citadeldb_strands_agents import CitadelSessionManager
|
|
13
|
+
|
|
14
|
+
SID = "s1"
|
|
15
|
+
AID = "a1"
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def path(name="s.cdl"):
|
|
19
|
+
return os.path.join(tempfile.mkdtemp(), name)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def mgr(p=None, session_id=SID):
|
|
23
|
+
return CitadelSessionManager(session_id, p or path(), "pw")
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def ref(session_id=SID):
|
|
27
|
+
return FileSessionManager(session_id, storage_dir=tempfile.mkdtemp())
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def msg(i, text="hello"):
|
|
31
|
+
return SessionMessage(
|
|
32
|
+
message={"role": "user", "content": [{"text": text}]}, message_id=i
|
|
33
|
+
)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def agent(aid=AID):
|
|
37
|
+
return SessionAgent(agent_id=aid, state={}, conversation_manager_state={})
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def seeded(m, n=3):
|
|
41
|
+
"""A session with an agent and n messages."""
|
|
42
|
+
m.create_agent(SID, agent())
|
|
43
|
+
for i in range(n):
|
|
44
|
+
m.create_message(SID, AID, msg(i, f"turn {i}"))
|
|
45
|
+
return m
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
# ---- construction mirrors the shipped managers ---------------------------
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def test_it_is_both_a_repository_and_a_manager():
|
|
52
|
+
from strands.session.session_manager import SessionManager
|
|
53
|
+
from strands.session.session_repository import SessionRepository
|
|
54
|
+
|
|
55
|
+
m = mgr()
|
|
56
|
+
assert isinstance(m, SessionRepository)
|
|
57
|
+
assert isinstance(m, SessionManager)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def test_constructing_it_creates_the_session_like_the_reference():
|
|
61
|
+
"""RepositorySessionManager creates the session on construction."""
|
|
62
|
+
assert mgr().read_session(SID) is not None
|
|
63
|
+
assert ref().read_session(SID) is not None
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def test_a_passphrase_is_required():
|
|
67
|
+
with pytest.raises(ValueError, match="passphrase"):
|
|
68
|
+
CitadelSessionManager(SID, path(), "")
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
# ---- session / agent / message contracts vs the reference ----------------
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def test_creating_a_duplicate_session_raises_like_the_reference():
|
|
75
|
+
m, r = mgr(), ref()
|
|
76
|
+
dup = Session(session_id=SID, session_type=SessionType.AGENT)
|
|
77
|
+
for store in (m, r):
|
|
78
|
+
with pytest.raises(SessionException, match="already exists"):
|
|
79
|
+
store.create_session(dup)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def test_reading_an_unknown_session_is_none_like_the_reference():
|
|
83
|
+
assert mgr().read_session("nope") is None
|
|
84
|
+
assert ref().read_session("nope") is None
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def test_reading_an_unknown_agent_is_none_like_the_reference():
|
|
88
|
+
assert mgr().read_agent(SID, "nope") is None
|
|
89
|
+
assert ref().read_agent(SID, "nope") is None
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def test_updating_an_unknown_agent_raises_like_the_reference():
|
|
93
|
+
m, r = mgr(), ref()
|
|
94
|
+
for store in (m, r):
|
|
95
|
+
with pytest.raises(SessionException, match="does not exist"):
|
|
96
|
+
store.update_agent(SID, agent("ghost"))
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def test_reading_an_unknown_message_is_none_like_the_reference():
|
|
100
|
+
m, r = seeded(mgr()), seeded(ref())
|
|
101
|
+
assert m.read_message(SID, AID, 99) is None
|
|
102
|
+
assert r.read_message(SID, AID, 99) is None
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def test_updating_an_unknown_message_raises_like_the_reference():
|
|
106
|
+
m, r = seeded(mgr()), seeded(ref())
|
|
107
|
+
for store in (m, r):
|
|
108
|
+
with pytest.raises(SessionException, match="does not exist"):
|
|
109
|
+
store.update_message(SID, AID, msg(99))
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def test_listing_messages_for_an_unknown_agent_raises_like_the_reference():
|
|
113
|
+
m, r = mgr(), ref()
|
|
114
|
+
for store in (m, r):
|
|
115
|
+
with pytest.raises(SessionException):
|
|
116
|
+
store.list_messages(SID, "ghost")
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
# ---- ordering and pagination --------------------------------------------
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def test_messages_come_back_in_conversation_order():
|
|
123
|
+
m = seeded(mgr(), 5)
|
|
124
|
+
assert [x.message_id for x in m.list_messages(SID, AID)] == [0, 1, 2, 3, 4]
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
def test_order_follows_the_index_not_the_write_time():
|
|
128
|
+
"""Updating an early message must not move it to the end."""
|
|
129
|
+
m = seeded(mgr(), 5)
|
|
130
|
+
m.update_message(SID, AID, msg(1, "edited"))
|
|
131
|
+
assert [x.message_id for x in m.list_messages(SID, AID)] == [0, 1, 2, 3, 4]
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
def test_pagination_matches_the_reference():
|
|
135
|
+
m, r = seeded(mgr(), 6), seeded(ref(), 6)
|
|
136
|
+
for limit, offset in ((2, 0), (2, 2), (None, 3), (10, 0), (2, 99)):
|
|
137
|
+
ours = [x.message_id for x in m.list_messages(SID, AID, limit, offset)]
|
|
138
|
+
theirs = [x.message_id for x in r.list_messages(SID, AID, limit, offset)]
|
|
139
|
+
assert ours == theirs, (limit, offset, ours, theirs)
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
# ---- no duplicates on repeat writes -------------------------------------
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def test_updating_a_message_does_not_duplicate_it():
|
|
146
|
+
m = seeded(mgr(), 3)
|
|
147
|
+
for text in ("one", "two", "three"):
|
|
148
|
+
m.update_message(SID, AID, msg(1, text))
|
|
149
|
+
assert len(m.list_messages(SID, AID)) == 3
|
|
150
|
+
assert m.read_message(SID, AID, 1).message["content"][0]["text"] == "three"
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def test_updating_an_agent_does_not_duplicate_it():
|
|
154
|
+
m = seeded(mgr())
|
|
155
|
+
for _ in range(5):
|
|
156
|
+
a = agent()
|
|
157
|
+
a.state = {"n": 1}
|
|
158
|
+
m.update_agent(SID, a)
|
|
159
|
+
assert m.read_agent(SID, AID).state == {"n": 1}
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
# ---- redaction: the reason to use an encrypted store --------------------
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def test_the_reference_keeps_the_redacted_text_on_disk():
|
|
166
|
+
"""Established as the baseline this integration improves on."""
|
|
167
|
+
r = seeded(ref(), 1)
|
|
168
|
+
original = "my card number is 4111 1111 1111 1111"
|
|
169
|
+
r.create_message(SID, AID, msg(1, original))
|
|
170
|
+
m = r.read_message(SID, AID, 1)
|
|
171
|
+
m.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
172
|
+
r.update_message(SID, AID, m)
|
|
173
|
+
|
|
174
|
+
stored = r.read_message(SID, AID, 1)
|
|
175
|
+
assert stored.redact_message is not None
|
|
176
|
+
assert original in str(stored.message), "the reference keeps the original beside it"
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def test_a_redaction_removes_the_original_from_the_store():
|
|
180
|
+
"""Keeping the original would leave the redacted content in the store."""
|
|
181
|
+
m = seeded(mgr(), 0)
|
|
182
|
+
secret = "my card number is 4111 1111 1111 1111"
|
|
183
|
+
m.create_message(SID, AID, msg(0, secret))
|
|
184
|
+
|
|
185
|
+
stored = m.read_message(SID, AID, 0)
|
|
186
|
+
stored.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
187
|
+
m.update_message(SID, AID, stored)
|
|
188
|
+
|
|
189
|
+
after = m.read_message(SID, AID, 0)
|
|
190
|
+
assert after.to_message() == {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
191
|
+
assert secret not in str(after.message), "the original survived the redaction"
|
|
192
|
+
assert secret not in str(after.to_dict()), "the original survived the redaction"
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def test_a_redaction_at_creation_time_also_drops_the_original():
|
|
196
|
+
m = seeded(mgr(), 0)
|
|
197
|
+
secret = "another secret"
|
|
198
|
+
sm = msg(0, secret)
|
|
199
|
+
sm.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
200
|
+
m.create_message(SID, AID, sm)
|
|
201
|
+
assert secret not in str(m.read_message(SID, AID, 0).to_dict())
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def test_an_unredacted_message_keeps_its_content():
|
|
205
|
+
"""Only a redaction drops the original; an ordinary update must not."""
|
|
206
|
+
m = seeded(mgr(), 0)
|
|
207
|
+
m.create_message(SID, AID, msg(0, "ordinary content"))
|
|
208
|
+
m.update_message(SID, AID, msg(0, "edited content"))
|
|
209
|
+
after = m.read_message(SID, AID, 0)
|
|
210
|
+
assert after.to_message()["content"][0]["text"] == "edited content"
|
|
211
|
+
assert after.redact_message is None
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def test_a_redaction_replaces_the_stored_record():
|
|
215
|
+
m = seeded(mgr(), 1)
|
|
216
|
+
original = "my card number is 4111 1111 1111 1111"
|
|
217
|
+
m.create_message(SID, AID, msg(1, original))
|
|
218
|
+
|
|
219
|
+
stored = m.read_message(SID, AID, 1)
|
|
220
|
+
stored.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
221
|
+
m.update_message(SID, AID, stored)
|
|
222
|
+
|
|
223
|
+
# One record for the index, carrying the redaction.
|
|
224
|
+
assert len([x for x in m.list_messages(SID, AID) if x.message_id == 1]) == 1
|
|
225
|
+
assert m.read_message(SID, AID, 1).redact_message is not None
|
|
226
|
+
|
|
227
|
+
|
|
228
|
+
def test_created_at_survives_an_update_like_the_reference():
|
|
229
|
+
m, r = seeded(mgr(), 1), seeded(ref(), 1)
|
|
230
|
+
for store in (m, r):
|
|
231
|
+
first = store.read_message(SID, AID, 0)
|
|
232
|
+
edited = msg(0, "edited")
|
|
233
|
+
store.update_message(SID, AID, edited)
|
|
234
|
+
assert store.read_message(SID, AID, 0).created_at == first.created_at
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
# ---- erasure -------------------------------------------------------------
|
|
238
|
+
|
|
239
|
+
|
|
240
|
+
def test_forget_session_destroys_everything_it_owns():
|
|
241
|
+
p = path()
|
|
242
|
+
m = seeded(mgr(p), 3)
|
|
243
|
+
assert m.forget_session(SID) == 5 # session + agent + 3 messages
|
|
244
|
+
assert m.read_session(SID) is None
|
|
245
|
+
assert m.read_agent(SID, AID) is None
|
|
246
|
+
|
|
247
|
+
|
|
248
|
+
def test_forgetting_an_unknown_session_is_zero():
|
|
249
|
+
assert mgr().forget_session("never") == 0
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
# ---- persistence ---------------------------------------------------------
|
|
253
|
+
|
|
254
|
+
|
|
255
|
+
def test_it_survives_a_reopen():
|
|
256
|
+
p = path()
|
|
257
|
+
first = seeded(mgr(p), 3)
|
|
258
|
+
first.update_agent(SID, SessionAgent(agent_id=AID, state={"k": "v"},
|
|
259
|
+
conversation_manager_state={}))
|
|
260
|
+
del first
|
|
261
|
+
gc.collect()
|
|
262
|
+
|
|
263
|
+
again = CitadelSessionManager(SID, p, "pw")
|
|
264
|
+
assert again.read_session(SID) is not None
|
|
265
|
+
assert again.read_agent(SID, AID).state == {"k": "v"}
|
|
266
|
+
assert [x.message_id for x in again.list_messages(SID, AID)] == [0, 1, 2]
|
|
267
|
+
|
|
268
|
+
|
|
269
|
+
def test_a_wrong_passphrase_cannot_reopen():
|
|
270
|
+
import citadeldb
|
|
271
|
+
|
|
272
|
+
p = path()
|
|
273
|
+
first = seeded(mgr(p))
|
|
274
|
+
del first
|
|
275
|
+
gc.collect()
|
|
276
|
+
with pytest.raises(citadeldb.EncryptionError):
|
|
277
|
+
CitadelSessionManager(SID, p, "wrong")
|
|
278
|
+
|
|
279
|
+
|
|
280
|
+
# ---- concurrency ---------------------------------------------------------
|
|
281
|
+
|
|
282
|
+
|
|
283
|
+
def test_concurrent_message_writes_all_land():
|
|
284
|
+
m = seeded(mgr(), 0)
|
|
285
|
+
with cf.ThreadPoolExecutor(max_workers=4) as ex:
|
|
286
|
+
list(ex.map(lambda i: m.create_message(SID, AID, msg(i, f"t{i}")), range(40)))
|
|
287
|
+
assert len(m.list_messages(SID, AID)) == 40
|
|
288
|
+
|
|
289
|
+
|
|
290
|
+
def test_a_message_id_is_stored_once_however_often_it_is_written():
|
|
291
|
+
"""message_id is the conversation index, so a second write of one id
|
|
292
|
+
replaces it. Two request handlers over one store is the normal shape."""
|
|
293
|
+
m = seeded(mgr(), 0)
|
|
294
|
+
m.create_message(SID, AID, msg(7, "from handler one"))
|
|
295
|
+
m.create_message(SID, AID, msg(7, "from handler two"))
|
|
296
|
+
stored = m.list_messages(SID, AID)
|
|
297
|
+
assert [s.message_id for s in stored] == [7]
|
|
298
|
+
assert m.read_message(SID, AID, 7).message["content"][0]["text"] == "from handler two"
|
|
299
|
+
|
|
300
|
+
|
|
301
|
+
def test_concurrent_writes_of_one_message_id_leave_one_record():
|
|
302
|
+
m = seeded(mgr(), 0)
|
|
303
|
+
with cf.ThreadPoolExecutor(max_workers=8) as ex:
|
|
304
|
+
list(ex.map(lambda i: m.create_message(SID, AID, msg(3, f"t{i}")), range(32)))
|
|
305
|
+
assert [s.message_id for s in m.list_messages(SID, AID)] == [3]
|
|
306
|
+
|
|
307
|
+
|
|
308
|
+
# ---- scoping and content shapes -----------------------------------------
|
|
309
|
+
|
|
310
|
+
|
|
311
|
+
def test_sessions_and_agents_are_scoped_exactly():
|
|
312
|
+
p = path()
|
|
313
|
+
m = CitadelSessionManager("sess", p, "pw")
|
|
314
|
+
other = CitadelSessionManager("sess-2", p, "pw")
|
|
315
|
+
m.create_agent("sess", agent())
|
|
316
|
+
other.create_agent("sess-2", agent())
|
|
317
|
+
m.create_message("sess", AID, msg(0, "mine"))
|
|
318
|
+
assert len(m.list_messages("sess", AID)) == 1
|
|
319
|
+
assert other.list_messages("sess-2", AID) == []
|
|
320
|
+
|
|
321
|
+
|
|
322
|
+
def test_agents_within_a_session_are_scoped_exactly():
|
|
323
|
+
m = mgr()
|
|
324
|
+
m.create_agent(SID, agent("a"))
|
|
325
|
+
m.create_agent(SID, agent("a-2"))
|
|
326
|
+
m.create_message(SID, "a", msg(0, "mine"))
|
|
327
|
+
assert len(m.list_messages(SID, "a")) == 1
|
|
328
|
+
assert m.list_messages(SID, "a-2") == []
|
|
329
|
+
|
|
330
|
+
|
|
331
|
+
def test_unicode_round_trips():
|
|
332
|
+
m = seeded(mgr(), 0)
|
|
333
|
+
text = "ünïcode — 中文 — emoji 🎉"
|
|
334
|
+
m.create_message(SID, AID, msg(0, text))
|
|
335
|
+
assert m.read_message(SID, AID, 0).message["content"][0]["text"] == text
|
|
336
|
+
|
|
337
|
+
|
|
338
|
+
def test_agent_state_named_like_internal_fields_round_trips():
|
|
339
|
+
m = seeded(mgr(), 0)
|
|
340
|
+
hostile = {"sid": "x", "aid": "x", "mid": "x", "message": "x", "agent": "x"}
|
|
341
|
+
a = agent()
|
|
342
|
+
a.state = dict(hostile)
|
|
343
|
+
m.update_agent(SID, a)
|
|
344
|
+
assert m.read_agent(SID, AID).state == hostile
|
|
345
|
+
|
|
346
|
+
|
|
347
|
+
def test_a_message_with_no_text_is_storable():
|
|
348
|
+
m = seeded(mgr(), 0)
|
|
349
|
+
m.create_message(SID, AID, SessionMessage(
|
|
350
|
+
message={"role": "assistant", "content": []}, message_id=0
|
|
351
|
+
))
|
|
352
|
+
assert len(m.list_messages(SID, AID)) == 1
|
|
353
|
+
|
|
354
|
+
|
|
355
|
+
def test_two_managers_share_one_database_file():
|
|
356
|
+
p = path()
|
|
357
|
+
a = CitadelSessionManager("sa", p, "pw")
|
|
358
|
+
b = CitadelSessionManager("sb", p, "pw")
|
|
359
|
+
a.create_agent("sa", agent())
|
|
360
|
+
a.create_message("sa", AID, msg(0, "in a"))
|
|
361
|
+
assert a.read_session("sa") is not None and b.read_session("sb") is not None
|
|
362
|
+
assert b.read_session("sa") is not None, "one file, both sessions visible"
|