memwal 0.2.0.dev0__tar.gz → 0.2.1__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.
- {memwal-0.2.0.dev0 → memwal-0.2.1}/CHANGELOG.md +13 -1
- {memwal-0.2.0.dev0 → memwal-0.2.1}/PKG-INFO +1 -1
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/__init__.py +1 -1
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/client.py +102 -43
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/middleware.py +161 -23
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/mock.py +5 -3
- {memwal-0.2.0.dev0 → memwal-0.2.1}/pyproject.toml +1 -1
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_middleware.py +309 -1
- memwal-0.2.1/tests/test_nonfinite_validation.py +238 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/.gitignore +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/README.md +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/examples/.env.example +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/examples/.gitignore +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/examples/async_remember_demo.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/examples/interactive_demo.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/examples/verify_credentials.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/compatibility.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/types.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/memwal/utils.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/notebooks/walrus_memory_python_sdk.ipynb +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/run_tests.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/__init__.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_auth_rejected_message.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_blank_text.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_client.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_empty_namespace.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_env_presets.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_integration.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_list_namespaces.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_mock.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_normalize_server_url.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_polling_delay.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_recall_weights.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_remember_derived_key.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_retry_after_http_date.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_signing.py +0 -0
- {memwal-0.2.0.dev0 → memwal-0.2.1}/tests/test_transport_retry.py +0 -0
|
@@ -1,9 +1,21 @@
|
|
|
1
1
|
# memwal
|
|
2
2
|
|
|
3
|
-
## 0.2.
|
|
3
|
+
## 0.2.1
|
|
4
4
|
|
|
5
5
|
Unreleased.
|
|
6
6
|
|
|
7
|
+
### Fixed
|
|
8
|
+
|
|
9
|
+
- `with_memwal_*` gains `save_mode` and confirmable saves: `memwal_wait_for_saves` (async and sync) polls every job the middleware enqueued, and restores the job ids if that poll raises so a retry can still confirm them. (WALM-307, #410, #411, #412)
|
|
10
|
+
- `wait_for_remember_jobs` chunks its status reads to the relayer's 20-id `MAX_BULK_ITEMS` cap. Above it the relayer answers a non-transient 400, so a wait over 21+ jobs raised on every attempt and could never confirm — reachable from `save_mode="remember"` after 21 turns, or from two fat `analyze` extracts. (WALM-307)
|
|
11
|
+
|
|
12
|
+
- Save confirmation preserves unconfirmed job IDs when an async wait is cancelled or a synchronous wait is interrupted, then propagates the cancellation or interrupt. (#1303)
|
|
13
|
+
- Reject explicitly empty namespaces instead of silently writing to the default namespace. (#1265)
|
|
14
|
+
- Reject non-finite vectors and invalid limits before signing requests. (#1264)
|
|
15
|
+
|
|
16
|
+
## 0.2.0
|
|
17
|
+
|
|
18
|
+
|
|
7
19
|
### Changed
|
|
8
20
|
|
|
9
21
|
- Python 3.11 or later is required. Python 3.9 and 3.10 have reached end of life; recreate existing virtual environments on a supported patched interpreter before upgrading.
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: memwal
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.1
|
|
4
4
|
Summary: Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing
|
|
5
5
|
Project-URL: Homepage, https://memory.walrus.xyz
|
|
6
6
|
Project-URL: Documentation, https://memory.walrus.xyz
|
|
@@ -98,6 +98,10 @@ PUBLIC_SUI_GRAPHQL_NETWORKS = ("mainnet", "testnet", "devnet")
|
|
|
98
98
|
# Sui read endpoint is unresponsive.
|
|
99
99
|
PACKAGE_VERSION_CHECK_TIMEOUT_S = 10.0
|
|
100
100
|
SEAL_SESSION_SAFETY_MARGIN_MS = 30_000
|
|
101
|
+
|
|
102
|
+
# The relayer's `MAX_BULK_ITEMS` (services/server/src/types.rs). A bulk read
|
|
103
|
+
# above it is rejected with a non-transient 400, so callers must chunk.
|
|
104
|
+
MAX_BULK_STATUS_IDS = 20
|
|
101
105
|
AUTH_REJECTED_MESSAGE = (
|
|
102
106
|
"401 from relayer: typically wrong private key, key not registered on this "
|
|
103
107
|
"account, account ID mismatch (a wallet address is not a MemWal account ID), "
|
|
@@ -410,6 +414,29 @@ def _occurred_at_to_wire(
|
|
|
410
414
|
)
|
|
411
415
|
|
|
412
416
|
|
|
417
|
+
|
|
418
|
+
def _check_limit(limit: Any) -> None:
|
|
419
|
+
"""Reject a limit that is not a positive integer before it is signed (#1218)."""
|
|
420
|
+
if isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0:
|
|
421
|
+
raise ValueError(f"limit must be a positive integer, got {limit!r}")
|
|
422
|
+
|
|
423
|
+
|
|
424
|
+
def _check_vector(vector: Any) -> None:
|
|
425
|
+
"""Reject a vector that is not a non-empty list of finite numbers (#1218).
|
|
426
|
+
|
|
427
|
+
json.dumps would otherwise emit bare NaN/Infinity tokens, which are not
|
|
428
|
+
JSON, and the relayer would answer with an opaque parse error.
|
|
429
|
+
"""
|
|
430
|
+
if (
|
|
431
|
+
not isinstance(vector, (list, tuple))
|
|
432
|
+
or len(vector) == 0
|
|
433
|
+
or not all(
|
|
434
|
+
isinstance(v, (int, float)) and not isinstance(v, bool) and math.isfinite(v)
|
|
435
|
+
for v in vector
|
|
436
|
+
)
|
|
437
|
+
):
|
|
438
|
+
raise ValueError("vector must be a non-empty list of finite numbers")
|
|
439
|
+
|
|
413
440
|
class MemWal:
|
|
414
441
|
"""Async-native Walrus Memory client.
|
|
415
442
|
|
|
@@ -907,60 +934,81 @@ class MemWal:
|
|
|
907
934
|
attempt += 1
|
|
908
935
|
retry_after_ms = 0
|
|
909
936
|
|
|
910
|
-
|
|
911
|
-
|
|
937
|
+
# The relayer rejects a bulk-status read of more than
|
|
938
|
+
# MAX_BULK_STATUS_IDS ids with a 400, which is NOT transient: an
|
|
939
|
+
# unchunked poll of 21+ jobs raises out of the whole wait and
|
|
940
|
+
# every retry hits the same 400. Chunk here so `memwal_wait_for_saves`
|
|
941
|
+
# and `remember_bulk_and_wait` are both fixed.
|
|
942
|
+
terminal_ids: Set[str] = set()
|
|
943
|
+
refused_ids: List[str] = []
|
|
944
|
+
for start in range(0, len(pending), MAX_BULK_STATUS_IDS):
|
|
945
|
+
chunk = pending[start : start + MAX_BULK_STATUS_IDS]
|
|
912
946
|
try:
|
|
913
|
-
|
|
914
|
-
|
|
915
|
-
|
|
916
|
-
|
|
917
|
-
|
|
918
|
-
|
|
919
|
-
|
|
920
|
-
|
|
921
|
-
|
|
922
|
-
if err.status
|
|
923
|
-
|
|
924
|
-
|
|
947
|
+
poll_timeout = self._poll_deadline_s(deadline_ms)
|
|
948
|
+
try:
|
|
949
|
+
batch = await self.get_remember_bulk_status(chunk, timeout=poll_timeout)
|
|
950
|
+
except TypeError as err:
|
|
951
|
+
if "unexpected keyword argument" in str(err) and "timeout" in str(err):
|
|
952
|
+
batch = await self.get_remember_bulk_status(chunk)
|
|
953
|
+
else:
|
|
954
|
+
raise
|
|
955
|
+
except _HttpStatusError as err:
|
|
956
|
+
if _is_transient_polling_status(err.status):
|
|
957
|
+
last_refusal = err.status
|
|
958
|
+
if err.status == 429:
|
|
959
|
+
saw_rate_limit = True
|
|
960
|
+
retry_after_ms = _clamped_retry_after_ms(err, deadline_ms)
|
|
961
|
+
# Refused, not answered: keep this chunk pending.
|
|
962
|
+
refused_ids.extend(chunk)
|
|
963
|
+
continue
|
|
964
|
+
raise
|
|
965
|
+
except httpx.TransportError:
|
|
966
|
+
# Network drops retry as status 0; continuing next loop iteration
|
|
967
|
+
# increments `attempt` so exponential backoff delay continues to grow.
|
|
968
|
+
last_refusal = 0
|
|
969
|
+
refused_ids.extend(chunk)
|
|
925
970
|
continue
|
|
926
|
-
|
|
927
|
-
|
|
928
|
-
|
|
929
|
-
|
|
930
|
-
|
|
931
|
-
|
|
932
|
-
|
|
971
|
+
answered = True
|
|
972
|
+
chunk_pending = set(chunk)
|
|
973
|
+
terminal_ids |= {
|
|
974
|
+
item.job_id
|
|
975
|
+
for item in batch.results
|
|
976
|
+
if item.job_id in chunk_pending and settle(item)
|
|
977
|
+
}
|
|
933
978
|
|
|
934
|
-
|
|
935
|
-
|
|
936
|
-
|
|
937
|
-
if item.job_id in pending and settle(item)
|
|
938
|
-
}
|
|
979
|
+
if not answered and refused_ids:
|
|
980
|
+
# Nothing got through this tick; back off and retry.
|
|
981
|
+
continue
|
|
939
982
|
pending = [jid for jid in pending if jid not in terminal_ids]
|
|
940
983
|
|
|
941
984
|
if opts.timeout_ms > 0 and len(pending) == len(job_ids) and not last_seen:
|
|
942
985
|
# Nothing settled and nothing was seen: the polls may all have been
|
|
943
986
|
# refused. Confirm with one direct read before reporting anything.
|
|
944
|
-
|
|
945
|
-
|
|
946
|
-
|
|
947
|
-
|
|
948
|
-
|
|
949
|
-
|
|
950
|
-
|
|
951
|
-
|
|
952
|
-
|
|
953
|
-
|
|
954
|
-
|
|
955
|
-
|
|
956
|
-
|
|
987
|
+
settled: Set[str] = set()
|
|
988
|
+
all_ids = list(job_ids)
|
|
989
|
+
for start in range(0, len(all_ids), MAX_BULK_STATUS_IDS):
|
|
990
|
+
chunk = all_ids[start : start + MAX_BULK_STATUS_IDS]
|
|
991
|
+
try:
|
|
992
|
+
probed = await self.get_remember_bulk_status(chunk)
|
|
993
|
+
except _HttpStatusError as err:
|
|
994
|
+
if err.status == 429:
|
|
995
|
+
# Header first, else the body's retry_after_seconds, so a
|
|
996
|
+
# proxy that strips Retry-After does not lose the hint.
|
|
997
|
+
wait_ms = _retry_after_ms(err)
|
|
998
|
+
retry_after = f"{wait_ms / 1000:g}" if wait_ms else None
|
|
999
|
+
raise MemWalRateLimited(job_ids, retry_after) from err
|
|
1000
|
+
# Any other probe failure is not evidence about the jobs.
|
|
1001
|
+
continue
|
|
1002
|
+
except httpx.TransportError:
|
|
1003
|
+
continue
|
|
957
1004
|
answered = True
|
|
958
|
-
|
|
1005
|
+
pending_now = set(pending)
|
|
1006
|
+
settled |= {
|
|
959
1007
|
item.job_id
|
|
960
1008
|
for item in probed.results
|
|
961
|
-
if item.job_id in
|
|
1009
|
+
if item.job_id in pending_now and settle(item)
|
|
962
1010
|
}
|
|
963
|
-
|
|
1011
|
+
pending = [jid for jid in pending if jid not in settled]
|
|
964
1012
|
|
|
965
1013
|
for job_id in pending:
|
|
966
1014
|
seen = last_seen.get(job_id)
|
|
@@ -1057,6 +1105,7 @@ class MemWal:
|
|
|
1057
1105
|
sort = params.sort
|
|
1058
1106
|
else:
|
|
1059
1107
|
query_text = query
|
|
1108
|
+
_check_limit(limit)
|
|
1060
1109
|
body = {
|
|
1061
1110
|
"query": query_text,
|
|
1062
1111
|
"limit": limit,
|
|
@@ -1299,6 +1348,7 @@ class MemWal:
|
|
|
1299
1348
|
downloads in parallel, then 3 SEAL decrypts in parallel, then
|
|
1300
1349
|
embeddings. Expect seconds-per-blob on cold caches.
|
|
1301
1350
|
"""
|
|
1351
|
+
_check_limit(limit)
|
|
1302
1352
|
data = await self._signed_request("POST", "/api/restore", {
|
|
1303
1353
|
"namespace": namespace,
|
|
1304
1354
|
"limit": limit,
|
|
@@ -1354,6 +1404,10 @@ class MemWal:
|
|
|
1354
1404
|
Returns:
|
|
1355
1405
|
:class:`NamespacesResult`.
|
|
1356
1406
|
"""
|
|
1407
|
+
# Validate before resolving the owner: on a fresh client that lookup is
|
|
1408
|
+
# a signed POST /api/stats, so a bad limit must not reach it (#1218).
|
|
1409
|
+
if limit is not None:
|
|
1410
|
+
_check_limit(limit)
|
|
1357
1411
|
owner = await self._resolve_owner()
|
|
1358
1412
|
|
|
1359
1413
|
params: Dict[str, str] = {}
|
|
@@ -1428,6 +1482,7 @@ class MemWal:
|
|
|
1428
1482
|
Returns:
|
|
1429
1483
|
:class:`RememberManualResult` with id, blob_id, owner, namespace.
|
|
1430
1484
|
"""
|
|
1485
|
+
_check_vector(opts.vector)
|
|
1431
1486
|
data = await self._signed_request(
|
|
1432
1487
|
"POST",
|
|
1433
1488
|
"/api/remember/manual",
|
|
@@ -1457,6 +1512,8 @@ class MemWal:
|
|
|
1457
1512
|
Returns:
|
|
1458
1513
|
:class:`RecallManualResult` with blob_id + distance pairs.
|
|
1459
1514
|
"""
|
|
1515
|
+
_check_vector(opts.vector)
|
|
1516
|
+
_check_limit(opts.limit)
|
|
1460
1517
|
body: Dict[str, Any] = {
|
|
1461
1518
|
"vector": opts.vector,
|
|
1462
1519
|
"limit": opts.limit,
|
|
@@ -1751,7 +1808,9 @@ class MemWal:
|
|
|
1751
1808
|
|
|
1752
1809
|
method_upper = method.upper()
|
|
1753
1810
|
timestamp = str(int(time.time()))
|
|
1754
|
-
body_str = "" if method_upper == "GET" else json.dumps(
|
|
1811
|
+
body_str = "" if method_upper == "GET" else json.dumps(
|
|
1812
|
+
body, separators=(",", ":"), allow_nan=False
|
|
1813
|
+
)
|
|
1755
1814
|
body_hash = sha256_hex(body_str)
|
|
1756
1815
|
# MED-1 / LOW-23: nonce + account_id are part of the canonical signed
|
|
1757
1816
|
# message. Server rejects the request as "unsupported legacy SDK"
|
|
@@ -55,7 +55,7 @@ from typing import (
|
|
|
55
55
|
)
|
|
56
56
|
|
|
57
57
|
from .client import MemWal, _with_fresh_http_client
|
|
58
|
-
from .types import RecallMemory
|
|
58
|
+
from .types import RecallMemory, RememberBulkOptions, RememberBulkResult
|
|
59
59
|
|
|
60
60
|
if TYPE_CHECKING:
|
|
61
61
|
from langchain_core.language_models.chat_models import BaseChatModel
|
|
@@ -178,6 +178,28 @@ def _format_memories(
|
|
|
178
178
|
])
|
|
179
179
|
|
|
180
180
|
|
|
181
|
+
SAVE_MODES = ("analyze", "remember")
|
|
182
|
+
"""Valid ``save_mode`` values for the auto-save path.
|
|
183
|
+
|
|
184
|
+
``"analyze"`` extracts spoken-fact-style statements via the server LLM and
|
|
185
|
+
stores one memory per fact — it is lossy by design, and non-sentence content
|
|
186
|
+
(code snippets, structured data) yields zero facts. ``"remember"`` stores the
|
|
187
|
+
user message verbatim as a single memory, which is what callers saving code or
|
|
188
|
+
other non-conversational content want.
|
|
189
|
+
"""
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def _validate_save_mode(save_mode: str) -> str:
|
|
193
|
+
if save_mode not in SAVE_MODES:
|
|
194
|
+
raise ValueError(
|
|
195
|
+
f"save_mode must be one of {SAVE_MODES}, got {save_mode!r}"
|
|
196
|
+
)
|
|
197
|
+
return save_mode
|
|
198
|
+
|
|
199
|
+
|
|
200
|
+
_MAX_TRACKED_JOB_IDS = 1000
|
|
201
|
+
|
|
202
|
+
|
|
181
203
|
class _PendingSaves:
|
|
182
204
|
"""Tracks in-flight fire-and-forget auto-save work (asyncio Tasks and
|
|
183
205
|
background Threads) so a caller can drain it deterministically via
|
|
@@ -197,6 +219,7 @@ class _PendingSaves:
|
|
|
197
219
|
self._failures: List[tuple[str, str]] = []
|
|
198
220
|
self._dropped = 0
|
|
199
221
|
self._on_save_error = on_save_error
|
|
222
|
+
self._job_ids: List[str] = []
|
|
200
223
|
self._lock = threading.Lock()
|
|
201
224
|
self._callback_threads: set[int] = set()
|
|
202
225
|
|
|
@@ -250,6 +273,23 @@ class _PendingSaves:
|
|
|
250
273
|
f"Walrus Memory auto-save failed ({total} call(s)): {detail}"
|
|
251
274
|
)
|
|
252
275
|
|
|
276
|
+
def track_job_ids(self, job_ids: List[str]) -> None:
|
|
277
|
+
"""Record the background remember jobs an auto-save enqueued, so a
|
|
278
|
+
caller can later confirm they finished and read their blob IDs."""
|
|
279
|
+
if not job_ids:
|
|
280
|
+
return
|
|
281
|
+
with self._lock:
|
|
282
|
+
self._job_ids.extend(job_ids)
|
|
283
|
+
if len(self._job_ids) > _MAX_TRACKED_JOB_IDS:
|
|
284
|
+
del self._job_ids[:-_MAX_TRACKED_JOB_IDS]
|
|
285
|
+
|
|
286
|
+
def drain_job_ids(self) -> List[str]:
|
|
287
|
+
"""Return every job ID recorded so far and forget them, so a second
|
|
288
|
+
wait does not re-poll jobs that already settled."""
|
|
289
|
+
with self._lock:
|
|
290
|
+
job_ids, self._job_ids = self._job_ids, []
|
|
291
|
+
return job_ids
|
|
292
|
+
|
|
253
293
|
def track_task(self, task: "asyncio.Task[Any]") -> None:
|
|
254
294
|
self._tasks.add(task)
|
|
255
295
|
task.add_done_callback(self._tasks.discard)
|
|
@@ -333,6 +373,46 @@ class _PendingSaves:
|
|
|
333
373
|
thread.join()
|
|
334
374
|
|
|
335
375
|
|
|
376
|
+
async def _run_auto_save(
|
|
377
|
+
memwal: MemWal,
|
|
378
|
+
text: str,
|
|
379
|
+
namespace: str,
|
|
380
|
+
save_mode: str,
|
|
381
|
+
pending: _PendingSaves,
|
|
382
|
+
log: Callable[..., Any],
|
|
383
|
+
) -> None:
|
|
384
|
+
"""Persist `text` via the configured save mode and record its job IDs.
|
|
385
|
+
|
|
386
|
+
``analyze()`` returns cleanly with zero facts for content it cannot read
|
|
387
|
+
as spoken facts (code, structured data), so a plain call there looks
|
|
388
|
+
identical to a successful save. Log in that case and point at
|
|
389
|
+
``save_mode="remember"`` rather than dropping the write silently.
|
|
390
|
+
"""
|
|
391
|
+
try:
|
|
392
|
+
if save_mode == "remember":
|
|
393
|
+
accepted = await memwal.remember(text, namespace)
|
|
394
|
+
if accepted is not None and getattr(accepted, "job_id", None):
|
|
395
|
+
pending.track_job_ids([accepted.job_id])
|
|
396
|
+
return
|
|
397
|
+
|
|
398
|
+
result = await memwal.analyze(text, namespace)
|
|
399
|
+
if result is None:
|
|
400
|
+
return
|
|
401
|
+
job_ids = list(getattr(result, "job_ids", None) or [])
|
|
402
|
+
facts = getattr(result, "facts", None)
|
|
403
|
+
pending.track_job_ids(job_ids)
|
|
404
|
+
if not facts and not job_ids:
|
|
405
|
+
log(
|
|
406
|
+
"[Walrus Memory] Auto-save extracted 0 facts from a "
|
|
407
|
+
f"{len(text)}-character message and stored nothing. analyze() "
|
|
408
|
+
"only extracts spoken-fact-style statements — pass "
|
|
409
|
+
'save_mode="remember" to store the message verbatim instead.'
|
|
410
|
+
)
|
|
411
|
+
except Exception as e: # noqa: BLE001 -- auto-save must never break the LLM call
|
|
412
|
+
log(f"[Walrus Memory] Auto-save failed: {e}")
|
|
413
|
+
pending.note_error(e)
|
|
414
|
+
|
|
415
|
+
|
|
336
416
|
def _expose_memwal_controls(obj: Any, memwal: MemWal, pending: _PendingSaves) -> None:
|
|
337
417
|
"""Attach the underlying client and a way to drain pending auto-saves —
|
|
338
418
|
without this, a short-lived process has no way to avoid silently
|
|
@@ -349,9 +429,48 @@ def _expose_memwal_controls(obj: Any, memwal: MemWal, pending: _PendingSaves) ->
|
|
|
349
429
|
(obj.memwal_flush) still finds it — OpenAI clients aren't Pydantic
|
|
350
430
|
models and work the same way either way.
|
|
351
431
|
"""
|
|
432
|
+
async def memwal_wait_for_saves(
|
|
433
|
+
opts: Optional[RememberBulkOptions] = None,
|
|
434
|
+
) -> RememberBulkResult:
|
|
435
|
+
"""Drain pending auto-saves, then poll their remember jobs to a
|
|
436
|
+
terminal state and return per-job status + blob IDs."""
|
|
437
|
+
await pending.flush()
|
|
438
|
+
job_ids = pending.drain_job_ids()
|
|
439
|
+
try:
|
|
440
|
+
return await memwal.wait_for_remember_jobs(job_ids, opts)
|
|
441
|
+
except BaseException:
|
|
442
|
+
# Cancellation and interrupts must leave unconfirmed jobs available
|
|
443
|
+
# to the next wait, while still propagating to the caller.
|
|
444
|
+
pending.track_job_ids(job_ids)
|
|
445
|
+
raise
|
|
446
|
+
|
|
447
|
+
def memwal_wait_for_saves_sync(
|
|
448
|
+
opts: Optional[RememberBulkOptions] = None,
|
|
449
|
+
) -> RememberBulkResult:
|
|
450
|
+
"""Sync counterpart of ``memwal_wait_for_saves`` for callers on the
|
|
451
|
+
blocking entry points (``OpenAI``, ``llm.invoke``)."""
|
|
452
|
+
pending.flush_sync()
|
|
453
|
+
job_ids = pending.drain_job_ids()
|
|
454
|
+
try:
|
|
455
|
+
return _run_blocking(
|
|
456
|
+
lambda: _with_fresh_http_client(
|
|
457
|
+
memwal, memwal.wait_for_remember_jobs(job_ids, opts)
|
|
458
|
+
)
|
|
459
|
+
)
|
|
460
|
+
except BaseException:
|
|
461
|
+
# Cancellation and interrupts must leave unconfirmed jobs available
|
|
462
|
+
# to the next wait, while still propagating to the caller.
|
|
463
|
+
pending.track_job_ids(job_ids)
|
|
464
|
+
raise
|
|
465
|
+
|
|
352
466
|
object.__setattr__(obj, "_memwal", memwal)
|
|
467
|
+
# Public alias: reaching wait_for_remember_jobs() and the rest of the
|
|
468
|
+
# low-level client shouldn't require touching a private attribute.
|
|
469
|
+
object.__setattr__(obj, "memwal", memwal)
|
|
353
470
|
object.__setattr__(obj, "memwal_flush", pending.flush)
|
|
354
471
|
object.__setattr__(obj, "memwal_flush_sync", pending.flush_sync)
|
|
472
|
+
object.__setattr__(obj, "memwal_wait_for_saves", memwal_wait_for_saves)
|
|
473
|
+
object.__setattr__(obj, "memwal_wait_for_saves_sync", memwal_wait_for_saves_sync)
|
|
355
474
|
|
|
356
475
|
|
|
357
476
|
async def _warn_if_cancelled(coro: Any, label: str) -> None:
|
|
@@ -416,6 +535,7 @@ def with_memwal_langchain(
|
|
|
416
535
|
namespace: str = "default",
|
|
417
536
|
max_memories: int = 5,
|
|
418
537
|
auto_save: bool = True,
|
|
538
|
+
save_mode: str = "analyze",
|
|
419
539
|
min_relevance: float = 0.3,
|
|
420
540
|
debug: bool = False,
|
|
421
541
|
env: Optional[str] = None,
|
|
@@ -438,6 +558,11 @@ def with_memwal_langchain(
|
|
|
438
558
|
namespace: Default namespace.
|
|
439
559
|
max_memories: Max memories to inject per request.
|
|
440
560
|
auto_save: Auto-save new facts from conversation.
|
|
561
|
+
save_mode: How auto_save persists the user message. ``"analyze"``
|
|
562
|
+
(default) extracts spoken-fact-style statements and stores one
|
|
563
|
+
memory per fact — non-sentence content such as code yields zero
|
|
564
|
+
facts and stores nothing. ``"remember"`` stores the message
|
|
565
|
+
verbatim as a single memory.
|
|
441
566
|
min_relevance: Minimum similarity score (0-1) to include a memory.
|
|
442
567
|
debug: Enable debug logging.
|
|
443
568
|
env: Optional relayer preset (``"prod"`` / ``"dev"`` / ``"staging"`` /
|
|
@@ -458,6 +583,8 @@ def with_memwal_langchain(
|
|
|
458
583
|
"Install with: pip install memwal[langchain]"
|
|
459
584
|
) from e
|
|
460
585
|
|
|
586
|
+
_validate_save_mode(save_mode)
|
|
587
|
+
|
|
461
588
|
memwal = MemWal.create(
|
|
462
589
|
key=key,
|
|
463
590
|
account_id=account_id,
|
|
@@ -552,16 +679,12 @@ def with_memwal_langchain(
|
|
|
552
679
|
return messages
|
|
553
680
|
|
|
554
681
|
async def _post_analyze(messages: List[BaseMessage]) -> None:
|
|
555
|
-
"""
|
|
682
|
+
"""Persist the user message via the configured save mode."""
|
|
556
683
|
if not auto_save or not _is_last_message_user(messages):
|
|
557
684
|
return
|
|
558
685
|
user_text = _find_last_user_message(messages)
|
|
559
686
|
if user_text:
|
|
560
|
-
|
|
561
|
-
await memwal.analyze(user_text, namespace)
|
|
562
|
-
except Exception as e:
|
|
563
|
-
log(f"[Walrus Memory] Auto-save failed: {e}")
|
|
564
|
-
pending.note_error(e)
|
|
687
|
+
await _run_auto_save(memwal, user_text, namespace, save_mode, pending, log)
|
|
565
688
|
|
|
566
689
|
async def patched_agenerate(
|
|
567
690
|
messages: List[BaseMessage], *args: Any, **kwargs: Any
|
|
@@ -604,6 +727,7 @@ def with_memwal_openai(
|
|
|
604
727
|
namespace: str = "default",
|
|
605
728
|
max_memories: int = 5,
|
|
606
729
|
auto_save: bool = True,
|
|
730
|
+
save_mode: str = "analyze",
|
|
607
731
|
min_relevance: float = 0.3,
|
|
608
732
|
debug: bool = False,
|
|
609
733
|
env: Optional[str] = None,
|
|
@@ -628,6 +752,11 @@ def with_memwal_openai(
|
|
|
628
752
|
namespace: Default namespace.
|
|
629
753
|
max_memories: Max memories to inject per request.
|
|
630
754
|
auto_save: Auto-save new facts from conversation.
|
|
755
|
+
save_mode: How auto_save persists the user message. ``"analyze"``
|
|
756
|
+
(default) extracts spoken-fact-style statements and stores one
|
|
757
|
+
memory per fact — non-sentence content such as code yields zero
|
|
758
|
+
facts and stores nothing. ``"remember"`` stores the message
|
|
759
|
+
verbatim as a single memory.
|
|
631
760
|
min_relevance: Minimum similarity score (0-1) to include a memory.
|
|
632
761
|
debug: Enable debug logging.
|
|
633
762
|
env: Optional relayer preset (``"prod"`` / ``"dev"`` / ``"staging"`` /
|
|
@@ -639,6 +768,8 @@ def with_memwal_openai(
|
|
|
639
768
|
Returns:
|
|
640
769
|
The same client, with ``chat.completions.create`` wrapped to use Walrus Memory.
|
|
641
770
|
"""
|
|
771
|
+
_validate_save_mode(save_mode)
|
|
772
|
+
|
|
642
773
|
memwal = MemWal.create(
|
|
643
774
|
key=key,
|
|
644
775
|
account_id=account_id,
|
|
@@ -654,11 +785,13 @@ def with_memwal_openai(
|
|
|
654
785
|
|
|
655
786
|
if is_async:
|
|
656
787
|
_wrap_async_openai(
|
|
657
|
-
client, memwal, namespace, max_memories, auto_save,
|
|
788
|
+
client, memwal, namespace, max_memories, auto_save, save_mode,
|
|
789
|
+
min_relevance, log, pending,
|
|
658
790
|
)
|
|
659
791
|
else:
|
|
660
792
|
_wrap_sync_openai(
|
|
661
|
-
client, memwal, namespace, max_memories, auto_save,
|
|
793
|
+
client, memwal, namespace, max_memories, auto_save, save_mode,
|
|
794
|
+
min_relevance, log, pending,
|
|
662
795
|
)
|
|
663
796
|
|
|
664
797
|
_expose_memwal_controls(client, memwal, pending)
|
|
@@ -672,6 +805,7 @@ def _wrap_async_openai(
|
|
|
672
805
|
namespace: str,
|
|
673
806
|
max_memories: int,
|
|
674
807
|
auto_save: bool,
|
|
808
|
+
save_mode: str,
|
|
675
809
|
min_relevance: float,
|
|
676
810
|
log: Callable[..., Any],
|
|
677
811
|
pending: _PendingSaves,
|
|
@@ -725,18 +859,14 @@ def _wrap_async_openai(
|
|
|
725
859
|
|
|
726
860
|
result = await original_create(*args, **kwargs)
|
|
727
861
|
|
|
728
|
-
# Fire-and-forget
|
|
862
|
+
# Fire-and-forget auto-save
|
|
729
863
|
if auto_save and _is_last_message_user(messages):
|
|
730
864
|
user_text = _find_last_user_message(messages)
|
|
731
865
|
if user_text:
|
|
732
|
-
|
|
733
|
-
|
|
734
|
-
|
|
735
|
-
|
|
736
|
-
log(f"[Walrus Memory] Auto-save failed: {e}")
|
|
737
|
-
pending.note_error(e)
|
|
738
|
-
|
|
739
|
-
_fire_and_forget(_analyze(), pending)
|
|
866
|
+
_fire_and_forget(
|
|
867
|
+
_run_auto_save(memwal, user_text, namespace, save_mode, pending, log),
|
|
868
|
+
pending,
|
|
869
|
+
)
|
|
740
870
|
|
|
741
871
|
return result
|
|
742
872
|
|
|
@@ -749,6 +879,7 @@ def _wrap_sync_openai(
|
|
|
749
879
|
namespace: str,
|
|
750
880
|
max_memories: int,
|
|
751
881
|
auto_save: bool,
|
|
882
|
+
save_mode: str,
|
|
752
883
|
min_relevance: float,
|
|
753
884
|
log: Callable[..., Any],
|
|
754
885
|
pending: _PendingSaves,
|
|
@@ -807,18 +938,25 @@ def _wrap_sync_openai(
|
|
|
807
938
|
|
|
808
939
|
result = original_create(*args, **kwargs)
|
|
809
940
|
|
|
810
|
-
# Fire-and-forget
|
|
941
|
+
# Fire-and-forget auto-save
|
|
811
942
|
if auto_save and _is_last_message_user(messages):
|
|
812
943
|
user_text = _find_last_user_message(messages)
|
|
813
944
|
if user_text:
|
|
814
|
-
def
|
|
945
|
+
def _save() -> None:
|
|
946
|
+
# _run_auto_save swallows request failures itself; this guards
|
|
947
|
+
# the loop/http-client plumbing around it so a failure there
|
|
948
|
+
# can't take down the background thread.
|
|
815
949
|
try:
|
|
816
|
-
_run_memwal(
|
|
817
|
-
|
|
950
|
+
_run_memwal(
|
|
951
|
+
lambda: _run_auto_save(
|
|
952
|
+
memwal, user_text, namespace, save_mode, pending, log
|
|
953
|
+
)
|
|
954
|
+
)
|
|
955
|
+
except Exception as e: # noqa: BLE001
|
|
818
956
|
log(f"[Walrus Memory] Auto-save failed: {e}")
|
|
819
957
|
pending.note_error(e)
|
|
820
958
|
|
|
821
|
-
pending.spawn_thread(
|
|
959
|
+
pending.spawn_thread(_save)
|
|
822
960
|
|
|
823
961
|
return result
|
|
824
962
|
|
|
@@ -12,6 +12,7 @@ from dataclasses import dataclass
|
|
|
12
12
|
from datetime import datetime, timedelta, timezone
|
|
13
13
|
from typing import Any, Dict, List, Optional, Sequence, Union
|
|
14
14
|
|
|
15
|
+
from .client import _check_limit
|
|
15
16
|
from .types import (
|
|
16
17
|
AnalyzedFact,
|
|
17
18
|
AnalyzeResult,
|
|
@@ -279,8 +280,7 @@ class MemWalMock:
|
|
|
279
280
|
sort = None
|
|
280
281
|
query_text = query
|
|
281
282
|
_validate_text(query_text, "query")
|
|
282
|
-
|
|
283
|
-
raise ValueError("limit must be a non-negative integer")
|
|
283
|
+
_check_limit(limit)
|
|
284
284
|
if sort not in (None, "relevance", "recent"):
|
|
285
285
|
raise ValueError("sort must be 'relevance' or 'recent'")
|
|
286
286
|
resolved_namespace = self._namespace if namespace is None else namespace
|
|
@@ -383,7 +383,7 @@ class MemWalMock:
|
|
|
383
383
|
)
|
|
384
384
|
|
|
385
385
|
async def restore(self, namespace: str, limit: int = 10) -> RestoreResult:
|
|
386
|
-
|
|
386
|
+
_check_limit(limit)
|
|
387
387
|
return RestoreResult(
|
|
388
388
|
restored=0,
|
|
389
389
|
skipped=0,
|
|
@@ -399,6 +399,8 @@ class MemWalMock:
|
|
|
399
399
|
cursor: Optional[str] = None,
|
|
400
400
|
limit: Optional[int] = None,
|
|
401
401
|
) -> NamespacesResult:
|
|
402
|
+
if limit is not None:
|
|
403
|
+
_check_limit(limit)
|
|
402
404
|
grouped: Dict[str, List[_Memory]] = {}
|
|
403
405
|
for memory in self._memories:
|
|
404
406
|
grouped.setdefault(memory.namespace, []).append(memory)
|
|
@@ -18,22 +18,27 @@ from __future__ import annotations
|
|
|
18
18
|
import asyncio
|
|
19
19
|
import json
|
|
20
20
|
import threading
|
|
21
|
+
from types import SimpleNamespace
|
|
21
22
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
22
23
|
|
|
23
24
|
import httpx
|
|
24
25
|
import nacl.signing
|
|
26
|
+
import pytest
|
|
25
27
|
import respx
|
|
26
28
|
|
|
27
29
|
from memwal.client import MemWal
|
|
28
30
|
from memwal.middleware import (
|
|
31
|
+
_MAX_TRACKED_JOB_IDS,
|
|
32
|
+
_expose_memwal_controls,
|
|
29
33
|
_find_last_user_message,
|
|
30
34
|
_format_memories,
|
|
31
35
|
_inject_openai_memory,
|
|
32
36
|
_PendingSaves,
|
|
37
|
+
_run_auto_save,
|
|
33
38
|
with_memwal_langchain,
|
|
34
39
|
with_memwal_openai,
|
|
35
40
|
)
|
|
36
|
-
from memwal.types import RecallMemory
|
|
41
|
+
from memwal.types import RecallMemory, RememberBulkOptions
|
|
37
42
|
from memwal.utils import bytes_to_hex
|
|
38
43
|
|
|
39
44
|
# ============================================================
|
|
@@ -131,6 +136,27 @@ def _mock_analyze() -> httpx.Response:
|
|
|
131
136
|
)
|
|
132
137
|
|
|
133
138
|
|
|
139
|
+
_REMEMBER_URL = f"{_SERVER}/api/remember"
|
|
140
|
+
_BULK_STATUS_URL = f"{_SERVER}/api/remember/bulk/status"
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def _mock_analyze_with_facts(*job_ids: str) -> httpx.Response:
|
|
144
|
+
return httpx.Response(
|
|
145
|
+
202,
|
|
146
|
+
json={
|
|
147
|
+
"facts": [{"text": f"fact {i}", "id": job_id} for i, job_id in enumerate(job_ids)],
|
|
148
|
+
"job_ids": list(job_ids),
|
|
149
|
+
"fact_count": len(job_ids),
|
|
150
|
+
"status": "pending",
|
|
151
|
+
"owner": "0xowner",
|
|
152
|
+
},
|
|
153
|
+
)
|
|
154
|
+
|
|
155
|
+
|
|
156
|
+
def _mock_remember_accepted(job_id: str = "job-remember") -> httpx.Response:
|
|
157
|
+
return httpx.Response(202, json={"job_id": job_id, "status": "pending"})
|
|
158
|
+
|
|
159
|
+
|
|
134
160
|
# ============================================================
|
|
135
161
|
# Pure function tests (no mocking needed)
|
|
136
162
|
# ============================================================
|
|
@@ -1363,3 +1389,285 @@ class TestPendingSaves:
|
|
|
1363
1389
|
else:
|
|
1364
1390
|
raise AssertionError("flush_sync did not raise")
|
|
1365
1391
|
|
|
1392
|
+
|
|
1393
|
+
# ============================================================
|
|
1394
|
+
# WALM-307: auto_save save_mode + confirmable saves
|
|
1395
|
+
# ============================================================
|
|
1396
|
+
|
|
1397
|
+
|
|
1398
|
+
class TestAutoSaveMode:
|
|
1399
|
+
"""analyze() silently stores nothing for non-fact content (GH #411);
|
|
1400
|
+
save_mode="remember" must store it verbatim instead."""
|
|
1401
|
+
|
|
1402
|
+
def _memwal(self) -> MemWal:
|
|
1403
|
+
return MemWal.create(
|
|
1404
|
+
key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER, namespace="default"
|
|
1405
|
+
)
|
|
1406
|
+
|
|
1407
|
+
@respx.mock
|
|
1408
|
+
async def test_analyze_mode_reports_zero_facts_through_wrapper_log(self) -> None:
|
|
1409
|
+
_mock_seal_session_prereqs()
|
|
1410
|
+
respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze())
|
|
1411
|
+
|
|
1412
|
+
pending = _PendingSaves()
|
|
1413
|
+
messages: list = []
|
|
1414
|
+
with patch("memwal.middleware.logger") as mock_logger:
|
|
1415
|
+
await _run_auto_save(
|
|
1416
|
+
self._memwal(), "def f():\n return 1", "default", "analyze",
|
|
1417
|
+
pending, messages.append,
|
|
1418
|
+
)
|
|
1419
|
+
|
|
1420
|
+
assert not mock_logger.warning.called
|
|
1421
|
+
assert len(messages) == 1
|
|
1422
|
+
assert "save_mode" in messages[0]
|
|
1423
|
+
assert pending.drain_job_ids() == []
|
|
1424
|
+
|
|
1425
|
+
@respx.mock
|
|
1426
|
+
async def test_analyze_mode_tracks_job_ids_without_warning(self) -> None:
|
|
1427
|
+
_mock_seal_session_prereqs()
|
|
1428
|
+
respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze_with_facts("j1", "j2"))
|
|
1429
|
+
|
|
1430
|
+
pending = _PendingSaves()
|
|
1431
|
+
with patch("memwal.middleware.logger") as mock_logger:
|
|
1432
|
+
await _run_auto_save(
|
|
1433
|
+
self._memwal(), "I am allergic to peanuts", "default", "analyze",
|
|
1434
|
+
pending, lambda *_: None,
|
|
1435
|
+
)
|
|
1436
|
+
|
|
1437
|
+
assert not mock_logger.warning.called
|
|
1438
|
+
assert pending.drain_job_ids() == ["j1", "j2"]
|
|
1439
|
+
|
|
1440
|
+
@respx.mock
|
|
1441
|
+
async def test_remember_mode_stores_verbatim_and_skips_analyze(self) -> None:
|
|
1442
|
+
_mock_seal_session_prereqs()
|
|
1443
|
+
analyze_route = respx.post(_ANALYZE_URL).mock(return_value=_mock_analyze())
|
|
1444
|
+
remember_route = respx.post(_REMEMBER_URL).mock(
|
|
1445
|
+
return_value=_mock_remember_accepted("job-1")
|
|
1446
|
+
)
|
|
1447
|
+
|
|
1448
|
+
pending = _PendingSaves()
|
|
1449
|
+
await _run_auto_save(
|
|
1450
|
+
self._memwal(), "def f():\n return 1", "default", "remember",
|
|
1451
|
+
pending, lambda *_: None,
|
|
1452
|
+
)
|
|
1453
|
+
|
|
1454
|
+
assert remember_route.called
|
|
1455
|
+
assert not analyze_route.called
|
|
1456
|
+
assert pending.drain_job_ids() == ["job-1"]
|
|
1457
|
+
|
|
1458
|
+
def test_invalid_save_mode_rejected_at_wrap_time(self) -> None:
|
|
1459
|
+
client = MagicMock()
|
|
1460
|
+
try:
|
|
1461
|
+
with_memwal_openai(
|
|
1462
|
+
client, key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER,
|
|
1463
|
+
save_mode="verbatim",
|
|
1464
|
+
)
|
|
1465
|
+
except ValueError as e:
|
|
1466
|
+
assert "save_mode" in str(e)
|
|
1467
|
+
else:
|
|
1468
|
+
raise AssertionError("expected ValueError for an unknown save_mode")
|
|
1469
|
+
|
|
1470
|
+
def test_drain_job_ids_is_one_shot(self) -> None:
|
|
1471
|
+
pending = _PendingSaves()
|
|
1472
|
+
pending.track_job_ids(["a", "b"])
|
|
1473
|
+
assert pending.drain_job_ids() == ["a", "b"]
|
|
1474
|
+
assert pending.drain_job_ids() == []
|
|
1475
|
+
|
|
1476
|
+
def test_tracked_job_ids_are_capped_keeping_the_newest(self) -> None:
|
|
1477
|
+
pending = _PendingSaves()
|
|
1478
|
+
for i in range(_MAX_TRACKED_JOB_IDS + 5):
|
|
1479
|
+
pending.track_job_ids([f"job-{i}"])
|
|
1480
|
+
|
|
1481
|
+
job_ids = pending.drain_job_ids()
|
|
1482
|
+
assert len(job_ids) == _MAX_TRACKED_JOB_IDS
|
|
1483
|
+
assert job_ids[0] == "job-5"
|
|
1484
|
+
assert job_ids[-1] == f"job-{_MAX_TRACKED_JOB_IDS + 4}"
|
|
1485
|
+
|
|
1486
|
+
|
|
1487
|
+
class TestConfirmableSaves:
|
|
1488
|
+
"""GH #410/#412: the easy-setup path had no way to confirm a save
|
|
1489
|
+
landed or read its blob ID without dropping to the low-level client."""
|
|
1490
|
+
|
|
1491
|
+
def _make_async_client(self) -> MagicMock:
|
|
1492
|
+
client = MagicMock()
|
|
1493
|
+
client._async_client = MagicMock()
|
|
1494
|
+
mock_response = MagicMock()
|
|
1495
|
+
mock_response.choices = [MagicMock()]
|
|
1496
|
+
mock_response.choices[0].message.content = "answer"
|
|
1497
|
+
client.chat = MagicMock()
|
|
1498
|
+
client.chat.completions = MagicMock()
|
|
1499
|
+
client.chat.completions.create = AsyncMock(return_value=mock_response)
|
|
1500
|
+
return client
|
|
1501
|
+
|
|
1502
|
+
@respx.mock
|
|
1503
|
+
async def test_wait_for_saves_returns_blob_ids(self) -> None:
|
|
1504
|
+
_mock_seal_session_prereqs()
|
|
1505
|
+
respx.post(_RECALL_URL).mock(return_value=_mock_recall([]))
|
|
1506
|
+
respx.post(_REMEMBER_URL).mock(return_value=_mock_remember_accepted("job-1"))
|
|
1507
|
+
respx.post(_BULK_STATUS_URL).mock(
|
|
1508
|
+
return_value=httpx.Response(
|
|
1509
|
+
200,
|
|
1510
|
+
json={
|
|
1511
|
+
"results": [
|
|
1512
|
+
{"job_id": "job-1", "status": "done", "blob_id": "blob-abc"}
|
|
1513
|
+
]
|
|
1514
|
+
},
|
|
1515
|
+
)
|
|
1516
|
+
)
|
|
1517
|
+
|
|
1518
|
+
smart = with_memwal_openai(
|
|
1519
|
+
self._make_async_client(), key=_KEY_HEX, account_id=_ACCOUNT_ID,
|
|
1520
|
+
server_url=_SERVER, auto_save=True, save_mode="remember",
|
|
1521
|
+
)
|
|
1522
|
+
await smart.chat.completions.create(
|
|
1523
|
+
model="gpt-4o",
|
|
1524
|
+
messages=[{"role": "user", "content": "def f(): return 1"}],
|
|
1525
|
+
)
|
|
1526
|
+
|
|
1527
|
+
result = await smart.memwal_wait_for_saves(
|
|
1528
|
+
RememberBulkOptions(poll_interval_ms=1, timeout_ms=5000)
|
|
1529
|
+
)
|
|
1530
|
+
|
|
1531
|
+
assert result.total == 1
|
|
1532
|
+
assert result.succeeded == 1
|
|
1533
|
+
assert result.results[0].blob_id == "blob-abc"
|
|
1534
|
+
|
|
1535
|
+
@respx.mock
|
|
1536
|
+
async def test_wrapper_exposes_low_level_client_publicly(self) -> None:
|
|
1537
|
+
_mock_seal_session_prereqs()
|
|
1538
|
+
smart = with_memwal_openai(
|
|
1539
|
+
self._make_async_client(), key=_KEY_HEX, account_id=_ACCOUNT_ID,
|
|
1540
|
+
server_url=_SERVER, auto_save=False,
|
|
1541
|
+
)
|
|
1542
|
+
|
|
1543
|
+
assert isinstance(smart.memwal, MemWal)
|
|
1544
|
+
assert callable(smart.memwal.wait_for_remember_jobs)
|
|
1545
|
+
assert callable(smart.memwal_wait_for_saves)
|
|
1546
|
+
assert callable(smart.memwal_wait_for_saves_sync)
|
|
1547
|
+
|
|
1548
|
+
@respx.mock
|
|
1549
|
+
async def test_failed_wait_leaves_job_ids_for_the_next_wait(self) -> None:
|
|
1550
|
+
_mock_seal_session_prereqs()
|
|
1551
|
+
respx.post(_RECALL_URL).mock(return_value=_mock_recall([]))
|
|
1552
|
+
respx.post(_REMEMBER_URL).mock(return_value=_mock_remember_accepted("job-1"))
|
|
1553
|
+
status_route = respx.post(_BULK_STATUS_URL).mock(
|
|
1554
|
+
return_value=httpx.Response(403, json={"error": "forbidden"})
|
|
1555
|
+
)
|
|
1556
|
+
|
|
1557
|
+
smart = with_memwal_openai(
|
|
1558
|
+
self._make_async_client(), key=_KEY_HEX, account_id=_ACCOUNT_ID,
|
|
1559
|
+
server_url=_SERVER, auto_save=True, save_mode="remember",
|
|
1560
|
+
)
|
|
1561
|
+
await smart.chat.completions.create(
|
|
1562
|
+
model="gpt-4o",
|
|
1563
|
+
messages=[{"role": "user", "content": "def f(): return 1"}],
|
|
1564
|
+
)
|
|
1565
|
+
|
|
1566
|
+
opts = RememberBulkOptions(poll_interval_ms=1, timeout_ms=5000)
|
|
1567
|
+
try:
|
|
1568
|
+
await smart.memwal_wait_for_saves(opts)
|
|
1569
|
+
except Exception:
|
|
1570
|
+
pass
|
|
1571
|
+
else:
|
|
1572
|
+
raise AssertionError("expected the failing poll to raise")
|
|
1573
|
+
|
|
1574
|
+
status_route.mock(
|
|
1575
|
+
return_value=httpx.Response(
|
|
1576
|
+
200,
|
|
1577
|
+
json={
|
|
1578
|
+
"results": [
|
|
1579
|
+
{"job_id": "job-1", "status": "done", "blob_id": "blob-abc"}
|
|
1580
|
+
]
|
|
1581
|
+
},
|
|
1582
|
+
)
|
|
1583
|
+
)
|
|
1584
|
+
result = await smart.memwal_wait_for_saves(opts)
|
|
1585
|
+
|
|
1586
|
+
assert result.total == 1
|
|
1587
|
+
assert result.succeeded == 1
|
|
1588
|
+
assert result.results[0].blob_id == "blob-abc"
|
|
1589
|
+
|
|
1590
|
+
async def test_cancelled_wait_preserves_jobs_and_propagates_cancellation(self) -> None:
|
|
1591
|
+
pending = _PendingSaves()
|
|
1592
|
+
pending.track_job_ids(["job-1"])
|
|
1593
|
+
entered = asyncio.Event()
|
|
1594
|
+
|
|
1595
|
+
async def poll_until_cancelled(job_ids, opts):
|
|
1596
|
+
entered.set()
|
|
1597
|
+
await asyncio.Event().wait()
|
|
1598
|
+
|
|
1599
|
+
memwal = MagicMock()
|
|
1600
|
+
memwal.wait_for_remember_jobs = AsyncMock(side_effect=poll_until_cancelled)
|
|
1601
|
+
wrapper = SimpleNamespace()
|
|
1602
|
+
_expose_memwal_controls(wrapper, memwal, pending)
|
|
1603
|
+
task = asyncio.create_task(wrapper.memwal_wait_for_saves())
|
|
1604
|
+
await entered.wait()
|
|
1605
|
+
# A save accepted while polling must survive alongside the restored batch.
|
|
1606
|
+
pending.track_job_ids(["job-2"])
|
|
1607
|
+
task.cancel()
|
|
1608
|
+
with pytest.raises(asyncio.CancelledError):
|
|
1609
|
+
_ = await task
|
|
1610
|
+
|
|
1611
|
+
memwal.wait_for_remember_jobs.side_effect = None
|
|
1612
|
+
memwal.wait_for_remember_jobs.return_value = "confirmed"
|
|
1613
|
+
assert await wrapper.memwal_wait_for_saves() == "confirmed"
|
|
1614
|
+
memwal.wait_for_remember_jobs.assert_awaited_with(["job-2", "job-1"], None)
|
|
1615
|
+
assert pending.drain_job_ids() == []
|
|
1616
|
+
|
|
1617
|
+
def test_interrupted_sync_wait_preserves_jobs_and_propagates_interrupt(self) -> None:
|
|
1618
|
+
pending = _PendingSaves()
|
|
1619
|
+
pending.track_job_ids(["job-1"])
|
|
1620
|
+
wrapper = SimpleNamespace()
|
|
1621
|
+
_expose_memwal_controls(wrapper, MagicMock(), pending)
|
|
1622
|
+
|
|
1623
|
+
with patch("memwal.middleware._run_blocking", side_effect=KeyboardInterrupt):
|
|
1624
|
+
with pytest.raises(KeyboardInterrupt):
|
|
1625
|
+
wrapper.memwal_wait_for_saves_sync()
|
|
1626
|
+
|
|
1627
|
+
assert pending.drain_job_ids() == ["job-1"]
|
|
1628
|
+
|
|
1629
|
+
@respx.mock
|
|
1630
|
+
async def test_wait_confirms_more_jobs_than_the_relayer_bulk_cap(self) -> None:
|
|
1631
|
+
"""21 tracked jobs must confirm, not 400-loop.
|
|
1632
|
+
|
|
1633
|
+
`/api/remember/bulk/status` rejects more than `MAX_BULK_ITEMS` (20)
|
|
1634
|
+
ids with a 400. 400 is not transient, so an unchunked poll raised,
|
|
1635
|
+
the ids were restored for the retry, and every retry hit the same
|
|
1636
|
+
400 — a wait that could never confirm. `save_mode="remember"` is one
|
|
1637
|
+
job per turn, so 21 turns reach this.
|
|
1638
|
+
"""
|
|
1639
|
+
_mock_seal_session_prereqs()
|
|
1640
|
+
job_ids = [f"job-{i}" for i in range(21)]
|
|
1641
|
+
seen_chunk_sizes: list[int] = []
|
|
1642
|
+
|
|
1643
|
+
def _status(request: httpx.Request) -> httpx.Response:
|
|
1644
|
+
body = json.loads(request.content)
|
|
1645
|
+
ids = body["job_ids"]
|
|
1646
|
+
seen_chunk_sizes.append(len(ids))
|
|
1647
|
+
if len(ids) > 20:
|
|
1648
|
+
return httpx.Response(
|
|
1649
|
+
400, json={"error": "job_ids exceeds MAX_BULK_ITEMS"}
|
|
1650
|
+
)
|
|
1651
|
+
return httpx.Response(
|
|
1652
|
+
200,
|
|
1653
|
+
json={
|
|
1654
|
+
"results": [
|
|
1655
|
+
{"job_id": jid, "status": "done", "blob_id": f"blob-{jid}"}
|
|
1656
|
+
for jid in ids
|
|
1657
|
+
]
|
|
1658
|
+
},
|
|
1659
|
+
)
|
|
1660
|
+
|
|
1661
|
+
respx.post(_BULK_STATUS_URL).mock(side_effect=_status)
|
|
1662
|
+
|
|
1663
|
+
client = MemWal.create(key=_KEY_HEX, account_id=_ACCOUNT_ID, server_url=_SERVER)
|
|
1664
|
+
result = await client.wait_for_remember_jobs(
|
|
1665
|
+
job_ids, RememberBulkOptions(poll_interval_ms=1, timeout_ms=5000)
|
|
1666
|
+
)
|
|
1667
|
+
|
|
1668
|
+
assert max(seen_chunk_sizes) <= 20, seen_chunk_sizes
|
|
1669
|
+
assert result.total == 21
|
|
1670
|
+
assert result.succeeded == 21
|
|
1671
|
+
assert result.timed_out == 0
|
|
1672
|
+
# Order still follows the caller's input, chunking notwithstanding.
|
|
1673
|
+
assert [r.id for r in result.results] == job_ids
|
|
@@ -0,0 +1,238 @@
|
|
|
1
|
+
"""#1218: non-finite vectors and bad limits must fail client-side.
|
|
2
|
+
|
|
3
|
+
json.dumps emits bare NaN/Infinity tokens, which are not JSON, so the
|
|
4
|
+
relayer rejected the body with an opaque parse error. Limits went out as
|
|
5
|
+
NaN, -5 or ?limit=nan.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import math
|
|
11
|
+
|
|
12
|
+
import httpx
|
|
13
|
+
import pytest
|
|
14
|
+
|
|
15
|
+
from memwal import (
|
|
16
|
+
MemWal,
|
|
17
|
+
MemWalError,
|
|
18
|
+
MemWalMock,
|
|
19
|
+
RecallManualOptions,
|
|
20
|
+
RecallParams,
|
|
21
|
+
RememberManualOptions,
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class CaptureMemWal(MemWal):
|
|
26
|
+
def __init__(self, *args, **kwargs):
|
|
27
|
+
super().__init__(*args, **kwargs)
|
|
28
|
+
self.captured = []
|
|
29
|
+
|
|
30
|
+
async def _ensure_compatible_relayer(self):
|
|
31
|
+
return {
|
|
32
|
+
"apiVersion": "1.0.0",
|
|
33
|
+
"relayerVersion": "capture",
|
|
34
|
+
"minSupportedSdk": {"python": "0.0.0"},
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
async def _resolve_owner(self):
|
|
38
|
+
return "0xowner"
|
|
39
|
+
|
|
40
|
+
async def _signed_request(
|
|
41
|
+
self,
|
|
42
|
+
method,
|
|
43
|
+
path,
|
|
44
|
+
body,
|
|
45
|
+
accepted_statuses=(200,),
|
|
46
|
+
include_seal_session=True,
|
|
47
|
+
timeout=None,
|
|
48
|
+
):
|
|
49
|
+
self.captured.append((method, path, body))
|
|
50
|
+
if path.startswith("/v1/owners"):
|
|
51
|
+
return {"namespaces": [], "has_more": False, "snapshot_version": 1}
|
|
52
|
+
if path == "/api/remember/manual":
|
|
53
|
+
return {"id": "1", "blob_id": "b", "owner": "0x1", "namespace": "default"}
|
|
54
|
+
if path == "/api/restore":
|
|
55
|
+
return {"restored": 0, "skipped": 0, "total": 0, "namespace": "ns", "owner": "0x1"}
|
|
56
|
+
return {"results": [], "total": 0}
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class FreshOwnerMemWal(MemWal):
|
|
60
|
+
"""Keeps the real ``_resolve_owner``, so owner lookup goes through signing.
|
|
61
|
+
|
|
62
|
+
Every signed request is recorded and then fails, like a relayer outage:
|
|
63
|
+
a bad limit must surface as ValueError, not as this error.
|
|
64
|
+
"""
|
|
65
|
+
|
|
66
|
+
def __init__(self, *args, **kwargs):
|
|
67
|
+
super().__init__(*args, **kwargs)
|
|
68
|
+
self.signed = []
|
|
69
|
+
|
|
70
|
+
async def _ensure_compatible_relayer(self):
|
|
71
|
+
return {
|
|
72
|
+
"apiVersion": "1.0.0",
|
|
73
|
+
"relayerVersion": "capture",
|
|
74
|
+
"minSupportedSdk": {"python": "0.0.0"},
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
async def _signed_request(
|
|
78
|
+
self,
|
|
79
|
+
method,
|
|
80
|
+
path,
|
|
81
|
+
body,
|
|
82
|
+
accepted_statuses=(200,),
|
|
83
|
+
include_seal_session=True,
|
|
84
|
+
timeout=None,
|
|
85
|
+
):
|
|
86
|
+
self.signed.append((method, path))
|
|
87
|
+
raise MemWalError("relayer unavailable")
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def _client():
|
|
91
|
+
return CaptureMemWal.create(
|
|
92
|
+
key="11" * 32, account_id="0x" + "ab" * 32, server_url="https://relayer.example"
|
|
93
|
+
)
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
BAD_VECTORS = [
|
|
97
|
+
[0.1, math.nan],
|
|
98
|
+
[math.inf],
|
|
99
|
+
[0.1, -math.inf],
|
|
100
|
+
[],
|
|
101
|
+
[0.1, "0.2"],
|
|
102
|
+
[0.1, None],
|
|
103
|
+
[True, 0.2],
|
|
104
|
+
"0.1,0.2",
|
|
105
|
+
None,
|
|
106
|
+
]
|
|
107
|
+
BAD_LIMITS = [math.nan, math.inf, -5, -1, 0, 2.5, "5", True]
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
@pytest.mark.parametrize("vector", BAD_VECTORS, ids=repr)
|
|
111
|
+
async def test_recall_manual_rejects_bad_vector(vector):
|
|
112
|
+
c = _client()
|
|
113
|
+
with pytest.raises(ValueError, match="vector"):
|
|
114
|
+
await c.recall_manual(RecallManualOptions(vector=vector))
|
|
115
|
+
assert c.captured == []
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
@pytest.mark.parametrize("vector", BAD_VECTORS, ids=repr)
|
|
119
|
+
async def test_remember_manual_rejects_bad_vector(vector):
|
|
120
|
+
c = _client()
|
|
121
|
+
with pytest.raises(ValueError, match="vector"):
|
|
122
|
+
await c.remember_manual(RememberManualOptions(encrypted_data="AAAA", vector=vector))
|
|
123
|
+
assert c.captured == []
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
@pytest.mark.parametrize("limit", BAD_LIMITS, ids=repr)
|
|
127
|
+
@pytest.mark.parametrize(
|
|
128
|
+
"call",
|
|
129
|
+
[
|
|
130
|
+
("recall", lambda c, n: c.recall("q", limit=n)),
|
|
131
|
+
("recall_params", lambda c, n: c.recall(RecallParams(query="q", limit=n))),
|
|
132
|
+
("recall_manual", lambda c, n: c.recall_manual(RecallManualOptions(vector=[0.1], limit=n))),
|
|
133
|
+
("restore", lambda c, n: c.restore("ns", limit=n)),
|
|
134
|
+
("list_namespaces", lambda c, n: c.list_namespaces(limit=n)),
|
|
135
|
+
],
|
|
136
|
+
ids=lambda x: x[0],
|
|
137
|
+
)
|
|
138
|
+
async def test_bad_limit_is_rejected_before_signing(call, limit):
|
|
139
|
+
c = _client()
|
|
140
|
+
with pytest.raises(ValueError, match="limit"):
|
|
141
|
+
await call[1](c, limit)
|
|
142
|
+
assert c.captured == []
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
async def test_valid_inputs_still_go_out_unchanged():
|
|
146
|
+
c = _client()
|
|
147
|
+
await c.recall_manual(RecallManualOptions(vector=[0.1, -2, 0], limit=3))
|
|
148
|
+
await c.remember_manual(RememberManualOptions(encrypted_data="AAAA", vector=[1, 0.5]))
|
|
149
|
+
await c.recall("q", limit=1)
|
|
150
|
+
await c.restore("ns", limit=50)
|
|
151
|
+
await c.list_namespaces(limit=500)
|
|
152
|
+
await c.list_namespaces()
|
|
153
|
+
bodies = [b for _, _, b in c.captured]
|
|
154
|
+
assert bodies[0]["vector"] == [0.1, -2, 0] and bodies[0]["limit"] == 3
|
|
155
|
+
assert bodies[1]["vector"] == [1, 0.5]
|
|
156
|
+
assert bodies[2]["limit"] == 1
|
|
157
|
+
assert bodies[3]["limit"] == 50
|
|
158
|
+
assert c.captured[4][1].endswith("?limit=500")
|
|
159
|
+
assert "limit" not in c.captured[5][1]
|
|
160
|
+
|
|
161
|
+
|
|
162
|
+
async def test_signed_request_never_emits_nan_tokens(monkeypatch):
|
|
163
|
+
"""Backstop: even if a NaN reaches the signer, it must not be serialised."""
|
|
164
|
+
client = MemWal.create(
|
|
165
|
+
key="11" * 32, account_id="0x" + "ab" * 32, server_url="https://relayer.example"
|
|
166
|
+
)
|
|
167
|
+
|
|
168
|
+
async def compatible():
|
|
169
|
+
return {
|
|
170
|
+
"apiVersion": "1.0.0",
|
|
171
|
+
"relayerVersion": "t",
|
|
172
|
+
"minSupportedSdk": {"python": "0.0.0"},
|
|
173
|
+
}
|
|
174
|
+
|
|
175
|
+
sent = []
|
|
176
|
+
|
|
177
|
+
async def fake_post(*args, **kwargs): # pragma: no cover - must not be reached
|
|
178
|
+
sent.append(kwargs)
|
|
179
|
+
raise AssertionError("request sent with a non-JSON body")
|
|
180
|
+
|
|
181
|
+
monkeypatch.setattr(client, "_ensure_compatible_relayer", compatible)
|
|
182
|
+
monkeypatch.setattr(httpx.AsyncClient, "send", fake_post)
|
|
183
|
+
with pytest.raises(ValueError):
|
|
184
|
+
await client._signed_request(
|
|
185
|
+
"POST", "/api/recall", {"limit": math.nan}, include_seal_session=False
|
|
186
|
+
)
|
|
187
|
+
assert sent == []
|
|
188
|
+
await client.close()
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
@pytest.mark.parametrize("limit", BAD_LIMITS, ids=repr)
|
|
192
|
+
async def test_list_namespaces_bad_limit_raises_before_owner_lookup(limit):
|
|
193
|
+
"""The owner lookup is a signed POST /api/stats; a bad limit must not reach it."""
|
|
194
|
+
c = FreshOwnerMemWal.create(
|
|
195
|
+
key="11" * 32, account_id="0x" + "ab" * 32, server_url="https://relayer.example"
|
|
196
|
+
)
|
|
197
|
+
with pytest.raises(ValueError, match="limit"):
|
|
198
|
+
await c.list_namespaces(limit=limit)
|
|
199
|
+
assert c.signed == []
|
|
200
|
+
await c.close()
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
async def test_list_namespaces_valid_limit_still_resolves_owner():
|
|
204
|
+
c = FreshOwnerMemWal.create(
|
|
205
|
+
key="11" * 32, account_id="0x" + "ab" * 32, server_url="https://relayer.example"
|
|
206
|
+
)
|
|
207
|
+
with pytest.raises(MemWalError, match="relayer unavailable"):
|
|
208
|
+
await c.list_namespaces(limit=5)
|
|
209
|
+
assert c.signed == [("POST", "/api/stats")]
|
|
210
|
+
await c.close()
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
@pytest.mark.parametrize("limit", BAD_LIMITS, ids=repr)
|
|
214
|
+
@pytest.mark.parametrize(
|
|
215
|
+
"call",
|
|
216
|
+
[
|
|
217
|
+
("recall", lambda m, n: m.recall("q", limit=n)),
|
|
218
|
+
("recall_params", lambda m, n: m.recall(RecallParams(query="q", limit=n))),
|
|
219
|
+
("ask", lambda m, n: m.ask("q", limit=n)),
|
|
220
|
+
("restore", lambda m, n: m.restore("ns", limit=n)),
|
|
221
|
+
("list_namespaces", lambda m, n: m.list_namespaces(limit=n)),
|
|
222
|
+
],
|
|
223
|
+
ids=lambda x: x[0],
|
|
224
|
+
)
|
|
225
|
+
async def test_mock_rejects_the_same_bad_limits_as_the_client(call, limit):
|
|
226
|
+
mock = MemWalMock.create()
|
|
227
|
+
await mock.remember("a fact")
|
|
228
|
+
with pytest.raises(ValueError, match="limit"):
|
|
229
|
+
await call[1](mock, limit)
|
|
230
|
+
|
|
231
|
+
|
|
232
|
+
async def test_mock_accepts_positive_integer_limits():
|
|
233
|
+
mock = MemWalMock.create()
|
|
234
|
+
await mock.remember("a fact")
|
|
235
|
+
assert len((await mock.recall("fact", limit=1)).results) == 1
|
|
236
|
+
assert (await mock.restore("ns", limit=1)).restored == 0
|
|
237
|
+
assert len((await mock.list_namespaces(limit=1)).namespaces) == 1
|
|
238
|
+
assert len((await mock.list_namespaces()).namespaces) == 1
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|