mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 02:46:02 +00:00
* feat(memory): add guaranteed injection for correction facts with graceful fallback When the token budget is tight, high-value facts (e.g. user corrections) can be silently evicted by lower-priority regular facts. This change: - Introduces configurable 'guaranteed_categories' (default: [correction]) whose facts draw from a separate 'guaranteed_token_budget', ensuring they are never dropped due to budget pressure. - Adds a graceful fallback to confidence-only ranking when the guaranteed-category path raises an unexpected exception. - Refactors fact selection into a header-agnostic helper (_select_fact_lines) with explicit token accounting in the caller, eliminating double-counting of separators. - Emits a single 'Facts:' header regardless of whether both guaranteed and regular facts are present. - Extends the final safety truncation limit to account for the additional guaranteed budget so guaranteed facts survive end-to-end. * refactor(memory): address review feedback on guaranteed injection - Restore strict break-on-overflow in `_select_fact_lines` to preserve the caller's confidence-ordered ranking; add a regression test locking in the invariant that a shorter lower-confidence fact never slips ahead of a skipped higher-confidence one. - Account for the inter-group `\n` separator between guaranteed and regular fact blocks in the regular budget (1-token precision fix). - Clarify docstrings on `format_memory_for_injection` and `MemoryConfig.guaranteed_token_budget` to distinguish the common *displacement* case (total stays within `max_tokens`) from the rarer *additive* case (safety-truncation ceiling raised when guaranteed lines alone would overflow). * fix(memory): address P1 safety truncation + P2s from review - Structure-aware safety truncation: Facts block is now a protected suffix so guaranteed-category facts can never be silently discarded by a prefix-cut on overflow. Only the preceding (user/history) sections are eligible for truncation. - Extend the same protected-suffix treatment to the except/fallback path by returning fact lines alongside the formatted section from _fallback_format_facts, avoiding string parsing. - Single inter-section separator: facts section no longer embeds its own leading \n\n; the final "\n\n".join(sections) is the single source of truth for section-to-section spacing. - Bare string for guaranteed_categories now raises TypeError instead of silently iterating single characters. - Category-less / malformed facts no longer default-promote into the guaranteed "context" pool — only facts with an explicit category field qualify. - Lift valid_facts pre-filter outside the try so the fallback path reuses it instead of re-doing validation work. - MemoryConfigResponse + DeerFlowClient.get_memory_config now expose guaranteed_categories / guaranteed_token_budget. - config.example.yaml: document the two new fields and bump config_version from 12 to 13. - Add regression tests for every finding. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
729 lines
33 KiB
Python
729 lines
33 KiB
Python
"""Prompt templates for memory update and injection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import math
|
|
import re
|
|
import threading
|
|
import time
|
|
from typing import Any, cast
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
import tiktoken
|
|
|
|
TIKTOKEN_AVAILABLE = True
|
|
except ImportError:
|
|
TIKTOKEN_AVAILABLE = False
|
|
|
|
# Prompt template for updating memory based on conversation
|
|
MEMORY_UPDATE_PROMPT = """You are a memory management system. Your task is to analyze a conversation and update the user's memory profile.
|
|
|
|
Current Memory State:
|
|
<current_memory>
|
|
{current_memory}
|
|
</current_memory>
|
|
|
|
New Conversation to Process:
|
|
<conversation>
|
|
{conversation}
|
|
</conversation>
|
|
|
|
Instructions:
|
|
1. Analyze the conversation for important information about the user
|
|
2. Extract relevant facts, preferences, and context with specific details (numbers, names, technologies)
|
|
3. Update the memory sections as needed following the detailed length guidelines below
|
|
|
|
Before extracting facts, perform a structured reflection on the conversation:
|
|
1. Error/Retry Detection: Did the agent encounter errors, require retries, or produce incorrect results?
|
|
If yes, record the root cause and correct approach as a high-confidence fact with category "correction".
|
|
2. User Correction Detection: Did the user correct the agent's direction, understanding, or output?
|
|
If yes, record the correct interpretation or approach as a high-confidence fact with category "correction".
|
|
Include what went wrong in "sourceError" only when category is "correction" and the mistake is explicit in the conversation.
|
|
3. Project Constraint Discovery: Were any project-specific constraints discovered during the conversation?
|
|
If yes, record them as facts with the most appropriate category and confidence.
|
|
|
|
{correction_hint}
|
|
|
|
Memory Section Guidelines:
|
|
|
|
**User Context** (Current state - concise summaries):
|
|
- workContext: Professional role, company, key projects, main technologies (2-3 sentences)
|
|
Example: Core contributor, project names with metrics (16k+ stars), technical stack
|
|
- personalContext: Languages, communication preferences, key interests (1-2 sentences)
|
|
Example: Bilingual capabilities, specific interest areas, expertise domains
|
|
- topOfMind: Multiple ongoing focus areas and priorities (3-5 sentences, detailed paragraph)
|
|
Example: Primary project work, parallel technical investigations, ongoing learning/tracking
|
|
Include: Active implementation work, troubleshooting issues, market/research interests
|
|
Note: This captures SEVERAL concurrent focus areas, not just one task
|
|
|
|
**History** (Temporal context - rich paragraphs):
|
|
- recentMonths: Detailed summary of recent activities (4-6 sentences or 1-2 paragraphs)
|
|
Timeline: Last 1-3 months of interactions
|
|
Include: Technologies explored, projects worked on, problems solved, interests demonstrated
|
|
- earlierContext: Important historical patterns (3-5 sentences or 1 paragraph)
|
|
Timeline: 3-12 months ago
|
|
Include: Past projects, learning journeys, established patterns
|
|
- longTermBackground: Persistent background and foundational context (2-4 sentences)
|
|
Timeline: Overall/foundational information
|
|
Include: Core expertise, longstanding interests, fundamental working style
|
|
|
|
**Facts Extraction**:
|
|
- Extract specific, quantifiable details (e.g., "16k+ GitHub stars", "200+ datasets")
|
|
- Include proper nouns (company names, project names, technology names)
|
|
- Preserve technical terminology and version numbers
|
|
- Categories:
|
|
* preference: Tools, styles, approaches user prefers/dislikes
|
|
* knowledge: Specific expertise, technologies mastered, domain knowledge
|
|
* context: Background facts (job title, projects, locations, languages)
|
|
* behavior: Working patterns, communication habits, problem-solving approaches
|
|
* goal: Stated objectives, learning targets, project ambitions
|
|
* correction: Explicit agent mistakes or user corrections, including the correct approach
|
|
- Confidence levels:
|
|
* 0.9-1.0: Explicitly stated facts ("I work on X", "My role is Y")
|
|
* 0.7-0.8: Strongly implied from actions/discussions
|
|
* 0.5-0.6: Inferred patterns (use sparingly, only for clear patterns)
|
|
|
|
**What Goes Where**:
|
|
- workContext: Current job, active projects, primary tech stack
|
|
- personalContext: Languages, personality, interests outside direct work tasks
|
|
- topOfMind: Multiple ongoing priorities and focus areas user cares about recently (gets updated most frequently)
|
|
Should capture 3-5 concurrent themes: main work, side explorations, learning/tracking interests
|
|
- recentMonths: Detailed account of recent technical explorations and work
|
|
- earlierContext: Patterns from slightly older interactions still relevant
|
|
- longTermBackground: Unchanging foundational facts about the user
|
|
|
|
**Multilingual Content**:
|
|
- Preserve original language for proper nouns and company names
|
|
- Keep technical terms in their original form (DeepSeek, LangGraph, etc.)
|
|
- Note language capabilities in personalContext
|
|
|
|
Output Format (JSON):
|
|
{{
|
|
"user": {{
|
|
"workContext": {{ "summary": "...", "shouldUpdate": true/false }},
|
|
"personalContext": {{ "summary": "...", "shouldUpdate": true/false }},
|
|
"topOfMind": {{ "summary": "...", "shouldUpdate": true/false }}
|
|
}},
|
|
"history": {{
|
|
"recentMonths": {{ "summary": "...", "shouldUpdate": true/false }},
|
|
"earlierContext": {{ "summary": "...", "shouldUpdate": true/false }},
|
|
"longTermBackground": {{ "summary": "...", "shouldUpdate": true/false }}
|
|
}},
|
|
"newFacts": [
|
|
{{ "content": "...", "category": "preference|knowledge|context|behavior|goal|correction", "confidence": 0.0-1.0 }}
|
|
],
|
|
"factsToRemove": ["fact_id_1", "fact_id_2"]
|
|
}}
|
|
|
|
Important Rules:
|
|
- Only set shouldUpdate=true if there's meaningful new information
|
|
- Follow length guidelines: workContext/personalContext are concise (1-3 sentences), topOfMind and history sections are detailed (paragraphs)
|
|
- Include specific metrics, version numbers, and proper nouns in facts
|
|
- Only add facts that are clearly stated (0.9+) or strongly implied (0.7+)
|
|
- Use category "correction" for explicit agent mistakes or user corrections; assign confidence >= 0.95 when the correction is explicit
|
|
- Include "sourceError" only for explicit correction facts when the prior mistake or wrong approach is clearly stated; omit it otherwise
|
|
- Remove facts that are contradicted by new information
|
|
- When updating topOfMind, integrate new focus areas while removing completed/abandoned ones
|
|
Keep 3-5 concurrent focus themes that are still active and relevant
|
|
- For history sections, integrate new information chronologically into appropriate time period
|
|
- Preserve technical accuracy - keep exact names of technologies, companies, projects
|
|
- Focus on information useful for future interactions and personalization
|
|
- IMPORTANT: Do NOT record file upload events in memory. Uploaded files are
|
|
session-specific and ephemeral — they will not be accessible in future sessions.
|
|
Recording upload events causes confusion in subsequent conversations.
|
|
|
|
Return ONLY valid JSON, no explanation or markdown."""
|
|
|
|
|
|
# Prompt template for extracting facts from a single message
|
|
FACT_EXTRACTION_PROMPT = """Extract factual information about the user from this message.
|
|
|
|
Message:
|
|
{message}
|
|
|
|
Extract facts in this JSON format:
|
|
{{
|
|
"facts": [
|
|
{{ "content": "...", "category": "preference|knowledge|context|behavior|goal|correction", "confidence": 0.0-1.0 }}
|
|
]
|
|
}}
|
|
|
|
Categories:
|
|
- preference: User preferences (likes/dislikes, styles, tools)
|
|
- knowledge: User's expertise or knowledge areas
|
|
- context: Background context (location, job, projects)
|
|
- behavior: Behavioral patterns
|
|
- goal: User's goals or objectives
|
|
- correction: Explicit corrections or mistakes to avoid repeating
|
|
|
|
Rules:
|
|
- Only extract clear, specific facts
|
|
- Confidence should reflect certainty (explicit statement = 0.9+, implied = 0.6-0.8)
|
|
- Skip vague or temporary information
|
|
|
|
Return ONLY valid JSON."""
|
|
|
|
|
|
# Module-level tiktoken encoding cache. Populated lazily on first use;
|
|
# subsequent calls are a dict lookup (no network I/O). Pre-warming at
|
|
# startup via :func:`warm_tiktoken_cache` avoids blocking a request on the
|
|
# (potentially slow) first ``get_encoding`` call.
|
|
#
|
|
# A *failed* load is cached as a ``(None, monotonic_timestamp)`` tuple so that
|
|
# a network-restricted environment does not re-attempt the blocking BPE
|
|
# download on every subsequent call. After ``_TIKTOKEN_RETRY_COOLDOWN_S`` the
|
|
# failure is allowed to expire so a transient network outage can self-heal back
|
|
# to accurate tiktoken counting without a process restart. A load already in
|
|
# progress is cached as ``_TIKTOKEN_ENCODING_LOADING`` so concurrent callers
|
|
# fall back immediately instead of spawning more blocking
|
|
# ``tiktoken.get_encoding`` threads. Use the ``memory.token_counting: char``
|
|
# config to skip tiktoken entirely.
|
|
_TIKTOKEN_ENCODING_MISSING = object()
|
|
_TIKTOKEN_ENCODING_LOADING = object()
|
|
# Cooldown before a *failed* tiktoken load is re-attempted. This is an internal
|
|
# tuning constant rather than a user-facing config: it only affects how quickly
|
|
# the default ``tiktoken`` mode self-heals after a transient network outage.
|
|
# Deployments that want to avoid tiktoken's network dependency entirely should
|
|
# set ``memory.token_counting: char`` instead of tuning this value.
|
|
_TIKTOKEN_RETRY_COOLDOWN_S = 600.0
|
|
_tiktoken_encoding_cache: dict[str, Any] = {}
|
|
_tiktoken_encoding_cache_lock = threading.Lock()
|
|
|
|
|
|
def _get_tiktoken_encoding(encoding_name: str = "cl100k_base") -> tiktoken.Encoding | None:
|
|
"""Return a cached tiktoken encoding, or ``None`` on failure / unavailability.
|
|
|
|
On the very first call for a given *encoding_name*, tiktoken may need to
|
|
download the BPE data from ``openaipublic.blob.core.windows.net``. In
|
|
network-restricted environments (e.g. deployments behind the GFW) this
|
|
download can block for tens of minutes before the OS TCP timeout kicks in.
|
|
The caller must therefore be prepared for this to block and should run it
|
|
off the event loop (e.g. via ``asyncio.to_thread``).
|
|
|
|
A failed load is remembered (with a timestamp) so subsequent calls fall
|
|
back immediately to character-based estimation instead of re-triggering the
|
|
blocking download. The failure expires after ``_TIKTOKEN_RETRY_COOLDOWN_S``
|
|
so a transient outage can self-heal without a restart. A load already in
|
|
progress is also remembered so that a timed-out caller does not leave a
|
|
window where later requests start more blocking ``get_encoding`` calls.
|
|
"""
|
|
if not TIKTOKEN_AVAILABLE:
|
|
return None
|
|
|
|
with _tiktoken_encoding_cache_lock:
|
|
cached = _tiktoken_encoding_cache.get(encoding_name, _TIKTOKEN_ENCODING_MISSING)
|
|
if cached is _TIKTOKEN_ENCODING_LOADING:
|
|
return None
|
|
if isinstance(cached, tuple):
|
|
# Cached failure: (None, failed_at). Retry only after cooldown.
|
|
_, failed_at = cached
|
|
if time.monotonic() - failed_at < _TIKTOKEN_RETRY_COOLDOWN_S:
|
|
return None
|
|
cached = _TIKTOKEN_ENCODING_MISSING
|
|
if cached is not _TIKTOKEN_ENCODING_MISSING:
|
|
return cast("tiktoken.Encoding", cached)
|
|
_tiktoken_encoding_cache[encoding_name] = _TIKTOKEN_ENCODING_LOADING
|
|
|
|
try:
|
|
encoding = tiktoken.get_encoding(encoding_name)
|
|
except Exception:
|
|
logger.warning("Failed to load tiktoken encoding %r; falling back to char-based estimation", encoding_name, exc_info=True)
|
|
with _tiktoken_encoding_cache_lock:
|
|
_tiktoken_encoding_cache[encoding_name] = (None, time.monotonic())
|
|
return None
|
|
|
|
with _tiktoken_encoding_cache_lock:
|
|
_tiktoken_encoding_cache[encoding_name] = encoding
|
|
return encoding
|
|
|
|
|
|
def _char_based_token_estimate(text: str) -> int:
|
|
"""Network-free token estimate that accounts for CJK density.
|
|
|
|
The plain ``len(text) // 4`` heuristic is reasonable for English/code
|
|
(~4 chars per token) but significantly under-estimates token counts for
|
|
Chinese, Japanese, and Korean text, where the ratio is closer to 1.5-2
|
|
characters per token. Counting CJK characters separately (~2 chars per
|
|
token) avoids over-filling the injection budget for CJK-heavy memory
|
|
content.
|
|
"""
|
|
cjk = sum(
|
|
1
|
|
for ch in text
|
|
if "\u4e00" <= ch <= "\u9fff" # CJK Unified Ideographs
|
|
or "\u3040" <= ch <= "\u30ff" # Hiragana + Katakana
|
|
or "\uac00" <= ch <= "\ud7a3" # Hangul syllables
|
|
)
|
|
return (len(text) - cjk) // 4 + cjk // 2
|
|
|
|
|
|
def _count_tokens(text: str, encoding_name: str = "cl100k_base", *, use_tiktoken: bool = True) -> int:
|
|
"""Count tokens in text using tiktoken.
|
|
|
|
Args:
|
|
text: The text to count tokens for.
|
|
encoding_name: The encoding to use (default: cl100k_base for GPT-4/3.5).
|
|
use_tiktoken: When ``False``, skip tiktoken entirely and use the
|
|
network-free character-based estimate. This guarantees no BPE
|
|
download is attempted (see ``memory.token_counting`` config).
|
|
|
|
Returns:
|
|
The number of tokens in the text.
|
|
"""
|
|
if not use_tiktoken:
|
|
return _char_based_token_estimate(text)
|
|
|
|
encoding = _get_tiktoken_encoding(encoding_name)
|
|
if encoding is None:
|
|
# Fallback to CJK-aware character estimation if tiktoken is not
|
|
# available or the encoding failed to load.
|
|
return _char_based_token_estimate(text)
|
|
|
|
try:
|
|
return len(encoding.encode(text))
|
|
except Exception:
|
|
# Fallback to CJK-aware character estimation on error.
|
|
return _char_based_token_estimate(text)
|
|
|
|
|
|
def warm_tiktoken_cache() -> bool:
|
|
"""Pre-warm the tiktoken encoding cache.
|
|
|
|
Call at startup (off the event loop) so the first request never blocks
|
|
on the BPE download. Returns ``True`` if the encoding was loaded
|
|
successfully (or was already cached), ``False`` if tiktoken is
|
|
unavailable or the download failed.
|
|
"""
|
|
return _get_tiktoken_encoding("cl100k_base") is not None
|
|
|
|
|
|
def _coerce_confidence(value: Any, default: float = 0.0) -> float:
|
|
"""Coerce a confidence-like value to a bounded float in [0, 1].
|
|
|
|
Non-finite values (NaN, inf, -inf) are treated as invalid and fall back
|
|
to the default before clamping, preventing them from dominating ranking.
|
|
The ``default`` parameter is assumed to be a finite value.
|
|
"""
|
|
try:
|
|
confidence = float(value)
|
|
except (TypeError, ValueError):
|
|
return max(0.0, min(1.0, default))
|
|
if not math.isfinite(confidence):
|
|
return max(0.0, min(1.0, default))
|
|
return max(0.0, min(1.0, confidence))
|
|
|
|
|
|
def _format_fact_line(fact: dict[str, Any]) -> str | None:
|
|
"""Build a single formatted fact line, or return ``None`` for invalid facts.
|
|
|
|
Extracted as a shared helper so the guaranteed-injection and regular-injection
|
|
paths produce identical line formatting.
|
|
"""
|
|
content_value = fact.get("content")
|
|
if not isinstance(content_value, str):
|
|
return None
|
|
content = content_value.strip()
|
|
if not content:
|
|
return None
|
|
category = str(fact.get("category", "context")).strip() or "context"
|
|
confidence = _coerce_confidence(fact.get("confidence"), default=0.0)
|
|
source_error = fact.get("sourceError")
|
|
if category == "correction" and isinstance(source_error, str) and source_error.strip():
|
|
return f"- [{category} | {confidence:.2f}] {content} (avoid: {source_error.strip()})"
|
|
return f"- [{category} | {confidence:.2f}] {content}"
|
|
|
|
|
|
def _select_fact_lines(
|
|
ranked_facts: list[dict[str, Any]],
|
|
*,
|
|
token_budget: int,
|
|
use_tiktoken: bool,
|
|
) -> tuple[list[str], int]:
|
|
"""Greedily select formatted fact lines within a *line-only* token budget.
|
|
|
|
This function is intentionally **header-agnostic**: it counts only the
|
|
fact lines themselves (including ``\\n`` separators between lines). The
|
|
caller is responsible for reserving tokens for the ``"Facts:\\n"`` header
|
|
and any inter-section ``"\\n\\n"`` separator *before* calling this
|
|
function, and passing the remaining capacity as *token_budget*.
|
|
|
|
Stops at the first fact that would exceed the budget so the caller's
|
|
pre-sorted order (typically confidence-descending) is preserved strictly:
|
|
a shorter lower-ranked fact can never slip ahead of a skipped
|
|
higher-ranked one.
|
|
|
|
Args:
|
|
ranked_facts: Facts pre-sorted by the caller's preferred ranking.
|
|
token_budget: Maximum tokens available for fact lines only.
|
|
use_tiktoken: Whether to use tiktoken for counting.
|
|
|
|
Returns:
|
|
``(selected_lines, consumed_tokens)`` — *consumed_tokens* is the
|
|
exact token cost of the returned lines (including inter-line
|
|
``\\n`` separators, but *not* a leading header).
|
|
"""
|
|
lines: list[str] = []
|
|
consumed = 0
|
|
for fact in ranked_facts:
|
|
formatted = _format_fact_line(fact)
|
|
if formatted is None:
|
|
continue
|
|
line_text = ("\n" + formatted) if lines else formatted
|
|
line_tokens = _count_tokens(line_text, use_tiktoken=use_tiktoken)
|
|
if consumed + line_tokens > token_budget:
|
|
break
|
|
lines.append(formatted)
|
|
consumed += line_tokens
|
|
return lines, consumed
|
|
|
|
|
|
def _fallback_format_facts(
|
|
valid_facts: list[dict[str, Any]],
|
|
*,
|
|
preceding_section_cost: int,
|
|
max_tokens: int,
|
|
use_tiktoken: bool,
|
|
) -> tuple[str, list[str]] | tuple[None, None]:
|
|
"""Confidence-only ranking used when the primary path raises an exception.
|
|
|
|
Returns a tuple ``(section_text, fact_lines)`` where ``section_text`` is the
|
|
formatted ``"Facts:\\n..."`` section string (without any leading inter-section
|
|
separator — the caller owns that), and ``fact_lines`` are the individual lines
|
|
that make up the facts block. Both elements are ``None`` if no facts survive.
|
|
|
|
Returning the lines separately lets the caller track them for the
|
|
structure-aware safety truncation so fallback facts enjoy the same
|
|
protected-suffix treatment as facts emitted by the primary path.
|
|
|
|
*valid_facts* is the already-filtered fact list built by the primary path so
|
|
the fallback does not redo validation work. *preceding_section_cost* is the
|
|
tokens already consumed by user-context / history sections (used to derive
|
|
the remaining budget).
|
|
"""
|
|
ranked = sorted(valid_facts, key=lambda f: _coerce_confidence(f.get("confidence"), default=0.0), reverse=True)
|
|
|
|
header = "Facts:\n"
|
|
overhead = _count_tokens(header, use_tiktoken=use_tiktoken)
|
|
line_budget = max_tokens - preceding_section_cost - overhead
|
|
if line_budget <= 0:
|
|
return None, None
|
|
|
|
lines, _ = _select_fact_lines(ranked, token_budget=line_budget, use_tiktoken=use_tiktoken)
|
|
if not lines:
|
|
return None, None
|
|
return header + "\n".join(lines), lines
|
|
|
|
|
|
def format_memory_for_injection(
|
|
memory_data: dict[str, Any],
|
|
max_tokens: int = 2000,
|
|
*,
|
|
use_tiktoken: bool = True,
|
|
guaranteed_categories: list[str] | None = None,
|
|
guaranteed_token_budget: int = 500,
|
|
) -> str:
|
|
"""Format memory data for injection into system prompt.
|
|
|
|
Args:
|
|
memory_data: The memory data dictionary.
|
|
max_tokens: Maximum tokens to use (counted via tiktoken for accuracy).
|
|
use_tiktoken: When ``False``, all token counting uses the network-free
|
|
character-based estimate instead of tiktoken (see
|
|
``memory.token_counting`` config). Defaults to ``True``.
|
|
guaranteed_categories: Fact categories that must always be injected
|
|
regardless of the regular token budget. These facts draw from a
|
|
separate *guaranteed_token_budget*. When ``None`` or empty, all
|
|
facts compete for the same budget (original behaviour).
|
|
guaranteed_token_budget: Token ceiling for the guaranteed section.
|
|
In the common case the guaranteed lines *displace* regular lines
|
|
within *max_tokens* (the total output stays ≤ ``max_tokens``);
|
|
the budget becomes truly additive only when the guaranteed lines
|
|
alone would push the assembled output past *max_tokens*, at which
|
|
point the safety-truncation ceiling is raised to
|
|
``max_tokens + guaranteed_actual_usage`` to protect them.
|
|
Ignored when *guaranteed_categories* is ``None`` or empty.
|
|
|
|
Returns:
|
|
Formatted memory string for system prompt injection.
|
|
"""
|
|
if not memory_data:
|
|
return ""
|
|
|
|
# Reject a bare string explicitly: iterating a ``str`` yields single
|
|
# characters, which would silently produce a meaningless frozenset of
|
|
# letters and turn the guarantee off without any warning. Config-layer
|
|
# callers go through Pydantic (which enforces ``list[str]``), so this
|
|
# only guards the public helper surface.
|
|
if isinstance(guaranteed_categories, str):
|
|
raise TypeError("guaranteed_categories must be an iterable of strings, not a bare str")
|
|
effective_guaranteed: frozenset[str] = frozenset(c.strip() for c in guaranteed_categories if isinstance(c, str) and c.strip()) if guaranteed_categories else frozenset()
|
|
|
|
sections: list[str] = []
|
|
|
|
# Format user context
|
|
user_data = memory_data.get("user", {})
|
|
if user_data:
|
|
user_sections = []
|
|
|
|
work_ctx = user_data.get("workContext", {})
|
|
if work_ctx.get("summary"):
|
|
user_sections.append(f"Work: {work_ctx['summary']}")
|
|
|
|
personal_ctx = user_data.get("personalContext", {})
|
|
if personal_ctx.get("summary"):
|
|
user_sections.append(f"Personal: {personal_ctx['summary']}")
|
|
|
|
top_of_mind = user_data.get("topOfMind", {})
|
|
if top_of_mind.get("summary"):
|
|
user_sections.append(f"Current Focus: {top_of_mind['summary']}")
|
|
|
|
if user_sections:
|
|
sections.append("User Context:\n" + "\n".join(f"- {s}" for s in user_sections))
|
|
|
|
# Format history
|
|
history_data = memory_data.get("history", {})
|
|
if history_data:
|
|
history_sections = []
|
|
|
|
recent = history_data.get("recentMonths", {})
|
|
if recent.get("summary"):
|
|
history_sections.append(f"Recent: {recent['summary']}")
|
|
|
|
earlier = history_data.get("earlierContext", {})
|
|
if earlier.get("summary"):
|
|
history_sections.append(f"Earlier: {earlier['summary']}")
|
|
|
|
background = history_data.get("longTermBackground", {})
|
|
if background.get("summary"):
|
|
history_sections.append(f"Background: {background['summary']}")
|
|
|
|
if history_sections:
|
|
sections.append("History:\n" + "\n".join(f"- {s}" for s in history_sections))
|
|
|
|
# ── Facts ────────────────────────────────────────────────────────────────
|
|
#
|
|
# Design notes
|
|
# ~~~~~~~~~~~~
|
|
# • A single ``"Facts:\\n"`` header is emitted at most once.
|
|
# • Guaranteed-category facts are selected first from their own
|
|
# *guaranteed_token_budget* and placed at the front of the Facts block,
|
|
# so they cannot be evicted by regular facts. In the common case the
|
|
# total output still fits within *max_tokens* (guaranteed lines displace
|
|
# regular ones); the budget becomes truly additive only when the
|
|
# guaranteed lines alone push the output past *max_tokens*, in which
|
|
# case the safety-truncation ceiling is raised accordingly.
|
|
# • Regular facts draw from *max_tokens* only.
|
|
# • All token accounting (header, separators, lines) is performed here
|
|
# in the caller; the ``_select_fact_lines`` helper is header-agnostic.
|
|
# • When the primary path raises any exception, ``_fallback_format_facts``
|
|
# performs a single-pass confidence-only ranking.
|
|
facts_data = memory_data.get("facts", [])
|
|
guaranteed_line_tokens = 0 # used later for the effective truncation limit
|
|
if isinstance(facts_data, list) and facts_data:
|
|
# Token cost of sections built above (user context, history).
|
|
base_text = "\n\n".join(sections)
|
|
base_tokens = _count_tokens(base_text, use_tiktoken=use_tiktoken) if base_text else 0
|
|
|
|
# Pre-filter valid facts *before* entering the try so the except
|
|
# path can pass the same list straight into the fallback without
|
|
# redoing validation work on the hot prompt-injection path.
|
|
valid_facts = [f for f in facts_data if isinstance(f, dict) and isinstance(f.get("content"), str) and f.get("content", "").strip()]
|
|
|
|
# Initialise the facts-block markers *before* the try so the
|
|
# structure-aware truncation at the bottom of the function can
|
|
# reason about them regardless of whether the primary path or
|
|
# the except/fallback path produced the final Facts section.
|
|
facts_header = "Facts:\n"
|
|
all_fact_lines: list[str] = []
|
|
|
|
try:
|
|
# 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
|
|
# a guaranteed pool whose operator configured
|
|
# ``guaranteed_categories=["context"]``. Missing-category facts
|
|
# always fall through to the regular path.
|
|
def _confidence_key(fact: dict[str, Any]) -> float:
|
|
return _coerce_confidence(fact.get("confidence"), default=0.0)
|
|
|
|
if effective_guaranteed:
|
|
|
|
def _category_match(fact: dict[str, Any]) -> bool:
|
|
raw = fact.get("category")
|
|
if not isinstance(raw, str):
|
|
return False
|
|
cat = raw.strip()
|
|
return bool(cat) and cat in effective_guaranteed
|
|
|
|
guaranteed = sorted(
|
|
[f for f in valid_facts if _category_match(f)],
|
|
key=_confidence_key,
|
|
reverse=True,
|
|
)
|
|
regular = sorted(
|
|
[f for f in valid_facts if not _category_match(f)],
|
|
key=_confidence_key,
|
|
reverse=True,
|
|
)
|
|
else:
|
|
guaranteed = []
|
|
regular = sorted(valid_facts, key=_confidence_key, reverse=True)
|
|
|
|
# ── Phase 1: select guaranteed lines ──────────────────────────
|
|
header_cost = _count_tokens(facts_header, use_tiktoken=use_tiktoken)
|
|
|
|
guaranteed_lines: list[str] = []
|
|
if guaranteed:
|
|
guaranteed_line_budget = guaranteed_token_budget
|
|
guaranteed_lines, guaranteed_line_tokens = _select_fact_lines(
|
|
guaranteed,
|
|
token_budget=guaranteed_line_budget,
|
|
use_tiktoken=use_tiktoken,
|
|
)
|
|
|
|
# ── Phase 2: select regular lines ────────────────────────────
|
|
# Regular facts compete for *max_tokens* (the main budget).
|
|
# Subtract everything already accounted for:
|
|
# base sections + inter-section separator + header
|
|
# + guaranteed lines + the inter-group ``\n`` that joins the
|
|
# regular block to the guaranteed block (when both are present).
|
|
regular_lines: list[str] = []
|
|
if regular:
|
|
inter_group_newline_tokens = _count_tokens("\n", use_tiktoken=use_tiktoken) if guaranteed_lines else 0
|
|
used_before_regular = base_tokens + header_cost + guaranteed_line_tokens + inter_group_newline_tokens
|
|
regular_line_budget = max_tokens - used_before_regular
|
|
if regular_line_budget > 0:
|
|
regular_lines, _ = _select_fact_lines(
|
|
regular,
|
|
token_budget=regular_line_budget,
|
|
use_tiktoken=use_tiktoken,
|
|
)
|
|
|
|
# ── Emit a single "Facts:" section ───────────────────────────
|
|
# Leading inter-section separator is NOT embedded here; the
|
|
# final ``"\n\n".join(sections)`` is the single source of truth
|
|
# for section-to-section spacing, preventing the prior
|
|
# double-``\n\n`` bug.
|
|
all_fact_lines = guaranteed_lines + regular_lines
|
|
if all_fact_lines:
|
|
section_text = facts_header + "\n".join(all_fact_lines)
|
|
sections.append(section_text)
|
|
|
|
except Exception:
|
|
# ── Fallback: confidence-only ranking, single budget ─────────
|
|
# Any unexpected error in the partition / guaranteed path must
|
|
# not prevent memory injection entirely. Fall back to the
|
|
# original single-pass confidence ranking. Re-use the
|
|
# pre-filtered ``valid_facts`` so we don't redo validation work
|
|
# on the hot fallback path.
|
|
logger.warning(
|
|
"Memory injection: guaranteed-category path failed, falling back to confidence-only ranking",
|
|
exc_info=True,
|
|
)
|
|
fallback, fallback_lines = _fallback_format_facts(
|
|
valid_facts,
|
|
preceding_section_cost=base_tokens,
|
|
max_tokens=max_tokens,
|
|
use_tiktoken=use_tiktoken,
|
|
)
|
|
if fallback:
|
|
sections.append(fallback)
|
|
# Surface the fallback's lines to ``all_fact_lines`` so the
|
|
# structure-aware truncation below treats fallback facts as a
|
|
# protected suffix too. Without this, a large user-context
|
|
# prefix could silently clip fallback facts via the original
|
|
# prefix-cut.
|
|
all_fact_lines = fallback_lines
|
|
|
|
if not sections:
|
|
return ""
|
|
|
|
result = "\n\n".join(sections)
|
|
|
|
token_count = _count_tokens(result, use_tiktoken=use_tiktoken)
|
|
effective_limit = max_tokens + guaranteed_line_tokens
|
|
if token_count > effective_limit:
|
|
# Structure-aware truncation: the ``Facts:\n...`` block is treated as
|
|
# a *protected suffix* so guaranteed-category facts — the very facts
|
|
# this PR exists to preserve — can never be silently discarded by a
|
|
# prefix-cut on overflow. Only the preceding (user-context / history)
|
|
# sections are eligible for truncation; if they alone exceed the
|
|
# budget available after reserving the Facts block, they are clipped
|
|
# from the tail. When *guaranteed_line_tokens* is zero (no
|
|
# guaranteed categories configured or no facts survived), the
|
|
# equation collapses to the original prefix-truncation against
|
|
# ``max_tokens``, so backward compatibility is preserved.
|
|
facts_block = (facts_header + "\n".join(all_fact_lines)) if all_fact_lines else ""
|
|
facts_block_tokens = _count_tokens(facts_block, use_tiktoken=use_tiktoken)
|
|
separator_tokens = _count_tokens("\n\n", use_tiktoken=use_tiktoken)
|
|
budget_for_non_facts = max(
|
|
0,
|
|
effective_limit - facts_block_tokens - (separator_tokens if facts_block else 0),
|
|
)
|
|
|
|
# Build the preceding (non-facts) portion from *sections* excluding
|
|
# the trailing Facts block.
|
|
preceding_sections = sections[:-1] if all_fact_lines else sections
|
|
preceding = "\n\n".join(preceding_sections)
|
|
|
|
if preceding:
|
|
preceding_tokens = _count_tokens(preceding, use_tiktoken=use_tiktoken)
|
|
if preceding_tokens > budget_for_non_facts:
|
|
char_per_token = len(preceding) / max(preceding_tokens, 1)
|
|
target_chars = int(budget_for_non_facts * char_per_token * 0.95)
|
|
preceding = preceding[:target_chars].rstrip() + "\n..."
|
|
result = (preceding + "\n\n" + facts_block) if facts_block else preceding
|
|
else:
|
|
result = facts_block
|
|
|
|
return result
|
|
|
|
|
|
def format_conversation_for_update(messages: list[Any]) -> str:
|
|
"""Format conversation messages for memory update prompt.
|
|
|
|
Args:
|
|
messages: List of conversation messages.
|
|
|
|
Returns:
|
|
Formatted conversation string.
|
|
"""
|
|
lines = []
|
|
for msg in messages:
|
|
role = getattr(msg, "type", "unknown")
|
|
content = getattr(msg, "content", str(msg))
|
|
|
|
# Handle content that might be a list (multimodal)
|
|
if isinstance(content, list):
|
|
text_parts = []
|
|
for p in content:
|
|
if isinstance(p, str):
|
|
text_parts.append(p)
|
|
elif isinstance(p, dict):
|
|
text_val = p.get("text")
|
|
if isinstance(text_val, str):
|
|
text_parts.append(text_val)
|
|
content = " ".join(text_parts) if text_parts else str(content)
|
|
|
|
# Strip uploaded_files tags from human messages to avoid persisting
|
|
# ephemeral file path info into long-term memory. Skip the turn entirely
|
|
# when nothing remains after stripping (upload-only message).
|
|
if role == "human":
|
|
content = re.sub(r"<uploaded_files>[\s\S]*?</uploaded_files>\n*", "", str(content)).strip()
|
|
if not content:
|
|
continue
|
|
|
|
# Truncate very long messages
|
|
if len(str(content)) > 1000:
|
|
content = str(content)[:1000] + "..."
|
|
|
|
if role == "human":
|
|
lines.append(f"User: {content}")
|
|
elif role == "ai":
|
|
lines.append(f"Assistant: {content}")
|
|
|
|
return "\n\n".join(lines)
|