strands_session_mongodb 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_session_mongodb-0.1.0/.gitignore +7 -0
- strands_session_mongodb-0.1.0/PKG-INFO +82 -0
- strands_session_mongodb-0.1.0/README.md +64 -0
- strands_session_mongodb-0.1.0/pyproject.toml +33 -0
- strands_session_mongodb-0.1.0/src/strands_session_mongodb/__init__.py +5 -0
- strands_session_mongodb-0.1.0/src/strands_session_mongodb/session_manager.py +102 -0
- strands_session_mongodb-0.1.0/tests/test_agent_llm.py +81 -0
- strands_session_mongodb-0.1.0/tests/test_repository.py +149 -0
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: strands_session_mongodb
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: MongoDB session manager (storage backend) for Strands Agents — persists sessions, agent state, and messages
|
|
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,mongodb,persistence,session,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: pymongo
|
|
15
|
+
Requires-Dist: strands-agents
|
|
16
|
+
Requires-Dist: strands-agents-session>=0.1.0
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
|
|
19
|
+
# strands-session-mongodb
|
|
20
|
+
|
|
21
|
+
MongoDB **session manager** (storage backend) for [Strands Agents](https://strandsagents.com). Persists an agent's sessions, agent state, conversation-manager state, and messages to MongoDB so conversations resume across runs.
|
|
22
|
+
|
|
23
|
+
Part of the [`strands-agents-session`](https://pypi.org/project/strands-agents-session/) family — it implements the base `SessionStorage` interface, so all the Strands session-repository logic (message indexing, restore, `removed_message_count` offsetting, tool-use repair) comes from the core.
|
|
24
|
+
|
|
25
|
+
> **Storage only, by design.** Message *pruning* in Strands is a `ConversationManager` concern, decoupled from storage. This package does not prune.
|
|
26
|
+
|
|
27
|
+
## Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
pip install strands-session-mongodb
|
|
31
|
+
# or, via the family's extra:
|
|
32
|
+
pip install "strands-agents-session[mongodb]"
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
**Requires Python 3.10+** and a reachable MongoDB (local, Atlas, or self-hosted).
|
|
36
|
+
|
|
37
|
+
## Quick start
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
from strands import Agent
|
|
41
|
+
from strands_session_mongodb import MongoDBSessionManager
|
|
42
|
+
|
|
43
|
+
session_manager = MongoDBSessionManager(
|
|
44
|
+
session_id="user-123",
|
|
45
|
+
connection_string="mongodb://127.0.0.1:27017/",
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
agent = Agent(session_manager=session_manager)
|
|
49
|
+
agent("Hi, I'm Kamal")
|
|
50
|
+
agent("What's my name?") # remembers within the session
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
Next run, same `session_id` → the agent restores its full history and state from MongoDB.
|
|
54
|
+
|
|
55
|
+
## API
|
|
56
|
+
|
|
57
|
+
### `MongoDBSessionManager(session_id, *, connection_string="mongodb://127.0.0.1:27017/", database_name="strands_sessions", collection_name="sessions", client=None, ttl_seconds=None)`
|
|
58
|
+
|
|
59
|
+
| Parameter | Description |
|
|
60
|
+
|---|---|
|
|
61
|
+
| `session_id` | Session identifier |
|
|
62
|
+
| `connection_string` | MongoDB URI (local, Atlas `mongodb+srv://…`, etc.) |
|
|
63
|
+
| `database_name` | Database (default `strands_sessions`) |
|
|
64
|
+
| `collection_name` | Collection (default `sessions`) |
|
|
65
|
+
| `client` | Optional pre-built `pymongo.MongoClient` to reuse |
|
|
66
|
+
| `ttl_seconds` | If set, adds a TTL index so documents expire automatically |
|
|
67
|
+
|
|
68
|
+
## Data model
|
|
69
|
+
|
|
70
|
+
A single collection; each document has a `pk` + `sk` (compound-indexed, unique) and a `data` field with the serialized payload:
|
|
71
|
+
|
|
72
|
+
| Item | pk | sk |
|
|
73
|
+
|---|---|---|
|
|
74
|
+
| Session | `SESSION#<session_id>` | `META` |
|
|
75
|
+
| Agent | `SESSION#<session_id>` | `AGENT#<agent_id>` |
|
|
76
|
+
| Message | `SESSION#<session_id>#AGENT#<agent_id>` | `MSG#<zero-padded id>` |
|
|
77
|
+
|
|
78
|
+
The `(pk, sk)` index gives ordered range scans, so `list_messages(offset, limit)` is a simple indexed `find(...).sort(sk).limit(...)` — matching the `removed_message_count` offset semantics Strands relies on.
|
|
79
|
+
|
|
80
|
+
## License
|
|
81
|
+
|
|
82
|
+
MIT
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
# strands-session-mongodb
|
|
2
|
+
|
|
3
|
+
MongoDB **session manager** (storage backend) for [Strands Agents](https://strandsagents.com). Persists an agent's sessions, agent state, conversation-manager state, and messages to MongoDB so conversations resume across runs.
|
|
4
|
+
|
|
5
|
+
Part of the [`strands-agents-session`](https://pypi.org/project/strands-agents-session/) family — it implements the base `SessionStorage` interface, so all the Strands session-repository logic (message indexing, restore, `removed_message_count` offsetting, tool-use repair) comes from the core.
|
|
6
|
+
|
|
7
|
+
> **Storage only, by design.** Message *pruning* in Strands is a `ConversationManager` concern, decoupled from storage. This package does not prune.
|
|
8
|
+
|
|
9
|
+
## Installation
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
pip install strands-session-mongodb
|
|
13
|
+
# or, via the family's extra:
|
|
14
|
+
pip install "strands-agents-session[mongodb]"
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
**Requires Python 3.10+** and a reachable MongoDB (local, Atlas, or self-hosted).
|
|
18
|
+
|
|
19
|
+
## Quick start
|
|
20
|
+
|
|
21
|
+
```python
|
|
22
|
+
from strands import Agent
|
|
23
|
+
from strands_session_mongodb import MongoDBSessionManager
|
|
24
|
+
|
|
25
|
+
session_manager = MongoDBSessionManager(
|
|
26
|
+
session_id="user-123",
|
|
27
|
+
connection_string="mongodb://127.0.0.1:27017/",
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
agent = Agent(session_manager=session_manager)
|
|
31
|
+
agent("Hi, I'm Kamal")
|
|
32
|
+
agent("What's my name?") # remembers within the session
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Next run, same `session_id` → the agent restores its full history and state from MongoDB.
|
|
36
|
+
|
|
37
|
+
## API
|
|
38
|
+
|
|
39
|
+
### `MongoDBSessionManager(session_id, *, connection_string="mongodb://127.0.0.1:27017/", database_name="strands_sessions", collection_name="sessions", client=None, ttl_seconds=None)`
|
|
40
|
+
|
|
41
|
+
| Parameter | Description |
|
|
42
|
+
|---|---|
|
|
43
|
+
| `session_id` | Session identifier |
|
|
44
|
+
| `connection_string` | MongoDB URI (local, Atlas `mongodb+srv://…`, etc.) |
|
|
45
|
+
| `database_name` | Database (default `strands_sessions`) |
|
|
46
|
+
| `collection_name` | Collection (default `sessions`) |
|
|
47
|
+
| `client` | Optional pre-built `pymongo.MongoClient` to reuse |
|
|
48
|
+
| `ttl_seconds` | If set, adds a TTL index so documents expire automatically |
|
|
49
|
+
|
|
50
|
+
## Data model
|
|
51
|
+
|
|
52
|
+
A single collection; each document has a `pk` + `sk` (compound-indexed, unique) and a `data` field with the serialized payload:
|
|
53
|
+
|
|
54
|
+
| Item | pk | sk |
|
|
55
|
+
|---|---|---|
|
|
56
|
+
| Session | `SESSION#<session_id>` | `META` |
|
|
57
|
+
| Agent | `SESSION#<session_id>` | `AGENT#<agent_id>` |
|
|
58
|
+
| Message | `SESSION#<session_id>#AGENT#<agent_id>` | `MSG#<zero-padded id>` |
|
|
59
|
+
|
|
60
|
+
The `(pk, sk)` index gives ordered range scans, so `list_messages(offset, limit)` is a simple indexed `find(...).sort(sk).limit(...)` — matching the `removed_message_count` offset semantics Strands relies on.
|
|
61
|
+
|
|
62
|
+
## License
|
|
63
|
+
|
|
64
|
+
MIT
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "strands_session_mongodb"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "MongoDB session manager (storage backend) for Strands Agents — persists sessions, agent state, and messages"
|
|
9
|
+
authors = [{name = "Kamal", email = "skamalj@github.com"}]
|
|
10
|
+
readme = "README.md"
|
|
11
|
+
requires-python = ">=3.10"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"strands-agents-session>=0.1.0",
|
|
14
|
+
"strands-agents",
|
|
15
|
+
"pymongo",
|
|
16
|
+
]
|
|
17
|
+
keywords = [
|
|
18
|
+
"strands", "strands-agents", "session", "mongodb",
|
|
19
|
+
"agent-state", "memory", "persistence",
|
|
20
|
+
]
|
|
21
|
+
classifiers = [
|
|
22
|
+
"Programming Language :: Python :: 3",
|
|
23
|
+
"License :: OSI Approved :: MIT License",
|
|
24
|
+
"Operating System :: OS Independent",
|
|
25
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
26
|
+
]
|
|
27
|
+
|
|
28
|
+
[project.urls]
|
|
29
|
+
Homepage = "https://github.com/skamalj/strands-agents-session"
|
|
30
|
+
Repository = "https://github.com/skamalj/strands-agents-session.git"
|
|
31
|
+
|
|
32
|
+
[tool.hatch.build.targets.wheel]
|
|
33
|
+
packages = ["src/strands_session_mongodb"]
|
|
@@ -0,0 +1,102 @@
|
|
|
1
|
+
"""MongoDB session backend for Strands Agents.
|
|
2
|
+
|
|
3
|
+
Provides a ``MongoDBSessionStorage`` (implementing the ``SessionStorage``
|
|
4
|
+
interface from ``strands-agents-session``) and a thin ``MongoDBSessionManager``
|
|
5
|
+
that wires it into ``KeyValueSessionManager``. All the session-repository logic
|
|
6
|
+
lives in the base package; this module is just the storage layer.
|
|
7
|
+
|
|
8
|
+
Records are stored in a single collection with a compound ``_id`` of the form
|
|
9
|
+
``{pk, sk}``, and a ``data`` field holding the serialized payload. A compound
|
|
10
|
+
index on ``(pk, sk)`` gives ordered range scans for ``list_messages``.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from typing import Any, Dict, List, Optional
|
|
14
|
+
|
|
15
|
+
from pymongo import ASCENDING, MongoClient
|
|
16
|
+
|
|
17
|
+
from strands_agents_session import KeyValueSessionManager, SessionStorage
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class MongoDBSessionStorage(SessionStorage):
|
|
21
|
+
"""MongoDB implementation of the session ``SessionStorage`` interface."""
|
|
22
|
+
|
|
23
|
+
def __init__(
|
|
24
|
+
self,
|
|
25
|
+
*,
|
|
26
|
+
connection_string: str = "mongodb://127.0.0.1:27017/",
|
|
27
|
+
database_name: str = "strands_sessions",
|
|
28
|
+
collection_name: str = "sessions",
|
|
29
|
+
client: Optional[MongoClient] = None,
|
|
30
|
+
ttl_seconds: Optional[int] = None,
|
|
31
|
+
) -> None:
|
|
32
|
+
self.client = client or MongoClient(connection_string)
|
|
33
|
+
self.collection = self.client[database_name][collection_name]
|
|
34
|
+
# Compound index on (pk, sk) — unique identity + ordered range scans.
|
|
35
|
+
self.collection.create_index([("pk", ASCENDING), ("sk", ASCENDING)], unique=True)
|
|
36
|
+
self.ttl_seconds = ttl_seconds
|
|
37
|
+
if ttl_seconds:
|
|
38
|
+
# MongoDB TTL index expires documents `ttl_seconds` after `created_at`.
|
|
39
|
+
self.collection.create_index("created_at", expireAfterSeconds=ttl_seconds)
|
|
40
|
+
|
|
41
|
+
# SessionStorage ----------------------------------------------------
|
|
42
|
+
|
|
43
|
+
def put(self, item: Dict[str, Any]) -> None:
|
|
44
|
+
doc = {"pk": item["pk"], "sk": item["sk"], "data": item["data"]}
|
|
45
|
+
if self.ttl_seconds:
|
|
46
|
+
import datetime
|
|
47
|
+
|
|
48
|
+
doc["created_at"] = datetime.datetime.now(datetime.timezone.utc)
|
|
49
|
+
self.collection.replace_one({"pk": item["pk"], "sk": item["sk"]}, doc, upsert=True)
|
|
50
|
+
|
|
51
|
+
def get(self, pk: str, sk: str) -> Optional[Dict[str, Any]]:
|
|
52
|
+
doc = self.collection.find_one({"pk": pk, "sk": sk})
|
|
53
|
+
if doc is None:
|
|
54
|
+
return None
|
|
55
|
+
return {"pk": doc["pk"], "sk": doc["sk"], "data": doc["data"]}
|
|
56
|
+
|
|
57
|
+
def query(
|
|
58
|
+
self, pk: str, sk_gte: Optional[str] = None, limit: Optional[int] = None
|
|
59
|
+
) -> List[Dict[str, Any]]:
|
|
60
|
+
query: Dict[str, Any] = {"pk": pk}
|
|
61
|
+
if sk_gte is not None:
|
|
62
|
+
query["sk"] = {"$gte": sk_gte}
|
|
63
|
+
cursor = self.collection.find(query).sort("sk", ASCENDING)
|
|
64
|
+
if limit is not None:
|
|
65
|
+
cursor = cursor.limit(limit)
|
|
66
|
+
return [{"pk": d["pk"], "sk": d["sk"], "data": d["data"]} for d in cursor]
|
|
67
|
+
|
|
68
|
+
def delete(self, pk: str, sk: str) -> None:
|
|
69
|
+
self.collection.delete_one({"pk": pk, "sk": sk})
|
|
70
|
+
|
|
71
|
+
def delete_partition(self, pk: str) -> None:
|
|
72
|
+
self.collection.delete_many({"pk": pk})
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
class MongoDBSessionManager(KeyValueSessionManager):
|
|
76
|
+
"""MongoDB-backed session manager for Strands Agents.
|
|
77
|
+
|
|
78
|
+
agent = Agent(session_manager=MongoDBSessionManager(
|
|
79
|
+
session_id="user-123",
|
|
80
|
+
connection_string="mongodb://127.0.0.1:27017/",
|
|
81
|
+
))
|
|
82
|
+
"""
|
|
83
|
+
|
|
84
|
+
def __init__(
|
|
85
|
+
self,
|
|
86
|
+
session_id: str,
|
|
87
|
+
*,
|
|
88
|
+
connection_string: str = "mongodb://127.0.0.1:27017/",
|
|
89
|
+
database_name: str = "strands_sessions",
|
|
90
|
+
collection_name: str = "sessions",
|
|
91
|
+
client: Optional[MongoClient] = None,
|
|
92
|
+
ttl_seconds: Optional[int] = None,
|
|
93
|
+
**kwargs: Any,
|
|
94
|
+
) -> None:
|
|
95
|
+
storage = MongoDBSessionStorage(
|
|
96
|
+
connection_string=connection_string,
|
|
97
|
+
database_name=database_name,
|
|
98
|
+
collection_name=collection_name,
|
|
99
|
+
client=client,
|
|
100
|
+
ttl_seconds=ttl_seconds,
|
|
101
|
+
)
|
|
102
|
+
super().__init__(session_id=session_id, storage=storage)
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
"""
|
|
2
|
+
End-to-end tests: a real Strands Agent (OpenAI gpt-4o-mini) persisting to a real
|
|
3
|
+
MongoDB via MongoDBSessionManager, then restoring in a fresh manager.
|
|
4
|
+
|
|
5
|
+
Requires:
|
|
6
|
+
- A reachable MongoDB (MONGO_URI, default mongodb://127.0.0.1:27017/)
|
|
7
|
+
- OPENAI_API_KEY
|
|
8
|
+
"""
|
|
9
|
+
import os
|
|
10
|
+
import uuid
|
|
11
|
+
|
|
12
|
+
import pytest
|
|
13
|
+
|
|
14
|
+
from strands import Agent
|
|
15
|
+
from strands.models.openai import OpenAIModel
|
|
16
|
+
|
|
17
|
+
from strands_session_mongodb import MongoDBSessionManager
|
|
18
|
+
|
|
19
|
+
MONGO_URI = os.environ.get("MONGO_URI", "mongodb://127.0.0.1:27017/")
|
|
20
|
+
DB = os.environ.get("MONGO_TEST_DB", "strands_sessions_test")
|
|
21
|
+
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
|
|
22
|
+
|
|
23
|
+
pytestmark = pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY not set")
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def make_model():
|
|
27
|
+
return OpenAIModel(
|
|
28
|
+
client_args={"api_key": OPENAI_API_KEY},
|
|
29
|
+
model_id="gpt-4o-mini",
|
|
30
|
+
params={"temperature": 0},
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def make_manager(session_id):
|
|
35
|
+
return MongoDBSessionManager(
|
|
36
|
+
session_id=session_id, connection_string=MONGO_URI, database_name=DB
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def text_of(result):
|
|
41
|
+
try:
|
|
42
|
+
return " ".join(
|
|
43
|
+
b.get("text", "") for b in result.message["content"] if isinstance(b, dict)
|
|
44
|
+
)
|
|
45
|
+
except Exception:
|
|
46
|
+
return str(result)
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def test_conversation_persists_and_restores_across_managers():
|
|
50
|
+
session_id = f"llm-{uuid.uuid4()}"
|
|
51
|
+
|
|
52
|
+
a1 = Agent(model=make_model(), agent_id="assistant", session_manager=make_manager(session_id))
|
|
53
|
+
a1("Hi, my name is Kamal and I live in Pune.")
|
|
54
|
+
|
|
55
|
+
a2 = Agent(model=make_model(), agent_id="assistant", session_manager=make_manager(session_id))
|
|
56
|
+
reply = text_of(a2("What is my name?"))
|
|
57
|
+
assert "kamal" in reply.lower(), f"expected restored recall of 'kamal', got: {reply}"
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def test_messages_written_to_repository():
|
|
61
|
+
session_id = f"llm-{uuid.uuid4()}"
|
|
62
|
+
m1 = make_manager(session_id)
|
|
63
|
+
a1 = Agent(model=make_model(), agent_id="assistant", session_manager=m1)
|
|
64
|
+
a1("Say hello in one word.")
|
|
65
|
+
|
|
66
|
+
msgs = m1.list_messages(session_id, "assistant")
|
|
67
|
+
assert len(msgs) >= 2
|
|
68
|
+
assert [mm.message_id for mm in msgs] == sorted(mm.message_id for mm in msgs)
|
|
69
|
+
assert msgs[0].message["role"] == "user"
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def test_agent_record_persisted():
|
|
73
|
+
session_id = f"llm-{uuid.uuid4()}"
|
|
74
|
+
m1 = make_manager(session_id)
|
|
75
|
+
a1 = Agent(model=make_model(), agent_id="assistant", session_manager=m1)
|
|
76
|
+
a1("Hello there.")
|
|
77
|
+
|
|
78
|
+
got = m1.read_agent(session_id, "assistant")
|
|
79
|
+
assert got is not None
|
|
80
|
+
assert got.agent_id == "assistant"
|
|
81
|
+
assert isinstance(got.conversation_manager_state, dict)
|
|
@@ -0,0 +1,149 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Comprehensive repository tests for MongoDBSessionManager against a real MongoDB.
|
|
3
|
+
|
|
4
|
+
Deterministic (no LLM): exercises the SessionRepository CRUD surface — sessions,
|
|
5
|
+
agents (with agent state), messages, pagination, updates/redaction, isolation.
|
|
6
|
+
|
|
7
|
+
Requires a reachable MongoDB (default mongodb://127.0.0.1:27017/, override with
|
|
8
|
+
MONGO_URI).
|
|
9
|
+
"""
|
|
10
|
+
import os
|
|
11
|
+
import uuid
|
|
12
|
+
|
|
13
|
+
import pytest
|
|
14
|
+
|
|
15
|
+
from strands.types.exceptions import SessionException
|
|
16
|
+
from strands.types.session import SessionAgent, SessionMessage
|
|
17
|
+
from strands_session_mongodb import MongoDBSessionManager
|
|
18
|
+
|
|
19
|
+
MONGO_URI = os.environ.get("MONGO_URI", "mongodb://127.0.0.1:27017/")
|
|
20
|
+
DB = os.environ.get("MONGO_TEST_DB", "strands_sessions_test")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def make_manager(session_id=None):
|
|
24
|
+
return MongoDBSessionManager(
|
|
25
|
+
session_id=session_id or f"sess-{uuid.uuid4()}",
|
|
26
|
+
connection_string=MONGO_URI,
|
|
27
|
+
database_name=DB,
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def msg(i, role="user", text=None):
|
|
32
|
+
return SessionMessage.from_message(
|
|
33
|
+
{"role": role, "content": [{"text": text or f"m{i}"}]}, i
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
# sessions ------------------------------------------------------------------
|
|
38
|
+
|
|
39
|
+
def test_session_created_on_init():
|
|
40
|
+
mgr = make_manager()
|
|
41
|
+
got = mgr.read_session(mgr.session_id)
|
|
42
|
+
assert got is not None and got.session_id == mgr.session_id
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def test_create_duplicate_session_raises():
|
|
46
|
+
mgr = make_manager()
|
|
47
|
+
with pytest.raises(SessionException):
|
|
48
|
+
mgr.create_session(mgr.session)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def test_read_missing_session_returns_none():
|
|
52
|
+
assert make_manager().read_session(f"nope-{uuid.uuid4()}") is None
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def test_delete_session_removes_everything():
|
|
56
|
+
mgr = make_manager()
|
|
57
|
+
sid = mgr.session_id
|
|
58
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
59
|
+
mgr.create_message(sid, "a1", msg(0))
|
|
60
|
+
mgr.delete_session(sid)
|
|
61
|
+
assert mgr.read_session(sid) is None
|
|
62
|
+
assert mgr.read_agent(sid, "a1") is None
|
|
63
|
+
assert mgr.list_messages(sid, "a1") == []
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
# agents (+ state) ----------------------------------------------------------
|
|
67
|
+
|
|
68
|
+
def test_agent_state_round_trip():
|
|
69
|
+
mgr = make_manager()
|
|
70
|
+
mgr.create_agent(mgr.session_id, SessionAgent(
|
|
71
|
+
agent_id="assistant", state={"tier": "gold", "lang": "en"},
|
|
72
|
+
conversation_manager_state={"removed_message_count": 0}))
|
|
73
|
+
got = mgr.read_agent(mgr.session_id, "assistant")
|
|
74
|
+
assert got.state == {"tier": "gold", "lang": "en"}
|
|
75
|
+
assert got.conversation_manager_state == {"removed_message_count": 0}
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def test_read_missing_agent_returns_none():
|
|
79
|
+
mgr = make_manager()
|
|
80
|
+
assert mgr.read_agent(mgr.session_id, "ghost") is None
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def test_update_agent_preserves_created_at():
|
|
84
|
+
mgr = make_manager()
|
|
85
|
+
sid = mgr.session_id
|
|
86
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={"v": 1}, conversation_manager_state={}))
|
|
87
|
+
created_at = mgr.read_agent(sid, "a1").created_at
|
|
88
|
+
mgr.update_agent(sid, SessionAgent(agent_id="a1", state={"v": 2}, conversation_manager_state={}))
|
|
89
|
+
got = mgr.read_agent(sid, "a1")
|
|
90
|
+
assert got.state == {"v": 2}
|
|
91
|
+
assert got.created_at == created_at
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def test_update_missing_agent_raises():
|
|
95
|
+
mgr = make_manager()
|
|
96
|
+
with pytest.raises(SessionException):
|
|
97
|
+
mgr.update_agent(mgr.session_id, SessionAgent(agent_id="x", state={}, conversation_manager_state={}))
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
# messages + pagination -----------------------------------------------------
|
|
101
|
+
|
|
102
|
+
def test_message_round_trip():
|
|
103
|
+
mgr = make_manager()
|
|
104
|
+
sid = mgr.session_id
|
|
105
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
106
|
+
mgr.create_message(sid, "a1", msg(0, text="hello"))
|
|
107
|
+
got = mgr.read_message(sid, "a1", 0)
|
|
108
|
+
assert got.message["content"][0]["text"] == "hello"
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def test_list_messages_sorted():
|
|
112
|
+
mgr = make_manager()
|
|
113
|
+
sid = mgr.session_id
|
|
114
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
115
|
+
for i in [2, 0, 4, 1, 3]:
|
|
116
|
+
mgr.create_message(sid, "a1", msg(i))
|
|
117
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1")] == [0, 1, 2, 3, 4]
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def test_list_messages_offset_and_limit():
|
|
121
|
+
mgr = make_manager()
|
|
122
|
+
sid = mgr.session_id
|
|
123
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
124
|
+
for i in range(10):
|
|
125
|
+
mgr.create_message(sid, "a1", msg(i))
|
|
126
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1", offset=2)] == [2, 3, 4, 5, 6, 7, 8, 9]
|
|
127
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1", limit=3)] == [0, 1, 2]
|
|
128
|
+
assert [m.message_id for m in mgr.list_messages(sid, "a1", offset=3, limit=4)] == [3, 4, 5, 6]
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
def test_update_message_redaction():
|
|
132
|
+
mgr = make_manager()
|
|
133
|
+
sid = mgr.session_id
|
|
134
|
+
mgr.create_agent(sid, SessionAgent(agent_id="a1", state={}, conversation_manager_state={}))
|
|
135
|
+
m = msg(0, text="secret 4271")
|
|
136
|
+
mgr.create_message(sid, "a1", m)
|
|
137
|
+
m.redact_message = {"role": "user", "content": [{"text": "[REDACTED]"}]}
|
|
138
|
+
mgr.update_message(sid, "a1", m)
|
|
139
|
+
got = mgr.read_message(sid, "a1", 0)
|
|
140
|
+
assert got.redact_message["content"][0]["text"] == "[REDACTED]"
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def test_sessions_isolated():
|
|
144
|
+
m1 = make_manager()
|
|
145
|
+
m2 = make_manager()
|
|
146
|
+
m1.create_agent(m1.session_id, SessionAgent(agent_id="a", state={"who": "one"}, conversation_manager_state={}))
|
|
147
|
+
m2.create_agent(m2.session_id, SessionAgent(agent_id="a", state={"who": "two"}, conversation_manager_state={}))
|
|
148
|
+
assert m1.read_agent(m1.session_id, "a").state == {"who": "one"}
|
|
149
|
+
assert m2.read_agent(m2.session_id, "a").state == {"who": "two"}
|