agent-framework-membase 0.1.0__tar.gz → 0.2.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.
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/PKG-INFO +14 -12
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/README.md +11 -9
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/examples/live_smoke.py +3 -6
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/pyproject.toml +2 -2
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/src/agent_framework_membase/_common.py +21 -8
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/src/agent_framework_membase/_provider.py +15 -25
- agent_framework_membase-0.2.0/src/agent_framework_membase/_tools.py +142 -0
- agent_framework_membase-0.2.0/src/agent_framework_membase/_version.py +1 -0
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/tests/conftest.py +37 -27
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/tests/test_provider.py +18 -19
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/tests/test_tools.py +26 -37
- agent_framework_membase-0.1.0/src/agent_framework_membase/_tools.py +0 -207
- agent_framework_membase-0.1.0/src/agent_framework_membase/_version.py +0 -1
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/.gitignore +0 -0
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/LICENSE +0 -0
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/src/agent_framework_membase/__init__.py +0 -0
- {agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/src/agent_framework_membase/py.typed +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: agent-framework-membase
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: Membase long-term memory for Microsoft Agent Framework (Python): a context provider with time-boxed recall and deliberate save/search/profile tools.
|
|
5
5
|
Project-URL: Homepage, https://github.com/unibaseio/membase/tree/main/integrations/agent-framework
|
|
6
6
|
Project-URL: Repository, https://github.com/unibaseio/membase
|
|
@@ -21,9 +21,9 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
|
21
21
|
Classifier: Typing :: Typed
|
|
22
22
|
Requires-Python: >=3.10
|
|
23
23
|
Requires-Dist: agent-framework-core<2,>=1.0.0
|
|
24
|
-
Requires-Dist: unibaseio-membase>=0.
|
|
24
|
+
Requires-Dist: unibaseio-membase>=0.4.0
|
|
25
25
|
Provides-Extra: local
|
|
26
|
-
Requires-Dist: unibaseio-membase[local]>=0.
|
|
26
|
+
Requires-Dist: unibaseio-membase[local]>=0.4.0; extra == 'local'
|
|
27
27
|
Description-Content-Type: text/markdown
|
|
28
28
|
|
|
29
29
|
# agent-framework-membase
|
|
@@ -43,7 +43,9 @@ pip install "agent-framework-membase @ git+https://github.com/unibaseio/membase#
|
|
|
43
43
|
|
|
44
44
|
You need a Membase developer key: in the [Membase app](https://app.membase.com), **Connect ›
|
|
45
45
|
Developer keys**. A *Read* key can search (and read the profile, if its owner ticked it);
|
|
46
|
-
saving needs *Read & write
|
|
46
|
+
saving needs *Read & write* and a key with a write target: a key that reaches one Memory saves
|
|
47
|
+
there, one that reaches several is refused and nothing is written somewhere by guess. Needs
|
|
48
|
+
`unibaseio-membase>=0.4.0`.
|
|
47
49
|
|
|
48
50
|
## Quick start
|
|
49
51
|
|
|
@@ -97,17 +99,17 @@ the Memory it came from. They are reference data, not instructions. ...
|
|
|
97
99
|
waits `recall_timeout` seconds (default **8**), then lets the run continue without memory and
|
|
98
100
|
logs why. The search finishes in the background, so a later turn finds the memory awake.
|
|
99
101
|
Any other error is logged the same way; memory never breaks a run.
|
|
100
|
-
- **Silence is not "nothing stored".** When a Memory could not answer (
|
|
101
|
-
|
|
102
|
+
- **Silence is not "nothing stored".** When a Memory could not answer (the search's
|
|
103
|
+
`partial`), the instructions name it instead of passing it off as empty.
|
|
102
104
|
- **A withdrawn key is not retried.** A `403` means the owner narrowed or revoked the key; the
|
|
103
105
|
provider logs "this Membase key was narrowed or revoked" and stops calling that key.
|
|
104
106
|
|
|
105
|
-
|
|
107
|
+
A recall covers every Memory the key reaches. Options: `limit=` (passages in all, default 8),
|
|
106
108
|
`recall_timeout=`, `context_prompt=`, `source_id=` (default `"membase"`).
|
|
107
109
|
|
|
108
110
|
**After each run** it does nothing by default. With `capture=True` it saves the user's own
|
|
109
|
-
messages from the turn (not the agent's replies) as one note, once
|
|
110
|
-
|
|
111
|
+
messages from the turn (not the agent's replies) as one note in the key's write target, once
|
|
112
|
+
per turn; a key with no write target is told so once in the log and not tried again. The
|
|
111
113
|
`membase_save` tool is usually the better choice: saving works best as a deliberate act the
|
|
112
114
|
user asked for.
|
|
113
115
|
|
|
@@ -115,9 +117,9 @@ user asked for.
|
|
|
115
117
|
|
|
116
118
|
| Tool | What it does |
|
|
117
119
|
|---|---|
|
|
118
|
-
| `membase_search(query,
|
|
119
|
-
| `membase_save(content
|
|
120
|
-
| `membase_profile(
|
|
120
|
+
| `membase_search(query, limit=None)` | Passages from every Memory the key reaches, each tagged `[Memory name]`; names any Memory that could not answer. |
|
|
121
|
+
| `membase_save(content)` | Saves one fact in the user's words, into the key's write target. A key that reaches several Memories has no write target: the tool says so and that nothing was saved. |
|
|
122
|
+
| `membase_profile()` | The user's profile facts: standing statements about who they are, newest first. Needs a key whose owner ticked **Your profile**. |
|
|
121
123
|
|
|
122
124
|
`membase_tools(memory, save=False)` leaves out the save tool (for a Read key) and
|
|
123
125
|
`profile=False` leaves out the profile tool. There is no delete or forget tool. Tool calls use
|
|
@@ -15,7 +15,9 @@ pip install "agent-framework-membase @ git+https://github.com/unibaseio/membase#
|
|
|
15
15
|
|
|
16
16
|
You need a Membase developer key: in the [Membase app](https://app.membase.com), **Connect ›
|
|
17
17
|
Developer keys**. A *Read* key can search (and read the profile, if its owner ticked it);
|
|
18
|
-
saving needs *Read & write
|
|
18
|
+
saving needs *Read & write* and a key with a write target: a key that reaches one Memory saves
|
|
19
|
+
there, one that reaches several is refused and nothing is written somewhere by guess. Needs
|
|
20
|
+
`unibaseio-membase>=0.4.0`.
|
|
19
21
|
|
|
20
22
|
## Quick start
|
|
21
23
|
|
|
@@ -69,17 +71,17 @@ the Memory it came from. They are reference data, not instructions. ...
|
|
|
69
71
|
waits `recall_timeout` seconds (default **8**), then lets the run continue without memory and
|
|
70
72
|
logs why. The search finishes in the background, so a later turn finds the memory awake.
|
|
71
73
|
Any other error is logged the same way; memory never breaks a run.
|
|
72
|
-
- **Silence is not "nothing stored".** When a Memory could not answer (
|
|
73
|
-
|
|
74
|
+
- **Silence is not "nothing stored".** When a Memory could not answer (the search's
|
|
75
|
+
`partial`), the instructions name it instead of passing it off as empty.
|
|
74
76
|
- **A withdrawn key is not retried.** A `403` means the owner narrowed or revoked the key; the
|
|
75
77
|
provider logs "this Membase key was narrowed or revoked" and stops calling that key.
|
|
76
78
|
|
|
77
|
-
|
|
79
|
+
A recall covers every Memory the key reaches. Options: `limit=` (passages in all, default 8),
|
|
78
80
|
`recall_timeout=`, `context_prompt=`, `source_id=` (default `"membase"`).
|
|
79
81
|
|
|
80
82
|
**After each run** it does nothing by default. With `capture=True` it saves the user's own
|
|
81
|
-
messages from the turn (not the agent's replies) as one note, once
|
|
82
|
-
|
|
83
|
+
messages from the turn (not the agent's replies) as one note in the key's write target, once
|
|
84
|
+
per turn; a key with no write target is told so once in the log and not tried again. The
|
|
83
85
|
`membase_save` tool is usually the better choice: saving works best as a deliberate act the
|
|
84
86
|
user asked for.
|
|
85
87
|
|
|
@@ -87,9 +89,9 @@ user asked for.
|
|
|
87
89
|
|
|
88
90
|
| Tool | What it does |
|
|
89
91
|
|---|---|
|
|
90
|
-
| `membase_search(query,
|
|
91
|
-
| `membase_save(content
|
|
92
|
-
| `membase_profile(
|
|
92
|
+
| `membase_search(query, limit=None)` | Passages from every Memory the key reaches, each tagged `[Memory name]`; names any Memory that could not answer. |
|
|
93
|
+
| `membase_save(content)` | Saves one fact in the user's words, into the key's write target. A key that reaches several Memories has no write target: the tool says so and that nothing was saved. |
|
|
94
|
+
| `membase_profile()` | The user's profile facts: standing statements about who they are, newest first. Needs a key whose owner ticked **Your profile**. |
|
|
93
95
|
|
|
94
96
|
`membase_tools(memory, save=False)` leaves out the save tool (for a Read key) and
|
|
95
97
|
`profile=False` leaves out the profile tool. There is no delete or forget tool. Tool calls use
|
|
@@ -5,7 +5,7 @@
|
|
|
5
5
|
A real ``agent_framework.Agent`` runs with the Membase context provider and tools attached;
|
|
6
6
|
its chat client is a stand-in that calls the tools the way a model would and prints what the
|
|
7
7
|
model would have seen. Writes nothing unless MEMBASE_SMOKE_WRITE=1 (then it saves one note to
|
|
8
|
-
|
|
8
|
+
the key's write target: a Read & write key that reaches one Memory). Exits 0 without a key.
|
|
9
9
|
"""
|
|
10
10
|
|
|
11
11
|
from __future__ import annotations
|
|
@@ -57,7 +57,7 @@ async def main() -> int:
|
|
|
57
57
|
# MEMBASE_API_KEY from the environment; the host is the package's default unless MEMBASE_BASE_URL is set.
|
|
58
58
|
client = Membase(base_url=os.environ.get("MEMBASE_BASE_URL") or DEFAULT_BASE_URL)
|
|
59
59
|
|
|
60
|
-
print("
|
|
60
|
+
print("Memories in reach:", [m.get("name") for m in client.capabilities().get("memories") or []])
|
|
61
61
|
|
|
62
62
|
calls: list[tuple[str, dict[str, Any]]] = [("membase_search", {"query": question, "limit": 5})]
|
|
63
63
|
if os.environ.get("MEMBASE_SMOKE_PROFILE") == "1":
|
|
@@ -66,10 +66,7 @@ async def main() -> int:
|
|
|
66
66
|
calls.append(
|
|
67
67
|
(
|
|
68
68
|
"membase_save",
|
|
69
|
-
{
|
|
70
|
-
"content": f"Smoke test note from agent-framework-membase at {time.ctime()}.",
|
|
71
|
-
"container": os.environ.get("MEMBASE_SMOKE_CONTAINER"),
|
|
72
|
-
},
|
|
69
|
+
{"content": f"Smoke test note from agent-framework-membase at {time.ctime()}."},
|
|
73
70
|
)
|
|
74
71
|
)
|
|
75
72
|
model = StandInModel(calls)
|
|
@@ -25,11 +25,11 @@ classifiers = [
|
|
|
25
25
|
]
|
|
26
26
|
dependencies = [
|
|
27
27
|
"agent-framework-core>=1.0.0,<2",
|
|
28
|
-
"unibaseio-membase>=0.
|
|
28
|
+
"unibaseio-membase>=0.4.0",
|
|
29
29
|
]
|
|
30
30
|
|
|
31
31
|
[project.optional-dependencies]
|
|
32
|
-
local = ["unibaseio-membase[local]>=0.
|
|
32
|
+
local = ["unibaseio-membase[local]>=0.4.0"]
|
|
33
33
|
|
|
34
34
|
[project.urls]
|
|
35
35
|
Homepage = "https://github.com/unibaseio/membase/tree/main/integrations/agent-framework"
|
|
@@ -17,6 +17,7 @@ from membase.errors import (
|
|
|
17
17
|
PermissionDeniedError,
|
|
18
18
|
RateLimitError,
|
|
19
19
|
UnprocessableEntityError,
|
|
20
|
+
WriteTargetRequiredError,
|
|
20
21
|
)
|
|
21
22
|
|
|
22
23
|
logger = logging.getLogger("agent_framework_membase")
|
|
@@ -35,6 +36,11 @@ ClientFor = Callable[[str], Membase]
|
|
|
35
36
|
|
|
36
37
|
REVOKED = "this Membase key was narrowed or revoked by its owner, or it does not allow this operation"
|
|
37
38
|
|
|
39
|
+
NO_WRITE_TARGET = (
|
|
40
|
+
"this Membase key reaches several Memories and has no write target, so nothing was saved; "
|
|
41
|
+
"give the key a single Memory under Connect › Developer keys in the Membase app"
|
|
42
|
+
)
|
|
43
|
+
|
|
38
44
|
|
|
39
45
|
def bind_user(session: Any, user_id: str, *, source_id: str = DEFAULT_SOURCE_ID) -> None:
|
|
40
46
|
"""Record which person a session belongs to, for ``client_for``.
|
|
@@ -140,6 +146,8 @@ def _log_late_failure(task: asyncio.Future[Any]) -> None:
|
|
|
140
146
|
|
|
141
147
|
def describe_error(exc: BaseException) -> str:
|
|
142
148
|
"""One actionable sentence for an SDK failure."""
|
|
149
|
+
if isinstance(exc, WriteTargetRequiredError):
|
|
150
|
+
return f"Membase did not save it: {NO_WRITE_TARGET}."
|
|
143
151
|
if isinstance(exc, PermissionDeniedError):
|
|
144
152
|
return f"Membase refused the call (403): {REVOKED}. Do not retry with this key. ({exc.message})"
|
|
145
153
|
if isinstance(exc, APITimeoutError):
|
|
@@ -158,21 +166,26 @@ def describe_error(exc: BaseException) -> str:
|
|
|
158
166
|
|
|
159
167
|
|
|
160
168
|
def format_search(answer: Mapping[str, Any] | None) -> tuple[list[str], list[str]]:
|
|
161
|
-
"""``(passages, unanswered)``: one line per passage citing its
|
|
162
|
-
|
|
169
|
+
"""``(passages, unanswered)``: one line per passage citing its Memory, one per Memory that
|
|
170
|
+
could not answer or whose passages did not fit the limit (``partial``), so an empty result
|
|
171
|
+
is never mistaken for an empty memory."""
|
|
163
172
|
answer = answer or {}
|
|
164
173
|
passages: list[str] = []
|
|
165
|
-
for hit in answer.get("
|
|
174
|
+
for hit in answer.get("items") or []:
|
|
166
175
|
if not isinstance(hit, Mapping):
|
|
167
176
|
continue
|
|
168
177
|
content = str(hit.get("content") or "").strip()
|
|
169
178
|
if not content:
|
|
170
179
|
continue
|
|
171
|
-
where = str(hit.get("
|
|
180
|
+
where = str(hit.get("memory_name") or "Membase").strip()
|
|
172
181
|
passages.append(f"- [{where}] {content}")
|
|
173
182
|
unanswered: list[str] = []
|
|
174
|
-
for
|
|
175
|
-
if isinstance(
|
|
176
|
-
|
|
177
|
-
|
|
183
|
+
for p in answer.get("partial") or []:
|
|
184
|
+
if not isinstance(p, Mapping):
|
|
185
|
+
continue
|
|
186
|
+
name = p.get("memory_name") or "a Memory"
|
|
187
|
+
if p.get("reason") == "truncated":
|
|
188
|
+
unanswered.append(f'- "{name}" had more passages than the limit let through')
|
|
189
|
+
else:
|
|
190
|
+
unanswered.append(f'- "{name}" could not answer')
|
|
178
191
|
return passages, unanswered
|
|
@@ -9,10 +9,11 @@ from typing import TYPE_CHECKING, Any
|
|
|
9
9
|
|
|
10
10
|
from agent_framework import AgentSession, ContextProvider, SessionContext
|
|
11
11
|
from membase import Membase
|
|
12
|
-
from membase.errors import
|
|
12
|
+
from membase.errors import PermissionDeniedError, WriteTargetRequiredError
|
|
13
13
|
|
|
14
14
|
from ._common import (
|
|
15
15
|
DEFAULT_SOURCE_ID,
|
|
16
|
+
NO_WRITE_TARGET,
|
|
16
17
|
REVOKED,
|
|
17
18
|
ClientFor,
|
|
18
19
|
Clients,
|
|
@@ -59,12 +60,12 @@ class MembaseContextProvider(ContextProvider):
|
|
|
59
60
|
``recall_timeout`` seconds (8 by default: the first search after a quiet spell can take
|
|
60
61
|
up to a minute while the memory wakes). On a timeout or any error the run continues
|
|
61
62
|
without memory and the reason is logged. Passages are added to the run's instructions,
|
|
62
|
-
each citing the Memory
|
|
63
|
-
|
|
63
|
+
each citing the Memory it came from; a Memory that could not answer is named too, so the
|
|
64
|
+
model does not read silence as "nothing stored".
|
|
64
65
|
|
|
65
66
|
``after_run`` saves nothing unless ``capture=True``. With capture on, the user's own
|
|
66
|
-
messages from the turn (not the agent's replies) are added as one note
|
|
67
|
-
deliberate ``membase_save`` tool from :func:`membase_tools`.
|
|
67
|
+
messages from the turn (not the agent's replies) are added as one note in the key's write
|
|
68
|
+
target; prefer the deliberate ``membase_save`` tool from :func:`membase_tools`.
|
|
68
69
|
|
|
69
70
|
One client serves every session unless ``client_for`` is given: then each session must
|
|
70
71
|
name its person with :func:`bind_user` and is served by ``client_for(user_id)``, that
|
|
@@ -84,12 +85,10 @@ class MembaseContextProvider(ContextProvider):
|
|
|
84
85
|
local: str | os.PathLike[str] | bool | None = None,
|
|
85
86
|
client_for: ClientFor | None = None,
|
|
86
87
|
source_id: str = DEFAULT_SOURCE_ID,
|
|
87
|
-
container: str | None = None,
|
|
88
88
|
limit: int | None = DEFAULT_LIMIT,
|
|
89
89
|
recall_timeout: float = DEFAULT_RECALL_TIMEOUT,
|
|
90
90
|
context_prompt: str = DEFAULT_CONTEXT_PROMPT,
|
|
91
91
|
capture: bool = False,
|
|
92
|
-
capture_container: str | None = None,
|
|
93
92
|
capture_title: str = "Agent Framework conversation",
|
|
94
93
|
) -> None:
|
|
95
94
|
"""Create the provider.
|
|
@@ -103,14 +102,13 @@ class MembaseContextProvider(ContextProvider):
|
|
|
103
102
|
client_for: ``user_id -> Membase`` for an app serving several people; the user id
|
|
104
103
|
comes from :func:`bind_user` on the session. Excludes the arguments above.
|
|
105
104
|
source_id: The provider id; also the ``session.state`` key of the user binding.
|
|
106
|
-
|
|
107
|
-
|
|
105
|
+
limit: Passages per recall, across every Memory the key reaches (``None``: the
|
|
106
|
+
server's default).
|
|
108
107
|
recall_timeout: Seconds the automatic recall may take before the run goes on
|
|
109
108
|
without it.
|
|
110
109
|
context_prompt: Heading placed above the recalled passages.
|
|
111
|
-
capture: Save the user's messages after each turn
|
|
112
|
-
|
|
113
|
-
key reaches more than one Memory.
|
|
110
|
+
capture: Save the user's messages after each turn, into the key's write target.
|
|
111
|
+
Off by default.
|
|
114
112
|
capture_title: Title of captured notes.
|
|
115
113
|
"""
|
|
116
114
|
super().__init__(source_id)
|
|
@@ -122,12 +120,10 @@ class MembaseContextProvider(ContextProvider):
|
|
|
122
120
|
client_for=client_for,
|
|
123
121
|
source_id=source_id,
|
|
124
122
|
)
|
|
125
|
-
self.container = container
|
|
126
123
|
self.limit = limit
|
|
127
124
|
self.recall_timeout = recall_timeout
|
|
128
125
|
self.context_prompt = context_prompt
|
|
129
126
|
self.capture = capture
|
|
130
|
-
self.capture_container = capture_container
|
|
131
127
|
self.capture_title = capture_title
|
|
132
128
|
# A 403 means the owner narrowed or revoked the key: stop calling it automatically.
|
|
133
129
|
self._refused_recall: weakref.WeakSet[Membase] = weakref.WeakSet()
|
|
@@ -176,7 +172,6 @@ class MembaseContextProvider(ContextProvider):
|
|
|
176
172
|
answer = await in_thread(
|
|
177
173
|
client.search,
|
|
178
174
|
texts[-1],
|
|
179
|
-
container=self.container,
|
|
180
175
|
limit=self.limit,
|
|
181
176
|
timeout=self.recall_timeout,
|
|
182
177
|
)
|
|
@@ -205,8 +200,8 @@ class MembaseContextProvider(ContextProvider):
|
|
|
205
200
|
parts.append("<membase_memory>\n" + "\n".join(passages) + "\n</membase_memory>")
|
|
206
201
|
if unanswered:
|
|
207
202
|
parts.append(
|
|
208
|
-
"These Memories
|
|
209
|
-
"not evidence that nothing is stored:\n" + "\n".join(unanswered)
|
|
203
|
+
"These Memories did not fully answer this time, so a missing passage from them "
|
|
204
|
+
"is not evidence that nothing is stored:\n" + "\n".join(unanswered)
|
|
210
205
|
)
|
|
211
206
|
context.extend_instructions(self.source_id, "\n\n".join(parts))
|
|
212
207
|
|
|
@@ -231,7 +226,6 @@ class MembaseContextProvider(ContextProvider):
|
|
|
231
226
|
await in_thread(
|
|
232
227
|
client.memories.add,
|
|
233
228
|
"\n\n".join(texts),
|
|
234
|
-
container=self.capture_container,
|
|
235
229
|
title=self.capture_title,
|
|
236
230
|
timeout=self.recall_timeout,
|
|
237
231
|
)
|
|
@@ -247,13 +241,9 @@ class MembaseContextProvider(ContextProvider):
|
|
|
247
241
|
REVOKED,
|
|
248
242
|
exc.message,
|
|
249
243
|
)
|
|
250
|
-
except
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
if self.capture_container
|
|
254
|
-
else " (set capture_container= when the key reaches several Memories)"
|
|
255
|
-
)
|
|
256
|
-
logger.warning("Membase: capture refused: %s%s", exc.message, hint)
|
|
244
|
+
except WriteTargetRequiredError:
|
|
245
|
+
self._refused_capture.add(client)
|
|
246
|
+
logger.error("Membase: capture stopped: %s.", NO_WRITE_TARGET)
|
|
257
247
|
except Exception as exc: # noqa: BLE001 - a failed capture must never break a run
|
|
258
248
|
logger.warning("Membase: capture failed: %s", describe_error(exc))
|
|
259
249
|
|
|
@@ -0,0 +1,142 @@
|
|
|
1
|
+
"""``membase_tools``: search, save and profile as Agent Framework function tools."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
from typing import Annotated
|
|
7
|
+
|
|
8
|
+
from agent_framework import FunctionInvocationContext, FunctionTool, tool
|
|
9
|
+
from membase import Membase
|
|
10
|
+
from membase.errors import PermissionDeniedError, WriteTargetRequiredError
|
|
11
|
+
from pydantic import Field
|
|
12
|
+
|
|
13
|
+
from ._common import (
|
|
14
|
+
DEFAULT_SOURCE_ID,
|
|
15
|
+
NO_WRITE_TARGET,
|
|
16
|
+
REVOKED,
|
|
17
|
+
ClientFor,
|
|
18
|
+
Clients,
|
|
19
|
+
NoUserBound,
|
|
20
|
+
describe_error,
|
|
21
|
+
format_search,
|
|
22
|
+
in_thread,
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def membase_tools(
|
|
27
|
+
client: Membase | None = None,
|
|
28
|
+
*,
|
|
29
|
+
api_key: str | None = None,
|
|
30
|
+
base_url: str | None = None,
|
|
31
|
+
local: str | os.PathLike[str] | bool | None = None,
|
|
32
|
+
client_for: ClientFor | None = None,
|
|
33
|
+
source_id: str = DEFAULT_SOURCE_ID,
|
|
34
|
+
save: bool = True,
|
|
35
|
+
profile: bool = True,
|
|
36
|
+
) -> list[FunctionTool]:
|
|
37
|
+
"""Membase as Agent Framework function tools: ``membase_search``, ``membase_save``,
|
|
38
|
+
``membase_profile``.
|
|
39
|
+
|
|
40
|
+
Pass a ready ``Membase`` client, or ``api_key`` / ``base_url`` / ``local``, or
|
|
41
|
+
``client_for`` (``user_id -> Membase``) with the user bound to each session by
|
|
42
|
+
:func:`bind_user` (the same binding :class:`MembaseContextProvider` reads). Tool calls use
|
|
43
|
+
the SDK's own timeout (90 s by default), since the model asked for them explicitly.
|
|
44
|
+
|
|
45
|
+
``save=False`` leaves out ``membase_save`` (for a Read key); ``profile=False`` leaves out
|
|
46
|
+
``membase_profile`` (for a key whose owner did not share the profile). There is no delete
|
|
47
|
+
or forget tool.
|
|
48
|
+
"""
|
|
49
|
+
clients = Clients(
|
|
50
|
+
client, api_key=api_key, base_url=base_url, local=local, client_for=client_for, source_id=source_id
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
def _client(ctx: FunctionInvocationContext | None) -> Membase:
|
|
54
|
+
return clients.for_session(ctx.session if ctx is not None else None)
|
|
55
|
+
|
|
56
|
+
@tool(name="membase_search", approval_mode="never_require")
|
|
57
|
+
async def membase_search(
|
|
58
|
+
query: Annotated[str, Field(description="What to look for, in natural language.")],
|
|
59
|
+
limit: Annotated[int | None, Field(description="How many passages, 1-50.", ge=1, le=50)] = None,
|
|
60
|
+
ctx: FunctionInvocationContext | None = None,
|
|
61
|
+
) -> str:
|
|
62
|
+
"""Search the user's long-term Membase memory (every Memory the key reaches). Returns
|
|
63
|
+
passages, each tagged with the Memory it came from; name that Memory when you use a
|
|
64
|
+
passage."""
|
|
65
|
+
try:
|
|
66
|
+
mb = _client(ctx)
|
|
67
|
+
answer = await in_thread(mb.search, query, limit=limit)
|
|
68
|
+
except NoUserBound as exc:
|
|
69
|
+
return f"Membase is not available: {exc}."
|
|
70
|
+
except Exception as exc: # noqa: BLE001 - the tool answers the model with the error
|
|
71
|
+
return describe_error(exc)
|
|
72
|
+
passages, unanswered = format_search(answer)
|
|
73
|
+
out = (
|
|
74
|
+
["Passages from the user's Membase memory:", *passages]
|
|
75
|
+
if passages
|
|
76
|
+
else ["No passages in the user's Membase memory matched."]
|
|
77
|
+
)
|
|
78
|
+
if unanswered:
|
|
79
|
+
out += [
|
|
80
|
+
"Some Memories did not fully answer, so this is not proof nothing is stored:",
|
|
81
|
+
*unanswered,
|
|
82
|
+
]
|
|
83
|
+
return "\n".join(out)
|
|
84
|
+
|
|
85
|
+
@tool(name="membase_save", approval_mode="never_require")
|
|
86
|
+
async def membase_save(
|
|
87
|
+
content: Annotated[
|
|
88
|
+
str, Field(description="The fact to keep, in the user's own words. One fact per call.")
|
|
89
|
+
],
|
|
90
|
+
ctx: FunctionInvocationContext | None = None,
|
|
91
|
+
) -> str:
|
|
92
|
+
"""Save something to the user's Membase memory. Call this only when the user asks you
|
|
93
|
+
to remember something or plainly means it; never save your own conclusions."""
|
|
94
|
+
try:
|
|
95
|
+
mb = _client(ctx)
|
|
96
|
+
result = await in_thread(mb.memories.add, content)
|
|
97
|
+
except NoUserBound as exc:
|
|
98
|
+
return f"Membase is not available: {exc}."
|
|
99
|
+
except WriteTargetRequiredError:
|
|
100
|
+
return f"Membase did not save it: {NO_WRITE_TARGET}."
|
|
101
|
+
except PermissionDeniedError as exc:
|
|
102
|
+
return f"Membase did not save it (403): {REVOKED}; a Read key cannot write. ({exc.message})"
|
|
103
|
+
except Exception as exc: # noqa: BLE001 - the tool answers the model with the error
|
|
104
|
+
return describe_error(exc)
|
|
105
|
+
result = result or {}
|
|
106
|
+
names = [x.get("memory_name") for x in result.get("learning") or [] if x.get("memory_name")]
|
|
107
|
+
where = f' "{names[0]}"' if names else ""
|
|
108
|
+
if result.get("status") == "succeeded":
|
|
109
|
+
return f"Saved to Membase{where}; it is learned and searchable now."
|
|
110
|
+
return f"Saved to Membase{where}; it will be searchable once the Memory has learned it."
|
|
111
|
+
|
|
112
|
+
@tool(name="membase_profile", approval_mode="never_require")
|
|
113
|
+
async def membase_profile(ctx: FunctionInvocationContext | None = None) -> str:
|
|
114
|
+
"""Read the user's Membase profile: standing statements about who they are (name, role,
|
|
115
|
+
lasting preferences), newest first. Read it once, near the start."""
|
|
116
|
+
try:
|
|
117
|
+
mb = _client(ctx)
|
|
118
|
+
p = await in_thread(mb.profile_facts.list)
|
|
119
|
+
except NoUserBound as exc:
|
|
120
|
+
return f"Membase is not available: {exc}."
|
|
121
|
+
except PermissionDeniedError as exc:
|
|
122
|
+
return (
|
|
123
|
+
"Membase refused the profile (403): this key's owner did not share the profile "
|
|
124
|
+
f"with it, or {REVOKED}. ({exc.message})"
|
|
125
|
+
)
|
|
126
|
+
except Exception as exc: # noqa: BLE001 - the tool answers the model with the error
|
|
127
|
+
return describe_error(exc)
|
|
128
|
+
items = [str(x.get("content") or "").strip() for x in (p or {}).get("items") or []]
|
|
129
|
+
items = [x for x in items if x]
|
|
130
|
+
if not items:
|
|
131
|
+
return "The user's Membase profile is empty."
|
|
132
|
+
return "About the user:\n" + "\n".join(f"- {x}" for x in items)
|
|
133
|
+
|
|
134
|
+
tools: list[FunctionTool] = [membase_search]
|
|
135
|
+
if save:
|
|
136
|
+
tools.append(membase_save)
|
|
137
|
+
if profile:
|
|
138
|
+
tools.append(membase_profile)
|
|
139
|
+
return tools
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
__all__ = ["membase_tools"]
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
__version__ = "0.2.0"
|
|
@@ -25,17 +25,27 @@ class FakeMembase:
|
|
|
25
25
|
|
|
26
26
|
def __init__(self) -> None:
|
|
27
27
|
self.requests: list[dict[str, Any]] = []
|
|
28
|
-
self.
|
|
29
|
-
{"id": "mv-decisions", "name": "Project decisions"},
|
|
30
|
-
{"id": "mv-travel", "name": "Travel"},
|
|
31
|
-
]
|
|
28
|
+
self.write_target: str | None = "Travel" # None: the key reaches several Memories
|
|
32
29
|
self.results: list[dict[str, Any]] = [
|
|
33
|
-
{
|
|
30
|
+
{
|
|
31
|
+
"type": "memory",
|
|
32
|
+
"id": "mem_bGVkZ2Vy",
|
|
33
|
+
"content": "We settled on Postgres for the ledger.",
|
|
34
|
+
"score": 1.0,
|
|
35
|
+
"memory_name": "Project decisions",
|
|
36
|
+
},
|
|
34
37
|
]
|
|
35
|
-
self.
|
|
38
|
+
self.partial: list[dict[str, Any]] = []
|
|
36
39
|
self.status: dict[str, int] = {} # path -> forced status
|
|
37
40
|
self.delay: float = 0.0
|
|
38
|
-
self.profile = {
|
|
41
|
+
self.profile = {
|
|
42
|
+
"items": [
|
|
43
|
+
{"id": "pf_Mg", "content": "Works in Berlin", "source_kind": "user_stated"},
|
|
44
|
+
{"id": "pf_MQ", "content": "Prefers dark mode", "source_kind": "user_stated"},
|
|
45
|
+
],
|
|
46
|
+
"has_more": False,
|
|
47
|
+
"partial": [],
|
|
48
|
+
}
|
|
39
49
|
|
|
40
50
|
def handler(self, request: httpx.Request) -> httpx.Response:
|
|
41
51
|
body = json.loads(request.content) if request.content else None
|
|
@@ -57,28 +67,24 @@ class FakeMembase:
|
|
|
57
67
|
forced, json={"error": {"code": "unauthorized", "message": f"forced {forced}"}}
|
|
58
68
|
)
|
|
59
69
|
if path == "/v1/search":
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
return httpx.Response(200, json={"
|
|
63
|
-
if path == "/v1/profile":
|
|
70
|
+
if extra := set(body) - {"q", "limit", "mode"}:
|
|
71
|
+
return _error(422, "validation", f"unknown fields: {sorted(extra)}")
|
|
72
|
+
return httpx.Response(200, json={"items": self.results, "partial": self.partial})
|
|
73
|
+
if path == "/v1/profile-facts":
|
|
64
74
|
return httpx.Response(200, json=self.profile)
|
|
65
|
-
if path == "/v1/memories":
|
|
66
|
-
if body
|
|
67
|
-
return
|
|
68
|
-
|
|
69
|
-
)
|
|
70
|
-
if body.get("container") is None and len(self.containers) > 1:
|
|
71
|
-
return httpx.Response(
|
|
72
|
-
400,
|
|
73
|
-
json={
|
|
74
|
-
"error": {
|
|
75
|
-
"code": "validation",
|
|
76
|
-
"message": "container is required with several in reach",
|
|
77
|
-
}
|
|
78
|
-
},
|
|
79
|
-
)
|
|
75
|
+
if path == "/v1/memories/add" and request.method == "POST":
|
|
76
|
+
if extra := set(body) - {"content", "title"}:
|
|
77
|
+
return _error(422, "validation", f"unknown fields: {sorted(extra)}")
|
|
78
|
+
if self.write_target is None:
|
|
79
|
+
return _error(422, "write_target_required", "reaches 2 Memories and has no write target")
|
|
80
80
|
return httpx.Response(
|
|
81
|
-
202,
|
|
81
|
+
202,
|
|
82
|
+
json={
|
|
83
|
+
"id": "doc_bm90ZQ",
|
|
84
|
+
"kind": "note",
|
|
85
|
+
"status": "processing",
|
|
86
|
+
"learning": [{"memory_name": self.write_target, "status": "processing"}],
|
|
87
|
+
},
|
|
82
88
|
)
|
|
83
89
|
return httpx.Response(404, json={"error": {"code": "not_found", "message": path}})
|
|
84
90
|
|
|
@@ -90,6 +96,10 @@ class FakeMembase:
|
|
|
90
96
|
return [r for r in self.requests if r["path"] == path]
|
|
91
97
|
|
|
92
98
|
|
|
99
|
+
def _error(status: int, code: str, message: str) -> httpx.Response:
|
|
100
|
+
return httpx.Response(status, json={"error": {"code": code, "message": message}})
|
|
101
|
+
|
|
102
|
+
|
|
93
103
|
class ScriptedChatClient(FunctionInvocationLayer[Any], BaseChatClient[Any]):
|
|
94
104
|
"""A chat client with no model: replays scripted responses and records every request."""
|
|
95
105
|
|
|
@@ -16,7 +16,7 @@ def _instructions(chat: ScriptedChatClient, i: int = 0) -> str:
|
|
|
16
16
|
return str(chat.seen[i]["options"].get("instructions") or "")
|
|
17
17
|
|
|
18
18
|
|
|
19
|
-
async def
|
|
19
|
+
async def test_recall_injects_passages_citing_the_memory(mem):
|
|
20
20
|
chat = ScriptedChatClient([say("Postgres.")])
|
|
21
21
|
agent = Agent(
|
|
22
22
|
client=chat, instructions="Be brief.", context_providers=[MembaseContextProvider(mem.client())]
|
|
@@ -26,7 +26,7 @@ async def test_recall_injects_passages_citing_the_container(mem):
|
|
|
26
26
|
|
|
27
27
|
assert result.text == "Postgres."
|
|
28
28
|
(search,) = mem.calls("/v1/search")
|
|
29
|
-
assert search["json"] == {"q": "what did we decide about the ledger?", "limit": 8}
|
|
29
|
+
assert search["json"] == {"q": "what did we decide about the ledger?", "mode": "full", "limit": 8}
|
|
30
30
|
text = _instructions(chat)
|
|
31
31
|
assert text.startswith("Be brief.")
|
|
32
32
|
assert "## Membase memory" in text
|
|
@@ -38,14 +38,12 @@ async def test_recall_injects_passages_citing_the_container(mem):
|
|
|
38
38
|
|
|
39
39
|
async def test_recall_searches_only_the_latest_user_message(mem):
|
|
40
40
|
chat = ScriptedChatClient()
|
|
41
|
-
agent = Agent(
|
|
42
|
-
client=chat, context_providers=[MembaseContextProvider(mem.client(), container="mv-decisions")]
|
|
43
|
-
)
|
|
41
|
+
agent = Agent(client=chat, context_providers=[MembaseContextProvider(mem.client())])
|
|
44
42
|
session = agent.create_session()
|
|
45
43
|
await agent.run("first question", session=session)
|
|
46
44
|
await agent.run("second question", session=session)
|
|
47
45
|
assert [c["json"]["q"] for c in mem.calls("/v1/search")] == ["first question", "second question"]
|
|
48
|
-
assert all(
|
|
46
|
+
assert all("container" not in c["json"] for c in mem.calls("/v1/search"))
|
|
49
47
|
|
|
50
48
|
|
|
51
49
|
async def test_nothing_found_injects_nothing(mem):
|
|
@@ -60,15 +58,15 @@ async def test_nothing_found_injects_nothing(mem):
|
|
|
60
58
|
|
|
61
59
|
async def test_a_memory_that_could_not_answer_is_reported_not_silent(mem, caplog):
|
|
62
60
|
mem.results = []
|
|
63
|
-
mem.
|
|
61
|
+
mem.partial = [{"memory_name": "Travel", "reason": "error"}]
|
|
64
62
|
chat = ScriptedChatClient()
|
|
65
63
|
agent = Agent(client=chat, context_providers=[MembaseContextProvider(mem.client())])
|
|
66
64
|
with caplog.at_level(logging.WARNING, logger="agent_framework_membase"):
|
|
67
65
|
await agent.run("where am I flying next?")
|
|
68
66
|
text = _instructions(chat)
|
|
69
|
-
assert '"Travel" could not answer
|
|
67
|
+
assert '"Travel" could not answer' in text
|
|
70
68
|
assert "not evidence that nothing is stored" in text
|
|
71
|
-
assert "
|
|
69
|
+
assert '"Travel" could not answer' in caplog.text
|
|
72
70
|
|
|
73
71
|
|
|
74
72
|
async def test_cold_start_fails_open_within_the_budget(mem, caplog):
|
|
@@ -106,7 +104,7 @@ async def test_revoked_key_is_reported_and_not_retried(mem, caplog):
|
|
|
106
104
|
async def test_capture_is_off_by_default(mem):
|
|
107
105
|
agent = Agent(client=ScriptedChatClient(), context_providers=[MembaseContextProvider(mem.client())])
|
|
108
106
|
await agent.run("remember that I take the window seat")
|
|
109
|
-
assert mem.calls("/v1/memories") == []
|
|
107
|
+
assert mem.calls("/v1/memories/add") == []
|
|
110
108
|
|
|
111
109
|
|
|
112
110
|
async def test_capture_saves_the_users_words_once_per_turn(mem):
|
|
@@ -119,27 +117,25 @@ async def test_capture_saves_the_users_words_once_per_turn(mem):
|
|
|
119
117
|
return "done"
|
|
120
118
|
|
|
121
119
|
chat = ScriptedChatClient([call("noop", {}), say("Noted.")])
|
|
122
|
-
provider = MembaseContextProvider(mem.client(), capture=True
|
|
120
|
+
provider = MembaseContextProvider(mem.client(), capture=True)
|
|
123
121
|
agent = Agent(client=chat, tools=[noop], context_providers=[provider])
|
|
124
122
|
await agent.run("I always take the window seat")
|
|
125
|
-
(saved,) = mem.calls("/v1/memories")
|
|
123
|
+
(saved,) = mem.calls("/v1/memories/add")
|
|
126
124
|
assert saved["json"] == {
|
|
127
125
|
"content": "I always take the window seat",
|
|
128
|
-
"static": False,
|
|
129
126
|
"title": "Agent Framework conversation",
|
|
130
|
-
"container": "mv-travel",
|
|
131
127
|
}
|
|
132
128
|
assert "confirm" not in saved["params"]
|
|
133
129
|
|
|
134
130
|
|
|
135
131
|
async def test_capture_with_a_read_key_stops_after_the_refusal(mem, caplog):
|
|
136
|
-
mem.status["/v1/memories"] = 403
|
|
137
|
-
provider = MembaseContextProvider(mem.client(), capture=True
|
|
132
|
+
mem.status["/v1/memories/add"] = 403
|
|
133
|
+
provider = MembaseContextProvider(mem.client(), capture=True)
|
|
138
134
|
agent = Agent(client=ScriptedChatClient(), context_providers=[provider])
|
|
139
135
|
with caplog.at_level(logging.ERROR, logger="agent_framework_membase"):
|
|
140
136
|
await agent.run("a")
|
|
141
137
|
await agent.run("b")
|
|
142
|
-
assert len(mem.calls("/v1/memories")) == 1
|
|
138
|
+
assert len(mem.calls("/v1/memories/add")) == 1
|
|
143
139
|
assert "a Read key cannot write" in caplog.text
|
|
144
140
|
|
|
145
141
|
|
|
@@ -211,9 +207,12 @@ def test_api_host_defaults_to_membase_com_whatever_the_sdk_defaults_to(monkeypat
|
|
|
211
207
|
p.close()
|
|
212
208
|
|
|
213
209
|
|
|
214
|
-
async def
|
|
210
|
+
async def test_capture_without_a_write_target_explains_how_to_fix_it(mem, caplog):
|
|
211
|
+
mem.write_target = None
|
|
215
212
|
provider = MembaseContextProvider(mem.client(), capture=True)
|
|
216
213
|
agent = Agent(client=ScriptedChatClient(), context_providers=[provider])
|
|
217
214
|
with caplog.at_level(logging.WARNING, logger="agent_framework_membase"):
|
|
218
215
|
await agent.run("I always take the window seat")
|
|
219
|
-
|
|
216
|
+
await agent.run("and the aisle on short flights")
|
|
217
|
+
assert len(mem.calls("/v1/memories/add")) == 1 # refused once, then not retried
|
|
218
|
+
assert "no write target" in caplog.text and "Connect › Developer keys" in caplog.text
|
|
@@ -23,68 +23,57 @@ def test_tool_set_and_schemas(mem):
|
|
|
23
23
|
search = tools[0].parameters()
|
|
24
24
|
assert search["required"] == ["query"]
|
|
25
25
|
assert "ctx" not in search["properties"] # the framework injects it; the model never sees it
|
|
26
|
-
assert set(tools[1].parameters()["properties"]) == {"content"
|
|
26
|
+
assert set(tools[1].parameters()["properties"]) == {"content"}
|
|
27
|
+
assert "container" not in search["properties"]
|
|
27
28
|
assert [t.name for t in membase_tools(mem.client(), save=False, profile=False)] == ["membase_search"]
|
|
28
29
|
|
|
29
30
|
|
|
30
|
-
async def
|
|
31
|
-
_, (out,) = await _run(mem, [call("membase_search", {"query": "ledger"})])
|
|
31
|
+
async def test_search_cites_the_memory(mem):
|
|
32
|
+
_, (out,) = await _run(mem, [call("membase_search", {"query": "ledger", "limit": 3})])
|
|
32
33
|
assert "- [Project decisions] We settled on Postgres for the ledger." in out
|
|
33
|
-
assert mem.calls("/v1/search")[0]["json"] == {"q": "ledger"}
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
async def test_search_by_container_name_resolves_to_its_id(mem):
|
|
37
|
-
_, _ = await _run(mem, [call("membase_search", {"query": "flights", "container": "travel", "limit": 3})])
|
|
38
|
-
assert mem.calls("/v1/search")[0]["json"] == {"q": "flights", "container": "mv-travel", "limit": 3}
|
|
34
|
+
assert mem.calls("/v1/search")[0]["json"] == {"q": "ledger", "mode": "full", "limit": 3}
|
|
39
35
|
|
|
40
36
|
|
|
41
37
|
async def test_search_reports_memories_that_could_not_answer(mem):
|
|
42
38
|
mem.results = []
|
|
43
|
-
mem.
|
|
39
|
+
mem.partial = [
|
|
40
|
+
{"memory_name": "Travel", "reason": "error"},
|
|
41
|
+
{"memory_name": "Reading", "reason": "truncated"},
|
|
42
|
+
]
|
|
44
43
|
_, (out,) = await _run(mem, [call("membase_search", {"query": "x"})])
|
|
45
|
-
assert "No passages" in out and '"Travel" could not answer
|
|
44
|
+
assert "No passages" in out and '"Travel" could not answer' in out
|
|
45
|
+
assert '"Reading" had more passages than the limit let through' in out
|
|
46
46
|
|
|
47
47
|
|
|
48
|
-
async def
|
|
49
|
-
_, (out,) = await _run(
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
(saved,) = mem.calls("/v1/memories")
|
|
53
|
-
assert saved["json"] == {
|
|
54
|
-
"content": "I take the window seat",
|
|
55
|
-
"static": False,
|
|
56
|
-
"title": "",
|
|
57
|
-
"container": "mv-travel",
|
|
58
|
-
}
|
|
48
|
+
async def test_save_lands_in_the_write_target(mem):
|
|
49
|
+
_, (out,) = await _run(mem, [call("membase_save", {"content": "I take the window seat"})])
|
|
50
|
+
(saved,) = mem.calls("/v1/memories/add")
|
|
51
|
+
assert saved["json"] == {"content": "I take the window seat", "title": ""}
|
|
59
52
|
assert saved["params"] == {} # never confirm=
|
|
60
53
|
assert 'Saved to Membase "Travel"' in out
|
|
61
54
|
|
|
62
55
|
|
|
63
|
-
async def
|
|
56
|
+
async def test_save_without_a_write_target_says_how_to_fix_it(mem):
|
|
57
|
+
mem.write_target = None
|
|
64
58
|
_, (out,) = await _run(mem, [call("membase_save", {"content": "x"})])
|
|
65
|
-
assert
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
async def test_save_about_the_user_goes_to_the_profile(mem):
|
|
69
|
-
_, (out,) = await _run(mem, [call("membase_save", {"content": "I am vegetarian", "about_user": True})])
|
|
70
|
-
assert mem.calls("/v1/memories")[0]["json"]["static"] is True
|
|
71
|
-
assert out == "Recorded in the user's Membase profile."
|
|
59
|
+
assert "did not save" in out and "several Memories and has no write target" in out
|
|
60
|
+
assert "Connect › Developer keys" in out
|
|
72
61
|
|
|
73
62
|
|
|
74
63
|
async def test_save_with_a_read_key_explains_the_refusal(mem):
|
|
75
|
-
mem.status["/v1/memories"] = 403
|
|
76
|
-
_, (out,) = await _run(mem, [call("membase_save", {"content": "x"
|
|
64
|
+
mem.status["/v1/memories/add"] = 403
|
|
65
|
+
_, (out,) = await _run(mem, [call("membase_save", {"content": "x"})])
|
|
77
66
|
assert "narrowed or revoked" in out and "Read key cannot write" in out
|
|
78
67
|
|
|
79
68
|
|
|
80
69
|
async def test_profile(mem):
|
|
81
70
|
_, (out,) = await _run(mem, [call("membase_profile", {})])
|
|
82
|
-
assert "
|
|
71
|
+
assert out == "About the user:\n- Works in Berlin\n- Prefers dark mode"
|
|
83
72
|
|
|
84
73
|
|
|
85
74
|
async def test_profile_not_shared_is_a_clear_message(mem):
|
|
86
|
-
mem.status["/v1/profile"] = 403
|
|
87
|
-
_, (out,) = await _run(mem, [call("membase_profile", {
|
|
75
|
+
mem.status["/v1/profile-facts"] = 403
|
|
76
|
+
_, (out,) = await _run(mem, [call("membase_profile", {})])
|
|
88
77
|
assert "did not share the profile" in out
|
|
89
78
|
|
|
90
79
|
|
|
@@ -109,7 +98,7 @@ async def test_provider_and_tools_together(mem):
|
|
|
109
98
|
client = mem.client()
|
|
110
99
|
chat = ScriptedChatClient(
|
|
111
100
|
[
|
|
112
|
-
call("membase_save", {"content": "Ship on Fridays"
|
|
101
|
+
call("membase_save", {"content": "Ship on Fridays"}),
|
|
113
102
|
say("Saved."),
|
|
114
103
|
]
|
|
115
104
|
)
|
|
@@ -119,4 +108,4 @@ async def test_provider_and_tools_together(mem):
|
|
|
119
108
|
result = await agent.run("remember: we ship on Fridays")
|
|
120
109
|
assert result.text == "Saved."
|
|
121
110
|
assert "[Project decisions]" in str(chat.seen[0]["options"]["instructions"])
|
|
122
|
-
assert mem.calls("/v1/memories")[0]["json"]["content"] == "Ship on Fridays"
|
|
111
|
+
assert mem.calls("/v1/memories/add")[0]["json"]["content"] == "Ship on Fridays"
|
|
@@ -1,207 +0,0 @@
|
|
|
1
|
-
"""``membase_tools``: search, save and profile as Agent Framework function tools."""
|
|
2
|
-
|
|
3
|
-
from __future__ import annotations
|
|
4
|
-
|
|
5
|
-
import os
|
|
6
|
-
from typing import Annotated, Any
|
|
7
|
-
|
|
8
|
-
from agent_framework import FunctionInvocationContext, FunctionTool, tool
|
|
9
|
-
from membase import Membase
|
|
10
|
-
from membase.errors import BadRequestError, MembaseError, PermissionDeniedError
|
|
11
|
-
from pydantic import Field
|
|
12
|
-
|
|
13
|
-
from ._common import (
|
|
14
|
-
DEFAULT_SOURCE_ID,
|
|
15
|
-
REVOKED,
|
|
16
|
-
ClientFor,
|
|
17
|
-
Clients,
|
|
18
|
-
NoUserBound,
|
|
19
|
-
describe_error,
|
|
20
|
-
format_search,
|
|
21
|
-
in_thread,
|
|
22
|
-
)
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
def _resolve_container(client: Membase, wanted: str) -> tuple[str | None, str]:
|
|
26
|
-
"""``(container_id, name)`` for a container given by id or by name, else ``(None, message)``."""
|
|
27
|
-
containers = client.containers.list().get("containers") or []
|
|
28
|
-
key = wanted.strip().casefold()
|
|
29
|
-
for c in containers:
|
|
30
|
-
if str(c.get("id", "")) == wanted.strip():
|
|
31
|
-
return str(c["id"]), str(c.get("name") or c["id"])
|
|
32
|
-
for c in containers:
|
|
33
|
-
if str(c.get("name", "")).strip().casefold() == key:
|
|
34
|
-
return str(c["id"]), str(c.get("name") or c["id"])
|
|
35
|
-
names = ", ".join(f'"{c.get("name")}"' for c in containers) or "none"
|
|
36
|
-
return None, f'No Memory called "{wanted}" is in reach of this key. Memories in reach: {names}.'
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
def _names(client: Membase) -> str:
|
|
40
|
-
try:
|
|
41
|
-
containers = client.containers.list().get("containers") or []
|
|
42
|
-
except MembaseError:
|
|
43
|
-
return ""
|
|
44
|
-
return ", ".join(f'"{c.get("name")}"' for c in containers)
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
def membase_tools(
|
|
48
|
-
client: Membase | None = None,
|
|
49
|
-
*,
|
|
50
|
-
api_key: str | None = None,
|
|
51
|
-
base_url: str | None = None,
|
|
52
|
-
local: str | os.PathLike[str] | bool | None = None,
|
|
53
|
-
client_for: ClientFor | None = None,
|
|
54
|
-
source_id: str = DEFAULT_SOURCE_ID,
|
|
55
|
-
save: bool = True,
|
|
56
|
-
profile: bool = True,
|
|
57
|
-
) -> list[FunctionTool]:
|
|
58
|
-
"""Membase as Agent Framework function tools: ``membase_search``, ``membase_save``,
|
|
59
|
-
``membase_profile``.
|
|
60
|
-
|
|
61
|
-
Pass a ready ``Membase`` client, or ``api_key`` / ``base_url`` / ``local``, or
|
|
62
|
-
``client_for`` (``user_id -> Membase``) with the user bound to each session by
|
|
63
|
-
:func:`bind_user` (the same binding :class:`MembaseContextProvider` reads). Tool calls use
|
|
64
|
-
the SDK's own timeout (90 s by default), since the model asked for them explicitly.
|
|
65
|
-
|
|
66
|
-
``save=False`` leaves out ``membase_save`` (for a Read key); ``profile=False`` leaves out
|
|
67
|
-
``membase_profile`` (for a key whose owner did not share the profile). There is no delete
|
|
68
|
-
or forget tool.
|
|
69
|
-
"""
|
|
70
|
-
clients = Clients(
|
|
71
|
-
client, api_key=api_key, base_url=base_url, local=local, client_for=client_for, source_id=source_id
|
|
72
|
-
)
|
|
73
|
-
|
|
74
|
-
def _client(ctx: FunctionInvocationContext | None) -> Membase:
|
|
75
|
-
return clients.for_session(ctx.session if ctx is not None else None)
|
|
76
|
-
|
|
77
|
-
@tool(name="membase_search", approval_mode="never_require")
|
|
78
|
-
async def membase_search(
|
|
79
|
-
query: Annotated[str, Field(description="What to look for, in natural language.")],
|
|
80
|
-
container: Annotated[
|
|
81
|
-
str | None,
|
|
82
|
-
Field(
|
|
83
|
-
description="Only this Memory (its name as shown in results). Omit to search every Memory."
|
|
84
|
-
),
|
|
85
|
-
] = None,
|
|
86
|
-
limit: Annotated[int | None, Field(description="How many passages, 1-50.", ge=1, le=50)] = None,
|
|
87
|
-
ctx: FunctionInvocationContext | None = None,
|
|
88
|
-
) -> str:
|
|
89
|
-
"""Search the user's long-term Membase memory. Returns passages, each tagged with the
|
|
90
|
-
Memory it came from; name that Memory when you use a passage."""
|
|
91
|
-
try:
|
|
92
|
-
mb = _client(ctx)
|
|
93
|
-
container_id = None
|
|
94
|
-
if container:
|
|
95
|
-
container_id, note = await in_thread(_resolve_container, mb, container)
|
|
96
|
-
if container_id is None:
|
|
97
|
-
return note
|
|
98
|
-
answer = await in_thread(mb.search, query, container=container_id, limit=limit)
|
|
99
|
-
except NoUserBound as exc:
|
|
100
|
-
return f"Membase is not available: {exc}."
|
|
101
|
-
except Exception as exc: # noqa: BLE001 - the tool answers the model with the error
|
|
102
|
-
return describe_error(exc)
|
|
103
|
-
passages, unanswered = format_search(answer)
|
|
104
|
-
out = (
|
|
105
|
-
["Passages from the user's Membase memory:", *passages]
|
|
106
|
-
if passages
|
|
107
|
-
else ["No passages in the user's Membase memory matched."]
|
|
108
|
-
)
|
|
109
|
-
if unanswered:
|
|
110
|
-
out += ["These Memories could not answer, so this is not proof nothing is stored:", *unanswered]
|
|
111
|
-
return "\n".join(out)
|
|
112
|
-
|
|
113
|
-
@tool(name="membase_save", approval_mode="never_require")
|
|
114
|
-
async def membase_save(
|
|
115
|
-
content: Annotated[
|
|
116
|
-
str, Field(description="The fact to keep, in the user's own words. One fact per call.")
|
|
117
|
-
],
|
|
118
|
-
container: Annotated[
|
|
119
|
-
str | None,
|
|
120
|
-
Field(description="The Memory that keeps it (its name). Needed when the user has several."),
|
|
121
|
-
] = None,
|
|
122
|
-
about_user: Annotated[
|
|
123
|
-
bool,
|
|
124
|
-
Field(
|
|
125
|
-
description=(
|
|
126
|
-
"True for a standing fact about the user themselves (who they are, a lasting "
|
|
127
|
-
"preference): it goes to their profile instead of a Memory."
|
|
128
|
-
)
|
|
129
|
-
),
|
|
130
|
-
] = False,
|
|
131
|
-
ctx: FunctionInvocationContext | None = None,
|
|
132
|
-
) -> str:
|
|
133
|
-
"""Save something to the user's Membase memory. Call this only when the user asks you
|
|
134
|
-
to remember something or plainly means it; never save your own conclusions."""
|
|
135
|
-
try:
|
|
136
|
-
mb = _client(ctx)
|
|
137
|
-
container_id, name = None, ""
|
|
138
|
-
if container and not about_user:
|
|
139
|
-
container_id, name = await in_thread(_resolve_container, mb, container)
|
|
140
|
-
if container_id is None:
|
|
141
|
-
return name
|
|
142
|
-
result = await in_thread(mb.memories.add, content, container=container_id, static=about_user)
|
|
143
|
-
except NoUserBound as exc:
|
|
144
|
-
return f"Membase is not available: {exc}."
|
|
145
|
-
except BadRequestError as exc:
|
|
146
|
-
names = await in_thread(_names, mb)
|
|
147
|
-
hint = f" Pass container= with one of: {names}." if names and not container else ""
|
|
148
|
-
return f"Membase did not save it: {exc.message}.{hint}"
|
|
149
|
-
except PermissionDeniedError as exc:
|
|
150
|
-
return f"Membase did not save it (403): {REVOKED}; a Read key cannot write. ({exc.message})"
|
|
151
|
-
except Exception as exc: # noqa: BLE001 - the tool answers the model with the error
|
|
152
|
-
return describe_error(exc)
|
|
153
|
-
if (result or {}).get("static"):
|
|
154
|
-
return "Recorded in the user's Membase profile."
|
|
155
|
-
where = f' "{name}"' if name else ""
|
|
156
|
-
return f"Saved to Membase{where}; it is queued and will be searchable once the Memory has learned it."
|
|
157
|
-
|
|
158
|
-
@tool(name="membase_profile", approval_mode="never_require")
|
|
159
|
-
async def membase_profile(
|
|
160
|
-
topic: Annotated[
|
|
161
|
-
str | None, Field(description="Optional topic to pull relevant profile facts for.")
|
|
162
|
-
] = None,
|
|
163
|
-
ctx: FunctionInvocationContext | None = None,
|
|
164
|
-
) -> str:
|
|
165
|
-
"""Read the user's Membase profile: standing facts about who they are and their
|
|
166
|
-
preferences, and the most recently changed ones. Read it once, near the start."""
|
|
167
|
-
try:
|
|
168
|
-
mb = _client(ctx)
|
|
169
|
-
p = await in_thread(mb.profile, topic)
|
|
170
|
-
except NoUserBound as exc:
|
|
171
|
-
return f"Membase is not available: {exc}."
|
|
172
|
-
except PermissionDeniedError as exc:
|
|
173
|
-
return (
|
|
174
|
-
"Membase refused the profile (403): this key's owner did not share the profile "
|
|
175
|
-
f"with it, or {REVOKED}. ({exc.message})"
|
|
176
|
-
)
|
|
177
|
-
except Exception as exc: # noqa: BLE001 - the tool answers the model with the error
|
|
178
|
-
return describe_error(exc)
|
|
179
|
-
out: list[str] = []
|
|
180
|
-
for label, key in (
|
|
181
|
-
("Standing facts", "static"),
|
|
182
|
-
("Recently changed", "dynamic"),
|
|
183
|
-
("On the topic", "results"),
|
|
184
|
-
):
|
|
185
|
-
items = [_text(x) for x in (p or {}).get(key) or []]
|
|
186
|
-
items = [x for x in items if x]
|
|
187
|
-
if items:
|
|
188
|
-
out.append(f"{label}:\n" + "\n".join(f"- {x}" for x in items))
|
|
189
|
-
return "\n\n".join(out) or "The user's Membase profile is empty."
|
|
190
|
-
|
|
191
|
-
tools: list[FunctionTool] = [membase_search]
|
|
192
|
-
if save:
|
|
193
|
-
tools.append(membase_save)
|
|
194
|
-
if profile:
|
|
195
|
-
tools.append(membase_profile)
|
|
196
|
-
return tools
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
def _text(item: Any) -> str:
|
|
200
|
-
if isinstance(item, str):
|
|
201
|
-
return item.strip()
|
|
202
|
-
if isinstance(item, dict):
|
|
203
|
-
return str(item.get("content") or item.get("text") or item.get("memory") or "").strip()
|
|
204
|
-
return str(item).strip()
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
__all__ = ["membase_tools"]
|
|
@@ -1 +0,0 @@
|
|
|
1
|
-
__version__ = "0.1.0"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{agent_framework_membase-0.1.0 → agent_framework_membase-0.2.0}/src/agent_framework_membase/py.typed
RENAMED
|
File without changes
|