From 7d1f21735010955468eb41539714700ecd87b965 Mon Sep 17 00:00:00 2001 From: huanghuoguoguo <60681390+huanghuoguoguo@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:41:07 +0800 Subject: [PATCH] fix(LongTermMemory): unify bounded current and historical recall --- .../LongTermMemory/RECALL_POLICY.md | 13 + .../event_listener/memory_injector.py | 7 +- .../knowledge_engine/memory_engine.py | 120 ++------ .../knowledge_engine/memory_engine.yaml | 4 +- .../components/tools/recall_memory.py | 18 +- .../components/tools/recall_memory.yaml | 10 +- .../LongTermMemory/store/memory_store.py | 105 +++---- .../LongTermMemory/store/recall.py | 195 +++++++++++++ .../tests/store/test_recall_quality.py | 269 ++++++++++++++++++ 9 files changed, 559 insertions(+), 182 deletions(-) create mode 100644 plugins/KnowledgeEngine/LongTermMemory/RECALL_POLICY.md create mode 100644 plugins/KnowledgeEngine/LongTermMemory/store/recall.py create mode 100644 plugins/KnowledgeEngine/LongTermMemory/tests/store/test_recall_quality.py diff --git a/plugins/KnowledgeEngine/LongTermMemory/RECALL_POLICY.md b/plugins/KnowledgeEngine/LongTermMemory/RECALL_POLICY.md new file mode 100644 index 000000000..110e7d264 --- /dev/null +++ b/plugins/KnowledgeEngine/LongTermMemory/RECALL_POLICY.md @@ -0,0 +1,13 @@ +# Recall policy + +Both explicit `recall_memory` and KnowledgeEngine retrieval use the same bounded record selector. Pipeline automatic memory injection reaches that KnowledgeEngine; stable profile blocks are bounded separately. + +- Backend result order is the primary relevance signal. Backend scores and distances are preserved for transport but are not compared on an assumed universal scale. Importance and explicit event recency are small secondary signals; an old recorded timestamp alone does not decay a stable or legacy fact. +- Exact IDs, speaker identity, tags, and literal content matches can boost a result. This is not a semantic relevance classifier, and no universal similarity threshold or semantic abstention guarantee is claimed. +- Current recall honors `[valid_from, valid_until)`. Missing validity fields keep legacy records eligible. Explicit malformed bounds fail closed for current recall. Hidden lifecycle status queries and explicit `time_after`/`time_before` bounds allow historical inspection. +- `time_after`/`time_before` retain their old meaning: recorded `timestamp`, not event time. Scope and identity equality filters are sent to the backend and rechecked locally. String timestamps and string importance use local range checks, since Chroma range operators require numeric metadata. +- Retrieval considers a bounded window of 50–80 backend candidates and returns at most 20 complete records. Local time/status filtering may therefore return fewer results even when eligible records exist beyond that window. Exhaustive historical inspection remains available through the management/export paths. +- Repeated records with the same content, speaker, and time are collapsed. Different actors or event times remain distinct. Similar wording only receives a small diversity penalty; it never marks a memory replaced or deleted. A selected consolidation summary and its source records do not consume the same result budget twice. An overlapping summary is skipped as a whole, so selection does not guarantee coverage of every source represented in that summary. +- Complete serialized record candidates share an 8,000-character budget. Tool profile blocks each have a 2,000-character serialized budget. Oversized records/blocks are omitted, never sliced into partial facts or invalid JSON. Tool responses report omitted profile blocks in `profiles_omitted_for_budget`. The raw-record budget is conservative: it includes backend metadata that may not appear in the rendered answer. These are deterministic character limits, not exact tokenizer accounting. + +The deterministic regressions cover scope and time boundaries, score-scale invariance, retention of old preferences, secondary importance/recency signals, distinct speakers/events, complete-record budgets, and summary/source diversity. They do not replace real-model extraction or answer-quality evaluation. diff --git a/plugins/KnowledgeEngine/LongTermMemory/components/event_listener/memory_injector.py b/plugins/KnowledgeEngine/LongTermMemory/components/event_listener/memory_injector.py index d4efb1f20..b16927ca3 100644 --- a/plugins/KnowledgeEngine/LongTermMemory/components/event_listener/memory_injector.py +++ b/plugins/KnowledgeEngine/LongTermMemory/components/event_listener/memory_injector.py @@ -6,6 +6,7 @@ import re from langbot_plugin.api.definition.components.common.event_listener import EventListener from langbot_plugin.api.entities import events, context from langbot_plugin.api.entities.builtin.provider.message import Message +from store.recall import bounded_profile from langbot_plugin.api.proxies.query_based_api import QueryBasedAPIProxy logger = logging.getLogger(__name__) @@ -206,9 +207,11 @@ class MemoryInjector(EventListener): # Build injection parts blocks: list[str] = [] if session_profile_block.strip(): - blocks.append(session_profile_block) + if bounded_profile(session_profile_block): + blocks.append(session_profile_block) if speaker_profile_block.strip(): - blocks.append(speaker_profile_block) + if bounded_profile(speaker_profile_block): + blocks.append(speaker_profile_block) # Inject current speaker identity so LLM knows who is talking if sender_name: diff --git a/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.py b/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.py index 0cd63e328..42b927390 100644 --- a/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.py +++ b/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import logging -import math import time import uuid as uuid_mod from datetime import datetime, timezone @@ -22,6 +21,7 @@ from langbot_plugin.api.entities.builtin.rag.models import ( IngestionResult, ) from langbot_plugin.api.entities.builtin.rag.enums import DocumentStatus +from store.recall import matches_filters, select_memories from store.serialization import serialized_write logger = logging.getLogger(__name__) @@ -63,28 +63,6 @@ class LongTermMemoryEngine(KnowledgeEngine): # time-decay scoring # ================================================================ - @staticmethod - def _time_decay_score( - similarity: float, timestamp_str: str, half_life_days: float = 30.0, - ) -> float: - """Blend similarity with exponential time decay. - - final_score = similarity * exp(-ln2 * age_days / half_life_days) - - A memory exactly ``half_life_days`` old retains 50 % of its - similarity weight; very recent memories are nearly unaffected. - """ - if not timestamp_str: - return similarity * 0.5 - try: - ts = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00")) - now = datetime.now(timezone.utc) - age_days = max(0.0, (now - ts).total_seconds() / 86400.0) - decay = math.exp(-math.log(2) * age_days / half_life_days) - return similarity * decay - except (ValueError, TypeError): - return similarity * 0.5 - @staticmethod def _parse_importance(metadata: dict[str, Any]) -> int: raw = metadata.get("importance", "2") @@ -93,29 +71,6 @@ class LongTermMemoryEngine(KnowledgeEngine): except (TypeError, ValueError): return 2 - @staticmethod - def _importance_weight(importance: int) -> float: - # Importance should influence ranking, but not overpower semantic match. - # The weights keep importance as a secondary signal. - return { - 1: 0.9, - 2: 1.0, - 3: 1.08, - 4: 1.16, - 5: 1.24, - }.get(importance, 1.0) - - @staticmethod - def _speaker_match_weight( - metadata: dict[str, Any], - current_sender_id: str, - ) -> float: - if not current_sender_id: - return 1.0 - if str(metadata.get("sender_id", "") or "") == current_sender_id: - return 1.12 - return 1.0 - @staticmethod def _has_update_signal(metadata: dict[str, Any]) -> bool: raw_tags = str(metadata.get("tags", "") or "") @@ -169,7 +124,10 @@ class LongTermMemoryEngine(KnowledgeEngine): store = self.plugin.memory_store embedding_model_uuid = settings.get("embedding_model_uuid", "") - top_k = retrieval_settings.get("top_k", 5) + try: + top_k = max(1, min(20, int(retrieval_settings.get("top_k", 5)))) + except (ValueError, TypeError): + top_k = 5 logger.info( "[LongTermMemory] engine retrieve called: collection_id=%s top_k=%s session_name=%s sender_id=%s bot_uuid=%s query_len=%s", collection_id, @@ -200,7 +158,7 @@ class LongTermMemoryEngine(KnowledgeEngine): isolation = settings.get("isolation", "session") # Over-fetch to allow time-decay re-ranking to surface recent # memories that would otherwise be pushed out by pure similarity. - fetch_k = top_k * 3 + fetch_k = max(50, min(20, top_k) * 4) retrieval_strategy = store.normalize_retrieval_strategy( settings.get("retrieval_strategy", "auto") ) @@ -250,6 +208,8 @@ class LongTermMemoryEngine(KnowledgeEngine): item_id = item.get("id", "") if item_id and item_id in seen_ids: continue + if not matches_filters(item.get("metadata", {}), filters): + continue if not store.episode_status_included(item.get("metadata", {})): continue if item_id: @@ -263,11 +223,7 @@ class LongTermMemoryEngine(KnowledgeEngine): bot_uuid, session_name, isolation ) - if sender_id: - await extend_results({"$and": [{"user_key": scope_key}, {"sender_id": sender_id}]}) - - if len(results) < fetch_k: - await extend_results({"user_key": scope_key}) + await extend_results({"user_key": scope_key}) else: # Refuse to search without scope — returning unfiltered results # would leak memories across sessions / users. @@ -277,48 +233,10 @@ class LongTermMemoryEngine(KnowledgeEngine): ) return RetrievalResponse(results=[], total_found=0) - # Time-decay re-ranking: blend similarity with recency so that - # recent memories can outrank older ones with marginally higher - # vector similarity. - half_life = float(settings.get("recency_half_life_days", 30)) - for r in results: - sim = r.get("score") or (1.0 - r.get("distance", 1.0)) - metadata = r.get("metadata", {}) - ts = metadata.get("timestamp", "") - importance = self._parse_importance(metadata) - time_score = self._time_decay_score(float(sim), ts, half_life) - importance_weight = self._importance_weight(importance) - speaker_weight = self._speaker_match_weight(metadata, sender_id) - update_weight = 1.08 if self._has_update_signal(metadata) else 1.0 - exact_boost = 1.0 - if exact_match_boost: - exact_score = store._metadata_exact_match_score( - query, - r.get("id", ""), - metadata, - ) - exact_boost += min(exact_score, 2.0) * 0.25 - final_score = ( - time_score - * importance_weight - * speaker_weight - * update_weight - * exact_boost - ) - r["_final_score"] = max(0.0, min(1.0, final_score)) - r["_recency_hint"] = self._recency_hint(ts) - r["_importance"] = importance - r["_speaker_match"] = speaker_weight > 1.0 - r["_has_update_signal"] = update_weight > 1.0 - r["_exact_match"] = exact_boost > 1.0 - - results.sort(key=lambda r: r["_final_score"], reverse=True) - results = results[:top_k] - logger.info( - "[LongTermMemory] engine retrieve completed: collection_id=%s result_count=%s half_life_days=%s", - collection_id, - len(results), - half_life, + results = select_memories( + results, query, top_k=top_k, exact_match_boost=exact_match_boost, + current_sender_id=sender_id, + half_life_days=settings.get("recency_half_life_days", 30), ) entries: list[RetrievalResultEntry] = [] @@ -329,14 +247,14 @@ class LongTermMemoryEngine(KnowledgeEngine): importance = str(r.get("_importance", self._parse_importance(meta))) tags = meta.get("tags", "") hints = [f"importance:{importance}"] - recency_hint = r.get("_recency_hint", "unknown") + recency_hint = self._recency_hint(meta.get("event_at") or timestamp) if recency_hint != "unknown": hints.append(f"recency:{recency_hint}") - if r.get("_speaker_match"): + if sender_id and str(meta.get("sender_id", "")) == sender_id: hints.append("speaker:current") - if r.get("_has_update_signal"): + if self._has_update_signal(meta): hints.append("signal:update") - if r.get("_exact_match"): + if exact_match_boost and store._metadata_exact_match_score(query, r.get("id", ""), meta): hints.append("match:exact") display = f"[{timestamp}] ({', '.join(hints)})" @@ -350,8 +268,8 @@ class LongTermMemoryEngine(KnowledgeEngine): id=r.get("id", ""), content=[{"type": "text", "text": display}], metadata=meta, - score=r.get("_final_score", r.get("score")), - distance=1.0 - r.get("_final_score", r.get("score", 0.0)), + score=r.get("score") or 0.0, + distance=r.get("distance") or 0.0, ) ) diff --git a/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.yaml b/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.yaml index eb40c4f00..ad2de6998 100644 --- a/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.yaml +++ b/plugins/KnowledgeEngine/LongTermMemory/components/knowledge_engine/memory_engine.yaml @@ -46,8 +46,8 @@ spec: en_US: Recency Half-Life Days zh_Hans: 时效半衰期天数 description: - en_US: Controls how quickly older episodic memories lose ranking weight during automatic recall. - zh_Hans: 控制自动召回时旧情景记忆的时效衰减速度。 + en_US: Controls the small recency bonus for explicitly dated events. Undated stable and legacy facts do not decay. + zh_Hans: 控制明确事件时间的近期加分。无事件时间的稳定或旧记忆不衰减。 type: integer required: false default: 30 diff --git a/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.py b/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.py index 8f7da9d32..d76df0b78 100644 --- a/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.py +++ b/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.py @@ -4,6 +4,8 @@ import json import logging from typing import Any +from store.recall import bounded_profile + from langbot_plugin.api.definition.components.tool.tool import Tool from langbot_plugin.api.entities.builtin.provider import session as provider_session @@ -47,8 +49,8 @@ class RecallMemory(Tool): return "Error: query is required." top_k = params.get("top_k", 5) - if not isinstance(top_k, int) or top_k <= 0: - return "Error: top_k must be a positive integer." + if not isinstance(top_k, int) or isinstance(top_k, bool) or not 1 <= top_k <= 20: + return "Error: top_k must be an integer between 1 and 20." speaker_id = params.get("speaker_id", "") if speaker_id is None: @@ -147,8 +149,16 @@ class RecallMemory(Tool): kb_id, len(episodes), ) + profile_texts = { + "session": store.format_profile_text(session_profile), + "speaker": store.format_profile_text(speaker_profile), + } return json.dumps({ - "session_profile": store.format_profile_text(session_profile), - "speaker_profile": store.format_profile_text(speaker_profile), + "session_profile": bounded_profile(profile_texts["session"]), + "speaker_profile": bounded_profile(profile_texts["speaker"]), "episodes": episodes, + "profiles_omitted_for_budget": [ + name for name, text in profile_texts.items() + if text and not bounded_profile(text) + ], }, ensure_ascii=False) diff --git a/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.yaml b/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.yaml index adeb40d6a..b19d10b4e 100644 --- a/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.yaml +++ b/plugins/KnowledgeEngine/LongTermMemory/components/tools/recall_memory.yaml @@ -18,7 +18,9 @@ spec: description: Semantic query for the memory search. top_k: type: integer - description: Number of memories to return. Optional, default is 5. + minimum: 1 + maximum: 20 + description: Number of memories to return, between 1 and 20. Optional, default is 5. speaker_id: type: string description: Optional speaker ID filter. Use when you need memories from a specific participant. @@ -70,7 +72,11 @@ spec: Important behavior: - This is a read-only lookup tool. Do not call remember or update_profile just because recall returned no result. - - These filters narrow results rather than broaden them. + - Identity and scope filters always narrow results; they also apply to exact episode-ID lookup. + - Default current recall omits records before valid_from or at/after valid_until. + - time_after/time_before continue to filter recorded timestamp (not event time). Providing either bound requests historical inspection and permits expired or future-valid records within that recorded-time range. + - Hidden lifecycle status filters also permit historical inspection. + - Complete records are selected within a fixed context budget; oversized records and profile blocks are omitted rather than cut in half. profiles_omitted_for_budget names profile blocks omitted for this reason. - By default, recall only returns active memories. Use status=superseded or status=archived only when the user asks to inspect older hidden records. - When multiple filters are provided together, they are typically combined as restrictive conditions. - Do not add filters unless they are clearly needed. diff --git a/plugins/KnowledgeEngine/LongTermMemory/store/memory_store.py b/plugins/KnowledgeEngine/LongTermMemory/store/memory_store.py index 043c05d07..d1f860a48 100644 --- a/plugins/KnowledgeEngine/LongTermMemory/store/memory_store.py +++ b/plugins/KnowledgeEngine/LongTermMemory/store/memory_store.py @@ -8,6 +8,7 @@ import logging from datetime import datetime, timezone from typing import Any from langbot_plugin.entities.io.errors import ActionCallError +from store.recall import exact_match_score, matches_filters, select_memories from store.serialization import FENCE_KEY_PREFIX, BindingWrites, serialized_write logger = logging.getLogger(__name__) @@ -144,27 +145,7 @@ class MemoryStore: @staticmethod def _metadata_exact_match_score(query: str, item_id: str, metadata: dict[str, Any]) -> float: - needle = query.strip().lower() - if not needle: - return 0.0 - score = 0.0 - if needle == str(item_id or "").strip().lower(): - score += 2.0 - for field in ("sender_id", "sender_name"): - value = str(metadata.get(field, "") or "").strip().lower() - if value and needle == value: - score += 1.2 - tags = { - tag.strip().lower() - for tag in str(metadata.get("tags", "") or "").split(",") - if tag.strip() - } - if needle in tags: - score += 1.2 - content = str(metadata.get("content", "") or "").lower() - if needle and needle in content: - score += 0.6 - return score + return exact_match_score(query, item_id, metadata) # ======================== common helpers ======================== @@ -1389,9 +1370,10 @@ class MemoryStore: # string comparison works correctly for single-digit values in this range. filters["importance"] = {"$gte": str(importance_min)} - # ChromaDB requires $and wrapper when there are multiple filter conditions - if len(filters) > 1: - filters = {"$and": [{k: v} for k, v in filters.items()]} + # Legacy timestamps/importance are strings; Chroma range operators require numbers. + backend_filters = {key: value for key, value in filters.items() if not isinstance(value, dict)} + if len(backend_filters) > 1: + backend_filters = {"$and": [{key: value} for key, value in backend_filters.items()]} exact_episode = None if exact_match_boost and user_key: @@ -1401,14 +1383,14 @@ class MemoryStore: user_key=user_key, ) - fetch_k = max(top_k, top_k * 4, 50) + fetch_k = max(50, min(20, top_k) * 4) strategy = self.normalize_retrieval_strategy(retrieval_strategy) normalized_weight = self.normalize_vector_weight(vector_weight) search_kwargs = { "collection_id": collection_id, "query_vector": query_vector, "top_k": fetch_k, - "filters": filters if filters else None, + "filters": backend_filters or None, } if strategy in {"auto", "hybrid"}: try: @@ -1439,54 +1421,35 @@ class MemoryStore: filters if filters else None, ) - seen_ids: set[str] = set() - episodes = [] - if exact_episode and self.episode_status_included( - exact_episode.get("metadata", {}), - include_statuses, - ): - exact_episode["score"] = 1.0 - exact_episode["_exact_match_score"] = 2.0 - episodes.append(exact_episode) - seen_ids.add(exact_episode["id"]) - - for r in results: - rid = r.get("id", "") - if rid in seen_ids: + candidates = [] + if exact_episode: + candidates.append({"id": exact_episode["id"], "metadata": exact_episode.get("metadata", {}), "score": exact_episode.get("score")}) + candidates.extend(results) + seen_ids = set() + scoped = [] + for record in candidates: + rid = record.get("id", "") + metadata = record.get("metadata", {}) + if rid in seen_ids or not matches_filters(metadata, filters): continue - meta = r.get("metadata", {}) - if not self.episode_status_included(meta, include_statuses): + if not self.episode_status_included(metadata, include_statuses): continue - exact_score = self._metadata_exact_match_score(query, rid, meta) if exact_match_boost else 0.0 - episodes.append( - { - "id": rid, - "content": meta.get("content", ""), - "tags": meta.get("tags", "").split(",") if meta.get("tags") else [], - "importance": int(meta.get("importance", "2")), - "timestamp": meta.get("timestamp", ""), - "sender_id": meta.get("sender_id", ""), - "sender_name": meta.get("sender_name", ""), - "source": meta.get("source", ""), - "status": self.episode_status_from_metadata(meta), - "superseded_by": meta.get("superseded_by", ""), - "score": r.get("score"), - "_exact_match_score": exact_score, - } - ) seen_ids.add(rid) - - if exact_match_boost: - episodes.sort( - key=lambda ep: ( - float(ep.get("_exact_match_score", 0.0)), - float(ep.get("score") or 0.0), - ), - reverse=True, - ) - for ep in episodes: - ep.pop("_exact_match_score", None) - return episodes[:top_k] + scoped.append(record) + selected = select_memories( + scoped, query, top_k=top_k, exact_match_boost=exact_match_boost, + historical=bool(time_after or time_before) or self.normalize_episode_statuses(include_statuses) != {self.EPISODE_STATUS_ACTIVE}, + ) + episodes = [] + for record in selected: + meta = record.get("metadata", {}) + episode = self._episode_from_vector_item(record) + episode.pop("metadata", None) + episode["score"] = record.get("score") + for field in ("event_at", "valid_from", "valid_until", "source_ref", "claim_type", "consolidation_source_ids"): + episode[field] = meta.get(field, "unknown" if field == "claim_type" else "") + episodes.append(episode) + return episodes @serialized_write async def delete_episodes_by_user( diff --git a/plugins/KnowledgeEngine/LongTermMemory/store/recall.py b/plugins/KnowledgeEngine/LongTermMemory/store/recall.py new file mode 100644 index 000000000..3384cea70 --- /dev/null +++ b/plugins/KnowledgeEngine/LongTermMemory/store/recall.py @@ -0,0 +1,195 @@ +"""Bounded, backend-independent selection of complete memory records.""" + +from __future__ import annotations + +import json +import math +import re +from datetime import datetime, timezone +from typing import Any + +MAX_RECALL_CHARS = 8000 +MAX_PROFILE_CHARS = 2000 +MAX_RECALL_RESULTS = 20 + + +def parse_time(value: Any) -> datetime | None: + if not value: + return None + try: + parsed = datetime.fromisoformat(str(value).replace("Z", "+00:00")) + return parsed.astimezone(timezone.utc) if parsed.tzinfo else None + except (ValueError, TypeError): + return None + + +def currently_valid(metadata: dict, now: datetime) -> bool: + start = parse_time(metadata.get("valid_from")) + end = parse_time(metadata.get("valid_until")) + if (metadata.get("valid_from") and start is None) or ( + metadata.get("valid_until") and end is None + ): + return False + return (start is None or start <= now) and (end is None or now < end) + + +def matches_filters(metadata: dict, filters: dict | None) -> bool: + """Recheck Host-controlled filters for ID lookups and backend responses.""" + if not filters: + return True + if "$and" in filters: + return all(matches_filters(metadata, item) for item in filters["$and"]) + for key, expected in filters.items(): + actual = str(metadata.get(key, "")) + if isinstance(expected, dict): + if not actual or any( + (operator == "$gte" and actual < str(value)) + or (operator == "$lte" and actual > str(value)) + or operator not in {"$gte", "$lte"} + for operator, value in expected.items() + ): + return False + elif actual != str(expected): + return False + return True + + +def exact_match_score(query: str, item_id: str, metadata: dict) -> float: + needle = query.strip().casefold() + if not needle: + return 0.0 + score = 2.0 if needle == str(item_id).casefold() else 0.0 + for field in ("sender_id", "sender_name"): + if needle == str(metadata.get(field, "")).strip().casefold(): + score += 1.2 + if needle in { + tag.strip().casefold() for tag in str(metadata.get("tags", "")).split(",") + }: + score += 1.2 + if needle in str(metadata.get("content", "")).casefold(): + score += 0.6 + return score + + +def _terms(content: str) -> set[str]: + content = content.casefold() + words = set(re.findall(r"\w+", content.casefold())) + # Character shingles also work for CJK text without a tokenizer dependency. + words.update( + content[index : index + 3] for index in range(max(0, len(content) - 2)) + ) + return words + + +def _source_ids(metadata: dict) -> set[str]: + try: + sources = json.loads(metadata.get("consolidation_source_ids", "[]")) + except (ValueError, TypeError): + return set() + return ( + {item for item in sources if isinstance(item, str)} + if isinstance(sources, list) + else set() + ) + + +def select_memories( + records: list[dict], + query: str, + *, + top_k: int = 5, + exact_match_boost: bool = True, + current_sender_id: str = "", + historical: bool = False, + max_chars: int = MAX_RECALL_CHARS, + half_life_days: float = 30.0, + now: datetime | None = None, +) -> list[dict]: + """Use backend rank as relevance; scores/distances have no shared calibration. + + Complete serialized records count against the character budget. Oversized + records are omitted instead of truncating facts or producing partial JSON. + """ + now = now or datetime.now(timezone.utc) + try: + half_life_days = float(half_life_days) + if not math.isfinite(half_life_days) or half_life_days <= 0: + half_life_days = 30.0 + except (ValueError, TypeError): + half_life_days = 30.0 + candidates = [] + identities = set() + for rank, record in enumerate(records): + metadata = record.get("metadata", {}) + if not historical and not currently_valid(metadata, now): + continue + content = str(metadata.get("content", "") or "") + if not content or len(content) > max_chars: + continue + if len(json.dumps(record, ensure_ascii=False, default=str)) + 2 > max_chars: + continue + identity = ( + " ".join(content.casefold().split()), + str(metadata.get("sender_id", "")), + str(metadata.get("event_at") or metadata.get("timestamp", "")), + str(metadata.get("valid_from", "")), + str(metadata.get("valid_until", "")), + ) + if identity in identities: + continue + identities.add(identity) + try: + importance = min(5, max(1, int(metadata.get("importance", 2)))) + except (ValueError, TypeError): + importance = 2 + score = 0.1 / (rank + 1) + (importance - 2) * 0.0025 + event_at = parse_time(metadata.get("event_at")) + if event_at is not None: + age = max(0.0, (now - event_at).total_seconds() / 86400) + score += 0.008 * math.exp(-math.log(2) * age / half_life_days) + if ( + current_sender_id + and str(metadata.get("sender_id", "")) == current_sender_id + ): + score += 0.004 + if exact_match_boost: + score += exact_match_score(query, str(record.get("id", "")), metadata) + candidates.append((record, score, _terms(content))) + chosen = [] + chosen_terms: list[set[str]] = [] + covered_ids: set[str] = set() + used_chars = 2 + while candidates and len(chosen) < max(0, min(MAX_RECALL_RESULTS, top_k)): + + def diverse_score(candidate): + terms = candidate[2] + overlap = max( + ( + len(terms & prior) / max(1, len(terms | prior)) + for prior in chosen_terms + ), + default=0.0, + ) + return candidate[1] - overlap * 0.015 + + best = max( + range(len(candidates)), key=lambda index: diverse_score(candidates[index]) + ) + record, _score, terms = candidates.pop(best) + ids = {str(record.get("id", ""))} | _source_ids(record.get("metadata", {})) + if ids & covered_ids: + continue + size = len(json.dumps(record, ensure_ascii=False, default=str)) + 2 + if used_chars + size > max_chars: + continue + chosen.append(record) + chosen_terms.append(terms) + covered_ids.update(ids) + used_chars += size + return chosen + + +def bounded_profile(text: str) -> str: + return ( + text if len(json.dumps(text, ensure_ascii=False)) <= MAX_PROFILE_CHARS else "" + ) diff --git a/plugins/KnowledgeEngine/LongTermMemory/tests/store/test_recall_quality.py b/plugins/KnowledgeEngine/LongTermMemory/tests/store/test_recall_quality.py new file mode 100644 index 000000000..b9fe79837 --- /dev/null +++ b/plugins/KnowledgeEngine/LongTermMemory/tests/store/test_recall_quality.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import json +from datetime import datetime, timezone + +import pytest + +from components.knowledge_engine.memory_engine import LongTermMemoryEngine +from langbot_plugin.api.entities.builtin.rag.context import RetrievalContext +from store.memory_store import MemoryStore +from store.recall import bounded_profile, matches_filters, select_memories + +NOW = datetime(2026, 11, 2, tzinfo=timezone.utc) + + +def record(item_id, content=None, **metadata): + return { + "id": item_id, + "score": 0.001, + "metadata": { + "content": content or item_id, + "user_key": "bot:group_room", + "sender_id": "alice", + "timestamp": "2026-01-01T00:00:00Z", + "status": "active", + "importance": "2", + **metadata, + }, + } + + +class Plugin: + def __init__(self, records): + self.records = records + self.memory_store = MemoryStore(self) + self.calls = [] + + async def invoke_embedding(self, _model, _texts): + return [[1.0]] + + async def vector_search(self, **_kwargs): + self.calls.append(_kwargs) + # Deliberately ignore filters: the caller must still recheck them. + return self.records + + async def vector_list(self, collection_id, filters=None, **_kwargs): + items = [ + item for item in self.records if matches_filters(item["metadata"], filters) + ] + return {"items": items, "total": len(items)} + + +def test_backend_score_scale_and_zero_values_do_not_change_selection(): + records = [record("first"), record("second")] + baseline = [r["id"] for r in select_memories(records, "different wording", now=NOW)] + records[0]["score"] = 0 + records[1]["score"] = 100000 + assert [ + r["id"] for r in select_memories(records, "different wording", now=NOW) + ] == baseline + + +def test_current_validity_and_explicit_history_have_distinct_semantics(): + records = [ + record("expired", valid_until="2026-11-02T00:00:00Z"), + record("future", valid_from="2026-11-03T00:00:00Z"), + record("legacy"), + record("current", valid_from="2026-11-02T00:00:00Z"), + ] + assert {r["id"] for r in select_memories(records, "unrelated", now=NOW)} == { + "legacy", + "current", + } + assert len(select_memories(records, "unrelated", historical=True, now=NOW)) == 4 + + +def test_old_timestamp_does_not_age_out_stable_or_legacy_facts(): + records = [ + record("preference", "Alice prefers tea", timestamp="2000-01-01T00:00:00Z"), + record("recent", event_at="2026-11-02T00:00:00Z"), + ] + assert ( + select_memories(records, "preferences", top_k=1, now=NOW)[0]["id"] + == "preference" + ) + + +def test_importance_and_explicit_event_recency_are_secondary_signals(): + records = [record(str(i)) for i in range(10)] + records[-1]["metadata"].update(importance="5", event_at="2026-11-02T00:00:00Z") + selected = [r["id"] for r in select_memories(records, "other", top_k=10, now=NOW)] + assert selected[0] == "0" + assert selected.index("9") < selected.index("8") + + +def test_duplicate_facts_do_not_collapse_distinct_speakers_or_events(): + records = [ + record("a", "Meeting Friday"), + record("copy", "Meeting Friday"), + record("bob", "Meeting Friday", sender_id="bob"), + record("later", "Meeting Friday", event_at="2026-11-03T00:00:00Z"), + ] + assert { + r["id"] for r in select_memories(records, "meeting", top_k=10, now=NOW) + } == {"a", "bob", "later"} + + +def test_complete_records_and_json_remain_within_budget(): + huge = record("huge", "Fact " * 2000) + small = record("small", "A complete source fact") + assert select_memories([huge, small], "other", max_chars=1000) == [small] + assert select_memories([small], "other", max_chars=1) == [] + assert bounded_profile("Profile " * 1000) == "" + assert bounded_profile("A whole profile") == "A whole profile" + assert ( + len( + json.dumps( + select_memories([huge, small], "other", max_chars=1000), + ensure_ascii=False, + ) + ) + <= 1000 + ) + + +def test_summary_and_its_sources_do_not_repeat_the_same_evidence(): + summary = record( + "summary", "Agreed milestones", consolidation_source_ids='["source"]' + ) + selected = select_memories( + [summary, record("source", "Agreed milestone")], "other", top_k=5 + ) + assert len(selected) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "filter_name,filter_value", + [ + ("sender_id", "bob"), + ("sender_name", "Bob"), + ("source", "import"), + ("importance_min", 5), + ("time_after", "2027-01-01T00:00:00Z"), + ], +) +async def test_exact_id_cannot_bypass_other_filters(filter_name, filter_value): + plugin = Plugin([record("target", source="agent", sender_name="Alice")]) + result = await plugin.memory_store.search_episodes( + "kb", "emb", "target", user_key="bot:group_room", **{filter_name: filter_value} + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_time_filter_is_recorded_time_and_allows_expired_active_history(): + plugin = Plugin([record("expired", valid_until="2001-01-01T00:00:00Z")]) + assert ( + await plugin.memory_store.search_episodes( + "kb", + "emb", + "expired", + user_key="bot:group_room", + include_statuses=["active"], + ) + == [] + ) + historical = await plugin.memory_store.search_episodes( + "kb", + "emb", + "expired", + user_key="bot:group_room", + time_after="2026-01-01T00:00:00Z", + ) + assert historical[0]["id"] == "expired" + + +@pytest.mark.asyncio +async def test_engine_uses_same_validity_and_scope_selection(): + plugin = Plugin( + [ + record("wrong-scope", user_key="other"), + record("expired", valid_until="2001-01-01T00:00:00Z"), + record("valid"), + ] + ) + engine = LongTermMemoryEngine() + engine.plugin = plugin + response = await engine.retrieve( + RetrievalContext( + query="other", + knowledge_base_id="kb", + creation_settings={"embedding_model_uuid": "emb"}, + retrieval_settings={ + "session_name": "group_room", + "bot_uuid": "bot", + "sender_id": "alice", + }, + ) + ) + assert [entry.id for entry in response.results] == ["valid"] + + +@pytest.mark.asyncio +async def test_string_time_ranges_are_checked_locally_not_pushed_to_chroma(): + plugin = Plugin( + [ + record("old", timestamp="2000-01-01T00:00:00Z"), + record("new", timestamp="2026-01-01T00:00:00Z"), + ] + ) + result = await plugin.memory_store.search_episodes( + "kb", + "emb", + "other", + user_key="bot:group_room", + time_after="2020-01-01T00:00:00Z", + importance_min=2, + ) + assert [item["id"] for item in result] == ["new"] + assert plugin.calls[0]["filters"] == {"user_key": "bot:group_room"} + + +@pytest.mark.asyncio +async def test_tool_omits_whole_oversized_profile_without_corrupting_json(): + from types import SimpleNamespace + from unittest.mock import AsyncMock + from components.tools.recall_memory import RecallMemory + from langbot_plugin.api.entities.builtin.provider.session import ( + Session, + LauncherTypes, + ) + + store = SimpleNamespace( + resolve_user_context=AsyncMock( + return_value=( + "scope", + "scope", + "kb", + "session", + {"embedding_model_uuid": "emb"}, + ) + ), + search_episodes=AsyncMock(return_value=[]), + load_session_profile=AsyncMock(return_value={"text": "Complete fact " * 1000}), + load_speaker_profile=AsyncMock( + return_value={"text": "Alice prefers concise replies"} + ), + format_profile_text=lambda profile: profile["text"], + has_profile_data=bool, + EPISODE_STATUSES={"active"}, + ) + tool = RecallMemory() + tool.plugin = SimpleNamespace(memory_store=store) + result = json.loads( + await tool.call( + {"query": "preferences"}, + Session( + launcher_type=LauncherTypes.GROUP, + launcher_id="room", + sender_id="alice", + bot_uuid="bot", + ), + 1, + ) + ) + assert result["session_profile"] == "" + assert result["speaker_profile"] == "Alice prefers concise replies" + assert result["profiles_omitted_for_budget"] == ["session"]