strands-agents-session 0.1.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.
- strands_agents_session-0.1.0/.gitignore +7 -0
- strands_agents_session-0.1.0/PKG-INFO +76 -0
- strands_agents_session-0.1.0/README.md +60 -0
- strands_agents_session-0.1.0/pyproject.toml +37 -0
- strands_agents_session-0.1.0/src/strands_agents_session/__init__.py +18 -0
- strands_agents_session-0.1.0/src/strands_agents_session/keys.py +37 -0
- strands_agents_session-0.1.0/src/strands_agents_session/manager.py +146 -0
- strands_agents_session-0.1.0/src/strands_agents_session/storage.py +69 -0
- strands_agents_session-0.1.0/tests/test_keyvalue_manager.py +137 -0
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: strands-agents-session
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Backend-agnostic session management for Strands Agents — a SessionStorage interface + KeyValueSessionManager to build session backends (DynamoDB, Redis, MongoDB, …)
|
|
5
|
+
Project-URL: Homepage, https://github.com/skamalj/strands_agents_session
|
|
6
|
+
Project-URL: Repository, https://github.com/skamalj/strands_agents_session.git
|
|
7
|
+
Author-email: Kamal <skamalj@github.com>
|
|
8
|
+
Keywords: agent-state,memory,persistence,session,session-manager,strands,strands-agents
|
|
9
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
10
|
+
Classifier: Operating System :: OS Independent
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Requires-Python: >=3.10
|
|
14
|
+
Requires-Dist: strands-agents
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
|
|
17
|
+
# strands-agents-session
|
|
18
|
+
|
|
19
|
+
Backend-agnostic **session management** for [Strands Agents](https://strandsagents.com). Provides the shared machinery so you can build a session storage backend (DynamoDB, Redis, MongoDB, Postgres, …) by implementing just a handful of methods.
|
|
20
|
+
|
|
21
|
+
This is the base package of a family. Concrete backends are shipped separately, e.g. [`strands-session-dynamodb`](https://pypi.org/project/strands-session-dynamodb/).
|
|
22
|
+
|
|
23
|
+
## What it gives you
|
|
24
|
+
|
|
25
|
+
- **`KeyValueSessionManager`** — implements the full Strands `SessionRepository` (all 8 CRUD methods) and mixes in `RepositorySessionManager`, so the Strands session lifecycle (message indexing, restore, `removed_message_count` offsetting, tool-use repair, change detection) is reused unchanged.
|
|
26
|
+
- **`SessionStorage`** — a tiny ordered keyed-record interface (`put` / `get` / `query` / `delete` / `delete_partition`). Implement it and you have a working Strands session backend.
|
|
27
|
+
- **`InMemorySessionStorage`** — a ready in-memory backend for tests and local development.
|
|
28
|
+
- **`keys`** — shared key/serialization conventions (zero-padded message sort keys so `list_messages(offset, limit)` is a native ordered range scan).
|
|
29
|
+
|
|
30
|
+
> **Storage only, by design.** Message *pruning* in Strands is a `ConversationManager` concern, deliberately decoupled from storage. This package (and its backends) never prune — doing so at the storage layer would corrupt Strands' message-index/offset restore logic.
|
|
31
|
+
|
|
32
|
+
## Installation
|
|
33
|
+
|
|
34
|
+
```bash
|
|
35
|
+
pip install strands-agents-session
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
**Requires Python 3.10+.**
|
|
39
|
+
|
|
40
|
+
## Using the in-memory backend
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
from strands import Agent
|
|
44
|
+
from strands_agents_session import KeyValueSessionManager, InMemorySessionStorage
|
|
45
|
+
|
|
46
|
+
session_manager = KeyValueSessionManager(
|
|
47
|
+
session_id="user-123",
|
|
48
|
+
storage=InMemorySessionStorage(),
|
|
49
|
+
)
|
|
50
|
+
agent = Agent(session_manager=session_manager)
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
## Building a backend
|
|
54
|
+
|
|
55
|
+
Implement `SessionStorage` and hand it to `KeyValueSessionManager`:
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
from strands_agents_session import KeyValueSessionManager, SessionStorage
|
|
59
|
+
|
|
60
|
+
class MyStorage(SessionStorage):
|
|
61
|
+
def put(self, item): ... # item = {"pk", "sk", "data"}
|
|
62
|
+
def get(self, pk, sk): ... # -> item | None
|
|
63
|
+
def query(self, pk, sk_gte=None, limit=None): ... # ordered by sk asc
|
|
64
|
+
def delete(self, pk, sk): ...
|
|
65
|
+
def delete_partition(self, pk): ...
|
|
66
|
+
|
|
67
|
+
class MySessionManager(KeyValueSessionManager):
|
|
68
|
+
def __init__(self, session_id, **cfg):
|
|
69
|
+
super().__init__(session_id=session_id, storage=MyStorage(**cfg))
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
That's it — all the `SessionRepository` methods, restore logic, and pagination come from the base.
|
|
73
|
+
|
|
74
|
+
## License
|
|
75
|
+
|
|
76
|
+
MIT
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
# strands-agents-session
|
|
2
|
+
|
|
3
|
+
Backend-agnostic **session management** for [Strands Agents](https://strandsagents.com). Provides the shared machinery so you can build a session storage backend (DynamoDB, Redis, MongoDB, Postgres, …) by implementing just a handful of methods.
|
|
4
|
+
|
|
5
|
+
This is the base package of a family. Concrete backends are shipped separately, e.g. [`strands-session-dynamodb`](https://pypi.org/project/strands-session-dynamodb/).
|
|
6
|
+
|
|
7
|
+
## What it gives you
|
|
8
|
+
|
|
9
|
+
- **`KeyValueSessionManager`** — implements the full Strands `SessionRepository` (all 8 CRUD methods) and mixes in `RepositorySessionManager`, so the Strands session lifecycle (message indexing, restore, `removed_message_count` offsetting, tool-use repair, change detection) is reused unchanged.
|
|
10
|
+
- **`SessionStorage`** — a tiny ordered keyed-record interface (`put` / `get` / `query` / `delete` / `delete_partition`). Implement it and you have a working Strands session backend.
|
|
11
|
+
- **`InMemorySessionStorage`** — a ready in-memory backend for tests and local development.
|
|
12
|
+
- **`keys`** — shared key/serialization conventions (zero-padded message sort keys so `list_messages(offset, limit)` is a native ordered range scan).
|
|
13
|
+
|
|
14
|
+
> **Storage only, by design.** Message *pruning* in Strands is a `ConversationManager` concern, deliberately decoupled from storage. This package (and its backends) never prune — doing so at the storage layer would corrupt Strands' message-index/offset restore logic.
|
|
15
|
+
|
|
16
|
+
## Installation
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
pip install strands-agents-session
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
**Requires Python 3.10+.**
|
|
23
|
+
|
|
24
|
+
## Using the in-memory backend
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
from strands import Agent
|
|
28
|
+
from strands_agents_session import KeyValueSessionManager, InMemorySessionStorage
|
|
29
|
+
|
|
30
|
+
session_manager = KeyValueSessionManager(
|
|
31
|
+
session_id="user-123",
|
|
32
|
+
storage=InMemorySessionStorage(),
|
|
33
|
+
)
|
|
34
|
+
agent = Agent(session_manager=session_manager)
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
## Building a backend
|
|
38
|
+
|
|
39
|
+
Implement `SessionStorage` and hand it to `KeyValueSessionManager`:
|
|
40
|
+
|
|
41
|
+
```python
|
|
42
|
+
from strands_agents_session import KeyValueSessionManager, SessionStorage
|
|
43
|
+
|
|
44
|
+
class MyStorage(SessionStorage):
|
|
45
|
+
def put(self, item): ... # item = {"pk", "sk", "data"}
|
|
46
|
+
def get(self, pk, sk): ... # -> item | None
|
|
47
|
+
def query(self, pk, sk_gte=None, limit=None): ... # ordered by sk asc
|
|
48
|
+
def delete(self, pk, sk): ...
|
|
49
|
+
def delete_partition(self, pk): ...
|
|
50
|
+
|
|
51
|
+
class MySessionManager(KeyValueSessionManager):
|
|
52
|
+
def __init__(self, session_id, **cfg):
|
|
53
|
+
super().__init__(session_id=session_id, storage=MyStorage(**cfg))
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
That's it — all the `SessionRepository` methods, restore logic, and pagination come from the base.
|
|
57
|
+
|
|
58
|
+
## License
|
|
59
|
+
|
|
60
|
+
MIT
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "strands-agents-session"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Backend-agnostic session management for Strands Agents — a SessionStorage interface + KeyValueSessionManager to build session backends (DynamoDB, Redis, MongoDB, …)"
|
|
9
|
+
authors = [{name = "Kamal", email = "skamalj@github.com"}]
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"strands-agents",
|
|
14
|
+
]
|
|
15
|
+
keywords = [
|
|
16
|
+
"strands", "strands-agents", "session", "agent-state", "memory",
|
|
17
|
+
"persistence", "session-manager",
|
|
18
|
+
]
|
|
19
|
+
classifiers = [
|
|
20
|
+
"Programming Language :: Python :: 3",
|
|
21
|
+
"License :: OSI Approved :: MIT License",
|
|
22
|
+
"Operating System :: OS Independent",
|
|
23
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[project.urls]
|
|
27
|
+
Homepage = "https://github.com/skamalj/strands_agents_session"
|
|
28
|
+
Repository = "https://github.com/skamalj/strands_agents_session.git"
|
|
29
|
+
|
|
30
|
+
[dependency-groups]
|
|
31
|
+
dev = [
|
|
32
|
+
"pytest>=7.0",
|
|
33
|
+
"pytest-cov",
|
|
34
|
+
]
|
|
35
|
+
|
|
36
|
+
[tool.hatch.build.targets.wheel]
|
|
37
|
+
packages = ["src/strands_agents_session"]
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
"""Backend-agnostic session management for Strands Agents.
|
|
2
|
+
|
|
3
|
+
Provides a ``SessionStorage`` interface and a ``KeyValueSessionManager`` that
|
|
4
|
+
implements the full Strands ``SessionRepository`` over it, so new storage
|
|
5
|
+
backends (DynamoDB, Redis, MongoDB, …) only need to implement a handful of
|
|
6
|
+
keyed-record methods.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from . import keys
|
|
10
|
+
from .manager import KeyValueSessionManager
|
|
11
|
+
from .storage import InMemorySessionStorage, SessionStorage
|
|
12
|
+
|
|
13
|
+
__all__ = [
|
|
14
|
+
"KeyValueSessionManager",
|
|
15
|
+
"SessionStorage",
|
|
16
|
+
"InMemorySessionStorage",
|
|
17
|
+
"keys",
|
|
18
|
+
]
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
"""Key construction helpers shared across session storage backends.
|
|
2
|
+
|
|
3
|
+
Single-logical-table layout (keys are opaque strings, so any KV/range store fits):
|
|
4
|
+
|
|
5
|
+
Session : pk = "SESSION#<session_id>" sk = "META"
|
|
6
|
+
Agent : pk = "SESSION#<session_id>" sk = "AGENT#<agent_id>"
|
|
7
|
+
Message : pk = "SESSION#<session_id>#AGENT#<agent_id>" sk = "MSG#<zero-padded id>"
|
|
8
|
+
|
|
9
|
+
Message sort keys are zero-padded so that lexical ordering matches numeric
|
|
10
|
+
ordering — making ``list_messages(offset, limit)`` a simple ordered range scan,
|
|
11
|
+
matching the ``removed_message_count`` offset semantics Strands relies on.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
# Width for zero-padding message ids so lexical sort == numeric sort.
|
|
15
|
+
MESSAGE_ID_WIDTH = 20
|
|
16
|
+
|
|
17
|
+
SESSION_SK = "META"
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def session_pk(session_id: str) -> str:
|
|
21
|
+
return f"SESSION#{session_id}"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def agent_sk(agent_id: str) -> str:
|
|
25
|
+
return f"AGENT#{agent_id}"
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def agent_id_from_sk(sk: str) -> str:
|
|
29
|
+
return sk.split("AGENT#", 1)[1]
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def messages_pk(session_id: str, agent_id: str) -> str:
|
|
33
|
+
return f"SESSION#{session_id}#AGENT#{agent_id}"
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def message_sk(message_id: int) -> str:
|
|
37
|
+
return f"MSG#{int(message_id):0{MESSAGE_ID_WIDTH}d}"
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
"""Backend-agnostic Strands session manager over a ``SessionStorage``.
|
|
2
|
+
|
|
3
|
+
Implements the Strands ``SessionRepository`` (all 8 CRUD methods) in terms of a
|
|
4
|
+
small ``SessionStorage`` interface, and mixes in ``RepositorySessionManager`` so
|
|
5
|
+
the Strands session lifecycle (message indexing, restore, offsetting, tool-use
|
|
6
|
+
repair, change detection) is reused unchanged.
|
|
7
|
+
|
|
8
|
+
A concrete backend only needs to provide a ``SessionStorage`` — see
|
|
9
|
+
``strands_session_dynamodb`` for an example.
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
import json
|
|
13
|
+
from typing import Any, List, Optional
|
|
14
|
+
|
|
15
|
+
from strands.session.repository_session_manager import RepositorySessionManager
|
|
16
|
+
from strands.session.session_repository import SessionRepository
|
|
17
|
+
from strands.types.exceptions import SessionException
|
|
18
|
+
from strands.types.session import Session, SessionAgent, SessionMessage
|
|
19
|
+
|
|
20
|
+
from .keys import (
|
|
21
|
+
SESSION_SK,
|
|
22
|
+
agent_id_from_sk,
|
|
23
|
+
agent_sk,
|
|
24
|
+
message_sk,
|
|
25
|
+
messages_pk,
|
|
26
|
+
session_pk,
|
|
27
|
+
)
|
|
28
|
+
from .storage import SessionStorage
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class KeyValueSessionManager(RepositorySessionManager, SessionRepository):
|
|
32
|
+
"""Strands session manager backed by any ``SessionStorage`` implementation."""
|
|
33
|
+
|
|
34
|
+
def __init__(self, session_id: str, storage: SessionStorage, **kwargs: Any) -> None:
|
|
35
|
+
self.storage = storage
|
|
36
|
+
super().__init__(session_id=session_id, session_repository=self)
|
|
37
|
+
|
|
38
|
+
# ------------------------------------------------------------------
|
|
39
|
+
# helpers
|
|
40
|
+
# ------------------------------------------------------------------
|
|
41
|
+
|
|
42
|
+
@staticmethod
|
|
43
|
+
def _record(pk: str, sk: str, payload: dict) -> dict:
|
|
44
|
+
return {"pk": pk, "sk": sk, "data": json.dumps(payload)}
|
|
45
|
+
|
|
46
|
+
# ------------------------------------------------------------------
|
|
47
|
+
# sessions
|
|
48
|
+
# ------------------------------------------------------------------
|
|
49
|
+
|
|
50
|
+
def create_session(self, session: Session, **kwargs: Any) -> Session:
|
|
51
|
+
pk = session_pk(session.session_id)
|
|
52
|
+
if self.storage.get(pk, SESSION_SK) is not None:
|
|
53
|
+
raise SessionException(f"Session {session.session_id} already exists")
|
|
54
|
+
self.storage.put(self._record(pk, SESSION_SK, session.to_dict()))
|
|
55
|
+
return session
|
|
56
|
+
|
|
57
|
+
def read_session(self, session_id: str, **kwargs: Any) -> Optional[Session]:
|
|
58
|
+
rec = self.storage.get(session_pk(session_id), SESSION_SK)
|
|
59
|
+
if rec is None:
|
|
60
|
+
return None
|
|
61
|
+
return Session.from_dict(json.loads(rec["data"]))
|
|
62
|
+
|
|
63
|
+
def delete_session(self, session_id: str, **kwargs: Any) -> None:
|
|
64
|
+
pk = session_pk(session_id)
|
|
65
|
+
items = self.storage.query(pk)
|
|
66
|
+
agent_ids = [agent_id_from_sk(i["sk"]) for i in items if i["sk"].startswith("AGENT#")]
|
|
67
|
+
for agent_id in agent_ids:
|
|
68
|
+
self.storage.delete_partition(messages_pk(session_id, agent_id))
|
|
69
|
+
self.storage.delete_partition(pk)
|
|
70
|
+
|
|
71
|
+
# ------------------------------------------------------------------
|
|
72
|
+
# agents
|
|
73
|
+
# ------------------------------------------------------------------
|
|
74
|
+
|
|
75
|
+
def create_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: Any) -> None:
|
|
76
|
+
self.storage.put(
|
|
77
|
+
self._record(session_pk(session_id), agent_sk(session_agent.agent_id), session_agent.to_dict())
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
def read_agent(self, session_id: str, agent_id: str, **kwargs: Any) -> Optional[SessionAgent]:
|
|
81
|
+
rec = self.storage.get(session_pk(session_id), agent_sk(agent_id))
|
|
82
|
+
if rec is None:
|
|
83
|
+
return None
|
|
84
|
+
return SessionAgent.from_dict(json.loads(rec["data"]))
|
|
85
|
+
|
|
86
|
+
def update_agent(self, session_id: str, session_agent: SessionAgent, **kwargs: Any) -> None:
|
|
87
|
+
previous = self.read_agent(session_id, session_agent.agent_id)
|
|
88
|
+
if previous is None:
|
|
89
|
+
raise SessionException(
|
|
90
|
+
f"Agent {session_agent.agent_id} in session {session_id} does not exist"
|
|
91
|
+
)
|
|
92
|
+
session_agent.created_at = previous.created_at # preserve creation timestamp
|
|
93
|
+
self.storage.put(
|
|
94
|
+
self._record(session_pk(session_id), agent_sk(session_agent.agent_id), session_agent.to_dict())
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
# ------------------------------------------------------------------
|
|
98
|
+
# messages
|
|
99
|
+
# ------------------------------------------------------------------
|
|
100
|
+
|
|
101
|
+
def create_message(
|
|
102
|
+
self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any
|
|
103
|
+
) -> None:
|
|
104
|
+
self.storage.put(
|
|
105
|
+
self._record(
|
|
106
|
+
messages_pk(session_id, agent_id),
|
|
107
|
+
message_sk(session_message.message_id),
|
|
108
|
+
session_message.to_dict(),
|
|
109
|
+
)
|
|
110
|
+
)
|
|
111
|
+
|
|
112
|
+
def read_message(
|
|
113
|
+
self, session_id: str, agent_id: str, message_id: int, **kwargs: Any
|
|
114
|
+
) -> Optional[SessionMessage]:
|
|
115
|
+
rec = self.storage.get(messages_pk(session_id, agent_id), message_sk(message_id))
|
|
116
|
+
if rec is None:
|
|
117
|
+
return None
|
|
118
|
+
return SessionMessage.from_dict(json.loads(rec["data"]))
|
|
119
|
+
|
|
120
|
+
def update_message(
|
|
121
|
+
self, session_id: str, agent_id: str, session_message: SessionMessage, **kwargs: Any
|
|
122
|
+
) -> None:
|
|
123
|
+
previous = self.read_message(session_id, agent_id, session_message.message_id)
|
|
124
|
+
if previous is None:
|
|
125
|
+
raise SessionException(f"Message {session_message.message_id} does not exist")
|
|
126
|
+
session_message.created_at = previous.created_at
|
|
127
|
+
self.storage.put(
|
|
128
|
+
self._record(
|
|
129
|
+
messages_pk(session_id, agent_id),
|
|
130
|
+
message_sk(session_message.message_id),
|
|
131
|
+
session_message.to_dict(),
|
|
132
|
+
)
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
def list_messages(
|
|
136
|
+
self,
|
|
137
|
+
session_id: str,
|
|
138
|
+
agent_id: str,
|
|
139
|
+
limit: Optional[int] = None,
|
|
140
|
+
offset: int = 0,
|
|
141
|
+
**kwargs: Any,
|
|
142
|
+
) -> List[SessionMessage]:
|
|
143
|
+
records = self.storage.query(
|
|
144
|
+
messages_pk(session_id, agent_id), sk_gte=message_sk(offset), limit=limit
|
|
145
|
+
)
|
|
146
|
+
return [SessionMessage.from_dict(json.loads(r["data"])) for r in records]
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
"""Storage interface for session backends.
|
|
2
|
+
|
|
3
|
+
A backend implements this small keyed-record interface; the
|
|
4
|
+
``KeyValueSessionManager`` builds the full Strands ``SessionRepository`` on top
|
|
5
|
+
of it. Records are plain dicts with at least ``pk``, ``sk``, and ``data`` keys
|
|
6
|
+
(``data`` is a serialized payload string).
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from abc import ABC, abstractmethod
|
|
10
|
+
from typing import Any, Dict, List, Optional
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class SessionStorage(ABC):
|
|
14
|
+
"""Minimal ordered keyed-record store backing a session manager."""
|
|
15
|
+
|
|
16
|
+
@abstractmethod
|
|
17
|
+
def put(self, item: Dict[str, Any]) -> None:
|
|
18
|
+
"""Insert or overwrite a record. ``item`` has ``pk``, ``sk``, ``data``."""
|
|
19
|
+
|
|
20
|
+
@abstractmethod
|
|
21
|
+
def get(self, pk: str, sk: str) -> Optional[Dict[str, Any]]:
|
|
22
|
+
"""Return the record for (pk, sk), or None."""
|
|
23
|
+
|
|
24
|
+
@abstractmethod
|
|
25
|
+
def query(
|
|
26
|
+
self, pk: str, sk_gte: Optional[str] = None, limit: Optional[int] = None
|
|
27
|
+
) -> List[Dict[str, Any]]:
|
|
28
|
+
"""Return records under ``pk`` ordered by ``sk`` ascending.
|
|
29
|
+
|
|
30
|
+
If ``sk_gte`` is given, only records with ``sk >= sk_gte``. If ``limit``
|
|
31
|
+
is given, return at most that many.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
@abstractmethod
|
|
35
|
+
def delete(self, pk: str, sk: str) -> None:
|
|
36
|
+
"""Delete the record for (pk, sk) if present."""
|
|
37
|
+
|
|
38
|
+
@abstractmethod
|
|
39
|
+
def delete_partition(self, pk: str) -> None:
|
|
40
|
+
"""Delete every record under ``pk``."""
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class InMemorySessionStorage(SessionStorage):
|
|
44
|
+
"""In-memory storage — useful for tests and local development."""
|
|
45
|
+
|
|
46
|
+
def __init__(self) -> None:
|
|
47
|
+
self._data: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
|
48
|
+
|
|
49
|
+
def put(self, item: Dict[str, Any]) -> None:
|
|
50
|
+
self._data.setdefault(item["pk"], {})[item["sk"]] = dict(item)
|
|
51
|
+
|
|
52
|
+
def get(self, pk: str, sk: str) -> Optional[Dict[str, Any]]:
|
|
53
|
+
rec = self._data.get(pk, {}).get(sk)
|
|
54
|
+
return dict(rec) if rec is not None else None
|
|
55
|
+
|
|
56
|
+
def query(
|
|
57
|
+
self, pk: str, sk_gte: Optional[str] = None, limit: Optional[int] = None
|
|
58
|
+
) -> List[Dict[str, Any]]:
|
|
59
|
+
partition = self._data.get(pk, {})
|
|
60
|
+
keys = sorted(k for k in partition if sk_gte is None or k >= sk_gte)
|
|
61
|
+
if limit is not None:
|
|
62
|
+
keys = keys[:limit]
|
|
63
|
+
return [dict(partition[k]) for k in keys]
|
|
64
|
+
|
|
65
|
+
def delete(self, pk: str, sk: str) -> None:
|
|
66
|
+
self._data.get(pk, {}).pop(sk, None)
|
|
67
|
+
|
|
68
|
+
def delete_partition(self, pk: str) -> None:
|
|
69
|
+
self._data.pop(pk, None)
|
|
@@ -0,0 +1,137 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Offline tests for KeyValueSessionManager over InMemorySessionStorage.
|
|
3
|
+
|
|
4
|
+
Fully deterministic, no cloud — exercises the backend-agnostic SessionRepository
|
|
5
|
+
implementation (sessions, agents + agent state, messages, pagination,
|
|
6
|
+
update/redaction, isolation, delete). Every concrete backend inherits this
|
|
7
|
+
behavior, so these cover the shared logic once.
|
|
8
|
+
"""
|
|
9
|
+
import uuid
|
|
10
|
+
|
|
11
|
+
import pytest
|
|
12
|
+
|
|
13
|
+
from strands.types.exceptions import SessionException
|
|
14
|
+
from strands.types.session import SessionAgent, SessionMessage
|
|
15
|
+
from strands_agents_session import InMemorySessionStorage, KeyValueSessionManager
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def make_manager(storage=None, session_id=None):
|
|
19
|
+
return KeyValueSessionManager(
|
|
20
|
+
session_id=session_id or f"sess-{uuid.uuid4()}",
|
|
21
|
+
storage=storage or InMemorySessionStorage(),
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def msg(i, role="user", text=None):
|
|
26
|
+
return SessionMessage.from_message({"role": role, "content": [{"text": text or f"m{i}"}]}, i)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
# sessions ------------------------------------------------------------------
|
|
30
|
+
|
|
31
|
+
def test_session_created_on_init():
|
|
32
|
+
mgr = make_manager()
|
|
33
|
+
got = mgr.read_session(mgr.session_id)
|
|
34
|
+
assert got is not None and got.session_id == mgr.session_id
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_create_duplicate_session_raises():
|
|
38
|
+
mgr = make_manager()
|
|
39
|
+
with pytest.raises(SessionException):
|
|
40
|
+
mgr.create_session(mgr.session)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def test_read_missing_session_returns_none():
|
|
44
|
+
assert make_manager().read_session("nope") is None
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def test_delete_session_removes_everything():
|
|
48
|
+
storage = InMemorySessionStorage()
|
|
49
|
+
mgr = make_manager(storage)
|
|
50
|
+
sid = mgr.session_id
|
|
51
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
52
|
+
mgr.create_message(sid, "a1", msg(0))
|
|
53
|
+
mgr.delete_session(sid)
|
|
54
|
+
assert mgr.read_session(sid) is None
|
|
55
|
+
assert mgr.read_agent(sid, "a1") is None
|
|
56
|
+
assert mgr.list_messages(sid, "a1") == []
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
# agents (+ state) ----------------------------------------------------------
|
|
60
|
+
|
|
61
|
+
def test_agent_state_round_trip():
|
|
62
|
+
mgr = make_manager()
|
|
63
|
+
mgr.create_agent(mgr.session_id, SessionAgent(
|
|
64
|
+
agent_id="assistant", state={"tier": "gold"}, conversation_manager_state={"removed_message_count": 0}))
|
|
65
|
+
got = mgr.read_agent(mgr.session_id, "assistant")
|
|
66
|
+
assert got.state == {"tier": "gold"}
|
|
67
|
+
assert got.conversation_manager_state == {"removed_message_count": 0}
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def test_update_agent_preserves_created_at():
|
|
71
|
+
mgr = make_manager()
|
|
72
|
+
sid = mgr.session_id
|
|
73
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={"v": 1}, conversation_manager_state={}))
|
|
74
|
+
created_at = mgr.read_agent(sid, "a1").created_at
|
|
75
|
+
mgr.update_agent(sid, SessionAgent(agent_id="a1", state={"v": 2}, conversation_manager_state={}))
|
|
76
|
+
got = mgr.read_agent(sid, "a1")
|
|
77
|
+
assert got.state == {"v": 2}
|
|
78
|
+
assert got.created_at == created_at
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def test_update_missing_agent_raises():
|
|
82
|
+
mgr = make_manager()
|
|
83
|
+
with pytest.raises(SessionException):
|
|
84
|
+
mgr.update_agent(mgr.session_id, SessionAgent(agent_id="x", state={}, conversation_manager_state={}))
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
# messages + pagination -----------------------------------------------------
|
|
88
|
+
|
|
89
|
+
def test_message_round_trip():
|
|
90
|
+
mgr = make_manager()
|
|
91
|
+
sid = mgr.session_id
|
|
92
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
93
|
+
mgr.create_message(sid, "a1", msg(0, text="hello"))
|
|
94
|
+
got = mgr.read_message(sid, "a1", 0)
|
|
95
|
+
assert got.message["content"][0]["text"] == "hello"
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def test_list_messages_sorted():
|
|
99
|
+
mgr = make_manager()
|
|
100
|
+
sid = mgr.session_id
|
|
101
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
102
|
+
for i in [2, 0, 4, 1, 3]:
|
|
103
|
+
mgr.create_message(sid, "a1", msg(i))
|
|
104
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1")] == [0, 1, 2, 3, 4]
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_list_messages_offset_and_limit():
|
|
108
|
+
mgr = make_manager()
|
|
109
|
+
sid = mgr.session_id
|
|
110
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
111
|
+
for i in range(10):
|
|
112
|
+
mgr.create_message(sid, "a1", msg(i))
|
|
113
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1", offset=2)] == [2, 3, 4, 5, 6, 7, 8, 9]
|
|
114
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1", limit=3)] == [0, 1, 2]
|
|
115
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1", offset=3, limit=4)] == [3, 4, 5, 6]
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
def test_update_message_redaction():
|
|
119
|
+
mgr = make_manager()
|
|
120
|
+
sid = mgr.session_id
|
|
121
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
122
|
+
m = msg(0, text="secret")
|
|
123
|
+
mgr.create_message(sid, "a1", m)
|
|
124
|
+
m.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
125
|
+
mgr.update_message(sid, "a1", m)
|
|
126
|
+
got = mgr.read_message(sid, "a1", 0)
|
|
127
|
+
assert got.redact_message["content"][0]["text"] == "[REDACTED]"
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def test_sessions_isolated():
|
|
131
|
+
storage = InMemorySessionStorage()
|
|
132
|
+
m1 = make_manager(storage)
|
|
133
|
+
m2 = make_manager(storage)
|
|
134
|
+
m1.create_agent(m1.session_id, SessionAgent(agent_id="a", state={"who": "one"}, conversation_manager_state={}))
|
|
135
|
+
m2.create_agent(m2.session_id, SessionAgent(agent_id="a", state={"who": "two"}, conversation_manager_state={}))
|
|
136
|
+
assert m1.read_agent(m1.session_id, "a").state == {"who": "one"}
|
|
137
|
+
assert m2.read_agent(m2.session_id, "a").state == {"who": "two"}
|