mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 02:56:17 +00:00
feat(memory): add opt-in relevance-aware retrieval ranking (#5251)
* feat(memory): add opt-in relevance-aware retrieval ranking Add a deterministic, network-free lexical relevance strategy for DeerMem (issue #4495): memory_search ranks every fact in scope by idf-weighted token overlap combined with confidence, with optional greedy-MMR diversity against near-duplicate facts; prompt injection ranks facts against the current-turn query threaded from DynamicContextMiddleware through the new optional `query` keyword on MemoryManager.get_context/aget_context. Defaults preserve the legacy confidence-only behavior exactly; no prompt, storage-format, or vector/embedding-dependency changes. Refs #4495 Signed-off-by: pwd11 <fvdsrc@163.com> * fix(memory): bound relevance retrieval and apply review feedback Bound tokenization and index shared stems, preserve mixed CJK tokens, warm jieba, and align missing confidence with legacy injection. Cache MMR token sets and stop selection at result or injection budgets. Document retrieval-adapter precedence and add regression coverage. Refs #4495. Signed-off-by: pwd11 <fvdsrc@163.com> * fix(memory): preserve backend compatibility and normalize relevance Signed-off-by: pwd11 <fvdsrc@163.com> * fix(memory): omit absent query hints and share injection IDF Signed-off-by: pwd11 <fvdsrc@163.com> * test(memory): retain timeout mock until injection worker exits Signed-off-by: pwd11 <fvdsrc@163.com> * docs(agents): drop root guidance compaction Signed-off-by: pwd11 <fvdsrc@163.com> * fix(memory): validate token prefixes and preserve upload queries --------- Signed-off-by: pwd11 <fvdsrc@163.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
wd_pan
GitHub
parent
972020cf85
commit
cc27730348
14
README.md
14
README.md
@ -1618,6 +1618,20 @@ When a fact scope reaches `max_facts`, DeerMem still uses the historical confide
|
||||
|
||||
File-backed memory now separates global user context from agent facts. Each user has one `memory.json` containing only the project-independent `user` and `history` summaries; every fact is a canonical Markdown file below `agents/{agent_name}/facts/`. Existing lead-agent middleware, API, Settings, import/export, and embedded-client calls that omit `agent_name` resolve inside DeerMem to the reserved `__default__` bucket. That bucket is outside the valid custom-agent name grammar, so a real custom agent named `lead-agent` has a separate fact repository and deleting a custom agent cannot delete a memory-only directory without `config.yaml`. Public agent identifiers are case-insensitive and canonicalized to lowercase. Runtime/API readers still receive a compatibility `facts` array for the selected/default agent, so the frontend does not read agent facts from `memory.json`; structured Markdown `source` metadata is projected to the historical string field at the MemoryManager boundary. An unscoped Clear All first migrates facts from unread legacy per-agent JSON without adopting its soon-to-be-cleared summaries, then removes shared summaries and facts from every agent bucket while preserving agent configuration files, so a later read cannot resurrect skipped legacy facts; an explicitly agent-scoped clear removes only that agent's facts. On first normal read, old facts embedded in the user JSON are migrated automatically to `__default__`; facts written to the earlier implicit `lead-agent` bucket are also moved when that directory is not a real custom agent. Migration and normal writes notify the configured retrieval adapter only after durable storage locks are released. DeerMem uses a scope-aware SQLite FTS5/BM25 adapter by default, stores only rebuildable derived index data under `.retrieval/`, and rebuilds it in the background during Gateway startup or lazily on the first scoped search. A corrupt derived index is recreated automatically. Set `memory.backend_config.retrieval_adapter` to an empty string to disable it and use the local substring fallback. Chinese tokenization is optional; install the backend `memory-zh` extra (`uv sync --extra memory-zh`) for jieba-assisted sub-phrase search. Journaled writes, a shared user lock, and optimistic user-memory revisions prevent silent lost updates.
|
||||
|
||||
Set `memory.backend_config.retrieval_relevance_enabled: true` to opt into deterministic relevance/confidence ranking and query-aware injection. This **bypasses `retrieval_adapter` for search**, including the default FTS5/BM25 and custom adapters; indexing remains configured. `retrieval_relevance_weight` controls the blend and `retrieval_diversity_weight` enables near-duplicate penalties (default 0). Scoring uses at most the first 4096 characters and 128 tokens per query/fact. Search diversifies only up to `top_k`; injection diversifies guaranteed and regular facts independently until their token budgets are reached. Leave the feature disabled to retain the existing retrieval and injection behavior.
|
||||
|
||||
Lexical relevance measures IDF-weighted coverage of distinct query terms, so
|
||||
repeated partial matches cannot tie a complete match merely by saturating the
|
||||
score. Prefix matches require one complete token to prefix the other, not just
|
||||
four shared characters. For uploads, query-aware injection uses the preserved
|
||||
user request rather than the prepended file descriptions; attachment-only
|
||||
messages retain query-less injection. Older custom memory backends can keep their existing `get_context`
|
||||
signature: prompt injection and the inherited async wrapper pass `query` only
|
||||
when that callable supports the keyword and the hint is not `None`.
|
||||
Search and automatic injection use the same IDF weighting for the same
|
||||
user/agent candidate scope. Category-filtered search and injection's separate
|
||||
guaranteed/regular token budgets can still yield different final selections.
|
||||
|
||||
Memory injection follows the configured operation mode. In `middleware` mode, DeerMem injects the user-global summaries and the selected agent's facts. Custom-agent bootstrap conversations use that agent's fact bucket as well, so setup details do not leak into the default agent's memory. In `tool` mode, the automatic `<memory>` block contains only the global `user` and `history` summaries; agent facts are retrieved explicitly through `memory_search`, avoiding duplicate automatic and tool-returned fact context. Setting `memory.injection_enabled: false` still disables the entire block in either mode.
|
||||
|
||||
An individual Custom Agent can opt out of memory without changing the global setting. Add `memory_enabled: false` to that agent's `users/{user_id}/agents/{name}/config.yaml`. The agent still receives the current-date reminder, but DeerFlow does not inject recalled memory, queue passive or summarization-driven memory updates (including manual `/compact`), expose memory tools, or add memory-tool instructions for that agent. If an existing agent is switched off, its previously injected memory block is removed from checkpoint state before the next model call while its date reminder and conversation remain. Omitting the field (or setting it to `true`) preserves the existing global `memory` behavior.
|
||||
|
||||
@ -777,6 +777,7 @@ def _get_memory_context(
|
||||
*,
|
||||
app_config: AppConfig | None = None,
|
||||
user_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
"""Get memory context for injection into system prompt.
|
||||
|
||||
@ -786,6 +787,10 @@ def _get_memory_context(
|
||||
are read from this value instead of the global config singleton.
|
||||
user_id: Explicit user bucket. When omitted, resolves the current
|
||||
Gateway or standalone LangGraph Server identity.
|
||||
query: Optional current-turn query hint forwarded to the memory
|
||||
backend. Backends that enable query-aware ranking (DeerMem
|
||||
``retrieval_relevance_enabled``) rank injected facts against it;
|
||||
others ignore it.
|
||||
|
||||
Returns:
|
||||
Formatted memory context string wrapped in XML tags, or empty string if disabled.
|
||||
@ -795,6 +800,7 @@ def _get_memory_context(
|
||||
config = None
|
||||
try:
|
||||
from deerflow.agents.memory import get_memory_manager
|
||||
from deerflow.agents.memory.manager import context_query_kwargs
|
||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||
|
||||
if app_config is None:
|
||||
@ -807,9 +813,11 @@ def _get_memory_context(
|
||||
if not config.enabled or not config.injection_enabled:
|
||||
return ""
|
||||
|
||||
memory_content = get_memory_manager().get_context(
|
||||
manager = get_memory_manager()
|
||||
memory_content = manager.get_context(
|
||||
user_id=user_id or resolve_runtime_user_id(None),
|
||||
agent_name=agent_name,
|
||||
**context_query_kwargs(manager.get_context, query),
|
||||
)
|
||||
|
||||
if not memory_content.strip():
|
||||
|
||||
@ -300,3 +300,45 @@ Latin words and CJK bigrams both participate in mixed-script similarity.
|
||||
Whitespace-separated CJK runs retain adjacent-character ordering.
|
||||
INFO logs identify the target and proposal index without memory content and
|
||||
explicitly describe a proposed merge, not a completed persistence audit.
|
||||
|
||||
#### Relevance-aware retrieval (opt-in)
|
||||
|
||||
The deterministic lexical strategy behind issue #4495 lives in
|
||||
`deermem/core/relevance.py` (token overlap + idf weights + confidence blend +
|
||||
greedy MMR diversity). It never touches the persisted memory format and never
|
||||
runs by default.
|
||||
|
||||
- `retrieval_relevance_enabled: true` opts in. `memory_search` then ranks every
|
||||
fact in scope (not only literal substring matches) and prompt injection ranks
|
||||
facts against the current query before the token-budget selection.
|
||||
This takes precedence over `retrieval_adapter`: search bypasses FTS5/custom
|
||||
retrieval, while adapter indexing and warm-up remain configured.
|
||||
- Ranking reads at most 4096 characters and 128 tokens per query/fact. The
|
||||
no-jieba fallback emits both Latin words and CJK bigrams, including mixed text.
|
||||
`DeerMem.warm()` initializes optional jieba before serving requests, even
|
||||
with character-based token counting. Invalid/missing confidence defaults to 0.
|
||||
- Search stops MMR after `top_k` picks. Injection diversifies guaranteed and
|
||||
regular pools independently and lazily, stopping when each token budget is
|
||||
exhausted; it never truncates candidates before the guaranteed partition.
|
||||
MMR caches token sets and incrementally updates maximum similarity penalties.
|
||||
- `retrieval_relevance_weight` blends lexical relevance with confidence;
|
||||
`retrieval_diversity_weight` demotes near-duplicate facts. Defaults preserve
|
||||
legacy ordering. Relevance is distinct-query-token IDF coverage; repeated
|
||||
content cannot replace missing terms or saturate a partial match.
|
||||
Prefix matching requires one complete token to prefix the other; a shared
|
||||
four-character bucket alone is not a match (Postman is not PostgreSQL).
|
||||
- DeerMem injection builds IDF once from the selected user/agent fact scope,
|
||||
before guaranteed/regular partitioning, using the same bounded tokenizer as
|
||||
search. No IDF work runs without an active lexical query. Category-filtered
|
||||
search uses its filtered corpus; budgets and separate diversity pools can
|
||||
still produce different final selections. No IDF cache crosses calls/scopes.
|
||||
- The current-turn query flows from `DynamicContextMiddleware` (bounded,
|
||||
user-message text) through the optional `query` keyword on
|
||||
`MemoryManager.get_context` / `aget_context`. Shared signature inspection
|
||||
omits `query` when it is `None` or the backend is old/uninspectable, preserving
|
||||
forwarding wrappers' absent-hint contract; backend errors never cause retries.
|
||||
Query extraction prefers preserved `original_user_content` before applying
|
||||
the character cap, so upload descriptions never displace the user's request.
|
||||
Attachment-only messages with an empty preserved request stay query-less.
|
||||
- Ranking must be deterministic, network-free, and mutation-free: caller-owned
|
||||
fact dicts are read-only inputs.
|
||||
|
||||
@ -44,6 +44,7 @@ from .deermem.core.message_processing import (
|
||||
from .deermem.core.paths import DEFAULT_AGENT_BUCKET
|
||||
from .deermem.core.prompt import format_memory_for_injection, load_prompt, load_prompt_messages, warm_tiktoken_cache
|
||||
from .deermem.core.queue import MemoryUpdateQueue, QueueFull
|
||||
from .deermem.core.relevance import build_idf, order_facts_for_query, tokenize, warm_tokenizer
|
||||
from .deermem.core.storage import MemoryRevisionConflict, MemoryStorageCorruption, create_storage
|
||||
from .deermem.core.updater import MemoryUpdater, _coerce_source_confidence
|
||||
|
||||
@ -300,6 +301,7 @@ class DeerMem(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
"""Load memory and format it for injection (plain text, no wrap).
|
||||
|
||||
@ -308,6 +310,11 @@ class DeerMem(MemoryManager):
|
||||
facts stay behind ``memory_search`` so they are not duplicated in the
|
||||
prompt and a later retrieval result.
|
||||
|
||||
When ``query`` is provided and ``retrieval_relevance_enabled`` is
|
||||
true, facts are ranked against the query (lexical relevance combined
|
||||
with confidence, then optional diversity) before the token-budget
|
||||
selection. Without a query the legacy confidence ordering is kept.
|
||||
|
||||
Format parameters come from DeerMem's own ``DeerMemConfig`` (set at
|
||||
construction from ``backend_config``). The ``enabled``/
|
||||
``injection_enabled`` gate and the ``<memory>`` wrapping stay at the
|
||||
@ -315,12 +322,24 @@ class DeerMem(MemoryManager):
|
||||
"""
|
||||
injection_agent = None if self.mode == "tool" else _resolve_agent_name(agent_name)
|
||||
memory_data = _call_backend(lambda: self._updater.get_memory_data(agent_name=injection_agent, user_id=user_id))
|
||||
relevance_enabled = self._config.retrieval_relevance_enabled and bool(query and query.strip())
|
||||
corpus_idf = None
|
||||
if relevance_enabled and self._config.retrieval_relevance_weight > 0:
|
||||
# Use the selected user/agent corpus, before budget-pool partitioning,
|
||||
# with the same bounded tokenizer and IDF as unfiltered search.
|
||||
facts = memory_data.get("facts", [])
|
||||
if isinstance(facts, list) and facts:
|
||||
corpus_idf = build_idf([tokenize(fact["content"]) for fact in facts if isinstance(fact, dict) and isinstance(fact.get("content"), str)])
|
||||
return format_memory_for_injection(
|
||||
memory_data,
|
||||
max_tokens=self._config.max_injection_tokens,
|
||||
use_tiktoken=(self._config.token_counting == "tiktoken"),
|
||||
guaranteed_categories=self._config.guaranteed_categories,
|
||||
guaranteed_token_budget=self._config.guaranteed_token_budget,
|
||||
query=query if relevance_enabled else None,
|
||||
relevance_weight=self._config.retrieval_relevance_weight if relevance_enabled else None,
|
||||
diversity_weight=self._config.retrieval_diversity_weight if relevance_enabled else None,
|
||||
idf=corpus_idf,
|
||||
)
|
||||
|
||||
def search(
|
||||
@ -336,18 +355,31 @@ class DeerMem(MemoryManager):
|
||||
|
||||
Retrieval errors never make canonical memory unavailable: the existing
|
||||
case-insensitive substring path remains the last-resort fallback.
|
||||
|
||||
With ``retrieval_relevance_enabled``, the opt-in relevance-aware
|
||||
strategy ranks every fact in scope (adapter-free, deterministic) and
|
||||
may return related facts without a literal substring match.
|
||||
"""
|
||||
if not query or not query.strip() or top_k <= 0:
|
||||
return []
|
||||
resolved_agent_name = _resolve_agent_name(agent_name)
|
||||
indexed = self._fts5_search(query, top_k=top_k, user_id=user_id, agent_name=resolved_agent_name, category=category)
|
||||
results = indexed or self._substring_search(
|
||||
query,
|
||||
top_k=top_k,
|
||||
user_id=user_id,
|
||||
agent_name=resolved_agent_name,
|
||||
category=category,
|
||||
)
|
||||
if self._config.retrieval_relevance_enabled:
|
||||
results = self._relevance_search(
|
||||
query,
|
||||
top_k=top_k,
|
||||
user_id=user_id,
|
||||
agent_name=resolved_agent_name,
|
||||
category=category,
|
||||
)
|
||||
else:
|
||||
indexed = self._fts5_search(query, top_k=top_k, user_id=user_id, agent_name=resolved_agent_name, category=category)
|
||||
results = indexed or self._substring_search(
|
||||
query,
|
||||
top_k=top_k,
|
||||
user_id=user_id,
|
||||
agent_name=resolved_agent_name,
|
||||
category=category,
|
||||
)
|
||||
if results and (self._config.fact_eviction_policy == EVICTION_POLICY_HYBRID_V1 or self._config.fact_eviction_shadow_enabled):
|
||||
try:
|
||||
self._storage.record_fact_accesses(
|
||||
@ -410,6 +442,38 @@ class DeerMem(MemoryManager):
|
||||
matched.sort(key=_coerce_source_confidence, reverse=True)
|
||||
return _compat_document({"facts": matched[:top_k]})["facts"]
|
||||
|
||||
def _relevance_search(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
top_k: int,
|
||||
user_id: str | None,
|
||||
agent_name: str | None,
|
||||
category: str | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Opt-in deterministic ranking over every fact in scope (issue #4495).
|
||||
|
||||
Candidates are NOT limited to literal substring matches: all facts in
|
||||
the requested user/agent scope compete, ranked by lexical relevance
|
||||
combined with confidence, then diversified. The category filter still
|
||||
applies before the ``top_k`` slice. Caller-owned fact dicts are never
|
||||
mutated.
|
||||
"""
|
||||
memory_data = _call_backend(lambda: self._updater.get_memory_data(agent_name=agent_name, user_id=user_id))
|
||||
facts = [fact for fact in memory_data.get("facts", []) if isinstance(fact.get("content"), str) and (category is None or fact.get("category") == category)]
|
||||
if not facts:
|
||||
return []
|
||||
corpus_idf = build_idf([tokenize(fact["content"]) for fact in facts])
|
||||
ranked = order_facts_for_query(
|
||||
facts,
|
||||
query,
|
||||
relevance_weight=self._config.retrieval_relevance_weight,
|
||||
diversity_weight=self._config.retrieval_diversity_weight,
|
||||
idf=corpus_idf,
|
||||
limit=top_k,
|
||||
)
|
||||
return _compat_document({"facts": ranked})["facts"]
|
||||
|
||||
def _ensure_retrieval_scopes(self, scopes: list[dict[str, str | None]]) -> None:
|
||||
"""Lazily rebuild every requested scope when warm-up was skipped."""
|
||||
if not hasattr(self, "_retrieval_lock"):
|
||||
@ -533,6 +597,8 @@ class DeerMem(MemoryManager):
|
||||
or warming was unnecessary); False if tiktoken is unavailable or the
|
||||
download failed.
|
||||
"""
|
||||
if self._config.retrieval_relevance_enabled:
|
||||
warm_tokenizer()
|
||||
if self._config.token_counting == "char":
|
||||
logger.info("token_counting='char'; tiktoken not used, skipping warm-up")
|
||||
return True
|
||||
|
||||
@ -71,7 +71,10 @@ class DeerMemConfig(BaseModel):
|
||||
)
|
||||
retrieval_adapter: str = Field(
|
||||
default="fts5",
|
||||
description="Retrieval adapter factory: 'fts5' (default), an empty string to disable, or a dotted factory receiving DeerMemConfig and implementing RetrievalPort.",
|
||||
description=(
|
||||
"Retrieval adapter factory: 'fts5' (default), an empty string to disable, or a dotted factory receiving DeerMemConfig and implementing RetrievalPort. "
|
||||
"Search bypasses this adapter when retrieval_relevance_enabled is true; indexing remains configured."
|
||||
),
|
||||
)
|
||||
fact_dedup_enabled: bool = Field(
|
||||
default=False,
|
||||
@ -94,6 +97,30 @@ class DeerMemConfig(BaseModel):
|
||||
le=1.0,
|
||||
description=("Minimum bounded token-Jaccard similarity for the write-side near-duplicate merge gate. Used only when fact_dedup_enabled is true."),
|
||||
)
|
||||
retrieval_relevance_enabled: bool = Field(
|
||||
default=False,
|
||||
description=(
|
||||
"Opt-in relevance-aware retrieval (issue #4495). When true, "
|
||||
"search bypasses retrieval_adapter (including FTS5 and custom factories); "
|
||||
"memory_search ranks all facts in scope by deterministic lexical "
|
||||
"relevance combined with confidence, related facts are returned "
|
||||
"even without a literal substring match, and prompt injection "
|
||||
"ranks facts against the current query. False preserves the "
|
||||
"legacy confidence-based behavior exactly."
|
||||
),
|
||||
)
|
||||
retrieval_relevance_weight: float = Field(
|
||||
default=0.5,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Weight of lexical relevance vs confidence in the combined retrieval score. 0.0 = confidence only; 1.0 = relevance only. Used only when retrieval_relevance_enabled is true.",
|
||||
)
|
||||
retrieval_diversity_weight: float = Field(
|
||||
default=0.0,
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Greedy MMR similarity penalty that demotes near-duplicate facts during relevance-aware ranking. 0.0 (default) = no diversification. Used only when retrieval_relevance_enabled is true.",
|
||||
)
|
||||
# ── Queue ────────────────────────────────────────────────────────────
|
||||
debounce_seconds: int = Field(
|
||||
default=30,
|
||||
|
||||
@ -8,12 +8,15 @@ import math
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
import yaml
|
||||
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
|
||||
|
||||
from .relevance import iter_diversify, score_facts
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@ -383,7 +386,7 @@ def _escape_summary(value: Any) -> str:
|
||||
|
||||
|
||||
def _select_fact_lines(
|
||||
ranked_facts: list[dict[str, Any]],
|
||||
ranked_facts: Iterable[dict[str, Any]],
|
||||
*,
|
||||
token_budget: int,
|
||||
use_tiktoken: bool,
|
||||
@ -470,6 +473,10 @@ def format_memory_for_injection(
|
||||
use_tiktoken: bool = True,
|
||||
guaranteed_categories: list[str] | None = None,
|
||||
guaranteed_token_budget: int = 500,
|
||||
query: str | None = None,
|
||||
relevance_weight: float | None = None,
|
||||
diversity_weight: float | None = None,
|
||||
idf: dict[str, float] | None = None,
|
||||
) -> str:
|
||||
"""Format memory data for injection into system prompt.
|
||||
|
||||
@ -491,6 +498,20 @@ def format_memory_for_injection(
|
||||
point the safety-truncation ceiling is raised to
|
||||
``max_tokens + guaranteed_actual_usage`` to protect them.
|
||||
Ignored when *guaranteed_categories* is ``None`` or empty.
|
||||
query: Optional current-turn query. When provided together with
|
||||
``relevance_weight``, facts are first ranked by deterministic
|
||||
lexical relevance combined with confidence (issue #4495) before
|
||||
the guaranteed/regular partition and budget selection. ``None``
|
||||
preserves the legacy confidence-only ordering.
|
||||
relevance_weight: Weight of lexical relevance vs confidence for the
|
||||
query-aware ranking (0.0 = confidence only). Ignored when
|
||||
``query`` is ``None``.
|
||||
diversity_weight: Optional greedy-MMR similarity penalty that demotes
|
||||
near-duplicate facts during query-aware ranking. Ignored when
|
||||
``query`` is ``None``.
|
||||
idf: Optional scope-wide query-term weights, shared by both budget
|
||||
pools. DeerMem supplies the same corpus IDF as unfiltered search;
|
||||
direct callers that omit it retain uniform term weights.
|
||||
|
||||
Returns:
|
||||
Formatted memory string for system prompt injection.
|
||||
@ -586,6 +607,19 @@ def format_memory_for_injection(
|
||||
valid_facts = [f for f in facts_data if isinstance(f, dict) and isinstance(f.get("content"), str) and f.get("content", "").strip()]
|
||||
|
||||
try:
|
||||
# Score once, then lazily diversify each budget pool. Do not run
|
||||
# full-scope MMR before the token-budget consumer can stop it.
|
||||
relevance_ordered = bool(query and query.strip() and relevance_weight is not None)
|
||||
scores = {}
|
||||
if relevance_ordered:
|
||||
scores = {id(fact): score for score, fact in score_facts(valid_facts, query, relevance_weight=relevance_weight, idf=idf)}
|
||||
|
||||
def _rank_pool(pool: list[dict[str, Any]]) -> Iterable[dict[str, Any]]:
|
||||
if not relevance_ordered:
|
||||
return sorted(pool, key=_confidence_key, reverse=True)
|
||||
scored = sorted(((scores[id(fact)], fact) for fact in pool), key=lambda pair: pair[0], reverse=True)
|
||||
return iter_diversify(scored, similarity_weight=diversity_weight or 0.0)
|
||||
|
||||
# Partition valid facts into guaranteed vs regular groups.
|
||||
# Use the *raw* category field (no ``or "context"`` default) so
|
||||
# a category-less legacy fact is never silently promoted into
|
||||
@ -604,19 +638,12 @@ def format_memory_for_injection(
|
||||
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,
|
||||
)
|
||||
guaranteed_pool = [f for f in valid_facts if _category_match(f)]
|
||||
regular_pool = [f for f in valid_facts if not _category_match(f)]
|
||||
guaranteed, regular = _rank_pool(guaranteed_pool), _rank_pool(regular_pool)
|
||||
else:
|
||||
guaranteed = []
|
||||
regular = sorted(valid_facts, key=_confidence_key, reverse=True)
|
||||
regular = _rank_pool(valid_facts)
|
||||
|
||||
# ── Phase 1: select guaranteed lines ──────────────────────────
|
||||
header_cost = _count_tokens(facts_header, use_tiktoken=use_tiktoken)
|
||||
|
||||
@ -0,0 +1,254 @@
|
||||
"""Deterministic lexical relevance ranking for DeerMem retrieval.
|
||||
|
||||
Pure-Python and network-free helpers behind the optional relevance-aware
|
||||
retrieval strategy (issue #4495):
|
||||
|
||||
- ``lexical_relevance`` — idf-weighted token overlap between a query and a
|
||||
fact's content, plus a containment signal so unsegmented (CJK) text stays
|
||||
usable without jieba;
|
||||
- ``score_facts`` / ``rank_facts`` — combine lexical relevance with the
|
||||
existing fact confidence (``relevance_weight * relevance +
|
||||
(1 - relevance_weight) * confidence``);
|
||||
- ``diversify`` — greedy MMR selection that demotes near-duplicate facts.
|
||||
|
||||
All helpers treat caller-owned fact dicts as read-only; ranking returns new
|
||||
lists. Token matching is case-insensitive with prefix matching for shared
|
||||
stems (``database``/``databases``), mirroring the retrieval layer's
|
||||
dependency-free style.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Iterator
|
||||
from itertools import islice
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import jieba
|
||||
|
||||
_jieba_available = True
|
||||
except ImportError: # pragma: no cover - exercised via the tokenizer fallback
|
||||
_jieba_available = False
|
||||
|
||||
_WORD_RE = re.compile(r"[a-zA-Z0-9_]+|[\u3400-\u4dbf\u4e00-\u9fff]+")
|
||||
|
||||
#: A token pair overlaps when the tokens are equal or one is a prefix of the
|
||||
#: other (minimum 4 characters so short words do not over-match).
|
||||
_PREFIX_MATCH_MIN_CHARS = 4
|
||||
|
||||
#: Shared bound for lexical ranking and near-duplicate similarity.
|
||||
_SIMILARITY_TOKEN_BUDGET = 128
|
||||
_TEXT_CHAR_BUDGET = 4096
|
||||
|
||||
|
||||
def warm_tokenizer() -> None:
|
||||
"""Load the optional segmenter's dictionary off the first-request path."""
|
||||
if _jieba_available:
|
||||
jieba.initialize()
|
||||
|
||||
|
||||
def tokenize(text: str) -> list[str]:
|
||||
"""Tokenize at most 4096 characters into at most 128 relevance tokens.
|
||||
|
||||
Space-free CJK text without jieba falls back to character bigrams so
|
||||
Chinese queries still produce deterministic token overlap.
|
||||
"""
|
||||
if not text:
|
||||
return []
|
||||
lowered = text[:_TEXT_CHAR_BUDGET].strip().lower()
|
||||
if _jieba_available:
|
||||
return list(islice((token for token in jieba.cut(lowered) if token.strip()), _SIMILARITY_TOKEN_BUDGET))
|
||||
|
||||
def fallback_tokens() -> Iterator[str]:
|
||||
for match in _WORD_RE.finditer(lowered):
|
||||
part = match.group()
|
||||
if "\u3400" <= part[0] <= "\u9fff":
|
||||
if len(part) == 1:
|
||||
yield part
|
||||
else:
|
||||
for index in range(len(part) - 1):
|
||||
yield part[index : index + 2]
|
||||
else:
|
||||
yield part
|
||||
|
||||
return list(islice(fallback_tokens(), _SIMILARITY_TOKEN_BUDGET))
|
||||
|
||||
|
||||
def build_idf(corpus: list[list[str]]) -> dict[str, float]:
|
||||
"""Smoothed inverse document frequency over a token corpus.
|
||||
|
||||
Tokens shared by every document get the smallest weight (1.0); rarer
|
||||
tokens get larger weights. Tokens absent from the corpus are handled by
|
||||
``lexical_relevance`` with the default weight.
|
||||
"""
|
||||
document_count = len(corpus)
|
||||
if document_count == 0:
|
||||
return {}
|
||||
document_frequency: dict[str, int] = {}
|
||||
for tokens in corpus:
|
||||
for token in set(tokens):
|
||||
document_frequency[token] = document_frequency.get(token, 0) + 1
|
||||
return {token: math.log((document_count + 1) / (frequency + 1)) + 1.0 for token, frequency in document_frequency.items()}
|
||||
|
||||
|
||||
def lexical_relevance(
|
||||
query: str,
|
||||
content: str,
|
||||
*,
|
||||
idf: dict[str, float] | None = None,
|
||||
) -> float:
|
||||
"""IDF-weighted query coverage in ``[0, 1]``.
|
||||
|
||||
A containment signal (whole query inside the content, or vice versa)
|
||||
contributes one matched unit so unsegmented text such as CJK
|
||||
content still scores above zero without a segmenter.
|
||||
"""
|
||||
query_text = (query or "")[:_TEXT_CHAR_BUDGET].strip().lower()
|
||||
return _lexical_relevance(query_text, tokenize(query_text), content, idf=idf)
|
||||
|
||||
|
||||
def _lexical_relevance(query_text: str, query_tokens: list[str], content: str, *, idf: dict[str, float] | None) -> float:
|
||||
"""Score coverage of a prepared query against bounded content tokens."""
|
||||
content_text = (content or "")[:_TEXT_CHAR_BUDGET].strip().lower()
|
||||
if not query_text or not content_text:
|
||||
return 0.0
|
||||
|
||||
content_tokens = tokenize(content_text)
|
||||
containment = (query_text in content_text) or (content_text in query_text)
|
||||
if not query_tokens and not containment:
|
||||
return 0.0
|
||||
|
||||
weights = idf or {}
|
||||
content_set = set(content_tokens)
|
||||
prefix_buckets: dict[str, list[str]] = {}
|
||||
for content_token in content_set:
|
||||
if len(content_token) >= _PREFIX_MATCH_MIN_CHARS:
|
||||
prefix_buckets.setdefault(content_token[:_PREFIX_MATCH_MIN_CHARS], []).append(content_token)
|
||||
# Each distinct query token contributes its squared IDF at most once.
|
||||
# Compare the matched-query norm to the complete-query norm: repetition
|
||||
# cannot replace missing terms, and the result needs no clipping. The
|
||||
# norm ratio keeps useful partial matches competitive in confidence blends.
|
||||
matched_weight = total_weight = 1.0 if containment else 0.0
|
||||
for token in dict.fromkeys(query_tokens):
|
||||
weight = weights.get(token, 1.0) ** 2
|
||||
total_weight += weight
|
||||
# Sharing four characters only narrows the candidate set. Count a
|
||||
# match only when one complete token is a prefix of the other.
|
||||
if token in content_set or (len(token) >= _PREFIX_MATCH_MIN_CHARS and any(token.startswith(candidate) or candidate.startswith(token) for candidate in prefix_buckets.get(token[:_PREFIX_MATCH_MIN_CHARS], ()))):
|
||||
matched_weight += weight
|
||||
return math.sqrt(matched_weight / total_weight) if total_weight > 0.0 else 0.0
|
||||
|
||||
|
||||
def _coerce_confidence(fact: dict[str, Any]) -> float:
|
||||
try:
|
||||
value = float(fact.get("confidence"))
|
||||
if not math.isfinite(value):
|
||||
raise ValueError
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
return min(1.0, max(0.0, value))
|
||||
|
||||
|
||||
def score_facts(
|
||||
facts: list[dict[str, Any]],
|
||||
query: str,
|
||||
*,
|
||||
relevance_weight: float = 0.5,
|
||||
idf: dict[str, float] | None = None,
|
||||
) -> list[tuple[float, dict[str, Any]]]:
|
||||
"""Return ``(combined_score, fact)`` pairs sorted descending (no mutation).
|
||||
|
||||
``relevance_weight == 0`` short-circuits to the legacy confidence-only
|
||||
ordering without computing relevance.
|
||||
"""
|
||||
if relevance_weight <= 0.0:
|
||||
return [
|
||||
(confidence, fact)
|
||||
for confidence, fact in sorted(
|
||||
((_coerce_confidence(fact), fact) for fact in facts),
|
||||
key=lambda pair: pair[0],
|
||||
reverse=True,
|
||||
)
|
||||
]
|
||||
|
||||
query_text = (query or "")[:_TEXT_CHAR_BUDGET].strip().lower()
|
||||
query_tokens = tokenize(query_text)
|
||||
scored: list[tuple[float, dict[str, Any]]] = []
|
||||
for fact in facts:
|
||||
content = fact.get("content")
|
||||
relevance = _lexical_relevance(query_text, query_tokens, content, idf=idf) if isinstance(content, str) else 0.0
|
||||
confidence = _coerce_confidence(fact)
|
||||
combined = relevance_weight * relevance + (1.0 - relevance_weight) * confidence
|
||||
scored.append((combined, fact))
|
||||
scored.sort(key=lambda pair: pair[0], reverse=True)
|
||||
return scored
|
||||
|
||||
|
||||
def rank_facts(
|
||||
facts: list[dict[str, Any]],
|
||||
query: str,
|
||||
*,
|
||||
relevance_weight: float = 0.5,
|
||||
idf: dict[str, float] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Convenience wrapper: ``score_facts`` without the scores."""
|
||||
return [fact for _, fact in score_facts(facts, query, relevance_weight=relevance_weight, idf=idf)]
|
||||
|
||||
|
||||
def diversify(
|
||||
scored: list[tuple[float, dict[str, Any]]],
|
||||
*,
|
||||
similarity_weight: float,
|
||||
limit: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Greedy MMR over ``(score, fact)`` pairs; returns facts (no mutation).
|
||||
|
||||
``similarity_weight == 0`` returns the score order unchanged.
|
||||
"""
|
||||
ordered = iter_diversify(scored, similarity_weight=similarity_weight)
|
||||
return list(islice(ordered, max(0, limit))) if limit is not None else list(ordered)
|
||||
|
||||
|
||||
def iter_diversify(scored: list[tuple[float, dict[str, Any]]], *, similarity_weight: float) -> Iterator[dict[str, Any]]:
|
||||
"""Lazy MMR: tokenize each fact once and update penalties once per pick.
|
||||
|
||||
Consumers can stop at their result or token budget without ranking the rest.
|
||||
Equal adjusted scores retain their input order.
|
||||
"""
|
||||
if similarity_weight <= 0.0:
|
||||
yield from (fact for _, fact in scored)
|
||||
return
|
||||
token_sets = [set(tokenize(fact["content"])) if isinstance(fact.get("content"), str) else set() for _, fact in scored]
|
||||
remaining = list(range(len(scored)))
|
||||
penalties = [0.0] * len(scored)
|
||||
while remaining:
|
||||
best_index = 0
|
||||
best_adjusted = -math.inf
|
||||
for index, fact_index in enumerate(remaining):
|
||||
adjusted = scored[fact_index][0] - similarity_weight * penalties[fact_index]
|
||||
if adjusted > best_adjusted:
|
||||
best_adjusted = adjusted
|
||||
best_index = index
|
||||
picked_index = remaining.pop(best_index)
|
||||
yield scored[picked_index][1]
|
||||
picked_tokens = token_sets[picked_index]
|
||||
for fact_index in remaining:
|
||||
tokens = token_sets[fact_index]
|
||||
similarity = len(tokens & picked_tokens) / len(tokens | picked_tokens) if tokens and picked_tokens else 0.0
|
||||
penalties[fact_index] = max(penalties[fact_index], similarity)
|
||||
|
||||
|
||||
def order_facts_for_query(
|
||||
facts: list[dict[str, Any]],
|
||||
query: str,
|
||||
*,
|
||||
relevance_weight: float = 0.5,
|
||||
diversity_weight: float = 0.0,
|
||||
idf: dict[str, float] | None = None,
|
||||
limit: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Score by relevance+confidence, then diversify; used by search/injection."""
|
||||
scored = score_facts(facts, query, relevance_weight=relevance_weight, idf=idf)
|
||||
return diversify(scored, similarity_weight=diversity_weight, limit=limit)
|
||||
@ -205,6 +205,7 @@ class HonchoMemoryManager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
workspace = self._workspace(user_id)
|
||||
if workspace is None or not user_id:
|
||||
@ -290,8 +291,9 @@ class HonchoMemoryManager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
return await asyncio.to_thread(self.get_context, user_id, agent_name=agent_name, thread_id=thread_id)
|
||||
return await asyncio.to_thread(self.get_context, user_id, agent_name=agent_name, thread_id=thread_id, query=query)
|
||||
|
||||
async def asearch(
|
||||
self,
|
||||
|
||||
@ -187,10 +187,11 @@ class Mem0Manager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
"""Query-less recall: the contract passes no current query, so inject
|
||||
the bucket's most recent memories (top_k). Query-aware recall is
|
||||
available via search() in mode="tool"."""
|
||||
"""Query-less recall: this backend ignores the optional ``query``
|
||||
hint and injects the bucket's most recent memories (top_k).
|
||||
Query-aware recall is available via search() in mode="tool"."""
|
||||
filters = _build_filters(user_id=user_id, agent_name=agent_name, run_id=thread_id)
|
||||
if filters is None:
|
||||
return ""
|
||||
@ -246,12 +247,14 @@ class Mem0Manager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
return await asyncio.to_thread(
|
||||
self.get_context,
|
||||
user_id,
|
||||
agent_name=agent_name,
|
||||
thread_id=thread_id,
|
||||
query=query,
|
||||
)
|
||||
|
||||
# ── Tier 2: search ───────────────────────────────────────────────────
|
||||
|
||||
@ -122,6 +122,7 @@ class NoopMemoryManager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
@ -166,6 +166,7 @@ class OpenVikingMemoryManager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
if not self._begin_operation():
|
||||
return ""
|
||||
@ -204,12 +205,14 @@ class OpenVikingMemoryManager(MemoryManager):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
return await asyncio.to_thread(
|
||||
self.get_context,
|
||||
user_id,
|
||||
agent_name=agent_name,
|
||||
thread_id=thread_id,
|
||||
query=query,
|
||||
)
|
||||
|
||||
def search(
|
||||
|
||||
@ -17,11 +17,13 @@ do not impede pluggability.
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from abc import abstractmethod
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
from typing import Any, ClassVar, Literal
|
||||
@ -44,6 +46,23 @@ _backends_cache: dict[str, type[MemoryManager]] | None = None
|
||||
_manager_lock = threading.Lock()
|
||||
|
||||
|
||||
def context_query_kwargs(get_context: Callable[..., str], query: str | None) -> dict[str, str | None]:
|
||||
"""Pass the optional hint only when a backend accepts that keyword.
|
||||
|
||||
Older plugins need no signature change. An uninspectable callable keeps
|
||||
the old call contract; backend TypeErrors must never trigger a retry.
|
||||
"""
|
||||
if query is None:
|
||||
return {}
|
||||
try:
|
||||
parameters = inspect.signature(get_context).parameters.values()
|
||||
except (TypeError, ValueError):
|
||||
return {}
|
||||
if any(parameter.kind == inspect.Parameter.VAR_KEYWORD or (parameter.name == "query" and parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)) for parameter in parameters):
|
||||
return {"query": query}
|
||||
return {}
|
||||
|
||||
|
||||
class MemoryCallbacks:
|
||||
"""Observability hooks for memory backends. Default implementations are
|
||||
no-ops; override the ones you need. The pre-LLM-call hook
|
||||
@ -255,9 +274,15 @@ class MemoryManager(BaseModel):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
"""Return injection-ready memory text for the given bucket.
|
||||
|
||||
``query`` is an optional current-turn query hint. Backends that
|
||||
support query-aware ranking (DeerMem with
|
||||
``retrieval_relevance_enabled``) may rank injected facts against it;
|
||||
other backends ignore it. ``None`` must preserve legacy behavior.
|
||||
|
||||
Implementations load their memory and format it however they choose;
|
||||
the returned string is injected verbatim by call sites. Format
|
||||
parameters are the backend's own private config (received via
|
||||
@ -530,8 +555,9 @@ class MemoryManager(BaseModel):
|
||||
*,
|
||||
agent_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
query: str | None = None,
|
||||
) -> str:
|
||||
return self.get_context(user_id, agent_name=agent_name, thread_id=thread_id)
|
||||
return self.get_context(user_id, agent_name=agent_name, thread_id=thread_id, **context_query_kwargs(self.get_context, query))
|
||||
|
||||
async def asearch(
|
||||
self,
|
||||
|
||||
@ -47,7 +47,7 @@ import os
|
||||
import posixpath
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from datetime import datetime, tzinfo
|
||||
from typing import TYPE_CHECKING, override
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
@ -61,7 +61,7 @@ from langgraph.runtime import Runtime
|
||||
from deerflow.projects.context import build_project_context_message, is_project_context_message, pinned_project_snapshot, project_context_insertion_index, render_documents_block, render_project_block
|
||||
from deerflow.runtime.context_keys import CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY
|
||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||
from deerflow.utils.messages import INJECTED_USER_MESSAGE_ID_SUFFIX, strip_injected_user_message_id_suffix
|
||||
from deerflow.utils.messages import INJECTED_USER_MESSAGE_ID_SUFFIX, ORIGINAL_USER_CONTENT_KEY, strip_injected_user_message_id_suffix
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deerflow.config.app_config import AppConfig
|
||||
@ -74,6 +74,11 @@ logger = logging.getLogger(__name__)
|
||||
# This cap ensures the request degrades gracefully instead of hanging.
|
||||
_INJECT_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
#: Hard bound on the current-turn query forwarded to the memory backend for
|
||||
#: query-aware fact ranking. Keeps ranking cost deterministic regardless of
|
||||
#: message length.
|
||||
_INJECTION_QUERY_MAX_CHARS = 1000
|
||||
|
||||
_DATE_RE = re.compile(r"<current_date>([^<]+)</current_date>")
|
||||
_DYNAMIC_CONTEXT_REMINDER_KEY = "dynamic_context_reminder"
|
||||
# Authoritative injected date, carried in additional_kwargs of the date
|
||||
@ -320,6 +325,37 @@ class SubagentDateContextMiddleware(AgentMiddleware):
|
||||
return None
|
||||
|
||||
|
||||
def _derive_injection_query(message: object) -> str | None:
|
||||
"""Extract a bounded text query from the user message being injected on.
|
||||
|
||||
Prefer the original user text preserved by UploadsMiddleware so file
|
||||
descriptions cannot consume the query budget. Otherwise handle plain
|
||||
text and multimodal lists. An empty original request remains query-less.
|
||||
"""
|
||||
content = getattr(message, "content", None)
|
||||
additional_kwargs = getattr(message, "additional_kwargs", None)
|
||||
if isinstance(additional_kwargs, Mapping):
|
||||
original_content = additional_kwargs.get(ORIGINAL_USER_CONTENT_KEY)
|
||||
if isinstance(original_content, str):
|
||||
content = original_content
|
||||
if isinstance(content, str):
|
||||
text = content
|
||||
elif isinstance(content, list):
|
||||
parts = []
|
||||
for item in content:
|
||||
if isinstance(item, dict) and isinstance(item.get("text"), str):
|
||||
part = item["text"].strip()
|
||||
if part:
|
||||
parts.append(part)
|
||||
text = " ".join(parts)
|
||||
else:
|
||||
return None
|
||||
text = text.strip()
|
||||
if not text:
|
||||
return None
|
||||
return text[:_INJECTION_QUERY_MAX_CHARS]
|
||||
|
||||
|
||||
class DynamicContextMiddleware(AgentMiddleware):
|
||||
"""Inject memory and current date as a SystemMessage <system-reminder>.
|
||||
|
||||
@ -385,13 +421,17 @@ class DynamicContextMiddleware(AgentMiddleware):
|
||||
"shelf_index_max_bytes": max_bytes,
|
||||
}
|
||||
|
||||
def _build_full_reminder(self, runtime: Runtime | None = None) -> tuple[str, str | None]:
|
||||
def _build_full_reminder(self, runtime: Runtime | None = None, *, query: str | None = None) -> tuple[str, str | None]:
|
||||
"""Return (date_reminder, memory_block | None).
|
||||
|
||||
Framework-owned data (date) is separated from user-owned data (memory)
|
||||
so the downstream SystemMessage carries only framework authority and
|
||||
memory stays at role:user — preventing untrusted content from gaining
|
||||
system privilege (OWASP LLM01).
|
||||
|
||||
``query`` is the optional current-turn text forwarded to the memory
|
||||
backend for query-aware fact ranking (issue #4495); ``None`` keeps the
|
||||
legacy confidence-only ordering.
|
||||
"""
|
||||
from deerflow.agents.lead_agent.prompt import _get_memory_context
|
||||
|
||||
@ -401,6 +441,7 @@ class DynamicContextMiddleware(AgentMiddleware):
|
||||
self._agent_name,
|
||||
app_config=self._app_config,
|
||||
user_id=resolve_runtime_user_id(runtime),
|
||||
query=query,
|
||||
)
|
||||
if injection_enabled
|
||||
else ""
|
||||
@ -541,7 +582,7 @@ class DynamicContextMiddleware(AgentMiddleware):
|
||||
target_idx = next((i for i in reversed(range(len(messages))) if _is_user_injection_target(messages[i])), None)
|
||||
if target_idx is None:
|
||||
return {"messages": memory_removals} if memory_removals else None
|
||||
date_reminder, memory_block = self._build_full_reminder(runtime)
|
||||
date_reminder, memory_block = self._build_full_reminder(runtime, query=_derive_injection_query(messages[target_idx]))
|
||||
logger.info(
|
||||
"DynamicContextMiddleware: injecting full reminder (has_memory=%s) into last HumanMessage id=%r",
|
||||
memory_block is not None,
|
||||
|
||||
@ -99,11 +99,11 @@ async def test_abefore_agent_does_not_block_event_loop() -> None:
|
||||
# event-loop blocking visible to the Blockbuster gate.
|
||||
original_build = mw._build_full_reminder
|
||||
|
||||
def slow_build_reminder(runtime=None):
|
||||
def slow_build_reminder(runtime=None, *, query=None):
|
||||
import time
|
||||
|
||||
time.sleep(0.05) # 50ms sync sleep — blocks the thread it runs on
|
||||
return original_build(runtime)
|
||||
return original_build(runtime, query=query)
|
||||
|
||||
with (
|
||||
mock.patch.object(mw, "_build_full_reminder", slow_build_reminder),
|
||||
@ -294,11 +294,13 @@ async def test_abefore_agent_propagates_strict_memory_timeout(
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("slow_policy", [False, True], ids=["normal_policy", "slow_policy"])
|
||||
async def test_abefore_agent_policy_resolution_failure_does_not_replace_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
manager_class: str,
|
||||
backend_config: dict,
|
||||
api_key: str | None,
|
||||
slow_policy: bool,
|
||||
) -> None:
|
||||
"""An unresolved timeout policy must fail closed with the original cause."""
|
||||
if api_key is None:
|
||||
@ -317,26 +319,40 @@ async def test_abefore_agent_policy_resolution_failure_does_not_replace_timeout(
|
||||
release = threading.Event()
|
||||
finished = threading.Event()
|
||||
|
||||
if slow_policy:
|
||||
original_policy = mw._read_failures_are_fatal
|
||||
|
||||
def delayed_policy(*, allow_io=True):
|
||||
if not allow_io:
|
||||
return None
|
||||
# Exercise a cold worker still resolving policy after the 10ms timeout.
|
||||
threading.Event().wait(0.05)
|
||||
return original_policy(allow_io=allow_io)
|
||||
|
||||
monkeypatch.setattr(mw, "_read_failures_are_fatal", delayed_policy)
|
||||
|
||||
def blocking_inject(state, runtime=None):
|
||||
started.set()
|
||||
release.wait(timeout=2)
|
||||
finished.set()
|
||||
|
||||
try:
|
||||
with (
|
||||
mock.patch.object(mw, "_inject", blocking_inject),
|
||||
mock.patch(
|
||||
"deerflow.agents.middlewares.dynamic_context_middleware._INJECT_TIMEOUT_SECONDS",
|
||||
0.01,
|
||||
),
|
||||
):
|
||||
with (
|
||||
mock.patch.object(mw, "_inject", blocking_inject),
|
||||
mock.patch(
|
||||
"deerflow.agents.middlewares.dynamic_context_middleware._INJECT_TIMEOUT_SECONDS",
|
||||
0.01,
|
||||
),
|
||||
):
|
||||
try:
|
||||
state = {"messages": [HumanMessage(content="Hello", id="msg-1")]}
|
||||
runtime = SimpleNamespace(context={})
|
||||
with pytest.raises(MemoryReadError) as exc_info:
|
||||
await mw.abefore_agent(state, runtime)
|
||||
finally:
|
||||
release.set()
|
||||
assert await asyncio.to_thread(finished.wait, 1)
|
||||
finally:
|
||||
# The worker can reach self._inject only after cold policy resolution.
|
||||
# Keep its mock installed until the worker exits, including on timeout.
|
||||
release.set()
|
||||
assert await asyncio.to_thread(finished.wait, 1)
|
||||
|
||||
assert isinstance(exc_info.value.__cause__, TimeoutError)
|
||||
assert started.is_set()
|
||||
|
||||
@ -255,6 +255,7 @@ def test_memory_lookup_uses_runtime_user_id():
|
||||
None,
|
||||
app_config=None,
|
||||
user_id="runtime-user",
|
||||
query="Hi",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@ -376,7 +376,7 @@ def test_get_memory_context_uses_explicit_app_config_without_global_config(monke
|
||||
def fail_get_memory_config():
|
||||
raise AssertionError("ambient get_memory_config() must not be used when app_config is explicit")
|
||||
|
||||
def fake_get_context(user_id, *, agent_name=None, thread_id=None):
|
||||
def fake_get_context(user_id, *, agent_name=None, thread_id=None, query=None):
|
||||
captured["agent_name"] = agent_name
|
||||
captured["user_id"] = user_id
|
||||
return "remember this"
|
||||
@ -457,7 +457,7 @@ def test_get_memory_context_prefers_explicit_user_id(monkeypatch):
|
||||
def fail_resolve_runtime_user_id(runtime):
|
||||
raise AssertionError("explicit user_id must bypass ambient identity resolution")
|
||||
|
||||
def fake_get_context(user_id, *, agent_name=None, thread_id=None):
|
||||
def fake_get_context(user_id, *, agent_name=None, thread_id=None, query=None):
|
||||
captured["agent_name"] = agent_name
|
||||
captured["user_id"] = user_id
|
||||
return "remember this"
|
||||
|
||||
@ -44,7 +44,7 @@ class _MinimalBackend(MemoryManager):
|
||||
def add(self, thread_id, messages, *, agent_name=None, user_id=None, trace_id=None) -> None:
|
||||
self._adds.append((thread_id, user_id))
|
||||
|
||||
def get_context(self, user_id, *, agent_name=None, thread_id=None) -> str:
|
||||
def get_context(self, user_id, *, agent_name=None, thread_id=None, query=None) -> str:
|
||||
return f"ctx:{user_id}"
|
||||
|
||||
@classmethod
|
||||
|
||||
578
backend/tests/test_memory_relevance_retrieval.py
Normal file
578
backend/tests/test_memory_relevance_retrieval.py
Normal file
@ -0,0 +1,578 @@
|
||||
"""Tests for the optional relevance-aware retrieval strategy (issue #4495).
|
||||
|
||||
The strategy is opt-in via DeerMem-private config
|
||||
(``retrieval_relevance_enabled``) and must never change the default
|
||||
confidence-based behavior. Coverage:
|
||||
|
||||
- deterministic lexical relevance + confidence scoring;
|
||||
- greedy MMR diversity selection;
|
||||
- ``DeerMem.search`` relevance mode (including related facts without a
|
||||
literal substring match);
|
||||
- prompt-injection fact ordering under a query;
|
||||
- the DynamicContextMiddleware -> ``_get_memory_context`` query wiring.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from deerflow.agents.memory.backends.deermem.deer_mem import DeerMem
|
||||
from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.relevance import (
|
||||
build_idf,
|
||||
diversify,
|
||||
lexical_relevance,
|
||||
rank_facts,
|
||||
tokenize,
|
||||
)
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import DynamicContextMiddleware
|
||||
|
||||
|
||||
def _make_fact(content: str, category: str = "context", confidence: float = 0.7) -> dict:
|
||||
return {
|
||||
"id": f"fact_test_{hash(content) & 0xFFFFFFFF:08x}",
|
||||
"content": content,
|
||||
"category": category,
|
||||
"confidence": confidence,
|
||||
"createdAt": "2026-07-09T00:00:00Z",
|
||||
"source": "test",
|
||||
}
|
||||
|
||||
|
||||
def _deer_mem_with_facts(facts: list[dict], backend_config: dict | None = None) -> DeerMem:
|
||||
"""Build a DeerMem whose updater returns the given facts (no disk I/O)."""
|
||||
mgr = DeerMem(backend_config=backend_config)
|
||||
mgr._updater = SimpleNamespace(get_memory_data=lambda agent_name=None, *, user_id=None: {"facts": facts})
|
||||
return mgr
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lexical relevance scoring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLexicalRelevance:
|
||||
def test_missing_confidence_defaults_to_zero(self):
|
||||
missing = {"content": "unrelated first"}
|
||||
low = _make_fact("unrelated second", confidence=0.1)
|
||||
assert rank_facts([missing, low], "python")[0] is low
|
||||
|
||||
def test_optional_segmenter_receives_bounded_input(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core import relevance
|
||||
|
||||
seen = []
|
||||
|
||||
def cut(text):
|
||||
seen.append(len(text))
|
||||
yield from ("token" for _ in range(10000))
|
||||
|
||||
monkeypatch.setattr(relevance, "_jieba_available", True)
|
||||
monkeypatch.setattr(relevance, "jieba", SimpleNamespace(cut=cut), raising=False)
|
||||
assert len(tokenize("word" * 10000)) == 128
|
||||
assert seen == [4096]
|
||||
|
||||
def test_mixed_cjk_without_jieba(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core import relevance
|
||||
|
||||
monkeypatch.setattr(relevance, "_jieba_available", False)
|
||||
assert {"python", "我喜", "喜欢", "编程"} <= set(tokenize("我喜欢Python编程"))
|
||||
assert {"你好", "世界"} <= set(tokenize("你好 世界"))
|
||||
assert lexical_relevance("数据库升级", "Python数据库迁移") > 0
|
||||
|
||||
@pytest.mark.parametrize("confidence", [None, "invalid", float("nan"), float("inf")])
|
||||
def test_invalid_confidence_does_not_outrank_low_confidence(self, confidence):
|
||||
invalid = _make_fact("unrelated first", confidence=confidence)
|
||||
low = _make_fact("unrelated second", confidence=0.1)
|
||||
assert rank_facts([invalid, low], "python")[0] is low
|
||||
|
||||
def test_bounded_tokens(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core import relevance
|
||||
|
||||
monkeypatch.setattr(relevance, "_jieba_available", False)
|
||||
assert len(tokenize("word " * 10000)) <= 128
|
||||
assert len(tokenize("数据库迁移" * 10000)) <= 128
|
||||
|
||||
def test_query_tokenized_once_per_ranking(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core import relevance
|
||||
|
||||
original = relevance.tokenize
|
||||
queries = []
|
||||
|
||||
def counted(text):
|
||||
if text == "database migration":
|
||||
queries.append(text)
|
||||
return original(text)
|
||||
|
||||
monkeypatch.setattr(relevance, "tokenize", counted)
|
||||
rank_facts([_make_fact(f"python fact {i}") for i in range(100)], "database migration")
|
||||
assert len(queries) == 1
|
||||
|
||||
def test_overlapping_content_scores_higher_than_unrelated(self):
|
||||
query = "database migration"
|
||||
related = lexical_relevance(query, "Migrations are managed with alembic and a PostgreSQL database")
|
||||
unrelated = lexical_relevance(query, "User prefers cooking Italian food on weekends")
|
||||
assert related > unrelated
|
||||
|
||||
def test_zero_for_no_overlap(self):
|
||||
assert lexical_relevance("python", "User lives in Beijing") == 0.0
|
||||
|
||||
def test_case_insensitive(self):
|
||||
assert lexical_relevance("PYTHON", "User prefers Python") > 0.0
|
||||
|
||||
def test_substring_signal_without_word_boundaries(self):
|
||||
"""CJK / unsegmented content: containment still contributes relevance."""
|
||||
assert lexical_relevance("Python", "我喜欢Python编程") > 0.0
|
||||
|
||||
def test_empty_query_scores_zero(self):
|
||||
assert lexical_relevance("", "anything") == 0.0
|
||||
assert lexical_relevance(" ", "anything") == 0.0
|
||||
|
||||
|
||||
class TestIdf:
|
||||
def test_common_tokens_are_downweighted(self):
|
||||
corpus = [
|
||||
tokenize("database migration conventions"),
|
||||
tokenize("database backup schedule"),
|
||||
tokenize("database replica lag"),
|
||||
tokenize("the database is used everywhere"),
|
||||
]
|
||||
idf = build_idf(corpus)
|
||||
assert idf["migration"] > idf["database"]
|
||||
|
||||
|
||||
class TestRankFacts:
|
||||
def test_combines_relevance_and_confidence(self):
|
||||
facts = [
|
||||
_make_fact("User prefers concise answers", confidence=0.95),
|
||||
_make_fact("Migrations are managed with alembic", confidence=0.5),
|
||||
]
|
||||
ranked = rank_facts(facts, "database migration", relevance_weight=0.7)
|
||||
assert ranked[0]["content"] == "Migrations are managed with alembic"
|
||||
|
||||
def test_pure_confidence_when_relevance_weight_is_zero(self):
|
||||
facts = [
|
||||
_make_fact("Low", confidence=0.2),
|
||||
_make_fact("High", confidence=0.9),
|
||||
]
|
||||
ranked = rank_facts(facts, "high", relevance_weight=0.0)
|
||||
assert [f["content"] for f in ranked] == ["High", "Low"]
|
||||
|
||||
def test_does_not_mutate_input(self):
|
||||
facts = [
|
||||
_make_fact("Migrations are managed with alembic", confidence=0.5),
|
||||
_make_fact("User prefers concise answers", confidence=0.95),
|
||||
]
|
||||
snapshot = [dict(f) for f in facts]
|
||||
rank_facts(facts, "database migration", relevance_weight=0.7)
|
||||
assert facts == snapshot
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Diversity
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDiversify:
|
||||
def test_incremental_penalties_match_reference_mmr(self):
|
||||
scored = [(0.9 - (i % 4) * 0.1, _make_fact(f"database {i % 3} fact {i % 5}")) for i in range(20)]
|
||||
remaining = list(scored)
|
||||
expected = []
|
||||
|
||||
def penalty(fact):
|
||||
left = set(tokenize(fact["content"]))
|
||||
return max((len(left & set(tokenize(picked["content"]))) / len(left | set(tokenize(picked["content"]))) for picked in expected), default=0.0)
|
||||
|
||||
while remaining:
|
||||
index = max(range(len(remaining)), key=lambda i: remaining[i][0] - 0.5 * penalty(remaining[i][1]))
|
||||
expected.append(remaining.pop(index)[1])
|
||||
for limit in (0, 1, 5, len(scored), len(scored) + 1):
|
||||
assert diversify(scored, similarity_weight=0.5, limit=limit) == expected[:limit]
|
||||
|
||||
def test_limit_preserves_full_prefix(self):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.relevance import order_facts_for_query
|
||||
|
||||
facts = [_make_fact(text) for text in ["database migrations", "database migration", "python testing", "Italian cooking"]]
|
||||
full = order_facts_for_query(facts, "database", diversity_weight=0.5)
|
||||
assert order_facts_for_query(facts, "database", diversity_weight=0.5, limit=2) == full[:2]
|
||||
assert order_facts_for_query(facts, "database", diversity_weight=0.5, limit=0) == []
|
||||
|
||||
def test_tokenization_is_linear(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core import relevance
|
||||
|
||||
calls = []
|
||||
original = relevance.tokenize
|
||||
|
||||
def counted(text):
|
||||
calls.append(text)
|
||||
return original(text)
|
||||
|
||||
monkeypatch.setattr(relevance, "tokenize", counted)
|
||||
scored = [(0.7, _make_fact(f"database fact {i}")) for i in range(30)]
|
||||
diversify(scored, similarity_weight=0.5, limit=5)
|
||||
assert len(calls) <= len(scored)
|
||||
|
||||
def test_promotes_distinct_fact_over_near_duplicate(self):
|
||||
facts = [
|
||||
_make_fact("Use ruff for linting"),
|
||||
_make_fact("Use ruff for linting"),
|
||||
_make_fact("Deploys go through GitHub Actions"),
|
||||
]
|
||||
ranked = rank_facts(facts, "linting", relevance_weight=0.7)
|
||||
scored = [(1.0 - index * 0.1, fact) for index, fact in enumerate(ranked)]
|
||||
picked = diversify(scored, similarity_weight=0.5, limit=2)
|
||||
contents = [fact["content"] for fact in picked]
|
||||
assert contents[0] == "Use ruff for linting"
|
||||
assert "Deploys go through GitHub Actions" in contents
|
||||
assert len(contents) == 2
|
||||
|
||||
def test_identity_when_similarity_weight_is_zero(self):
|
||||
facts = [
|
||||
_make_fact("Use ruff for linting"),
|
||||
_make_fact("Deploys go through GitHub Actions"),
|
||||
]
|
||||
ranked = rank_facts(facts, "linting", relevance_weight=0.7)
|
||||
scored = [(1.0 - index * 0.1, fact) for index, fact in enumerate(ranked)]
|
||||
picked = diversify(scored, similarity_weight=0.0)
|
||||
assert [f["content"] for f in picked] == [f["content"] for f in ranked]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRelevanceConfig:
|
||||
def test_defaults_keep_legacy_behavior(self):
|
||||
config = DeerMemConfig()
|
||||
assert config.retrieval_relevance_enabled is False
|
||||
assert config.retrieval_relevance_weight == 0.5
|
||||
assert config.retrieval_diversity_weight == 0.0
|
||||
|
||||
def test_backend_config_accepts_new_knobs(self):
|
||||
config = DeerMemConfig.from_backend_config(
|
||||
{
|
||||
"retrieval_relevance_enabled": True,
|
||||
"retrieval_relevance_weight": 0.8,
|
||||
"retrieval_diversity_weight": 0.4,
|
||||
}
|
||||
)
|
||||
assert config.retrieval_relevance_enabled is True
|
||||
assert config.retrieval_relevance_weight == 0.8
|
||||
assert config.retrieval_diversity_weight == 0.4
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DeerMem.search with relevance mode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRelevanceSearch:
|
||||
def test_search_passes_top_k_to_mmr(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem import deer_mem
|
||||
|
||||
facts = [_make_fact(f"database fact {i}") for i in range(30)]
|
||||
original = deer_mem.order_facts_for_query
|
||||
limits = []
|
||||
|
||||
def ranked(*args, **kwargs):
|
||||
limits.append(kwargs.get("limit"))
|
||||
return original(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(deer_mem, "order_facts_for_query", ranked)
|
||||
mgr = _deer_mem_with_facts(facts, {"retrieval_relevance_enabled": True, "retrieval_diversity_weight": 0.5, "retrieval_adapter": ""})
|
||||
assert len(mgr.search("database", top_k=3)) == 3
|
||||
assert limits == [3]
|
||||
|
||||
@pytest.mark.parametrize("enabled", [False, True])
|
||||
@pytest.mark.parametrize("counting", ["char", "tiktoken"])
|
||||
def test_warms_segmenter_only_when_enabled(self, monkeypatch, enabled, counting):
|
||||
from deerflow.agents.memory.backends.deermem import deer_mem
|
||||
|
||||
calls = []
|
||||
monkeypatch.setattr(deer_mem, "warm_tokenizer", lambda: calls.append("jieba"))
|
||||
monkeypatch.setattr(deer_mem, "warm_tiktoken_cache", lambda: calls.append("tiktoken") or True)
|
||||
mgr = _deer_mem_with_facts([], {"retrieval_relevance_enabled": enabled, "token_counting": counting, "retrieval_adapter": ""})
|
||||
assert mgr.warm() is True
|
||||
assert calls == (["jieba"] if enabled else []) + (["tiktoken"] if counting == "tiktoken" else [])
|
||||
|
||||
def test_returns_related_fact_without_literal_substring(self):
|
||||
facts = [
|
||||
_make_fact("Database migrations are handled with alembic", "project", 0.4),
|
||||
_make_fact("User prefers concise answers", "preference", 0.9),
|
||||
]
|
||||
mgr = _deer_mem_with_facts(
|
||||
facts,
|
||||
backend_config={
|
||||
"retrieval_relevance_enabled": True,
|
||||
"retrieval_adapter": "",
|
||||
"retrieval_relevance_weight": 0.7,
|
||||
},
|
||||
)
|
||||
|
||||
results = mgr.search("how do I add a database migration", top_k=5)
|
||||
assert results[0]["content"] == "Database migrations are handled with alembic"
|
||||
assert len(results) == 2 # every fact in scope competes, not only substring matches
|
||||
|
||||
def test_relevance_outweighs_confidence(self):
|
||||
facts = [
|
||||
_make_fact("User prefers concise answers", "preference", 0.9),
|
||||
_make_fact("Migrations are managed with alembic", "project", 0.4),
|
||||
]
|
||||
mgr = _deer_mem_with_facts(
|
||||
facts,
|
||||
backend_config={
|
||||
"retrieval_relevance_enabled": True,
|
||||
"retrieval_adapter": "",
|
||||
"retrieval_relevance_weight": 0.7,
|
||||
},
|
||||
)
|
||||
|
||||
results = mgr.search("database migration", top_k=5)
|
||||
assert results[0]["content"] == "Migrations are managed with alembic"
|
||||
|
||||
def test_respects_category_filter_and_top_k(self):
|
||||
facts = [_make_fact(f"Database fact {index}", "project", 0.5) for index in range(6)] + [_make_fact("Unrelated preference", "preference", 0.9)]
|
||||
mgr = _deer_mem_with_facts(
|
||||
facts,
|
||||
backend_config={"retrieval_relevance_enabled": True, "retrieval_adapter": ""},
|
||||
)
|
||||
|
||||
results = mgr.search("database", top_k=3, category="project")
|
||||
assert len(results) == 3
|
||||
assert all(fact["category"] == "project" for fact in results)
|
||||
|
||||
def test_diversity_dedups_near_duplicates(self):
|
||||
facts = [
|
||||
_make_fact("Use ruff for linting", confidence=0.9),
|
||||
_make_fact("Use ruff for linting", confidence=0.8),
|
||||
_make_fact("CI lints on every pull request", confidence=0.7),
|
||||
]
|
||||
mgr = _deer_mem_with_facts(
|
||||
facts,
|
||||
backend_config={
|
||||
"retrieval_relevance_enabled": True,
|
||||
"retrieval_adapter": "",
|
||||
"retrieval_diversity_weight": 0.5,
|
||||
},
|
||||
)
|
||||
|
||||
# Both "lints" and "linting" extend this complete query token, so
|
||||
# the test isolates diversity rather than arbitrary shared stems.
|
||||
results = mgr.search("lint", top_k=2)
|
||||
assert len(results) == 2
|
||||
assert "CI lints on every pull request" in [fact["content"] for fact in results]
|
||||
|
||||
def test_legacy_behavior_unchanged_when_disabled(self):
|
||||
facts = [
|
||||
_make_fact("Fact A", confidence=0.3),
|
||||
_make_fact("Fact B", confidence=0.9),
|
||||
]
|
||||
mgr = _deer_mem_with_facts(facts) # default config
|
||||
|
||||
results = mgr.search("Fact", top_k=5)
|
||||
assert [fact["confidence"] for fact in results] == [0.9, 0.3]
|
||||
|
||||
def test_legacy_empty_result_without_substring_match_when_disabled(self):
|
||||
facts = [_make_fact("The project uses PostgreSQL for persistence")]
|
||||
mgr = _deer_mem_with_facts(facts)
|
||||
|
||||
assert mgr.search("database migration", top_k=5) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt injection with query-aware ranking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestInjectionRelevance:
|
||||
def test_diversification_stops_at_budget_and_preserves_guaranteed_pool(self, monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core import prompt
|
||||
|
||||
original = prompt.iter_diversify
|
||||
picked = []
|
||||
|
||||
def counted(*args, **kwargs):
|
||||
for fact in original(*args, **kwargs):
|
||||
picked.append(fact)
|
||||
yield fact
|
||||
|
||||
monkeypatch.setattr(prompt, "iter_diversify", counted)
|
||||
facts = [_make_fact(f"database fact {i}", confidence=0.9) for i in range(100)]
|
||||
facts.append(_make_fact("Always ask before deleting files", category="correction", confidence=0.1))
|
||||
result = prompt.format_memory_for_injection(
|
||||
{"facts": facts},
|
||||
query="database",
|
||||
relevance_weight=0.7,
|
||||
diversity_weight=0.5,
|
||||
**self._injection_args(max_tokens=40, guaranteed_categories=["correction"], guaranteed_token_budget=20),
|
||||
)
|
||||
assert "Always ask before deleting files" in result
|
||||
assert "database fact" in result
|
||||
assert len(picked) < 10
|
||||
|
||||
def _injection_args(self, **overrides):
|
||||
args = {
|
||||
"max_tokens": 300,
|
||||
"use_tiktoken": False,
|
||||
"guaranteed_categories": None,
|
||||
"guaranteed_token_budget": 500,
|
||||
}
|
||||
args.update(overrides)
|
||||
return args
|
||||
|
||||
def test_relevance_reranks_facts_under_token_budget(self):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.prompt import (
|
||||
format_memory_for_injection,
|
||||
)
|
||||
|
||||
facts = [
|
||||
_make_fact("User prefers concise answers", "preference", 0.95),
|
||||
_make_fact("Migrations are managed with alembic", "project", 0.4),
|
||||
]
|
||||
memory_data = {"facts": facts}
|
||||
|
||||
legacy = format_memory_for_injection(
|
||||
memory_data,
|
||||
**self._injection_args(max_tokens=20),
|
||||
)
|
||||
relevance = format_memory_for_injection(
|
||||
memory_data,
|
||||
query="how do I add a database migration",
|
||||
relevance_weight=0.7,
|
||||
**self._injection_args(max_tokens=20),
|
||||
)
|
||||
|
||||
assert "concise answers" in legacy
|
||||
assert "alembic" in relevance
|
||||
assert "alembic" not in legacy
|
||||
|
||||
def test_query_none_preserves_legacy_order(self):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.prompt import (
|
||||
format_memory_for_injection,
|
||||
)
|
||||
|
||||
facts = [
|
||||
_make_fact("User prefers concise answers", "preference", 0.95),
|
||||
_make_fact("Migrations are managed with alembic", "project", 0.4),
|
||||
_make_fact("User lives in Beijing", "personal", 0.8),
|
||||
]
|
||||
memory_data = {"facts": facts}
|
||||
|
||||
legacy = format_memory_for_injection(memory_data, **self._injection_args())
|
||||
with_query_none = format_memory_for_injection(memory_data, query=None, relevance_weight=0.7, **self._injection_args())
|
||||
assert legacy == with_query_none
|
||||
|
||||
|
||||
class TestGetContextQuery:
|
||||
def test_get_context_uses_query_when_enabled(self):
|
||||
facts = [
|
||||
_make_fact("User prefers concise answers", "preference", 0.95),
|
||||
_make_fact("Migrations are managed with alembic", "project", 0.4),
|
||||
]
|
||||
mgr = _deer_mem_with_facts(
|
||||
facts,
|
||||
backend_config={"retrieval_relevance_enabled": True, "retrieval_relevance_weight": 0.7},
|
||||
)
|
||||
|
||||
body = mgr.get_context("user-1", agent_name="assistant", query="how do I add a database migration")
|
||||
assert "alembic" in body
|
||||
|
||||
def test_get_context_without_query_keeps_confidence_order(self):
|
||||
facts = [
|
||||
_make_fact("User prefers concise answers", "preference", 0.95),
|
||||
_make_fact("Migrations are managed with alembic", "project", 0.4),
|
||||
]
|
||||
enabled = _deer_mem_with_facts(
|
||||
facts,
|
||||
backend_config={"retrieval_relevance_enabled": True},
|
||||
)
|
||||
disabled = _deer_mem_with_facts(facts)
|
||||
|
||||
assert enabled.get_context("user-1", agent_name="assistant") == disabled.get_context("user-1", agent_name="assistant")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Middleware wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMiddlewareQueryWiring:
|
||||
@pytest.mark.parametrize("multimodal", [False, True])
|
||||
@pytest.mark.parametrize("user_text", ["Use my PostgreSQL preferences to analyze these reports.", "", "PostgreSQL " * 200], ids=["request", "attachment_only", "bounded_request"])
|
||||
def test_upload_context_does_not_replace_original_query(self, monkeypatch, tmp_path, multimodal, user_text):
|
||||
from unittest import mock
|
||||
|
||||
from deerflow.agents.middlewares.uploads_middleware import UploadsMiddleware
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
uploads = UploadsMiddleware(base_dir=str(tmp_path))
|
||||
files = [{"filename": f"report-{i}.csv", "size": 1024, "path": f"/mnt/user-data/uploads/report-{i}.csv", "extension": ".csv"} for i in range(5)]
|
||||
monkeypatch.setattr(uploads, "_files_from_kwargs", lambda *_: files)
|
||||
content = [{"type": "text", "text": user_text}] if multimodal else user_text
|
||||
runtime = SimpleNamespace(context={})
|
||||
update = uploads.before_agent({"messages": [HumanMessage(content=content, id="msg-1")]}, runtime)
|
||||
uploaded_message = update["messages"][0]
|
||||
assert uploaded_message.additional_kwargs[ORIGINAL_USER_CONTENT_KEY] == user_text
|
||||
with mock.patch("deerflow.agents.lead_agent.prompt._get_memory_context", return_value="") as get_context:
|
||||
DynamicContextMiddleware().before_agent({"messages": [uploaded_message]}, runtime)
|
||||
get_context.assert_called_once()
|
||||
assert get_context.call_args.kwargs["query"] == (user_text.strip()[:1000] or None)
|
||||
|
||||
def test_invalid_original_content_metadata_uses_message_text(self):
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _derive_injection_query
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
|
||||
message = HumanMessage(content="database migration", additional_kwargs={ORIGINAL_USER_CONTENT_KEY: ["not a string"]})
|
||||
assert _derive_injection_query(message) == "database migration"
|
||||
|
||||
def test_first_turn_passes_current_query_to_memory_context(self):
|
||||
from unittest import mock
|
||||
|
||||
mw = DynamicContextMiddleware()
|
||||
state = {
|
||||
"messages": [
|
||||
HumanMessage(content="how do I add a database migration", id="msg-1"),
|
||||
]
|
||||
}
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"deerflow.agents.lead_agent.prompt._get_memory_context",
|
||||
return_value="",
|
||||
) as get_context,
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
mw.before_agent(state, SimpleNamespace(context={}))
|
||||
|
||||
get_context.assert_called_once()
|
||||
assert get_context.call_args.kwargs.get("query") == "how do I add a database migration"
|
||||
|
||||
def test_multimodal_content_yields_text_query(self):
|
||||
from unittest import mock
|
||||
|
||||
mw = DynamicContextMiddleware()
|
||||
state = {
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
content=[
|
||||
{"type": "text", "text": "how do I "},
|
||||
{"type": "text", "text": "add a database migration"},
|
||||
],
|
||||
id="msg-1",
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
with (
|
||||
mock.patch(
|
||||
"deerflow.agents.lead_agent.prompt._get_memory_context",
|
||||
return_value="",
|
||||
) as get_context,
|
||||
mock.patch("deerflow.agents.middlewares.dynamic_context_middleware.datetime") as mock_dt,
|
||||
):
|
||||
mock_dt.now.return_value.strftime.return_value = "2026-05-08, Friday"
|
||||
mw.before_agent(state, SimpleNamespace(context={}))
|
||||
|
||||
assert get_context.call_args.kwargs.get("query") == "how do I add a database migration"
|
||||
230
backend/tests/test_memory_relevance_review.py
Normal file
230
backend/tests/test_memory_relevance_review.py
Normal file
@ -0,0 +1,230 @@
|
||||
"""Regression coverage for PR 5251's backend compatibility and IDF review."""
|
||||
|
||||
import asyncio
|
||||
from copy import deepcopy
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.agents.lead_agent.prompt import _get_memory_context
|
||||
from deerflow.agents.memory import MemoryManager
|
||||
from deerflow.agents.memory.backends.deermem.deer_mem import DeerMem
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.relevance import build_idf, lexical_relevance, tokenize
|
||||
|
||||
|
||||
class _LegacyBackend(MemoryManager):
|
||||
@classmethod
|
||||
def from_config(cls, backend_config, *, mode="middleware", **host_hooks):
|
||||
return cls(backend_config=backend_config or {}, mode=mode)
|
||||
|
||||
def add(self, thread_id, messages, *, agent_name=None, user_id=None, trace_id=None):
|
||||
pass
|
||||
|
||||
def get_context(self, user_id, *, agent_name=None, thread_id=None):
|
||||
return f"memory:{user_id}:{agent_name}:{thread_id}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", [None, "", "database migration"])
|
||||
def test_old_backend_signature_keeps_prompt_memory(monkeypatch, query):
|
||||
manager = _LegacyBackend()
|
||||
monkeypatch.setattr("deerflow.agents.memory.get_memory_manager", lambda: manager)
|
||||
config = SimpleNamespace(memory=SimpleNamespace(enabled=True, injection_enabled=True, backend_config={}))
|
||||
context = _get_memory_context("agent-a", app_config=config, user_id="user-a", query=query)
|
||||
assert "<memory>" in context
|
||||
assert "memory:user-a:agent-a:None" in context
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", [None, "", "database migration"])
|
||||
def test_old_backend_signature_keeps_inherited_async_context(query):
|
||||
assert asyncio.run(_LegacyBackend().aget_context("user-a", agent_name="agent-a", thread_id="thread-a", query=query)) == "memory:user-a:agent-a:thread-a"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("accepts_kwargs", [False, True])
|
||||
def test_query_capable_backend_receives_hint_in_prompt_and_async(monkeypatch, accepts_kwargs):
|
||||
calls = []
|
||||
|
||||
def explicit(self, user_id, *, agent_name=None, thread_id=None, query=None):
|
||||
calls.append((user_id, agent_name, thread_id, query))
|
||||
return "query-aware memory"
|
||||
|
||||
def variadic(self, user_id, *, agent_name=None, thread_id=None, **kwargs):
|
||||
return explicit(self, user_id, agent_name=agent_name, thread_id=thread_id, query=kwargs.get("query"))
|
||||
|
||||
class QueryBackend(_LegacyBackend):
|
||||
get_context = variadic if accepts_kwargs else explicit
|
||||
|
||||
manager = QueryBackend()
|
||||
monkeypatch.setattr("deerflow.agents.memory.get_memory_manager", lambda: manager)
|
||||
config = SimpleNamespace(memory=SimpleNamespace(enabled=True, injection_enabled=True))
|
||||
assert "query-aware memory" in _get_memory_context("a", app_config=config, user_id="u", query="migration")
|
||||
assert asyncio.run(manager.aget_context("u", agent_name="a", thread_id="t", query="migration")) == "query-aware memory"
|
||||
assert calls == [("u", "a", None, "migration"), ("u", "a", "t", "migration")]
|
||||
|
||||
|
||||
def test_backend_typeerror_is_not_retried_as_legacy_signature(monkeypatch):
|
||||
calls = []
|
||||
|
||||
class BrokenBackend(_LegacyBackend):
|
||||
def get_context(self, user_id, *, agent_name=None, thread_id=None, query=None):
|
||||
calls.append(query)
|
||||
raise TypeError("backend implementation failed")
|
||||
|
||||
manager = BrokenBackend()
|
||||
with pytest.raises(TypeError, match="backend implementation failed"):
|
||||
asyncio.run(manager.aget_context("u", query="migration"))
|
||||
assert calls == ["migration"]
|
||||
monkeypatch.setattr("deerflow.agents.memory.get_memory_manager", lambda: manager)
|
||||
config = SimpleNamespace(memory=SimpleNamespace(enabled=True, injection_enabled=True))
|
||||
assert _get_memory_context(app_config=config, user_id="u", query="migration") == ""
|
||||
assert calls == ["migration", "migration"]
|
||||
|
||||
|
||||
def test_uninspectable_legacy_callable_keeps_prompt_memory(monkeypatch):
|
||||
class LegacyCallable:
|
||||
__signature__ = object()
|
||||
|
||||
def __call__(self, user_id, *, agent_name=None, thread_id=None):
|
||||
return "legacy callable memory"
|
||||
|
||||
monkeypatch.setattr("deerflow.agents.memory.get_memory_manager", lambda: SimpleNamespace(get_context=LegacyCallable()))
|
||||
config = SimpleNamespace(memory=SimpleNamespace(enabled=True, injection_enabled=True))
|
||||
assert "legacy callable memory" in _get_memory_context(app_config=config, user_id="u", query="migration")
|
||||
|
||||
|
||||
def _corpus():
|
||||
contents = ["Python coding conventions", "Python database migration uses Alembic"] + [f"Unrelated cooking recipe {index}" for index in range(8)]
|
||||
return [{"id": f"fact_{index}", "content": content, "confidence": 0.7, "category": "context", "createdAt": "2026-01-01T00:00:00Z"} for index, content in enumerate(contents)]
|
||||
|
||||
|
||||
def test_complete_query_coverage_outranks_partial_match_with_corpus_idf():
|
||||
facts = _corpus()
|
||||
idf = build_idf([tokenize(fact["content"]) for fact in facts])
|
||||
partial = lexical_relevance("python database migration", facts[0]["content"], idf=idf)
|
||||
complete = lexical_relevance("python database migration", facts[1]["content"], idf=idf)
|
||||
assert 0 < partial < complete <= 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("partial_first", [False, True])
|
||||
def test_search_top_one_prefers_complete_match_independent_of_input_order(partial_first):
|
||||
facts = _corpus()
|
||||
if not partial_first:
|
||||
facts[0], facts[1] = facts[1], facts[0]
|
||||
manager = DeerMem(backend_config={"retrieval_relevance_enabled": True, "retrieval_relevance_weight": 1.0})
|
||||
manager._updater = SimpleNamespace(get_memory_data=lambda agent_name=None, *, user_id=None: {"facts": facts})
|
||||
result = manager.search("python database migration", top_k=1)
|
||||
assert result[0]["content"] == "Python database migration uses Alembic"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"query,partial,complete",
|
||||
[
|
||||
("python database migration", "python " * 30, "migration database python additional detail"),
|
||||
("python python database migration", "Python coding conventions", "migration database python"),
|
||||
("python database migration", "python", "migration database python"),
|
||||
("python databases migrations", "python database " * 20, "python database migration"),
|
||||
],
|
||||
)
|
||||
def test_partial_matches_cannot_saturate_from_repetition_or_containment(query, partial, complete):
|
||||
idf = build_idf([tokenize(text) for text in [partial, complete, "unrelated"]])
|
||||
assert lexical_relevance(query, partial, idf=idf) < lexical_relevance(query, complete, idf=idf)
|
||||
|
||||
|
||||
def test_rare_query_terms_keep_more_weight_than_common_terms():
|
||||
idf = build_idf([tokenize(text) for text in ["python database migration", "python coding", "python testing", "python packaging"]])
|
||||
assert lexical_relevance("python database migration", "database", idf=idf) > lexical_relevance("python database migration", "python", idf=idf)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query,content", [("PostgreSQL", "Postman collections"), ("authorization", "authentication settings"), ("database", "dataframe columns")])
|
||||
def test_shared_four_character_prefix_is_not_a_lexical_match(query, content):
|
||||
assert lexical_relevance(query, content) == 0.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query,content", [("database", "databases"), ("databases", "database"), ("migration", "migrations"), ("migrations", "migration")])
|
||||
def test_complete_token_prefixes_still_match(query, content):
|
||||
assert lexical_relevance(query, content) == 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("unrelated_first", [False, True])
|
||||
def test_search_prefers_exact_token_over_shared_prefix(unrelated_first):
|
||||
facts = [
|
||||
{"id": "unrelated", "content": "Postman collections", "confidence": 0.7, "category": "context"},
|
||||
{"id": "exact", "content": "PostgreSQL database", "confidence": 0.7, "category": "context"},
|
||||
]
|
||||
if not unrelated_first:
|
||||
facts.reverse()
|
||||
manager = DeerMem(backend_config={"retrieval_relevance_enabled": True, "retrieval_relevance_weight": 1.0})
|
||||
manager._updater = SimpleNamespace(get_memory_data=lambda agent_name=None, *, user_id=None: {"facts": facts})
|
||||
assert manager.search("PostgreSQL", top_k=1)[0]["id"] == "exact"
|
||||
|
||||
|
||||
def test_absent_query_keeps_forwarding_wrapper_legacy_contract(monkeypatch):
|
||||
calls = []
|
||||
inner = _LegacyBackend()
|
||||
|
||||
class ForwardingBackend(_LegacyBackend):
|
||||
def get_context(self, user_id, **kwargs):
|
||||
calls.append(kwargs.copy())
|
||||
return inner.get_context(user_id, **kwargs)
|
||||
|
||||
manager = ForwardingBackend()
|
||||
monkeypatch.setattr("deerflow.agents.memory.get_memory_manager", lambda: manager)
|
||||
config = SimpleNamespace(memory=SimpleNamespace(enabled=True, injection_enabled=True))
|
||||
assert "memory:u:a:None" in _get_memory_context("a", app_config=config, user_id="u")
|
||||
assert asyncio.run(manager.aget_context("u", agent_name="a", thread_id="t")) == "memory:u:a:t"
|
||||
assert calls == [{"agent_name": "a"}, {"agent_name": "a", "thread_id": "t"}]
|
||||
|
||||
|
||||
def _idf_scope(common):
|
||||
contents = ["python conventions", "migration conventions"] + [f"{common} background {index}" for index in range(8)]
|
||||
return {"facts": [{"id": f"fact_{index}", "content": content, "category": "preference" if index < 2 else "context", "confidence": 0.7} for index, content in enumerate(contents)]}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("guaranteed", [[], ["preference"]])
|
||||
@pytest.mark.parametrize("reverse", [False, True])
|
||||
def test_search_and_injection_share_scope_idf_without_cross_scope_leakage(guaranteed, reverse):
|
||||
scopes = {("u1", "a"): _idf_scope("python"), ("u2", "a"): _idf_scope("migration"), ("u1", "b"): _idf_scope("migration")}
|
||||
if reverse:
|
||||
for data in scopes.values():
|
||||
data["facts"].reverse()
|
||||
original = deepcopy(scopes)
|
||||
manager = DeerMem(backend_config={"retrieval_relevance_enabled": True, "retrieval_relevance_weight": 1.0, "token_counting": "char", "guaranteed_categories": guaranteed})
|
||||
manager._updater = SimpleNamespace(get_memory_data=lambda agent_name=None, *, user_id=None: scopes[(user_id, agent_name)])
|
||||
# Return to the first scope to detect accidental reuse of another user's IDF.
|
||||
for user, agent in [("u1", "a"), ("u2", "a"), ("u1", "b"), ("u1", "a")]:
|
||||
search = manager.search("python migration", top_k=10, user_id=user, agent_name=agent)
|
||||
context = manager.get_context(user, agent_name=agent, query="python migration")
|
||||
expected, other = ("migration conventions", "python conventions") if (user, agent) == ("u1", "a") else ("python conventions", "migration conventions")
|
||||
assert search[0]["content"] == expected
|
||||
assert context.index(expected) < context.index(other)
|
||||
assert scopes == original
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled,query,weight", [(False, "python migration", 1.0), (True, None, 1.0), (True, "", 1.0), (True, " ", 1.0), (True, "python migration", 0.0)])
|
||||
def test_inactive_relevance_does_not_build_injection_idf(monkeypatch, enabled, query, weight):
|
||||
def unexpected_idf(corpus):
|
||||
pytest.fail("IDF must not be built when lexical relevance is unused")
|
||||
|
||||
monkeypatch.setattr("deerflow.agents.memory.backends.deermem.deer_mem.build_idf", unexpected_idf)
|
||||
manager = DeerMem(backend_config={"retrieval_relevance_enabled": enabled, "retrieval_relevance_weight": weight, "token_counting": "char"})
|
||||
manager._updater = SimpleNamespace(get_memory_data=lambda agent_name=None, *, user_id=None: _idf_scope("python"))
|
||||
assert manager.get_context("u", agent_name="a", query=query) == manager.get_context("u", agent_name="a")
|
||||
|
||||
|
||||
def test_large_scope_idf_is_bounded_and_keeps_rare_fact_within_budget(monkeypatch):
|
||||
from deerflow.agents.memory.backends.deermem.deermem.core.prompt import _count_tokens
|
||||
|
||||
corpora = []
|
||||
|
||||
def observed_idf(corpus):
|
||||
corpora.append((len(corpus), max(map(len, corpus))))
|
||||
return build_idf(corpus)
|
||||
|
||||
monkeypatch.setattr("deerflow.agents.memory.backends.deermem.deer_mem.build_idf", observed_idf)
|
||||
facts = [{"id": f"fact_{i}", "content": "python background " * 1000, "category": "context", "confidence": 0.7} for i in range(499)]
|
||||
facts.append({"id": "fact_rare", "content": "migration conventions", "category": "context", "confidence": 0.7})
|
||||
manager = DeerMem(backend_config={"retrieval_relevance_enabled": True, "retrieval_relevance_weight": 1.0, "token_counting": "char", "max_injection_tokens": 100, "guaranteed_categories": []})
|
||||
manager._updater = SimpleNamespace(get_memory_data=lambda agent_name=None, *, user_id=None: {"facts": facts})
|
||||
context = manager.get_context("u", agent_name="a", query="python migration")
|
||||
assert "migration conventions" in context
|
||||
assert _count_tokens(context, use_tiktoken=False) <= 100
|
||||
assert corpora == [(500, 128)]
|
||||
@ -113,7 +113,7 @@ class TestDynamicContextMemoryStamping:
|
||||
|
||||
from deerflow.agents.middlewares import dynamic_context_middleware as module
|
||||
|
||||
monkeypatch.setattr(module.DynamicContextMiddleware, "_build_full_reminder", lambda self, runtime=None: ("<system-reminder></system-reminder>", "some recalled memory"))
|
||||
monkeypatch.setattr(module.DynamicContextMiddleware, "_build_full_reminder", lambda self, runtime=None, *, query=None: ("<system-reminder></system-reminder>", "some recalled memory"))
|
||||
middleware = module.DynamicContextMiddleware()
|
||||
result = middleware._inject({"messages": [HumanMessage(content="hello", id="u1")]})
|
||||
memory_messages = [m for m in result["messages"] if str(m.id or "").endswith("__memory")]
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user