mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-14 09:01:22 +00:00
fix(auto-memory): persist same-ID transcript corrections (#668)
* fix(auto-memory): persist same-ID transcript corrections * test(auto-memory): satisfy sequence emptiness lint * Fix session transcript retry comparisons * Fix nested hint transcript retry stability
This commit is contained in:
Adarsh
GitHub
parent
fa6d8fcef7
commit
70877ce85b
3 changed files with 232 additions and 6 deletions
|
|
@ -147,6 +147,47 @@ def _sanitize_msg_for_save(msg: Msg) -> Msg:
|
|||
return msg.model_copy(update={"content": new_content})
|
||||
|
||||
|
||||
def _canonical_message_json(msg: Msg) -> str:
|
||||
"""Compare the JSON values written to disk without Python bool/int aliasing or key-order noise."""
|
||||
persisted = json.loads(msg.model_dump_json())
|
||||
return json.dumps(persisted, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
||||
|
||||
|
||||
def _stabilize_content_block(raw_block: dict, path: tuple[int, ...], message_timestamp: str | None) -> dict:
|
||||
"""Stabilize defaults generated while parsing a raw content block."""
|
||||
block = dict(raw_block)
|
||||
if isinstance(block.get("hint"), list):
|
||||
block["hint"] = [
|
||||
_stabilize_content_block(child, (*path, index), message_timestamp) if isinstance(child, dict) else child
|
||||
for index, child in enumerate(block["hint"])
|
||||
]
|
||||
|
||||
if "id" not in block:
|
||||
canonical = json.dumps(raw_block, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)
|
||||
position = str(path[0]) if len(path) == 1 else ".".join(map(str, path))
|
||||
block["id"] = hashlib.sha256(f"{position}:{canonical}".encode("utf-8")).hexdigest()
|
||||
if "created_at" not in block and message_timestamp:
|
||||
block["created_at"] = message_timestamp
|
||||
if block.get("type") == "hint" and "finished_at" not in block and message_timestamp:
|
||||
block["finished_at"] = message_timestamp
|
||||
return block
|
||||
|
||||
|
||||
def _stabilize_raw_content_ids(item: dict) -> dict:
|
||||
"""Stabilize generated IDs and timestamps in raw (including nested hint) blocks."""
|
||||
content = item.get("content")
|
||||
if not isinstance(content, list):
|
||||
return item
|
||||
|
||||
timestamp = item.get("created_at")
|
||||
stable_content = [
|
||||
_stabilize_content_block(block, (index,), timestamp) if isinstance(block, dict) else block
|
||||
for index, block in enumerate(content)
|
||||
]
|
||||
changed = stable_content != content
|
||||
return {**item, "content": stable_content} if changed else item
|
||||
|
||||
|
||||
def _normalize_msg_timestamp(item: dict) -> dict:
|
||||
"""Map common message timestamp aliases to AgentScope's ``created_at`` field."""
|
||||
if item.get("created_at"):
|
||||
|
|
@ -320,10 +361,13 @@ class AutoMemoryStep(BaseStep):
|
|||
by_id[msg.id] = msg
|
||||
for msg in messages:
|
||||
by_id[msg.id] = msg
|
||||
merged = sorted(by_id.values(), key=lambda m: m.created_at)
|
||||
merged = [_sanitize_msg_for_save(msg) for msg in sorted(by_id.values(), key=lambda m: m.created_at)]
|
||||
|
||||
# An unchanged ID does not imply unchanged persisted content. Compare
|
||||
# sanitized records so identical retries with omitted image/tool data
|
||||
# retain the append fast path, while corrections force a rewrite.
|
||||
can_append = 0 < len(existing) <= len(merged) and all(
|
||||
merged[i].id == existing[i].id for i in range(len(existing))
|
||||
_canonical_message_json(merged[i]) == _canonical_message_json(existing[i]) for i in range(len(existing))
|
||||
)
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
|
@ -333,7 +377,7 @@ class AutoMemoryStep(BaseStep):
|
|||
if new_msgs:
|
||||
async with aiofiles.open(path, "a", encoding="utf-8", newline="") as f:
|
||||
for msg in new_msgs:
|
||||
await f.write(_sanitize_msg_for_save(msg).model_dump_json() + "\n")
|
||||
await f.write(msg.model_dump_json() + "\n")
|
||||
self.logger.info(
|
||||
f"[{self.name}] save session appended session_id={session_id!r} "
|
||||
f"existing={len(existing)} appended={len(new_msgs)} total={len(merged)}",
|
||||
|
|
@ -346,7 +390,7 @@ class AutoMemoryStep(BaseStep):
|
|||
else:
|
||||
async with aiofiles.open(path, "w", encoding="utf-8", newline="") as f:
|
||||
for msg in merged:
|
||||
await f.write(_sanitize_msg_for_save(msg).model_dump_json() + "\n")
|
||||
await f.write(msg.model_dump_json() + "\n")
|
||||
self.logger.info(
|
||||
f"[{self.name}] save session rewrote session_id={session_id!r} "
|
||||
f"existing={len(existing)} total={len(merged)}",
|
||||
|
|
@ -360,6 +404,8 @@ class AutoMemoryStep(BaseStep):
|
|||
item = _normalize_msg_timestamp(item)
|
||||
if isinstance(item, dict) and isinstance(item.get("content"), str):
|
||||
item = {**item, "content": [{"type": "text", "text": item["content"]}]}
|
||||
if isinstance(item, dict):
|
||||
item = _stabilize_raw_content_ids(item)
|
||||
return Msg.model_validate(item)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -414,7 +414,7 @@ async def test_history_merge_keeps_main_source_contract(setup, enabled, history)
|
|||
elif history == "backfill":
|
||||
messages, expected = [older, original], _saved_line(older) + initial
|
||||
elif history == "same-id":
|
||||
messages, expected = [replacement], initial
|
||||
messages, expected = [replacement], _saved_line(replacement)
|
||||
else:
|
||||
messages, expected = [older, replacement], _saved_line(older) + _saved_line(replacement)
|
||||
before = [message.model_dump() for message in messages]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@
|
|||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
|
||||
from agentscope.message import Msg, TextBlock
|
||||
from agentscope.message import Base64Source, DataBlock, Msg, TextBlock
|
||||
import pytest
|
||||
|
||||
from reme.steps.evolve import auto_memory
|
||||
|
|
@ -197,3 +197,183 @@ async def test_concurrent_append_survives_rewrite_race(step, monkeypatch):
|
|||
|
||||
_assert_transcript(step._session_path(_SESSION), ["m0", "m_mid", "m1", "m2"])
|
||||
_assert_serialized(interleaving)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("append_new", [False, True])
|
||||
@pytest.mark.parametrize("field", ["content", "metadata", "created_at"])
|
||||
async def test_existing_message_updates_are_persisted(step, append_new, field):
|
||||
original = _msg("m0", "2026-10-01T10:00:00")
|
||||
await step._save_session_messages(_SESSION, [original])
|
||||
updates = {
|
||||
"content": [TextBlock(text="corrected observation")],
|
||||
"metadata": {"corrected": True},
|
||||
"created_at": "2026-10-01T10:00:30",
|
||||
}
|
||||
corrected = original.model_copy(update={field: updates[field]})
|
||||
messages = [corrected]
|
||||
if append_new:
|
||||
messages.append(_msg("m1", "2026-10-01T10:01:00"))
|
||||
|
||||
await step._save_session_messages(_SESSION, messages)
|
||||
|
||||
saved = [Msg.model_validate_json(line) for line in step._session_path(_SESSION).read_text().splitlines()]
|
||||
assert [message.model_dump() for message in saved] == [message.model_dump() for message in messages]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identical_retry_with_sanitized_content_does_not_rewrite(step, monkeypatch):
|
||||
message = _msg("m0", "2026-10-01T10:00:00")
|
||||
message.content.append(DataBlock(source=Base64Source(media_type="image/png", data="aW1hZ2U=")))
|
||||
await step._save_session_messages(_SESSION, [message])
|
||||
original = step._session_path(_SESSION).read_bytes()
|
||||
real_open = auto_memory.aiofiles.open
|
||||
write_modes = []
|
||||
|
||||
def record_open(path, mode="r", **kwargs):
|
||||
if mode in ("a", "w"):
|
||||
write_modes.append(mode)
|
||||
return real_open(path, mode, **kwargs)
|
||||
|
||||
monkeypatch.setattr(auto_memory.aiofiles, "open", record_open)
|
||||
await step._save_session_messages(_SESSION, [message])
|
||||
assert not write_modes
|
||||
assert step._session_path(_SESSION).read_bytes() == original
|
||||
await step._save_session_messages(_SESSION, [message, _msg("m1", "2026-10-01T10:01:00")])
|
||||
assert write_modes == ["a"]
|
||||
assert len(message.content) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("append_new", [False, True])
|
||||
async def test_metadata_bool_to_int_correction_is_persisted(step, append_new):
|
||||
"""JSON true and 1 differ on disk even though Python considers them equal."""
|
||||
original = _msg("m0", "2026-10-01T10:00:00").model_copy(update={"metadata": {"value": True}})
|
||||
await step._save_session_messages(_SESSION, [original])
|
||||
corrected = original.model_copy(update={"metadata": {"value": 1}})
|
||||
messages = [corrected]
|
||||
if append_new:
|
||||
messages.append(_msg("m1", "2026-10-01T10:01:00"))
|
||||
|
||||
await step._save_session_messages(_SESSION, messages)
|
||||
|
||||
saved = [Msg.model_validate_json(line) for line in step._session_path(_SESSION).read_text().splitlines()]
|
||||
assert saved[0].metadata["value"] == 1
|
||||
assert type(saved[0].metadata["value"]) is int
|
||||
assert [message.id for message in saved] == [message.id for message in messages]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
[
|
||||
[{"type": "text", "text": "same raw payload"}],
|
||||
[{"type": "hint", "hint": "notification"}],
|
||||
[
|
||||
{
|
||||
"type": "hint",
|
||||
"hint": [{"type": "text", "text": "notification"}],
|
||||
"finished_at": "2026-10-01T10:00:00",
|
||||
},
|
||||
],
|
||||
],
|
||||
ids=["text", "hint", "nested hint"],
|
||||
)
|
||||
async def test_identical_raw_retry_with_generated_content_defaults_does_not_rewrite(step, monkeypatch, content):
|
||||
"""Parsing identical raw messages twice should not churn generated block defaults."""
|
||||
raw = {
|
||||
"id": "m0",
|
||||
"name": "Assistant",
|
||||
"role": "assistant",
|
||||
"created_at": "2026-10-01T10:00:00",
|
||||
"content": content,
|
||||
}
|
||||
first = step._build_messages([raw])
|
||||
await step._save_session_messages(_SESSION, first)
|
||||
original = step._session_path(_SESSION).read_bytes()
|
||||
real_open = auto_memory.aiofiles.open
|
||||
write_modes = []
|
||||
|
||||
def record_open(path, mode="r", **kwargs):
|
||||
if mode in ("a", "w"):
|
||||
write_modes.append(mode)
|
||||
return real_open(path, mode, **kwargs)
|
||||
|
||||
monkeypatch.setattr(auto_memory.aiofiles, "open", record_open)
|
||||
second = step._build_messages([raw])
|
||||
await step._save_session_messages(_SESSION, second)
|
||||
|
||||
assert not write_modes
|
||||
assert step._session_path(_SESSION).read_bytes() == original
|
||||
|
||||
|
||||
def test_raw_hint_conversion_preserves_explicit_nested_ids_and_timestamps(step):
|
||||
"""Caller-provided IDs and timestamps remain authoritative at every hint depth."""
|
||||
message = step._to_msg(
|
||||
{
|
||||
"id": "m0",
|
||||
"name": "Alice",
|
||||
"role": "assistant",
|
||||
"created_at": "2026-10-01T10:00:00",
|
||||
"content": [
|
||||
{
|
||||
"type": "hint",
|
||||
"id": "outer-id",
|
||||
"created_at": "2026-10-01T10:00:01",
|
||||
"finished_at": "2026-10-01T10:00:02",
|
||||
"hint": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "notification",
|
||||
"id": "nested-id",
|
||||
"created_at": "2026-10-01T10:00:03",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
outer = message.content[0]
|
||||
nested = outer.hint[0]
|
||||
assert (outer.id, outer.created_at, outer.finished_at) == (
|
||||
"outer-id",
|
||||
"2026-10-01T10:00:01",
|
||||
"2026-10-01T10:00:02",
|
||||
)
|
||||
assert (nested.id, nested.created_at) == ("nested-id", "2026-10-01T10:00:03")
|
||||
|
||||
|
||||
def test_raw_message_conversion_preserves_explicit_content_block_ids(step):
|
||||
"""Caller-provided block IDs remain authoritative during normalization."""
|
||||
message = step._to_msg(
|
||||
{
|
||||
"id": "m0",
|
||||
"name": "Alice",
|
||||
"role": "user",
|
||||
"created_at": "2026-10-01T10:00:00",
|
||||
"content": [{"type": "text", "text": "same raw payload", "id": "caller-block-id"}],
|
||||
},
|
||||
)
|
||||
|
||||
assert message.content[0].id == "caller-block-id"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_identical_metadata_with_different_key_order_does_not_rewrite(step, monkeypatch):
|
||||
"""Canonical JSON comparison ignores object key order while preserving value types."""
|
||||
original = _msg("m0", "2026-10-01T10:00:00").model_copy(update={"metadata": {"a": 1, "b": 2}})
|
||||
await step._save_session_messages(_SESSION, [original])
|
||||
real_open = auto_memory.aiofiles.open
|
||||
write_modes = []
|
||||
|
||||
def record_open(path, mode="r", **kwargs):
|
||||
if mode in ("a", "w"):
|
||||
write_modes.append(mode)
|
||||
return real_open(path, mode, **kwargs)
|
||||
|
||||
monkeypatch.setattr(auto_memory.aiofiles, "open", record_open)
|
||||
reordered = original.model_copy(update={"metadata": {"b": 2, "a": 1}})
|
||||
await step._save_session_messages(_SESSION, [reordered])
|
||||
|
||||
assert not write_modes
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue