mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 19:06:01 +00:00
fix(memory): truncate mem0 context injection on entry boundaries (#4600)
* fix(memory): truncate mem0 context injection on entry boundaries Mem0Manager.get_context built the injection block and then hard-cut it at max_injection_chars, which could sever the last memory mid-line and leave a dangling partial entry in the agent prompt. Accumulate whole entries against the remaining budget instead: a memory that does not fit is skipped (a shorter later one may still fit). When not even the first memory fits, fall back to the previous hard-truncation behavior for that single entry rather than injecting nothing. Tests: entry-boundary truncation, skip-oversized-keep-later, and the oversized-first-entry fallback. * fix(memory): keep entry-boundary guarantee when no mem0 memory fits Follow-up to PR #4600 review: remove the hard-truncation fallback that could inject a partial memory when max_injection_chars is smaller than every recalled entry. Instead return empty context and log a warning that surfaced the undersized budget.
This commit is contained in:
parent
85c3909c2e
commit
c86071442c
@ -38,7 +38,8 @@ class Mem0Config:
|
||||
top_k: int = 8
|
||||
#: Minimum relevance score for search() results (mem0 `threshold`, 0-1).
|
||||
score_threshold: float = 0.1
|
||||
#: Hard cap on the injection text returned by get_context.
|
||||
#: Hard cap on the injection text returned by get_context; memories that
|
||||
#: do not fit whole are skipped (truncation happens on entry boundaries).
|
||||
max_injection_chars: int = 12000
|
||||
#: Per-request HTTP timeout in seconds.
|
||||
timeout_seconds: float = 10.0
|
||||
|
||||
@ -196,19 +196,41 @@ class Mem0Manager(MemoryManager):
|
||||
max_items=top_k,
|
||||
),
|
||||
)
|
||||
budget = self._config.max_injection_chars
|
||||
seen: set[str] = set()
|
||||
lines: list[str] = []
|
||||
used = 0
|
||||
shortest_line: int | None = None
|
||||
for record in records:
|
||||
rid = record.get("id")
|
||||
if rid in seen:
|
||||
continue
|
||||
seen.add(rid)
|
||||
text = str(record.get("memory") or "").strip()
|
||||
if text:
|
||||
lines.append(f"- {text}")
|
||||
if not text:
|
||||
continue
|
||||
line = f"- {text}"
|
||||
line_len = len(line)
|
||||
shortest_line = line_len if shortest_line is None else min(shortest_line, line_len)
|
||||
# Truncate on entry boundaries: keep only memories that fit whole
|
||||
# within the remaining budget (+1 for the joining newline), so the
|
||||
# injection never ends mid-entry with a dangling partial line. An
|
||||
# oversized entry is skipped -- a shorter later one may still fit.
|
||||
added = line_len if not lines else line_len + 1
|
||||
if used + added > budget:
|
||||
continue
|
||||
lines.append(line)
|
||||
used += added
|
||||
context = "\n".join(lines)
|
||||
if len(context) > self._config.max_injection_chars:
|
||||
context = context[: self._config.max_injection_chars]
|
||||
if not context and shortest_line is not None:
|
||||
# Every recalled memory was longer than the configured budget.
|
||||
# Keep the entry-boundary guarantee and surface the config problem
|
||||
# with a warning rather than injecting a partial fact.
|
||||
logger.warning(
|
||||
"max_injection_chars=%d is smaller than the shortest recalled memory (%d chars); returning empty context",
|
||||
budget,
|
||||
shortest_line,
|
||||
)
|
||||
return context
|
||||
|
||||
async def aget_context(
|
||||
|
||||
@ -476,6 +476,37 @@ class TestMem0ManagerGetContext:
|
||||
ctx = mgr.get_context("u1")
|
||||
assert len(ctx) <= 20
|
||||
|
||||
def test_truncates_on_entry_boundary(self) -> None:
|
||||
"""Budget truncation must keep whole entries, never cut a memory
|
||||
mid-line and leave a dangling partial entry in the prompt."""
|
||||
mgr, fake = _manager({"max_injection_chars": 20})
|
||||
fake.list_results = [
|
||||
{"id": "m1", "memory": "a" * 10}, # "- aaaaaaaaaa" = 12 chars, fits
|
||||
{"id": "m2", "memory": "b" * 10}, # + "\n" + 12 = 25 > 20, must be dropped whole
|
||||
]
|
||||
ctx = mgr.get_context("u1")
|
||||
assert ctx == "- " + "a" * 10
|
||||
|
||||
def test_skips_oversized_entry_and_keeps_later_fitting_one(self) -> None:
|
||||
"""An entry that does not fit whole is skipped; a shorter later entry
|
||||
may still fit within the remaining budget."""
|
||||
mgr, fake = _manager({"max_injection_chars": 20})
|
||||
fake.list_results = [
|
||||
{"id": "m1", "memory": "x" * 30}, # 32-char line, does not fit whole
|
||||
{"id": "m2", "memory": "short"}, # "- short" = 7 chars, fits
|
||||
]
|
||||
ctx = mgr.get_context("u1")
|
||||
assert ctx == "- short"
|
||||
|
||||
def test_oversized_entries_return_empty_with_warning(self, caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""When no memory fits the budget, keep the entry-boundary guarantee by
|
||||
returning empty context and logging a diagnosable warning."""
|
||||
mgr, fake = _manager({"max_injection_chars": 20})
|
||||
fake.list_results = [{"id": "m1", "memory": "x" * 30}]
|
||||
ctx = mgr.get_context("u1")
|
||||
assert ctx == ""
|
||||
assert any("max_injection_chars=20" in r.message and "shortest recalled memory" in r.message for r in caplog.records)
|
||||
|
||||
def test_async_get_context_offloads_sync_http_client(self) -> None:
|
||||
mgr, fake = _manager()
|
||||
event_loop_thread = threading.get_ident()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user