mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 02:46:02 +00:00
feat(memory): add guaranteed injection for correction facts with graceful fallback (#3592)
* feat(memory): add guaranteed injection for correction facts with graceful fallback When the token budget is tight, high-value facts (e.g. user corrections) can be silently evicted by lower-priority regular facts. This change: - Introduces configurable 'guaranteed_categories' (default: [correction]) whose facts draw from a separate 'guaranteed_token_budget', ensuring they are never dropped due to budget pressure. - Adds a graceful fallback to confidence-only ranking when the guaranteed-category path raises an unexpected exception. - Refactors fact selection into a header-agnostic helper (_select_fact_lines) with explicit token accounting in the caller, eliminating double-counting of separators. - Emits a single 'Facts:' header regardless of whether both guaranteed and regular facts are present. - Extends the final safety truncation limit to account for the additional guaranteed budget so guaranteed facts survive end-to-end. * refactor(memory): address review feedback on guaranteed injection - Restore strict break-on-overflow in `_select_fact_lines` to preserve the caller's confidence-ordered ranking; add a regression test locking in the invariant that a shorter lower-confidence fact never slips ahead of a skipped higher-confidence one. - Account for the inter-group `\n` separator between guaranteed and regular fact blocks in the regular budget (1-token precision fix). - Clarify docstrings on `format_memory_for_injection` and `MemoryConfig.guaranteed_token_budget` to distinguish the common *displacement* case (total stays within `max_tokens`) from the rarer *additive* case (safety-truncation ceiling raised when guaranteed lines alone would overflow). * fix(memory): address P1 safety truncation + P2s from review - Structure-aware safety truncation: Facts block is now a protected suffix so guaranteed-category facts can never be silently discarded by a prefix-cut on overflow. Only the preceding (user/history) sections are eligible for truncation. - Extend the same protected-suffix treatment to the except/fallback path by returning fact lines alongside the formatted section from _fallback_format_facts, avoiding string parsing. - Single inter-section separator: facts section no longer embeds its own leading \n\n; the final "\n\n".join(sections) is the single source of truth for section-to-section spacing. - Bare string for guaranteed_categories now raises TypeError instead of silently iterating single characters. - Category-less / malformed facts no longer default-promote into the guaranteed "context" pool — only facts with an explicit category field qualify. - Lift valid_facts pre-filter outside the try so the fallback path reuses it instead of re-doing validation work. - MemoryConfigResponse + DeerFlowClient.get_memory_config now expose guaranteed_categories / guaranteed_token_budget. - config.example.yaml: document the two new fields and bump config_version from 12 to 13. - Add regression tests for every finding. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
0ee35ca38f
commit
b990da785f
@ -123,6 +123,14 @@ class MemoryConfigResponse(BaseModel):
|
|||||||
injection_enabled: bool = Field(..., description="Whether memory injection is enabled")
|
injection_enabled: bool = Field(..., description="Whether memory injection is enabled")
|
||||||
max_injection_tokens: int = Field(..., description="Maximum tokens for memory injection")
|
max_injection_tokens: int = Field(..., description="Maximum tokens for memory injection")
|
||||||
token_counting: str = Field(..., description="Token counting strategy for memory injection ('tiktoken' or 'char')")
|
token_counting: str = Field(..., description="Token counting strategy for memory injection ('tiktoken' or 'char')")
|
||||||
|
guaranteed_categories: list[str] = Field(
|
||||||
|
...,
|
||||||
|
description="Fact categories that bypass the regular injection budget (always injected from a reserved allowance)",
|
||||||
|
)
|
||||||
|
guaranteed_token_budget: int = Field(
|
||||||
|
...,
|
||||||
|
description="Token ceiling for guaranteed-category facts (displaces regular lines in the common case; additive only when guaranteed alone overflows max_injection_tokens)",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class MemoryStatusResponse(BaseModel):
|
class MemoryStatusResponse(BaseModel):
|
||||||
@ -350,6 +358,8 @@ async def get_memory_config_endpoint() -> MemoryConfigResponse:
|
|||||||
injection_enabled=config.injection_enabled,
|
injection_enabled=config.injection_enabled,
|
||||||
max_injection_tokens=config.max_injection_tokens,
|
max_injection_tokens=config.max_injection_tokens,
|
||||||
token_counting=config.token_counting,
|
token_counting=config.token_counting,
|
||||||
|
guaranteed_categories=config.guaranteed_categories,
|
||||||
|
guaranteed_token_budget=config.guaranteed_token_budget,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@ -379,6 +389,8 @@ async def get_memory_status(http_request: Request) -> MemoryStatusResponse:
|
|||||||
injection_enabled=config.injection_enabled,
|
injection_enabled=config.injection_enabled,
|
||||||
max_injection_tokens=config.max_injection_tokens,
|
max_injection_tokens=config.max_injection_tokens,
|
||||||
token_counting=config.token_counting,
|
token_counting=config.token_counting,
|
||||||
|
guaranteed_categories=config.guaranteed_categories,
|
||||||
|
guaranteed_token_budget=config.guaranteed_token_budget,
|
||||||
),
|
),
|
||||||
data=MemoryResponse(**memory_data),
|
data=MemoryResponse(**memory_data),
|
||||||
)
|
)
|
||||||
|
|||||||
@ -610,6 +610,8 @@ def _get_memory_context(agent_name: str | None = None, *, app_config: AppConfig
|
|||||||
memory_data,
|
memory_data,
|
||||||
max_tokens=config.max_injection_tokens,
|
max_tokens=config.max_injection_tokens,
|
||||||
use_tiktoken=(config.token_counting == "tiktoken"),
|
use_tiktoken=(config.token_counting == "tiktoken"),
|
||||||
|
guaranteed_categories=getattr(config, "guaranteed_categories", None),
|
||||||
|
guaranteed_token_budget=getattr(config, "guaranteed_token_budget", 500),
|
||||||
)
|
)
|
||||||
|
|
||||||
if not memory_content.strip():
|
if not memory_content.strip():
|
||||||
|
|||||||
@ -316,7 +316,115 @@ def _coerce_confidence(value: Any, default: float = 0.0) -> float:
|
|||||||
return max(0.0, min(1.0, confidence))
|
return max(0.0, min(1.0, confidence))
|
||||||
|
|
||||||
|
|
||||||
def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2000, *, use_tiktoken: bool = True) -> str:
|
def _format_fact_line(fact: dict[str, Any]) -> str | None:
|
||||||
|
"""Build a single formatted fact line, or return ``None`` for invalid facts.
|
||||||
|
|
||||||
|
Extracted as a shared helper so the guaranteed-injection and regular-injection
|
||||||
|
paths produce identical line formatting.
|
||||||
|
"""
|
||||||
|
content_value = fact.get("content")
|
||||||
|
if not isinstance(content_value, str):
|
||||||
|
return None
|
||||||
|
content = content_value.strip()
|
||||||
|
if not content:
|
||||||
|
return None
|
||||||
|
category = str(fact.get("category", "context")).strip() or "context"
|
||||||
|
confidence = _coerce_confidence(fact.get("confidence"), default=0.0)
|
||||||
|
source_error = fact.get("sourceError")
|
||||||
|
if category == "correction" and isinstance(source_error, str) and source_error.strip():
|
||||||
|
return f"- [{category} | {confidence:.2f}] {content} (avoid: {source_error.strip()})"
|
||||||
|
return f"- [{category} | {confidence:.2f}] {content}"
|
||||||
|
|
||||||
|
|
||||||
|
def _select_fact_lines(
|
||||||
|
ranked_facts: list[dict[str, Any]],
|
||||||
|
*,
|
||||||
|
token_budget: int,
|
||||||
|
use_tiktoken: bool,
|
||||||
|
) -> tuple[list[str], int]:
|
||||||
|
"""Greedily select formatted fact lines within a *line-only* token budget.
|
||||||
|
|
||||||
|
This function is intentionally **header-agnostic**: it counts only the
|
||||||
|
fact lines themselves (including ``\\n`` separators between lines). The
|
||||||
|
caller is responsible for reserving tokens for the ``"Facts:\\n"`` header
|
||||||
|
and any inter-section ``"\\n\\n"`` separator *before* calling this
|
||||||
|
function, and passing the remaining capacity as *token_budget*.
|
||||||
|
|
||||||
|
Stops at the first fact that would exceed the budget so the caller's
|
||||||
|
pre-sorted order (typically confidence-descending) is preserved strictly:
|
||||||
|
a shorter lower-ranked fact can never slip ahead of a skipped
|
||||||
|
higher-ranked one.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
ranked_facts: Facts pre-sorted by the caller's preferred ranking.
|
||||||
|
token_budget: Maximum tokens available for fact lines only.
|
||||||
|
use_tiktoken: Whether to use tiktoken for counting.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
``(selected_lines, consumed_tokens)`` — *consumed_tokens* is the
|
||||||
|
exact token cost of the returned lines (including inter-line
|
||||||
|
``\\n`` separators, but *not* a leading header).
|
||||||
|
"""
|
||||||
|
lines: list[str] = []
|
||||||
|
consumed = 0
|
||||||
|
for fact in ranked_facts:
|
||||||
|
formatted = _format_fact_line(fact)
|
||||||
|
if formatted is None:
|
||||||
|
continue
|
||||||
|
line_text = ("\n" + formatted) if lines else formatted
|
||||||
|
line_tokens = _count_tokens(line_text, use_tiktoken=use_tiktoken)
|
||||||
|
if consumed + line_tokens > token_budget:
|
||||||
|
break
|
||||||
|
lines.append(formatted)
|
||||||
|
consumed += line_tokens
|
||||||
|
return lines, consumed
|
||||||
|
|
||||||
|
|
||||||
|
def _fallback_format_facts(
|
||||||
|
valid_facts: list[dict[str, Any]],
|
||||||
|
*,
|
||||||
|
preceding_section_cost: int,
|
||||||
|
max_tokens: int,
|
||||||
|
use_tiktoken: bool,
|
||||||
|
) -> tuple[str, list[str]] | tuple[None, None]:
|
||||||
|
"""Confidence-only ranking used when the primary path raises an exception.
|
||||||
|
|
||||||
|
Returns a tuple ``(section_text, fact_lines)`` where ``section_text`` is the
|
||||||
|
formatted ``"Facts:\\n..."`` section string (without any leading inter-section
|
||||||
|
separator — the caller owns that), and ``fact_lines`` are the individual lines
|
||||||
|
that make up the facts block. Both elements are ``None`` if no facts survive.
|
||||||
|
|
||||||
|
Returning the lines separately lets the caller track them for the
|
||||||
|
structure-aware safety truncation so fallback facts enjoy the same
|
||||||
|
protected-suffix treatment as facts emitted by the primary path.
|
||||||
|
|
||||||
|
*valid_facts* is the already-filtered fact list built by the primary path so
|
||||||
|
the fallback does not redo validation work. *preceding_section_cost* is the
|
||||||
|
tokens already consumed by user-context / history sections (used to derive
|
||||||
|
the remaining budget).
|
||||||
|
"""
|
||||||
|
ranked = sorted(valid_facts, key=lambda f: _coerce_confidence(f.get("confidence"), default=0.0), reverse=True)
|
||||||
|
|
||||||
|
header = "Facts:\n"
|
||||||
|
overhead = _count_tokens(header, use_tiktoken=use_tiktoken)
|
||||||
|
line_budget = max_tokens - preceding_section_cost - overhead
|
||||||
|
if line_budget <= 0:
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
lines, _ = _select_fact_lines(ranked, token_budget=line_budget, use_tiktoken=use_tiktoken)
|
||||||
|
if not lines:
|
||||||
|
return None, None
|
||||||
|
return header + "\n".join(lines), lines
|
||||||
|
|
||||||
|
|
||||||
|
def format_memory_for_injection(
|
||||||
|
memory_data: dict[str, Any],
|
||||||
|
max_tokens: int = 2000,
|
||||||
|
*,
|
||||||
|
use_tiktoken: bool = True,
|
||||||
|
guaranteed_categories: list[str] | None = None,
|
||||||
|
guaranteed_token_budget: int = 500,
|
||||||
|
) -> str:
|
||||||
"""Format memory data for injection into system prompt.
|
"""Format memory data for injection into system prompt.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
@ -325,6 +433,18 @@ def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2
|
|||||||
use_tiktoken: When ``False``, all token counting uses the network-free
|
use_tiktoken: When ``False``, all token counting uses the network-free
|
||||||
character-based estimate instead of tiktoken (see
|
character-based estimate instead of tiktoken (see
|
||||||
``memory.token_counting`` config). Defaults to ``True``.
|
``memory.token_counting`` config). Defaults to ``True``.
|
||||||
|
guaranteed_categories: Fact categories that must always be injected
|
||||||
|
regardless of the regular token budget. These facts draw from a
|
||||||
|
separate *guaranteed_token_budget*. When ``None`` or empty, all
|
||||||
|
facts compete for the same budget (original behaviour).
|
||||||
|
guaranteed_token_budget: Token ceiling for the guaranteed section.
|
||||||
|
In the common case the guaranteed lines *displace* regular lines
|
||||||
|
within *max_tokens* (the total output stays ≤ ``max_tokens``);
|
||||||
|
the budget becomes truly additive only when the guaranteed lines
|
||||||
|
alone would push the assembled output past *max_tokens*, at which
|
||||||
|
point the safety-truncation ceiling is raised to
|
||||||
|
``max_tokens + guaranteed_actual_usage`` to protect them.
|
||||||
|
Ignored when *guaranteed_categories* is ``None`` or empty.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Formatted memory string for system prompt injection.
|
Formatted memory string for system prompt injection.
|
||||||
@ -332,7 +452,16 @@ def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2
|
|||||||
if not memory_data:
|
if not memory_data:
|
||||||
return ""
|
return ""
|
||||||
|
|
||||||
sections = []
|
# Reject a bare string explicitly: iterating a ``str`` yields single
|
||||||
|
# characters, which would silently produce a meaningless frozenset of
|
||||||
|
# letters and turn the guarantee off without any warning. Config-layer
|
||||||
|
# callers go through Pydantic (which enforces ``list[str]``), so this
|
||||||
|
# only guards the public helper surface.
|
||||||
|
if isinstance(guaranteed_categories, str):
|
||||||
|
raise TypeError("guaranteed_categories must be an iterable of strings, not a bare str")
|
||||||
|
effective_guaranteed: frozenset[str] = frozenset(c.strip() for c in guaranteed_categories if isinstance(c, str) and c.strip()) if guaranteed_categories else frozenset()
|
||||||
|
|
||||||
|
sections: list[str] = []
|
||||||
|
|
||||||
# Format user context
|
# Format user context
|
||||||
user_data = memory_data.get("user", {})
|
user_data = memory_data.get("user", {})
|
||||||
@ -374,67 +503,181 @@ def format_memory_for_injection(memory_data: dict[str, Any], max_tokens: int = 2
|
|||||||
if history_sections:
|
if history_sections:
|
||||||
sections.append("History:\n" + "\n".join(f"- {s}" for s in history_sections))
|
sections.append("History:\n" + "\n".join(f"- {s}" for s in history_sections))
|
||||||
|
|
||||||
# Format facts (sorted by confidence; include as many as token budget allows)
|
# ── Facts ────────────────────────────────────────────────────────────────
|
||||||
|
#
|
||||||
|
# Design notes
|
||||||
|
# ~~~~~~~~~~~~
|
||||||
|
# • A single ``"Facts:\\n"`` header is emitted at most once.
|
||||||
|
# • Guaranteed-category facts are selected first from their own
|
||||||
|
# *guaranteed_token_budget* and placed at the front of the Facts block,
|
||||||
|
# so they cannot be evicted by regular facts. In the common case the
|
||||||
|
# total output still fits within *max_tokens* (guaranteed lines displace
|
||||||
|
# regular ones); the budget becomes truly additive only when the
|
||||||
|
# guaranteed lines alone push the output past *max_tokens*, in which
|
||||||
|
# case the safety-truncation ceiling is raised accordingly.
|
||||||
|
# • Regular facts draw from *max_tokens* only.
|
||||||
|
# • All token accounting (header, separators, lines) is performed here
|
||||||
|
# in the caller; the ``_select_fact_lines`` helper is header-agnostic.
|
||||||
|
# • When the primary path raises any exception, ``_fallback_format_facts``
|
||||||
|
# performs a single-pass confidence-only ranking.
|
||||||
facts_data = memory_data.get("facts", [])
|
facts_data = memory_data.get("facts", [])
|
||||||
|
guaranteed_line_tokens = 0 # used later for the effective truncation limit
|
||||||
if isinstance(facts_data, list) and facts_data:
|
if isinstance(facts_data, list) and facts_data:
|
||||||
ranked_facts = sorted(
|
# Token cost of sections built above (user context, history).
|
||||||
(f for f in facts_data if isinstance(f, dict) and isinstance(f.get("content"), str) and f.get("content").strip()),
|
|
||||||
key=lambda fact: _coerce_confidence(fact.get("confidence"), default=0.0),
|
|
||||||
reverse=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Compute token count for existing sections once, then account
|
|
||||||
# incrementally for each fact line to avoid full-string re-tokenization.
|
|
||||||
base_text = "\n\n".join(sections)
|
base_text = "\n\n".join(sections)
|
||||||
base_tokens = _count_tokens(base_text, use_tiktoken=use_tiktoken) if base_text else 0
|
base_tokens = _count_tokens(base_text, use_tiktoken=use_tiktoken) if base_text else 0
|
||||||
# Account for the separator between existing sections and the facts section.
|
|
||||||
|
# Pre-filter valid facts *before* entering the try so the except
|
||||||
|
# path can pass the same list straight into the fallback without
|
||||||
|
# redoing validation work on the hot prompt-injection path.
|
||||||
|
valid_facts = [f for f in facts_data if isinstance(f, dict) and isinstance(f.get("content"), str) and f.get("content", "").strip()]
|
||||||
|
|
||||||
|
# Initialise the facts-block markers *before* the try so the
|
||||||
|
# structure-aware truncation at the bottom of the function can
|
||||||
|
# reason about them regardless of whether the primary path or
|
||||||
|
# the except/fallback path produced the final Facts section.
|
||||||
facts_header = "Facts:\n"
|
facts_header = "Facts:\n"
|
||||||
separator_tokens = _count_tokens("\n\n" + facts_header, use_tiktoken=use_tiktoken) if base_text else _count_tokens(facts_header, use_tiktoken=use_tiktoken)
|
all_fact_lines: list[str] = []
|
||||||
running_tokens = base_tokens + separator_tokens
|
|
||||||
|
|
||||||
fact_lines: list[str] = []
|
try:
|
||||||
for fact in ranked_facts:
|
# Partition valid facts into guaranteed vs regular groups.
|
||||||
content_value = fact.get("content")
|
# Use the *raw* category field (no ``or "context"`` default) so
|
||||||
if not isinstance(content_value, str):
|
# a category-less legacy fact is never silently promoted into
|
||||||
continue
|
# a guaranteed pool whose operator configured
|
||||||
content = content_value.strip()
|
# ``guaranteed_categories=["context"]``. Missing-category facts
|
||||||
if not content:
|
# always fall through to the regular path.
|
||||||
continue
|
def _confidence_key(fact: dict[str, Any]) -> float:
|
||||||
category = str(fact.get("category", "context")).strip() or "context"
|
return _coerce_confidence(fact.get("confidence"), default=0.0)
|
||||||
confidence = _coerce_confidence(fact.get("confidence"), default=0.0)
|
|
||||||
source_error = fact.get("sourceError")
|
if effective_guaranteed:
|
||||||
if category == "correction" and isinstance(source_error, str) and source_error.strip():
|
|
||||||
line = f"- [{category} | {confidence:.2f}] {content} (avoid: {source_error.strip()})"
|
def _category_match(fact: dict[str, Any]) -> bool:
|
||||||
|
raw = fact.get("category")
|
||||||
|
if not isinstance(raw, str):
|
||||||
|
return False
|
||||||
|
cat = raw.strip()
|
||||||
|
return bool(cat) and cat in effective_guaranteed
|
||||||
|
|
||||||
|
guaranteed = sorted(
|
||||||
|
[f for f in valid_facts if _category_match(f)],
|
||||||
|
key=_confidence_key,
|
||||||
|
reverse=True,
|
||||||
|
)
|
||||||
|
regular = sorted(
|
||||||
|
[f for f in valid_facts if not _category_match(f)],
|
||||||
|
key=_confidence_key,
|
||||||
|
reverse=True,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
line = f"- [{category} | {confidence:.2f}] {content}"
|
guaranteed = []
|
||||||
|
regular = sorted(valid_facts, key=_confidence_key, reverse=True)
|
||||||
|
|
||||||
# Each additional line is preceded by a newline (except the first).
|
# ── Phase 1: select guaranteed lines ──────────────────────────
|
||||||
line_text = ("\n" + line) if fact_lines else line
|
header_cost = _count_tokens(facts_header, use_tiktoken=use_tiktoken)
|
||||||
line_tokens = _count_tokens(line_text, use_tiktoken=use_tiktoken)
|
|
||||||
|
|
||||||
if running_tokens + line_tokens <= max_tokens:
|
guaranteed_lines: list[str] = []
|
||||||
fact_lines.append(line)
|
if guaranteed:
|
||||||
running_tokens += line_tokens
|
guaranteed_line_budget = guaranteed_token_budget
|
||||||
else:
|
guaranteed_lines, guaranteed_line_tokens = _select_fact_lines(
|
||||||
break
|
guaranteed,
|
||||||
|
token_budget=guaranteed_line_budget,
|
||||||
|
use_tiktoken=use_tiktoken,
|
||||||
|
)
|
||||||
|
|
||||||
if fact_lines:
|
# ── Phase 2: select regular lines ────────────────────────────
|
||||||
sections.append("Facts:\n" + "\n".join(fact_lines))
|
# Regular facts compete for *max_tokens* (the main budget).
|
||||||
|
# Subtract everything already accounted for:
|
||||||
|
# base sections + inter-section separator + header
|
||||||
|
# + guaranteed lines + the inter-group ``\n`` that joins the
|
||||||
|
# regular block to the guaranteed block (when both are present).
|
||||||
|
regular_lines: list[str] = []
|
||||||
|
if regular:
|
||||||
|
inter_group_newline_tokens = _count_tokens("\n", use_tiktoken=use_tiktoken) if guaranteed_lines else 0
|
||||||
|
used_before_regular = base_tokens + header_cost + guaranteed_line_tokens + inter_group_newline_tokens
|
||||||
|
regular_line_budget = max_tokens - used_before_regular
|
||||||
|
if regular_line_budget > 0:
|
||||||
|
regular_lines, _ = _select_fact_lines(
|
||||||
|
regular,
|
||||||
|
token_budget=regular_line_budget,
|
||||||
|
use_tiktoken=use_tiktoken,
|
||||||
|
)
|
||||||
|
|
||||||
|
# ── Emit a single "Facts:" section ───────────────────────────
|
||||||
|
# Leading inter-section separator is NOT embedded here; the
|
||||||
|
# final ``"\n\n".join(sections)`` is the single source of truth
|
||||||
|
# for section-to-section spacing, preventing the prior
|
||||||
|
# double-``\n\n`` bug.
|
||||||
|
all_fact_lines = guaranteed_lines + regular_lines
|
||||||
|
if all_fact_lines:
|
||||||
|
section_text = facts_header + "\n".join(all_fact_lines)
|
||||||
|
sections.append(section_text)
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
# ── Fallback: confidence-only ranking, single budget ─────────
|
||||||
|
# Any unexpected error in the partition / guaranteed path must
|
||||||
|
# not prevent memory injection entirely. Fall back to the
|
||||||
|
# original single-pass confidence ranking. Re-use the
|
||||||
|
# pre-filtered ``valid_facts`` so we don't redo validation work
|
||||||
|
# on the hot fallback path.
|
||||||
|
logger.warning(
|
||||||
|
"Memory injection: guaranteed-category path failed, falling back to confidence-only ranking",
|
||||||
|
exc_info=True,
|
||||||
|
)
|
||||||
|
fallback, fallback_lines = _fallback_format_facts(
|
||||||
|
valid_facts,
|
||||||
|
preceding_section_cost=base_tokens,
|
||||||
|
max_tokens=max_tokens,
|
||||||
|
use_tiktoken=use_tiktoken,
|
||||||
|
)
|
||||||
|
if fallback:
|
||||||
|
sections.append(fallback)
|
||||||
|
# Surface the fallback's lines to ``all_fact_lines`` so the
|
||||||
|
# structure-aware truncation below treats fallback facts as a
|
||||||
|
# protected suffix too. Without this, a large user-context
|
||||||
|
# prefix could silently clip fallback facts via the original
|
||||||
|
# prefix-cut.
|
||||||
|
all_fact_lines = fallback_lines
|
||||||
|
|
||||||
if not sections:
|
if not sections:
|
||||||
return ""
|
return ""
|
||||||
|
|
||||||
result = "\n\n".join(sections)
|
result = "\n\n".join(sections)
|
||||||
|
|
||||||
# Use accurate token counting with tiktoken (or the char-based estimate
|
|
||||||
# when use_tiktoken is False).
|
|
||||||
token_count = _count_tokens(result, use_tiktoken=use_tiktoken)
|
token_count = _count_tokens(result, use_tiktoken=use_tiktoken)
|
||||||
if token_count > max_tokens:
|
effective_limit = max_tokens + guaranteed_line_tokens
|
||||||
# Truncate to fit within token limit
|
if token_count > effective_limit:
|
||||||
# Estimate characters to remove based on token ratio
|
# Structure-aware truncation: the ``Facts:\n...`` block is treated as
|
||||||
char_per_token = len(result) / token_count
|
# a *protected suffix* so guaranteed-category facts — the very facts
|
||||||
target_chars = int(max_tokens * char_per_token * 0.95) # 95% to leave margin
|
# this PR exists to preserve — can never be silently discarded by a
|
||||||
result = result[:target_chars] + "\n..."
|
# prefix-cut on overflow. Only the preceding (user-context / history)
|
||||||
|
# sections are eligible for truncation; if they alone exceed the
|
||||||
|
# budget available after reserving the Facts block, they are clipped
|
||||||
|
# from the tail. When *guaranteed_line_tokens* is zero (no
|
||||||
|
# guaranteed categories configured or no facts survived), the
|
||||||
|
# equation collapses to the original prefix-truncation against
|
||||||
|
# ``max_tokens``, so backward compatibility is preserved.
|
||||||
|
facts_block = (facts_header + "\n".join(all_fact_lines)) if all_fact_lines else ""
|
||||||
|
facts_block_tokens = _count_tokens(facts_block, use_tiktoken=use_tiktoken)
|
||||||
|
separator_tokens = _count_tokens("\n\n", use_tiktoken=use_tiktoken)
|
||||||
|
budget_for_non_facts = max(
|
||||||
|
0,
|
||||||
|
effective_limit - facts_block_tokens - (separator_tokens if facts_block else 0),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Build the preceding (non-facts) portion from *sections* excluding
|
||||||
|
# the trailing Facts block.
|
||||||
|
preceding_sections = sections[:-1] if all_fact_lines else sections
|
||||||
|
preceding = "\n\n".join(preceding_sections)
|
||||||
|
|
||||||
|
if preceding:
|
||||||
|
preceding_tokens = _count_tokens(preceding, use_tiktoken=use_tiktoken)
|
||||||
|
if preceding_tokens > budget_for_non_facts:
|
||||||
|
char_per_token = len(preceding) / max(preceding_tokens, 1)
|
||||||
|
target_chars = int(budget_for_non_facts * char_per_token * 0.95)
|
||||||
|
preceding = preceding[:target_chars].rstrip() + "\n..."
|
||||||
|
result = (preceding + "\n\n" + facts_block) if facts_block else preceding
|
||||||
|
else:
|
||||||
|
result = facts_block
|
||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|||||||
@ -1142,6 +1142,8 @@ class DeerFlowClient:
|
|||||||
"injection_enabled": config.injection_enabled,
|
"injection_enabled": config.injection_enabled,
|
||||||
"max_injection_tokens": config.max_injection_tokens,
|
"max_injection_tokens": config.max_injection_tokens,
|
||||||
"token_counting": config.token_counting,
|
"token_counting": config.token_counting,
|
||||||
|
"guaranteed_categories": config.guaranteed_categories,
|
||||||
|
"guaranteed_token_budget": config.guaranteed_token_budget,
|
||||||
}
|
}
|
||||||
|
|
||||||
def get_memory_status(self) -> dict:
|
def get_memory_status(self) -> dict:
|
||||||
|
|||||||
@ -73,6 +73,31 @@ class MemoryConfig(BaseModel):
|
|||||||
"CJK-aware character-based estimate and never touches tiktoken."
|
"CJK-aware character-based estimate and never touches tiktoken."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
guaranteed_categories: list[str] = Field(
|
||||||
|
default_factory=lambda: ["correction"],
|
||||||
|
description=(
|
||||||
|
"Fact categories that are always injected into the prompt regardless "
|
||||||
|
"of the regular token budget. These facts are allocated from a "
|
||||||
|
"separate reserved budget (``guaranteed_token_budget``). "
|
||||||
|
"This ensures high-value facts such as explicit user corrections "
|
||||||
|
"are never silently dropped when the token budget is tight."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
guaranteed_token_budget: int = Field(
|
||||||
|
default=500,
|
||||||
|
ge=50,
|
||||||
|
le=2000,
|
||||||
|
description=(
|
||||||
|
"Token ceiling for guaranteed-category facts. "
|
||||||
|
"Guaranteed facts are selected first from this budget and placed at "
|
||||||
|
"the front of the Facts block so they cannot be evicted by regular "
|
||||||
|
"facts. In the common case the total output still fits within "
|
||||||
|
"``max_injection_tokens`` (guaranteed lines displace regular ones); "
|
||||||
|
"the budget becomes additive only when guaranteed lines alone push "
|
||||||
|
"the output past ``max_injection_tokens``, in which case the "
|
||||||
|
"safety-truncation ceiling is raised accordingly."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# Global configuration instance
|
# Global configuration instance
|
||||||
|
|||||||
@ -204,7 +204,14 @@ def test_get_memory_context_uses_explicit_app_config_without_global_config(monke
|
|||||||
captured["user_id"] = user_id
|
captured["user_id"] = user_id
|
||||||
return {"facts": []}
|
return {"facts": []}
|
||||||
|
|
||||||
def fake_format_memory_for_injection(memory_data, *, max_tokens, use_tiktoken=True):
|
def fake_format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
*,
|
||||||
|
max_tokens,
|
||||||
|
use_tiktoken=True,
|
||||||
|
guaranteed_categories=None,
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
):
|
||||||
captured["memory_data"] = memory_data
|
captured["memory_data"] = memory_data
|
||||||
captured["max_tokens"] = max_tokens
|
captured["max_tokens"] = max_tokens
|
||||||
captured["use_tiktoken"] = use_tiktoken
|
captured["use_tiktoken"] = use_tiktoken
|
||||||
|
|||||||
@ -2,6 +2,8 @@
|
|||||||
|
|
||||||
import math
|
import math
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
from deerflow.agents.memory.prompt import _coerce_confidence, format_memory_for_injection
|
from deerflow.agents.memory.prompt import _coerce_confidence, format_memory_for_injection
|
||||||
|
|
||||||
|
|
||||||
@ -173,3 +175,482 @@ def test_format_memory_includes_long_term_background() -> None:
|
|||||||
assert "Background: Core expertise in distributed systems" in result
|
assert "Background: Core expertise in distributed systems" in result
|
||||||
assert "Recent: Recent activity summary" in result
|
assert "Recent: Recent activity summary" in result
|
||||||
assert "Earlier: Earlier context summary" in result
|
assert "Earlier: Earlier context summary" in result
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Guaranteed-category injection tests
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def test_guaranteed_correction_injected_when_budget_tight(monkeypatch) -> None:
|
||||||
|
"""Correction facts must be injected even when the regular budget is exhausted."""
|
||||||
|
# Deterministic char-based counting.
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._count_tokens",
|
||||||
|
lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text),
|
||||||
|
)
|
||||||
|
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
# Many high-confidence regular facts that will eat the budget.
|
||||||
|
{"content": "Regular fact A " * 20, "category": "knowledge", "confidence": 0.95},
|
||||||
|
{"content": "Regular fact B " * 20, "category": "knowledge", "confidence": 0.90},
|
||||||
|
{"content": "Regular fact C " * 20, "category": "knowledge", "confidence": 0.85},
|
||||||
|
# A correction fact with lower confidence.
|
||||||
|
{"content": "Use make dev, not npm start", "category": "correction", "confidence": 0.7},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
# Tight budget that cannot fit all facts.
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=200,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=100,
|
||||||
|
)
|
||||||
|
|
||||||
|
# The correction fact MUST appear regardless of budget pressure.
|
||||||
|
assert "Use make dev, not npm start" in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_guaranteed_facts_sorted_by_confidence() -> None:
|
||||||
|
"""Guaranteed facts should be sorted by confidence descending."""
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "Low conf correction", "category": "correction", "confidence": 0.6},
|
||||||
|
{"content": "High conf correction", "category": "correction", "confidence": 0.95},
|
||||||
|
{"content": "Regular fact", "category": "knowledge", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "High conf correction" in result
|
||||||
|
assert "Low conf correction" in result
|
||||||
|
assert result.index("High conf correction") < result.index("Low conf correction")
|
||||||
|
|
||||||
|
|
||||||
|
def test_guaranteed_budget_isolation() -> None:
|
||||||
|
"""Guaranteed facts draw from their own budget, not the regular budget."""
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "Correction one", "category": "correction", "confidence": 0.9},
|
||||||
|
{"content": "Regular knowledge", "category": "knowledge", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Both facts should appear (separate budgets).
|
||||||
|
assert "Correction one" in result
|
||||||
|
assert "Regular knowledge" in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_no_guaranteed_categories_backward_compatible() -> None:
|
||||||
|
"""When guaranteed_categories is None, behaviour matches the original."""
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "High conf", "category": "knowledge", "confidence": 0.95},
|
||||||
|
{"content": "Low conf", "category": "context", "confidence": 0.4},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
# No guaranteed_categories passed → original behaviour.
|
||||||
|
result = format_memory_for_injection(memory_data, max_tokens=2000)
|
||||||
|
|
||||||
|
assert "High conf" in result
|
||||||
|
assert result.index("High conf") < result.index("Low conf")
|
||||||
|
|
||||||
|
|
||||||
|
def test_empty_guaranteed_list_backward_compatible() -> None:
|
||||||
|
"""An empty guaranteed_categories list should behave like None."""
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "Correction fact", "category": "correction", "confidence": 0.9},
|
||||||
|
{"content": "Regular fact", "category": "knowledge", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "Correction fact" in result
|
||||||
|
assert "Regular fact" in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_fallback_on_ranking_error(monkeypatch) -> None:
|
||||||
|
"""If the guaranteed path raises, fall back to confidence-only ranking."""
|
||||||
|
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "Fact A", "category": "knowledge", "confidence": 0.9},
|
||||||
|
{"content": "Fact B", "category": "correction", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
# Force _select_fact_lines to raise on the *first* call (the guaranteed
|
||||||
|
# path) but succeed on subsequent calls (the fallback path).
|
||||||
|
call_count = {"n": 0}
|
||||||
|
prompt_module = __import__("deerflow.agents.memory.prompt", fromlist=["_select_fact_lines"])
|
||||||
|
original_select = prompt_module._select_fact_lines
|
||||||
|
|
||||||
|
def flaky_select(*args, **kwargs):
|
||||||
|
call_count["n"] += 1
|
||||||
|
if call_count["n"] == 1:
|
||||||
|
raise RuntimeError("simulated error in guaranteed path")
|
||||||
|
return original_select(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._select_fact_lines",
|
||||||
|
flaky_select,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Both facts should still appear via the fallback path.
|
||||||
|
assert "Fact A" in result
|
||||||
|
assert "Fact B" in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_guaranteed_respects_its_own_budget_limit(monkeypatch) -> None:
|
||||||
|
"""Even guaranteed facts are capped by guaranteed_token_budget."""
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._count_tokens",
|
||||||
|
lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Many correction facts that together exceed the guaranteed budget.
|
||||||
|
# Formatted line example: "- [correction | 0.95] CorrA xxxxxxxxxxxxxxxx"
|
||||||
|
# Each line is ~50 chars; with "Facts:\n" header (7 chars), two lines
|
||||||
|
# need ~107 chars, exceeding the 80-char guaranteed budget.
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "CorrA " + "x" * 20, "category": "correction", "confidence": 0.95},
|
||||||
|
{"content": "CorrB " + "x" * 20, "category": "correction", "confidence": 0.90},
|
||||||
|
{"content": "CorrC " + "x" * 20, "category": "correction", "confidence": 0.85},
|
||||||
|
{"content": "Short regular", "category": "knowledge", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=80, # Small guaranteed budget — fits 1 fact line only.
|
||||||
|
)
|
||||||
|
|
||||||
|
# At least the highest-confidence correction should appear.
|
||||||
|
assert "CorrA" in result
|
||||||
|
# The regular fact should also appear (it has its own budget).
|
||||||
|
assert "Short regular" in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_guaranteed_fact_with_source_error_rendered() -> None:
|
||||||
|
"""Guaranteed correction facts should still render sourceError."""
|
||||||
|
memory_data = {
|
||||||
|
"facts": [
|
||||||
|
{
|
||||||
|
"content": "Use uv, not pip.",
|
||||||
|
"category": "correction",
|
||||||
|
"confidence": 0.95,
|
||||||
|
"sourceError": "Agent suggested pip install.",
|
||||||
|
},
|
||||||
|
{"content": "Likes Python", "category": "preference", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "Use uv, not pip." in result
|
||||||
|
assert "avoid: Agent suggested pip install." in result
|
||||||
|
assert "Likes Python" in result
|
||||||
|
|
||||||
|
|
||||||
|
def test_single_facts_header_when_both_guaranteed_and_regular() -> None:
|
||||||
|
"""When both guaranteed and regular facts exist, emit exactly one 'Facts:' header."""
|
||||||
|
memory_data = {
|
||||||
|
"user": {"workContext": {"summary": "Dev"}}, # non-empty preceding section
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "Correction fact", "category": "correction", "confidence": 0.95},
|
||||||
|
{"content": "Knowledge fact", "category": "knowledge", "confidence": 0.80},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Exactly one "Facts:" header.
|
||||||
|
assert result.count("Facts:") == 1, f"Expected exactly one 'Facts:' header, got:\n{result}"
|
||||||
|
# Both facts appear under the single header.
|
||||||
|
assert "Correction fact" in result
|
||||||
|
assert "Knowledge fact" in result
|
||||||
|
# Guaranteed fact comes first (higher confidence + guaranteed).
|
||||||
|
assert result.index("Correction fact") < result.index("Knowledge fact")
|
||||||
|
|
||||||
|
|
||||||
|
def test_strict_confidence_order_when_high_confidence_fact_overflows(monkeypatch) -> None:
|
||||||
|
"""Within a single budget, a higher-confidence fact that exceeds the
|
||||||
|
remaining budget must NOT be skipped in favour of a shorter, lower-
|
||||||
|
confidence fact ranked after it.
|
||||||
|
|
||||||
|
This locks in the strict confidence-ordered selection semantics.
|
||||||
|
"""
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._count_tokens",
|
||||||
|
lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text),
|
||||||
|
)
|
||||||
|
|
||||||
|
memory_data = {
|
||||||
|
"user": {},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
# Higher-confidence but long enough to exceed the remaining budget.
|
||||||
|
{"content": "Long high-confidence fact " + "x" * 50, "category": "knowledge", "confidence": 0.95},
|
||||||
|
# Lower-confidence but short — would fit if we kept scanning past
|
||||||
|
# the over-budget high-confidence fact above.
|
||||||
|
{"content": "Short low", "category": "knowledge", "confidence": 0.50},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
# Budget large enough only for ~one short fact, not the long one.
|
||||||
|
result = format_memory_for_injection(memory_data, max_tokens=70, guaranteed_categories=None)
|
||||||
|
|
||||||
|
# The high-confidence fact does not fit, and the low-confidence fact
|
||||||
|
# MUST NOT slip in ahead of it.
|
||||||
|
assert "Short low" not in result, "Lower-confidence fact should not be selected when a higher-confidence fact ranked before it was skipped (strict ordering)."
|
||||||
|
|
||||||
|
|
||||||
|
# ── Regression tests for willem-bd's review on PR #3592 ──────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def test_structure_aware_truncation_preserves_guaranteed_on_overflow(monkeypatch) -> None:
|
||||||
|
"""[P1] When user context overflows, the trailing ``Facts:\\n...`` block
|
||||||
|
is treated as a protected suffix and only the preceding sections are
|
||||||
|
clipped — guaranteed-category facts can never be silently discarded by
|
||||||
|
a prefix-cut on overflow.
|
||||||
|
|
||||||
|
Locks in the fix for willem-bd's P1 finding on PR #3592.
|
||||||
|
"""
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._count_tokens",
|
||||||
|
lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text),
|
||||||
|
)
|
||||||
|
|
||||||
|
memory_data = {
|
||||||
|
# Oversized preceding section that would otherwise push Facts past the
|
||||||
|
# effective truncation ceiling.
|
||||||
|
"user": {"workContext": {"summary": "X" * 4000}},
|
||||||
|
"facts": [
|
||||||
|
{
|
||||||
|
"content": "CRITICAL: never use pip",
|
||||||
|
"category": "correction",
|
||||||
|
"confidence": 1.0,
|
||||||
|
"sourceError": "pip is deprecated",
|
||||||
|
},
|
||||||
|
{"content": "B", "category": "knowledge", "confidence": 0.5},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=200,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
guaranteed_token_budget=500,
|
||||||
|
use_tiktoken=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Guaranteed correction must survive even when preceding sections are huge.
|
||||||
|
assert "never use pip" in result, f"Guaranteed correction was silently truncated away:\n{result[-200:]}"
|
||||||
|
assert "pip is deprecated" in result
|
||||||
|
# The protected suffix shape: Facts block is at the tail.
|
||||||
|
assert result.rstrip().endswith("(avoid: pip is deprecated)")
|
||||||
|
|
||||||
|
|
||||||
|
def test_single_inter_section_separator_between_user_and_facts() -> None:
|
||||||
|
"""[P2] Exactly one ``\\n\\n`` separator between ``User Context:`` and
|
||||||
|
``Facts:`` — never four newlines.
|
||||||
|
|
||||||
|
Locks in the fix for willem-bd's P2 separator finding on PR #3592.
|
||||||
|
"""
|
||||||
|
memory_data = {
|
||||||
|
"user": {"workContext": {"summary": "Python developer"}},
|
||||||
|
"history": {},
|
||||||
|
"facts": [
|
||||||
|
{"content": "fact A", "category": "knowledge", "confidence": 0.9},
|
||||||
|
{
|
||||||
|
"content": "fact B",
|
||||||
|
"category": "correction",
|
||||||
|
"confidence": 0.8,
|
||||||
|
"sourceError": "avoid X",
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(memory_data, max_tokens=2000)
|
||||||
|
|
||||||
|
assert "\n\n\n\n" not in result, f"Found four consecutive newlines between sections:\n{result[:200]!r}"
|
||||||
|
# Exactly one \n\n between User Context: and Facts:.
|
||||||
|
idx_user = result.index("User Context:")
|
||||||
|
idx_facts = result.index("Facts:")
|
||||||
|
between = result[idx_user:idx_facts]
|
||||||
|
assert between.count("\n\n") == 1, f"Expected exactly one \\n\\n between sections, got:\n{between!r}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_bare_string_guaranteed_categories_raises_type_error() -> None:
|
||||||
|
"""[P2] Passing a bare ``str`` for *guaranteed_categories* must raise
|
||||||
|
``TypeError`` instead of silently iterating single characters and
|
||||||
|
disabling the guarantee.
|
||||||
|
|
||||||
|
Locks in the fix for willem-bd's P2 bare-string finding on PR #3592.
|
||||||
|
"""
|
||||||
|
memory_data = {
|
||||||
|
"facts": [
|
||||||
|
{"content": "CRITICAL", "category": "correction", "confidence": 0.8},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
with pytest.raises(TypeError, match="iterable"):
|
||||||
|
format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
guaranteed_categories="correction", # type: ignore[arg-type]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_categoryless_fact_not_promoted_into_guaranteed_context_pool(monkeypatch) -> None:
|
||||||
|
"""[P2] A fact with a missing/empty ``category`` field is *never*
|
||||||
|
silently promoted into a ``guaranteed_categories=["context"]`` pool —
|
||||||
|
only facts with an *explicit* ``category == "context"`` qualify.
|
||||||
|
|
||||||
|
Strategy: set a guaranteed budget tight enough to fit only the short
|
||||||
|
*explicit* ``context`` fact. If the legacy (no-category) fact were
|
||||||
|
silently promoted into the guaranteed pool, it would claim the budget
|
||||||
|
first (higher confidence) and push the explicit one out into the
|
||||||
|
regular pool where, under a tight ``max_tokens``, it would be lost.
|
||||||
|
If the fix holds, the explicit fact owns the guaranteed pool alone
|
||||||
|
and survives.
|
||||||
|
|
||||||
|
Locks in the fix for willem-bd's P2 category-less finding on PR #3592.
|
||||||
|
"""
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._count_tokens",
|
||||||
|
lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text),
|
||||||
|
)
|
||||||
|
|
||||||
|
memory_data = {
|
||||||
|
"facts": [
|
||||||
|
# Long legacy fact with NO category field.
|
||||||
|
{
|
||||||
|
"content": "legacy " + "x" * 80,
|
||||||
|
"confidence": 0.95,
|
||||||
|
},
|
||||||
|
# Short explicit context fact.
|
||||||
|
{
|
||||||
|
"content": "explicit ctx",
|
||||||
|
"category": "context",
|
||||||
|
"confidence": 0.9,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
# Guaranteed budget sized for the short explicit fact only.
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=200,
|
||||||
|
guaranteed_categories=["context"],
|
||||||
|
guaranteed_token_budget=40,
|
||||||
|
use_tiktoken=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# The explicit context fact must survive in the guaranteed pool.
|
||||||
|
assert "explicit ctx" in result, f"Explicit 'context' fact was evicted — legacy no-category fact was silently promoted into the guaranteed pool.\n{result!r}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_fallback_uses_prefiltered_valid_facts(monkeypatch) -> None:
|
||||||
|
"""[P2] When the primary path raises after ``valid_facts`` has been
|
||||||
|
built, the fallback operates on the pre-filtered list (no raw-content
|
||||||
|
facts leak through) and still produces a valid ``Facts:`` section.
|
||||||
|
|
||||||
|
Locks in the fix for willem-bd's P2 fallback-duplication finding on
|
||||||
|
PR #3592.
|
||||||
|
"""
|
||||||
|
monkeypatch.setattr(
|
||||||
|
"deerflow.agents.memory.prompt._count_tokens",
|
||||||
|
lambda text, encoding_name="cl100k_base", *, use_tiktoken=True: len(text),
|
||||||
|
)
|
||||||
|
|
||||||
|
call_count = {"select": 0}
|
||||||
|
original_select = __import__("deerflow.agents.memory.prompt", fromlist=["_select_fact_lines"])._select_fact_lines
|
||||||
|
|
||||||
|
def raising_select(*args, **kwargs):
|
||||||
|
call_count["select"] += 1
|
||||||
|
if call_count["select"] == 1:
|
||||||
|
raise RuntimeError("primary path failure")
|
||||||
|
return original_select(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr("deerflow.agents.memory.prompt._select_fact_lines", raising_select)
|
||||||
|
|
||||||
|
memory_data = {
|
||||||
|
"facts": [
|
||||||
|
{"content": "valid fact", "category": "knowledge", "confidence": 0.9},
|
||||||
|
# Malformed: no content field — should be pre-filtered and never
|
||||||
|
# reach the fallback's ranking.
|
||||||
|
{"category": "knowledge", "confidence": 0.95},
|
||||||
|
# Empty content — also pre-filtered.
|
||||||
|
{"content": " ", "category": "knowledge", "confidence": 0.9},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
result = format_memory_for_injection(
|
||||||
|
memory_data,
|
||||||
|
max_tokens=2000,
|
||||||
|
guaranteed_categories=["correction"],
|
||||||
|
use_tiktoken=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Fallback kicked in and still produced the Facts section.
|
||||||
|
assert "Facts:" in result
|
||||||
|
# The valid fact survived pre-filtering and fallback ranking.
|
||||||
|
assert "valid fact" in result
|
||||||
|
# Malformed facts were pre-filtered and never rendered.
|
||||||
|
assert result.count("- [") == 1
|
||||||
|
|||||||
@ -1191,6 +1191,20 @@ memory:
|
|||||||
# char - network-free CJK-aware character-based estimate; never touches
|
# char - network-free CJK-aware character-based estimate; never touches
|
||||||
# tiktoken. Slightly less precise budgeting, zero network I/O.
|
# tiktoken. Slightly less precise budgeting, zero network I/O.
|
||||||
token_counting: tiktoken
|
token_counting: tiktoken
|
||||||
|
# Guaranteed injection: fact categories that bypass the regular token budget
|
||||||
|
# and draw from a reserved allowance, so high-signal corrections (e.g.
|
||||||
|
# "don't use `pip`, use `uv`") survive even when the budget is tight.
|
||||||
|
# guaranteed_categories - list of fact categories to guarantee. Pass [] to
|
||||||
|
# disable; defaults to ["correction"].
|
||||||
|
# guaranteed_token_budget - token ceiling for guaranteed facts. In the
|
||||||
|
# common case the total injection stays within ``max_injection_tokens``
|
||||||
|
# (guaranteed lines displace regular ones); the allowance becomes
|
||||||
|
# additive only when guaranteed lines alone would overflow
|
||||||
|
# ``max_injection_tokens``, in which case the safety-truncation ceiling
|
||||||
|
# is raised accordingly.
|
||||||
|
guaranteed_categories:
|
||||||
|
- correction
|
||||||
|
guaranteed_token_budget: 500
|
||||||
|
|
||||||
# ============================================================================
|
# ============================================================================
|
||||||
# Custom Agent Management API
|
# Custom Agent Management API
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user