mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
fix(summarization): cover all selected history before compaction (#6918)
* fix(summarization): cover all selected history before compaction * fix(summarization): bound compaction work and recover oversized carries * docs: leave headroom in middleware guidance chain * fix(summarization): cool down repeated chunk-limit failures * test(runtime): complete callback cleanup graph fixture --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
b0a545ec90
commit
2464566942
11
README.md
11
README.md
@ -2051,6 +2051,17 @@ When your role lacks `runs:create`, the Web UI rejects a new task or `/goal <com
|
||||
|
||||
### Manual Context Compaction
|
||||
|
||||
Automatic and manual compaction process all selected history in chronological
|
||||
chunks before replacing it. `summarization.trim_tokens_to_summarize` bounds raw
|
||||
input per step; generated summaries use at most half to leave room for source.
|
||||
Oversized summaries retain marked head/tail text; raw history is never clipped.
|
||||
A pass is limited to 64 chunks and leaves history unchanged on failure or limit
|
||||
exhaustion. A limit breach starts a five-minute per-process cooldown for the
|
||||
same thread/history window and policy, including rebuilt middleware. A larger
|
||||
input budget or smaller source window permits an immediate retry. Model fallback
|
||||
and SDK retries can add calls. See the
|
||||
[summarization guide](backend/docs/summarization.md).
|
||||
|
||||
With `task_continuity.enabled`, `history_search` searches the current task's active
|
||||
and compacted history. Its optional `role` accepts `user`, `assistant`, or `tool`
|
||||
and filters before the eight-result limit; omitting it or passing `null` preserves
|
||||
|
||||
@ -14,6 +14,45 @@ The summarization feature uses LangChain's `SummarizationMiddleware` to monitor
|
||||
4. Maintains AI/Tool message pairs together for context continuity
|
||||
5. Stores the summary in `ThreadState.summary_text` and projects it ephemerally through durable context data
|
||||
|
||||
Selected history is processed from oldest to newest in bounded text chunks.
|
||||
Each call receives the preceding summary plus the next source fragment. Already
|
||||
generated summaries, including legacy stored summaries and final model output,
|
||||
use at most half the raw-input budget. Oversized summaries retain their head and
|
||||
tail with an explicit `[summary shortened]` marker; selected raw history is never
|
||||
shortened this way. Unused carry capacity is available to source text. A single oversized
|
||||
message is split into contiguous fragments, including its middle. Input is
|
||||
PII-redacted before splitting, then escaped before insertion into prompt blocks.
|
||||
No human-anchor filter or successful "too long" placeholder can discard selected
|
||||
source. Recent-message retention and AI/Tool partition boundaries stay unchanged.
|
||||
|
||||
Only after all chunks succeed are hooks, the compaction event, history capture
|
||||
and message replacement committed once. Any failed chunk leaves active history
|
||||
unchanged; manual `/compact` reports generation failure. A pass processes at most
|
||||
64 chunks, including cached chunks; hitting the limit also keeps original history.
|
||||
For a stable thread identity, a cap breach starts a five-minute cooldown keyed
|
||||
by user, thread, executor, opening source-message ID, prior summary and policy.
|
||||
Further automatic or manual attempts make no summary-model calls during that
|
||||
cooldown, even as history grows or middleware is rebuilt. A smaller source window,
|
||||
new summary/window or changed policy (including a larger input budget) allows an
|
||||
immediate retry. Warnings give the remaining delay and suggest increasing
|
||||
`trim_tokens_to_summarize` within model capacity or reducing selected history.
|
||||
The shared, locked cache keeps at most 128 hashed entries per process; eviction,
|
||||
process restart or another worker can permit an independent attempt. It is not a
|
||||
durable or distributed cost budget. Callers without a thread ID keep only the
|
||||
per-pass limit. Other generation failures retain their existing retry behavior.
|
||||
Configured/run-model fallback may make two candidate attempts per chunk, and SDK
|
||||
retries are separate. This bounds middleware steps, not a monetary bill or deadline.
|
||||
For legitimate histories beyond this limit, increase the existing per-call budget
|
||||
within the model's capacity or compact earlier. There is no new configuration.
|
||||
Impossibly small budgets that cannot fit a nonempty marked carry or source fail
|
||||
without eviction. `null` uses one unbounded input step and does not shorten carry.
|
||||
This guarantees new source input coverage, not lossless summaries: bounding
|
||||
already-compressed carry can omit earlier summarized facts.
|
||||
|
||||
Prefix selection expands a local window and then searches within it using the
|
||||
actual token counter, avoiding repeated processing of the entire remaining source.
|
||||
Warning diagnostics include numeric budget, carry and progress values, never text.
|
||||
|
||||
## Reuse and telemetry
|
||||
|
||||
The middleware keeps a bounded cache of exact rendered summary prompts and ordered
|
||||
@ -29,11 +68,11 @@ Each successful summary result emits a `middleware:summarize` journal event with
|
||||
`action="summary_result"`. Its `changes` include:
|
||||
|
||||
- `noop`: the result is byte-identical to the previous summary.
|
||||
- `llm_call_skipped`: this result came from the no-op cache.
|
||||
- `llm_call_skipped`: every chunk in this pass came from the no-op cache.
|
||||
- `summary_chars`: length of the new summary.
|
||||
- `previous_summary_chars`: length of the previous summary, or `null` when there was none.
|
||||
- `summarized_message_count`: messages folded into this summary.
|
||||
- `call_count`: completed summary results, including cache reuse and canned results.
|
||||
- `call_count`: completed compaction passes, including cache reuse.
|
||||
- `noop_count`: completed results identical to the previous summary.
|
||||
- `skip_count`: completed results reused from the no-op cache.
|
||||
|
||||
@ -43,8 +82,8 @@ Event counter snapshots include their own result and remain consistent during
|
||||
concurrent compactions. Counters are cumulative across runs sharing that middleware
|
||||
instance; concurrent journal events may arrive in a different order.
|
||||
|
||||
`call_count` is not an LLM request count: a compaction can try multiple candidates
|
||||
or produce a canned result without invoking a model. Use per-call LLM journal events
|
||||
`call_count` is not an LLM request count: a pass can process multiple chunks and
|
||||
try multiple candidates. Use per-call LLM journal events
|
||||
for model usage and cost, and `skip_count / call_count` for the cache reuse rate.
|
||||
|
||||
## Todo reminders
|
||||
@ -189,7 +228,7 @@ keep:
|
||||
#### `trim_tokens_to_summarize`
|
||||
- **Type**: Integer or null
|
||||
- **Default**: `4000`
|
||||
- **Description**: Token budget used to trim the raw input sections for the summarization call. Escaping, wrapper tags, and the summary prompt add overhead beyond this budget; it is not a hard limit on the final model request. When preserving the current user request leaves an assistant/tool-only summary window, trimming favors the most recent content in that window. If a mixed window still contains a human message but the human-anchored trim is empty, the existing final-message fallback is preserved. Set to `null` to skip trimming (not recommended for very long conversations).
|
||||
- **Description**: Per-step raw-text budget shared by generated carry (at most half) and the next chronological source chunk. Only generated or legacy summaries may be shortened, with a visible marker; raw source is covered without elision. A pass is limited to 64 chunks, retaining original history on failure or limit exhaustion. Escaping, wrapper tags, prompt overhead and SDK retries remain outside this budget. Set to `null` for one unbounded input step with no carry shortening (not recommended for very long conversations).
|
||||
|
||||
#### `summary_prompt`
|
||||
- **Type**: String or null
|
||||
|
||||
@ -9,12 +9,14 @@ Preserve the recognition sample and 5,000,000 UTF-8-byte guard. Count incrementa
|
||||
on parsing or field-limit errors, report an undetermined count rather than a
|
||||
physical-line or sample total.
|
||||
|
||||
Compaction keeps state `SystemMessage`s; transient instructions use request
|
||||
wrappers, and fully rescued partitions skip compaction. If latest-user rescue
|
||||
empties an AI/Tool-only window, use `_build_summary_input_text(strategy="last")`;
|
||||
mixed windows keep normal anchoring and final-message fallback. Budget raw text
|
||||
before escaping/wrapping; pass `trim_tokens_to_summarize=None` to avoid the
|
||||
LangChain default.
|
||||
Keep state `SystemMessage`s; inject transient instructions via request wrappers.
|
||||
Skip fully rescued partitions. Fold all raw source in order without human
|
||||
anchoring. PII-redact before splitting; budget before escaping. Generated carry:
|
||||
marked head/tail within half the raw budget; at most 64 chunks/pass. Commit
|
||||
hooks/history/removal only on success. Cap breaches use a bounded process-local
|
||||
5m cooldown across builds for the same user/thread/executor/window/policy;
|
||||
smaller source or changed policy retries immediately. See the linked contract.
|
||||
`trim_tokens_to_summarize=None` uses one unbounded step.
|
||||
|
||||
Todo compaction: [contract](../../../../../docs/summarization.md#todo-reminders).
|
||||
|
||||
@ -28,25 +30,23 @@ stay `in_progress`. Never infer status from reply text.
|
||||
|
||||
Assembly order: `tool_error_handling_middleware.py::_build_runtime_middlewares` (exposed as `build_lead_runtime_middlewares`), then `../lead_agent/agent.py::build_middlewares`. Optional entries require their config/runtime condition.
|
||||
|
||||
**Message provenance.** Stamp injected/rewritten messages with
|
||||
`provenance_kwargs()` from `deerflow_extension_api.provenance`: server-owned
|
||||
`deerflow_content_kind`, `deerflow_producer_kind`, and optional entity ID, even
|
||||
**Message provenance.** Injected/rewritten messages use
|
||||
`deerflow_extension_api.provenance.provenance_kwargs()`: server-owned
|
||||
`deerflow_content_kind`, `deerflow_producer_kind`, optional entity ID, even
|
||||
without observers. Producers: DynamicContext, DurableContext,
|
||||
SystemMessageCoalescing, ViewImage, SkillActivation. Summarization/Title use
|
||||
`SystemOperationKind.SUMMARIZATION`/`.TITLE`; summaries enter via DurableContext's
|
||||
`durable_context_data`, memory recall via DynamicContext's
|
||||
`dynamic_context_memory`. Memory only queues extraction.
|
||||
Rewrites that only modify the model-bound request view and never enter state do
|
||||
not require provenance stamps (precedents: InputSanitizationMiddleware,
|
||||
DanglingToolCallMiddleware, ModelContentCompatibilityMiddleware).
|
||||
`SystemOperationKind.SUMMARIZATION`/`.TITLE`; summaries enter through
|
||||
DurableContext's `durable_context_data`, recall through DynamicContext's
|
||||
`dynamic_context_memory`; Memory only queues extraction. Request-only rewrites
|
||||
need no stamps (InputSanitization, DanglingToolCall, ModelContentCompatibility).
|
||||
|
||||
**Self-description.** Configurable middleware exposes JSON-serialisable
|
||||
**Self-description.** Configurable middleware: JSON-safe
|
||||
`release_policy_parameters()` (`ReleasePolicyProvider`, duck typed). Update
|
||||
`collect_release_policies()` declarations with behavior; use `canonical_hash`
|
||||
for long text. Summarization declares enabled task-continuity retention or None;
|
||||
DurableContext declares normalized skills root, sorted read tools and continuity
|
||||
switch. History readers, including capture failures, validate persisted metadata
|
||||
so malformed values cannot abort compaction/model calls.
|
||||
`collect_release_policies()` with behavior; hash long text with `canonical_hash`.
|
||||
Summarization: enabled task-continuity retention or None. DurableContext:
|
||||
normalized skills root, sorted read tools, continuity switch. All history readers,
|
||||
including capture failures, validate persisted metadata; malformed values must
|
||||
not abort compaction/model calls.
|
||||
|
||||
**Removing tool calls.** Use `clone_ai_message_with_tool_calls`, not a bare
|
||||
`tool_calls` update: adapters resend stale `content` tool-call blocks, which
|
||||
@ -112,7 +112,7 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
17. **DynamicContextMiddleware** - Injects date and optional memory outside the static prompt. Opt-out removes its frozen server memory but retains date/user messages. Date follows server-local time unless `DEER_FLOW_DATE_TIMEZONE` names an IANA zone (invalid values fall back locally).
|
||||
18. **SkillActivationMiddleware** - Detects strict `/skill-name task` syntax on the latest real user message, resolves only enabled and runtime-allowed skills, injects the `SKILL.md` body as hidden current-turn context, and records a `middleware:skill_activation` audit event; a retried activation call replays the recorded reminder without re-auditing
|
||||
19. **SkillToolPolicyMiddleware** - Applies `allowed-tools` only after real activation; passive enabled skills and a custom agent's configured skill allowlist do not clamp the lead toolset. A run-scoped slash activation is authoritative and suppresses `skill_context` as a policy source, so reading another skill cannot widen the explicit skill's tools; without slash activation, skills captured after configured `read_file` loads retain the existing union semantics. The middleware filters model-visible schemas and blocks unauthorized execution, resolving canonical paths against the live enabled/agent-allowed registry on every model call, then stores a versioned, JSON-safe, middleware-token-bound decision signed by policy source plus active paths in run context for the resulting tool calls to reuse. The next model call always refreshes it, and malformed, foreign, stale, or unmatched decisions fall back to live resolution. `tool_search` and `describe_skill` remain framework-safe discovery tools under a restrictive policy; they may reveal or promote metadata, but a deferred business tool must still be declared by the active policy before its schema or execution can survive the policy middleware. The decision's owner token is authorization-sensitive, so its reserved context key is owned by `runtime.secret_context` and included in `REDACTED_CONTEXT_KEYS` for observable and persisted context copies. Registry load failures and a non-empty active set with no authorized skill fail closed to framework-safe tools; an individual stale path is skipped only when at least one valid active skill remains. This is best-effort behavioral scoping rather than a hard security boundary: alternate loads such as `bash cat` are not captured, and bounded autonomous `skill_context` can evict old entries. `task` is not framework-exempt, so a restricted skill cannot delegate around its policy. The middleware must remain immediately after `SkillActivationMiddleware` (which publishes the slash source through `runtime.secret_context`'s public path helpers authenticated by a required token shared only within the assembled middleware chain) and immediately before `DurableContextMiddleware`; assembly and compiled-graph tests pin ordering, token sharing, schema filtering, and execution blocking.
|
||||
20. **DurableContextMiddleware** - Before compaction, captures task dispatches/result summaries in `ThreadState.delegations` and loaded skill name/path/description references (never bodies) in `ThreadState.skill_context`; projects them into each model request. Static authority rules are injected as a `SystemMessage`; untrusted field values (`summary_text`, delegation results, skill descriptions) are injected separately as a hidden `HumanMessage` data block so compressed history, delegated work, and which skills are active stay visible without being stored as `messages` or promoted to system-role instructions. `build_subagent_runtime_middlewares` also attaches this middleware immediately before subagent summarization so a compacted `summary_text` is projected ahead of a preserved assistant/tool tail instead of leaving strict providers with an assistant-first request. See [tool artifacts](TOOL_ARTIFACTS.md) for projection, retry limits and delegation scope.
|
||||
20. **DurableContextMiddleware** - Before compaction, captures task dispatch/result summaries in `ThreadState.delegations` and loaded skill name/path/description (never bodies) in `ThreadState.skill_context`; projects them per request. Static authority uses `SystemMessage`; untrusted `summary_text`, delegation results and skill descriptions use a separate hidden `HumanMessage` data block, never persisted as messages or promoted to system instructions. Subagents attach it immediately before summarization: summary precedes preserved AI/tool history for strict providers. See [tool artifacts](TOOL_ARTIFACTS.md) for projection, retry limits and delegation scope.
|
||||
21. **SummarizationMiddleware** - *(optional)* Compacts near token limits. Fraction trigger/keep profiles belong to the active run model; an explicit summary model only generates. Memory flush/manual compaction follow runtime/checkpoint agent policy. Preserve the latest user request and DynamicContext reminders; a backward cutoff may retain old AI/tool turns or make first-turn compaction a no-op.
|
||||
22. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with `write_todos`. Skips reminders on `model_length_termination` so capped turns end cleanly. A raising model call requeues its reminder, uncounted, for the retry.
|
||||
23. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is read from terminal `ToolMessage.additional_kwargs` in the current run and merged back into the dispatching AIMessage by message position. The same state update marks the ToolMessage with `subagent_token_usage_attributed=true`, so checkpoint replay or middleware re-entry cannot add the cumulative snapshot twice; missing/malformed usage or a result with no matching dispatch remains unmarked and retryable.
|
||||
|
||||
@ -9,12 +9,13 @@ import threading
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, Protocol, override, runtime_checkable
|
||||
from time import monotonic
|
||||
from typing import Any, Protocol, override, runtime_checkable
|
||||
|
||||
from deerflow_extension_api import CompactionEvent, canonical_hash
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import SummarizationMiddleware
|
||||
from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, SystemMessage, get_buffer_string, trim_messages
|
||||
from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, SystemMessage, get_buffer_string
|
||||
from langgraph.config import get_config
|
||||
from langgraph.constants import TAG_NOSTREAM
|
||||
from langgraph.graph.message import REMOVE_ALL_MESSAGES
|
||||
@ -35,25 +36,21 @@ from deerflow.runtime.events.catalog import MIDDLEWARE_SUMMARIZE_TAG
|
||||
logger = logging.getLogger(__name__)
|
||||
_SUMMARY_TRIGGER_MESSAGE_NAME = "summary"
|
||||
_COMPACTION_TRANSFORM_KIND = "summarization"
|
||||
_COMPACTION_TRANSFORM_VERSION = "1"
|
||||
_COMPACTION_TRANSFORM_VERSION = "2"
|
||||
_UNSET = object()
|
||||
_NOOP_CACHE_MAX_ENTRIES = 128
|
||||
_MAX_SUMMARY_CHUNKS = 64
|
||||
_SUMMARY_CHUNK_RETRY_SECONDS = 300
|
||||
_SUMMARY_CHUNK_RETRY_MAX_ENTRIES = 128
|
||||
_SUMMARY_CHUNK_RETRY_CACHE: OrderedDict[str, tuple[float, int]] = OrderedDict()
|
||||
_SUMMARY_CHUNK_RETRY_LOCK = threading.Lock()
|
||||
_SHORTENED_SUMMARY_MARKER = "\n[summary shortened]\n"
|
||||
|
||||
|
||||
class _ReusedSummary(str):
|
||||
"""Marks a summary reused from an identical, previously no-op prompt (no LLM call)."""
|
||||
|
||||
|
||||
# Valid non-generated summaries for the empty / too-long-to-summarize edges; these
|
||||
# short-circuit model invocation (and must not be treated as generation failures).
|
||||
_CANNED_SUMMARIES = frozenset(
|
||||
{
|
||||
"No previous conversation history.",
|
||||
"Previous conversation was too long to summarize.",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class SummaryGenerationError(RuntimeError):
|
||||
"""Summary generation failed after exhausting the run-model fallback.
|
||||
|
||||
@ -201,6 +198,11 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
"trigger": plain_size(self.trigger),
|
||||
"keep": plain_size(self.keep),
|
||||
"trim_tokens_to_summarize": self.trim_tokens_to_summarize,
|
||||
"source_coverage": "chronological_chunks",
|
||||
"max_summary_chunks": _MAX_SUMMARY_CHUNKS,
|
||||
"chunk_limit_retry_seconds": _SUMMARY_CHUNK_RETRY_SECONDS,
|
||||
"chunk_limit_retry_max_entries": _SUMMARY_CHUNK_RETRY_MAX_ENTRIES,
|
||||
"summary_carry_budget_fraction": 0.5,
|
||||
"summary_prompt_hash": canonical_hash(self.summary_prompt),
|
||||
# Model objects are not JSON-serialisable. Keep generation and context
|
||||
# budget ownership distinct in the identity: changing either can alter
|
||||
@ -285,50 +287,84 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
async def _acreate_summary(self, messages_to_summarize: list[AnyMessage]) -> str | None:
|
||||
return await self._asummarize_with(messages_to_summarize)
|
||||
|
||||
def _prepare_summary_prompt(self, messages_to_summarize: list[AnyMessage], previous_summary: str | None) -> str | None:
|
||||
"""Return the formatted prompt, or a canned string for the empty/too-long edges.
|
||||
|
||||
A non-``None`` return that is not a real prompt (the two canned strings) is a
|
||||
valid summary and short-circuits generation; ``None`` means "build a prompt".
|
||||
"""
|
||||
if not messages_to_summarize:
|
||||
return "No previous conversation history."
|
||||
prompt = self._build_summary_prompt(messages_to_summarize, previous_summary=previous_summary)
|
||||
if prompt is None:
|
||||
return "Previous conversation was too long to summarize."
|
||||
return prompt
|
||||
|
||||
@staticmethod
|
||||
def _nonempty_summary(text: Any) -> str | None:
|
||||
"""Normalize a model response's text; a blank/whitespace-only body is a failure.
|
||||
|
||||
Committing ``""`` as a summary would fire the before_summarization hooks and
|
||||
remove all prior history for an empty replacement, so an empty body is treated
|
||||
as generation failure (try the fallback, or leave state unchanged) rather than
|
||||
a valid summary.
|
||||
"""
|
||||
"""A blank model response is a generation failure, never a replacement."""
|
||||
stripped = text.strip() if isinstance(text, str) else ""
|
||||
return stripped or None
|
||||
|
||||
def _summarize_with(self, messages_to_summarize: list[AnyMessage], previous_summary: str | None = None) -> str | None:
|
||||
"""Mirror the parent ``_create_summary`` but invoke the nostream-tagged model.
|
||||
def _summary_source_text(self, messages: list[AnyMessage]) -> str:
|
||||
# Redact the complete source before splitting: an identifier spanning a
|
||||
# chunk boundary must not bypass PII screening. Input remains plain text,
|
||||
# so the chat-message human-anchor filter is neither needed nor safe.
|
||||
return redact_text(get_buffer_string(messages), getattr(self._app_config, "pii_redaction", None))
|
||||
|
||||
We do not swap ``self.model`` at the instance level: the agent/middleware is
|
||||
cached and reused across concurrent runs, so a temporary swap would leak the
|
||||
``RunnableBinding`` to other coroutines during ``await`` and break parent logic
|
||||
that inspects the raw model (``profile`` / ``_get_ls_params``).
|
||||
def _summarize_with(self, messages_to_summarize: list[AnyMessage], previous_summary: str | None = None, *, runtime: Runtime | None = None) -> str | None:
|
||||
"""Fold every selected source fragment before committing any removal.
|
||||
|
||||
Generation uses the run's own model (``model_name: null``) or the explicitly
|
||||
configured summary model, falling back to the run model on failure so a broken
|
||||
summary provider cannot disable compaction while a working model is available.
|
||||
Chunk progress and summaries are local to this call; cached middleware
|
||||
instances can serve concurrent runs without sharing mutable cursors.
|
||||
The existing generation fallback and no-op cache apply to each prompt.
|
||||
"""
|
||||
prompt = self._prepare_summary_prompt(messages_to_summarize, previous_summary)
|
||||
if prompt is None or prompt in _CANNED_SUMMARIES:
|
||||
return prompt
|
||||
# Walk the ordered candidates; each attempt owns its full lifecycle (lazy
|
||||
# guarded construction -> invoke -> text extraction -> non-empty validation),
|
||||
# and any failure at any stage falls through to the next candidate. When all
|
||||
# candidates fail the caller leaves compaction state unchanged.
|
||||
if not messages_to_summarize:
|
||||
return "No previous conversation history."
|
||||
source = self._summary_source_text(messages_to_summarize)
|
||||
retry_key = self._summary_retry_key(runtime, messages_to_summarize, previous_summary)
|
||||
if self._summary_retry_blocked(retry_key, len(source)):
|
||||
return None
|
||||
summary = self._bound_summary_carry(previous_summary) if previous_summary else None
|
||||
if previous_summary and summary is None:
|
||||
return None
|
||||
cursor, chunks, all_reused = 0, 0, True
|
||||
while cursor < len(source):
|
||||
if self._summary_chunk_limit_reached(chunks, cursor, len(source), summary):
|
||||
self._remember_summary_chunk_breach(retry_key, len(source))
|
||||
return None
|
||||
step = self._next_summary_step(source, cursor, summary)
|
||||
if step is None:
|
||||
return None
|
||||
prompt, end = step
|
||||
summary = self._summarize_prompt(prompt, summary)
|
||||
if summary is None:
|
||||
return None
|
||||
all_reused = all_reused and isinstance(summary, _ReusedSummary)
|
||||
cursor = end
|
||||
chunks += 1
|
||||
if summary is None:
|
||||
return None
|
||||
return _ReusedSummary(summary) if all_reused else str(summary)
|
||||
|
||||
async def _asummarize_with(self, messages_to_summarize: list[AnyMessage], previous_summary: str | None = None, *, task_store=None, runtime: Runtime | None = None) -> str | None:
|
||||
"""Async fold; hooks, state replacement and success telemetry stay atomic."""
|
||||
if not messages_to_summarize:
|
||||
return "No previous conversation history."
|
||||
source = self._summary_source_text(messages_to_summarize)
|
||||
retry_key = self._summary_retry_key(runtime, messages_to_summarize, previous_summary)
|
||||
if self._summary_retry_blocked(retry_key, len(source)):
|
||||
return None
|
||||
summary = self._bound_summary_carry(previous_summary) if previous_summary else None
|
||||
if previous_summary and summary is None:
|
||||
return None
|
||||
cursor, chunks, all_reused = 0, 0, True
|
||||
while cursor < len(source):
|
||||
if self._summary_chunk_limit_reached(chunks, cursor, len(source), summary):
|
||||
self._remember_summary_chunk_breach(retry_key, len(source))
|
||||
return None
|
||||
step = self._next_summary_step(source, cursor, summary)
|
||||
if step is None:
|
||||
return None
|
||||
prompt, end = step
|
||||
summary = await self._asummarize_prompt(prompt, summary, task_store=task_store)
|
||||
if summary is None:
|
||||
return None
|
||||
all_reused = all_reused and isinstance(summary, _ReusedSummary)
|
||||
cursor = end
|
||||
chunks += 1
|
||||
if summary is None:
|
||||
return None
|
||||
return _ReusedSummary(summary) if all_reused else str(summary)
|
||||
|
||||
def _summarize_prompt(self, prompt: str, previous_summary: str | None) -> str | None:
|
||||
names = self._generation_candidate_names()
|
||||
noop_key = self._noop_cache_key(prompt, names)
|
||||
reused = self._reuse_noop_summary(noop_key, previous_summary)
|
||||
@ -337,23 +373,13 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
for index, name in enumerate(names):
|
||||
text = self._invoke_summary(self._model_for(name), prompt, last=index == len(names) - 1)
|
||||
if text is not None:
|
||||
# A fallback no-op must not suppress retrying a failed earlier candidate.
|
||||
# A fallback no-op must not suppress retrying a failed primary.
|
||||
if index == 0:
|
||||
self._remember_noop_summary(noop_key, text, previous_summary)
|
||||
return text
|
||||
return None
|
||||
|
||||
async def _asummarize_with(
|
||||
self,
|
||||
messages_to_summarize: list[AnyMessage],
|
||||
previous_summary: str | None = None,
|
||||
*,
|
||||
task_store=None,
|
||||
) -> str | None:
|
||||
"""Async counterpart of :meth:`_summarize_with` using the nostream model."""
|
||||
prompt = self._prepare_summary_prompt(messages_to_summarize, previous_summary)
|
||||
if prompt is None or prompt in _CANNED_SUMMARIES:
|
||||
return prompt
|
||||
async def _asummarize_prompt(self, prompt: str, previous_summary: str | None = None, *, task_store=None) -> str | None:
|
||||
names = self._generation_candidate_names()
|
||||
noop_key = self._noop_cache_key(prompt, names)
|
||||
reused = self._reuse_noop_summary(noop_key, previous_summary)
|
||||
@ -481,7 +507,10 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
summary = self._nonempty_summary(getattr(response, "text", None))
|
||||
if summary is None:
|
||||
self._log_summary_empty(last)
|
||||
return summary
|
||||
return None
|
||||
# Bound every generated carry, including the final output, before it
|
||||
# enters either the no-op cache or persisted summary_text.
|
||||
return self._bound_summary_carry(summary)
|
||||
|
||||
@staticmethod
|
||||
def _log_summary_error(last: bool) -> None:
|
||||
@ -506,136 +535,183 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
return messages
|
||||
return [*messages, self._summary_count_message(summary_text)]
|
||||
|
||||
@staticmethod
|
||||
def _bound_text(text: str, cap: int) -> str:
|
||||
if len(text) <= cap:
|
||||
return text
|
||||
if cap <= 0:
|
||||
return ""
|
||||
head = cap * 2 // 3
|
||||
omitted_marker = "\n...\n"
|
||||
if cap <= len(omitted_marker):
|
||||
return text[:cap]
|
||||
tail = max(0, cap - head - len(omitted_marker))
|
||||
if tail == 0:
|
||||
return text[:cap]
|
||||
return f"{text[:head]}{omitted_marker}{text[-tail:]}"
|
||||
|
||||
def _trim_summary_section_text(self, text: str, max_tokens: int, *, strategy: str) -> str:
|
||||
if not text.strip():
|
||||
return ""
|
||||
max_tokens = max(1, max_tokens)
|
||||
def _summary_text_tokens(self, text: str) -> int:
|
||||
if not text:
|
||||
return 0
|
||||
try:
|
||||
trimmed = trim_messages(
|
||||
[HumanMessage(content=text)],
|
||||
max_tokens=max_tokens,
|
||||
token_counter=self.token_counter,
|
||||
strategy=strategy,
|
||||
allow_partial=True,
|
||||
text_splitter=list,
|
||||
)
|
||||
if trimmed:
|
||||
content = trimmed[-1].content
|
||||
if isinstance(content, str) and content.strip():
|
||||
return content
|
||||
return max(0, int(self.token_counter([HumanMessage(content=text)])))
|
||||
except Exception:
|
||||
logger.debug("Failed to trim summary prompt section with token counter; falling back to deterministic text cap", exc_info=True)
|
||||
if strategy == "last":
|
||||
omitted_marker = "\n...\n"
|
||||
if len(text) > max_tokens and max_tokens > len(omitted_marker):
|
||||
return omitted_marker + text[-(max_tokens - len(omitted_marker)) :]
|
||||
return text[-max_tokens:]
|
||||
return self._bound_text(text, max_tokens)
|
||||
logger.debug("Failed to count summary input; using a character budget", exc_info=True)
|
||||
return len(text)
|
||||
|
||||
def _build_summary_input_text(self, formatted_messages: str, previous_summary: str | None = None, *, new_messages_strategy: Literal["first", "last"] = "first") -> str | None:
|
||||
"""Trim raw input sections before adding escaping and prompt overhead."""
|
||||
def _bound_summary_carry(self, summary: str) -> str | None:
|
||||
"""Reserve source capacity by bounding only already-generated summaries.
|
||||
|
||||
Never apply this to selected raw history. Redact before shortening so
|
||||
slicing an identifier cannot bypass PII handling. Short summaries keep
|
||||
their exact text; oversized ones retain both ends and a visible marker.
|
||||
"""
|
||||
text = redact_text(summary, getattr(self._app_config, "pii_redaction", None))
|
||||
if self.trim_tokens_to_summarize is None:
|
||||
trimmed_new_messages = formatted_messages
|
||||
trimmed_previous_summary = previous_summary.strip() if previous_summary else ""
|
||||
else:
|
||||
max_tokens = max(1, self.trim_tokens_to_summarize)
|
||||
if previous_summary:
|
||||
new_message_tokens = max(1, max_tokens // 2)
|
||||
previous_summary_tokens = max(1, max_tokens - new_message_tokens)
|
||||
trimmed_previous_summary = self._trim_summary_section_text(
|
||||
previous_summary.strip(),
|
||||
previous_summary_tokens,
|
||||
strategy="last",
|
||||
)
|
||||
trimmed_new_messages = self._trim_summary_section_text(
|
||||
formatted_messages,
|
||||
new_message_tokens,
|
||||
strategy=new_messages_strategy,
|
||||
)
|
||||
else:
|
||||
trimmed_previous_summary = ""
|
||||
trimmed_new_messages = self._trim_summary_section_text(
|
||||
formatted_messages,
|
||||
max_tokens,
|
||||
strategy=new_messages_strategy,
|
||||
)
|
||||
return text
|
||||
budget = max(1, self.trim_tokens_to_summarize)
|
||||
carry_budget = budget // 2
|
||||
tokens = self._summary_text_tokens(text)
|
||||
if tokens <= carry_budget:
|
||||
return text
|
||||
|
||||
# Escape < > & before embedding into the <existing_summary>/<new_messages>
|
||||
# blocks. new_messages is get_buffer_string over the raw state["messages"]
|
||||
# tail (InputSanitizationMiddleware only overrides the ModelRequest, never
|
||||
# state, so the summarizer sees genuine user text); existing_summary is the
|
||||
# prior turn's summary_text. An unescaped value like "</new_messages>..."
|
||||
# would close the block and forge an authority section for the extraction
|
||||
# LLM. Same block-breakout defense #4162 applied to the <conversation> block
|
||||
# and #4097 to the <memory> block. Escape after trimming so a trailing "..."
|
||||
# cannot split an entity; quote=False because content lands in element-text
|
||||
# position (never an attribute value).
|
||||
parts: list[str] = []
|
||||
if trimmed_previous_summary:
|
||||
parts.extend(
|
||||
[
|
||||
"<existing_summary>",
|
||||
html.escape(trimmed_previous_summary, quote=False),
|
||||
"</existing_summary>",
|
||||
"",
|
||||
]
|
||||
)
|
||||
if trimmed_new_messages:
|
||||
parts.extend(
|
||||
[
|
||||
"<new_messages>",
|
||||
html.escape(trimmed_new_messages, quote=False),
|
||||
"</new_messages>",
|
||||
]
|
||||
)
|
||||
if not parts:
|
||||
def shortened(keep: int) -> str:
|
||||
head, tail = (keep + 1) // 2, keep // 2
|
||||
return text[:head] + _SHORTENED_SUMMARY_MARKER + text[-tail:]
|
||||
|
||||
# The complete candidate (including its marker) is measured. An
|
||||
# impossibly small budget fails instead of producing an empty carry.
|
||||
low, high, bounded = 2, len(text) - 1, None
|
||||
while low <= high:
|
||||
middle = (low + high) // 2
|
||||
candidate = shortened(middle)
|
||||
if self._summary_text_tokens(candidate) <= carry_budget:
|
||||
bounded, low = candidate, middle + 1
|
||||
else:
|
||||
high = middle - 1
|
||||
if bounded is None:
|
||||
logger.warning("Summary carry cannot fit: carry_tokens=%d carry_budget=%d input_budget=%d", tokens, carry_budget, budget)
|
||||
else:
|
||||
logger.debug("Summary carry shortened: carry_tokens=%d bounded_tokens=%d carry_budget=%d input_budget=%d", tokens, self._summary_text_tokens(bounded), carry_budget, budget)
|
||||
return bounded
|
||||
|
||||
def _summary_retry_key(self, runtime: Runtime | None, messages: list[AnyMessage], previous_summary: str | None) -> str | None:
|
||||
"""Scope cap failures across builds without retaining history or identifiers.
|
||||
|
||||
The opening source ID survives appended turns, unlike a full-source hash.
|
||||
A new history window, prior summary or effective policy gets a fresh key.
|
||||
Callers without a stable thread identity retain the per-pass limit only.
|
||||
"""
|
||||
if runtime is None or not (thread_id := _resolve_thread_id(runtime)):
|
||||
return None
|
||||
return "\n".join(parts)
|
||||
from deerflow.runtime.user_context import resolve_runtime_user_id
|
||||
|
||||
context = runtime.context if isinstance(runtime.context, Mapping) else {}
|
||||
is_subagent = context.get("is_subagent") is True
|
||||
agent = context.get("agent_id") if is_subagent else (_resolve_agent_name(runtime) or "lead_agent")
|
||||
# Subagents sharing a parent thread still own separate executions.
|
||||
execution = context.get("sandbox_lease_owner_id") if is_subagent else None
|
||||
return canonical_hash([resolve_runtime_user_id(runtime), str(thread_id), agent, execution, messages[0].id, previous_summary, self.release_policy_parameters()])
|
||||
|
||||
def _summary_retry_blocked(self, key: str | None, source_chars: int) -> bool:
|
||||
if key is None:
|
||||
return False
|
||||
now = monotonic()
|
||||
with _SUMMARY_CHUNK_RETRY_LOCK:
|
||||
blocked = _SUMMARY_CHUNK_RETRY_CACHE.get(key)
|
||||
if blocked is None:
|
||||
return False
|
||||
deadline, failed_source_chars = blocked
|
||||
if now >= deadline or source_chars < failed_source_chars:
|
||||
del _SUMMARY_CHUNK_RETRY_CACHE[key]
|
||||
return False
|
||||
_SUMMARY_CHUNK_RETRY_CACHE.move_to_end(key)
|
||||
logger.warning(
|
||||
"Summary chunk-limit cooldown active; keeping history: retry_after_seconds=%.1f input_budget=%s. Increase summarization.trim_tokens_to_summarize within model capacity or reduce selected history.",
|
||||
deadline - now,
|
||||
self.trim_tokens_to_summarize,
|
||||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _remember_summary_chunk_breach(key: str | None, source_chars: int) -> None:
|
||||
if key is None:
|
||||
return
|
||||
with _SUMMARY_CHUNK_RETRY_LOCK:
|
||||
_SUMMARY_CHUNK_RETRY_CACHE[key] = (monotonic() + _SUMMARY_CHUNK_RETRY_SECONDS, source_chars)
|
||||
_SUMMARY_CHUNK_RETRY_CACHE.move_to_end(key)
|
||||
while len(_SUMMARY_CHUNK_RETRY_CACHE) > _SUMMARY_CHUNK_RETRY_MAX_ENTRIES:
|
||||
_SUMMARY_CHUNK_RETRY_CACHE.popitem(last=False)
|
||||
|
||||
def _summary_chunk_limit_reached(self, chunks: int, cursor: int, source_chars: int, summary: str | None) -> bool:
|
||||
if chunks < _MAX_SUMMARY_CHUNKS:
|
||||
return False
|
||||
logger.warning(
|
||||
"Summary chunk limit reached; keeping history: chunks=%d limit=%d processed_chars=%d source_chars=%d carry_tokens=%d input_budget=%s. "
|
||||
"Increase summarization.trim_tokens_to_summarize within model capacity or reduce selected history.",
|
||||
chunks,
|
||||
_MAX_SUMMARY_CHUNKS,
|
||||
cursor,
|
||||
source_chars,
|
||||
self._summary_text_tokens(summary or ""),
|
||||
self.trim_tokens_to_summarize,
|
||||
)
|
||||
return True
|
||||
|
||||
def _summary_prefix_length(self, source: str, cursor: int, available: int) -> int:
|
||||
"""Find a verified fitting prefix without repeatedly scanning the suffix.
|
||||
|
||||
Exponentially grow a local window, then bisect only its final interval.
|
||||
Character lengths guide probing, never stand in for the token counter.
|
||||
"""
|
||||
remaining = len(source) - cursor
|
||||
low, high = 0, min(remaining, max(1, available))
|
||||
while self._summary_text_tokens(source[cursor : cursor + high]) <= available:
|
||||
low = high
|
||||
if high == remaining:
|
||||
return high
|
||||
high = min(remaining, high * 2)
|
||||
high -= 1
|
||||
best = low
|
||||
low += 1
|
||||
while low <= high:
|
||||
middle = (low + high) // 2
|
||||
if self._summary_text_tokens(source[cursor : cursor + middle]) <= available:
|
||||
best, low = middle, middle + 1
|
||||
else:
|
||||
high = middle - 1
|
||||
return best
|
||||
|
||||
def _next_summary_step(self, source: str, cursor: int, previous_summary: str | None) -> tuple[str, int] | None:
|
||||
previous = redact_text(previous_summary, getattr(self._app_config, "pii_redaction", None)) if previous_summary else None
|
||||
end = len(source)
|
||||
if self.trim_tokens_to_summarize is not None:
|
||||
available = max(1, self.trim_tokens_to_summarize) - self._summary_text_tokens(previous or "")
|
||||
if available <= 0:
|
||||
logger.warning("Summary input cannot fit carry: carry_tokens=%d input_budget=%d", self._summary_text_tokens(previous or ""), self.trim_tokens_to_summarize)
|
||||
return None
|
||||
end = cursor + self._summary_prefix_length(source, cursor, available)
|
||||
if end < len(source):
|
||||
# Prefer complete serialized lines when possible. Oversized
|
||||
# individual messages still advance through contiguous slices.
|
||||
boundary = source.rfind("\n", cursor, end) + 1
|
||||
if boundary > cursor and self._summary_text_tokens(source[cursor:boundary]) <= available:
|
||||
end = boundary
|
||||
if end <= cursor:
|
||||
logger.warning("Summary input cannot fit source: source_budget=%d carry_tokens=%d input_budget=%d", available, self._summary_text_tokens(previous or ""), self.trim_tokens_to_summarize)
|
||||
return None
|
||||
body = self._build_summary_input_text(source[cursor:end], previous)
|
||||
if body is None:
|
||||
return None
|
||||
return self.summary_prompt.format(messages=body).rstrip(), end
|
||||
|
||||
def _build_summary_input_text(self, formatted_messages: str, previous_summary: str | None = None) -> str | None:
|
||||
"""Wrap complete raw sections; never trim material selected for removal.
|
||||
|
||||
The raw-text budget excludes escaping, wrapper tags and prompt overhead.
|
||||
A large carry or an impossibly small budget fails rather than replacing
|
||||
source history with a successful placeholder summary.
|
||||
"""
|
||||
if self.trim_tokens_to_summarize is not None and self._summary_text_tokens(formatted_messages) + self._summary_text_tokens(previous_summary or "") > max(1, self.trim_tokens_to_summarize):
|
||||
return None
|
||||
parts: list[str] = []
|
||||
if previous_summary:
|
||||
parts.extend(["<existing_summary>", html.escape(previous_summary, quote=False), "</existing_summary>", ""])
|
||||
if formatted_messages:
|
||||
parts.extend(["<new_messages>", html.escape(formatted_messages, quote=False), "</new_messages>"])
|
||||
return "\n".join(parts) if parts else None
|
||||
|
||||
def _build_summary_prompt(self, messages_to_summarize: list[AnyMessage], previous_summary: str | None = None) -> str | None:
|
||||
"""Build the summary prompt, returning ``None`` when trimming leaves nothing."""
|
||||
trimmed_messages = self._trim_messages_for_summary(messages_to_summarize)
|
||||
new_messages_strategy: Literal["first", "last"] = "first"
|
||||
if not trimmed_messages:
|
||||
if any(isinstance(message, HumanMessage) for message in messages_to_summarize):
|
||||
# The human anchor can fall outside the token-limited tail.
|
||||
# Preserve the existing final-message fallback for this case.
|
||||
trimmed_messages = messages_to_summarize[-1:]
|
||||
else:
|
||||
# Rescuing the current request can leave an AI/Tool-only window,
|
||||
# which the inherited human-anchored trimmer rejects even below
|
||||
# budget. Bound its raw text while favoring recent content.
|
||||
trimmed_messages = messages_to_summarize
|
||||
new_messages_strategy = "last"
|
||||
if not trimmed_messages:
|
||||
return None
|
||||
# Format messages to avoid token inflation from metadata when str() is called on
|
||||
# message objects.
|
||||
formatted_messages = get_buffer_string(trimmed_messages)
|
||||
formatted_messages = self._build_summary_input_text(formatted_messages, previous_summary=previous_summary, new_messages_strategy=new_messages_strategy)
|
||||
if not formatted_messages:
|
||||
return None
|
||||
# The summary model is invoked directly from before_model, outside
|
||||
# PiiRedactionMiddleware's wrap_model_call (#3190), so the compaction
|
||||
# input is redacted here; summaries then carry placeholders and the
|
||||
# summary_text DurableContextMiddleware reinjects stays clean.
|
||||
formatted_messages = redact_text(formatted_messages, getattr(self._app_config, "pii_redaction", None))
|
||||
return self.summary_prompt.format(messages=formatted_messages).rstrip()
|
||||
"""Build one complete prompt, or reject a window that needs chunking."""
|
||||
source = self._summary_source_text(messages_to_summarize)
|
||||
previous = redact_text(previous_summary, getattr(self._app_config, "pii_redaction", None)) if previous_summary else None
|
||||
body = self._build_summary_input_text(source, previous)
|
||||
return self.summary_prompt.format(messages=body).rstrip() if body else None
|
||||
|
||||
def before_model(self, state: AgentState, runtime: Runtime) -> dict | None:
|
||||
return self._maybe_summarize(state, runtime)
|
||||
@ -736,8 +812,9 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
def _record_summary_telemetry(self, runtime: Runtime, previous_summary: str | None, summary: str, summarized_count: int, *, llm_call_skipped: bool = False) -> None:
|
||||
"""Count successful summary results, unchanged results, and no-op cache reuse.
|
||||
|
||||
``call_count`` includes reuse and canned results, not individual LLM requests.
|
||||
Counters and event snapshots update atomically; telemetry failures stay soft.
|
||||
``call_count`` includes completed passes and reuse, not individual LLM requests.
|
||||
Counters and event snapshots update once per complete pass, regardless of
|
||||
its chunk count; telemetry failures stay soft.
|
||||
"""
|
||||
try:
|
||||
noop = previous_summary is not None and summary.encode("utf-8") == previous_summary.encode("utf-8")
|
||||
@ -794,7 +871,7 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
return None
|
||||
messages_to_summarize, preserved_messages, previous_summary, total_tokens = prepared
|
||||
source_content_hashes = self._freeze_compaction_sources(messages_to_summarize)
|
||||
summary = self._summarize_with(messages_to_summarize, previous_summary=previous_summary)
|
||||
summary = self._summarize_with(messages_to_summarize, previous_summary=previous_summary, runtime=runtime)
|
||||
if summary is None:
|
||||
if raise_on_failure:
|
||||
raise SummaryGenerationError("summary generation failed")
|
||||
@ -846,6 +923,7 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
||||
messages_to_summarize,
|
||||
previous_summary=previous_summary,
|
||||
task_store=task_store_from_runtime(runtime),
|
||||
runtime=runtime,
|
||||
)
|
||||
if summary is None:
|
||||
if raise_on_failure:
|
||||
|
||||
@ -85,7 +85,7 @@ class SummarizationConfig(BaseModel):
|
||||
trim_tokens_to_summarize: int | None = Field(
|
||||
default=4000,
|
||||
ge=1,
|
||||
description="Maximum tokens to keep when preparing messages for summarization. Pass null to skip trimming.",
|
||||
description="Raw-text budget per summary step. Generated carry uses at most half, reserving the rest for source; a pass has at most 64 chunks. Pass null for one unbounded step.",
|
||||
)
|
||||
|
||||
@field_validator("trim_tokens_to_summarize", mode="before")
|
||||
|
||||
@ -69,12 +69,12 @@ by lease/3. `tests/test_batch_lease_lifecycle.py` covers real SQLite interleavin
|
||||
threaded initialization, peer cancellation and renewal drain.
|
||||
|
||||
**Built-in Agents**: `general-purpose` (all tools except `task`) and `bash` (command specialist)
|
||||
**Registry and managed definitions**: Runtime resolution is built-in → `config.yaml custom_agents` → enabled administrator-managed definitions, followed by explicit `subagents.agents.<name>` overrides. Managed definitions are deployment-wide, persist through the same `agent_storage.backend` selection as Custom Agent definitions, and remain stored but are excluded from runtime when a built-in or later-added config definition owns the same name. The default Lead Agent sees the whole enabled catalog. A Custom Agent's `allowed_subagents` is snapshotted into run metadata (`None` = all, `[]` = hard deny, list = allowlist) and must filter both prompt discovery and `task` execution; never reload caller policy from mutable agent config inside the tool.
|
||||
**Benefit-based routing policy**: Direct execution is the default. Permit `task` when parallel latency, specialist skill, or context isolation clearly outweighs startup, duplicate discovery, synthesis, state conflicts, and side effects. Parallel scopes need independent outputs and disjoint mutable state; use the fewest useful agents and reassess each batch. A bounded sequential subagent is allowed for strong specialist/context benefit. With a per-response limit of 1, prompts omit parallel and multi-batch benefits and allow only material specialist/context benefit. Keep `lead_agent/prompt.py`, the `task` description, and built-in roles aligned; see `tests/test_subagent_routing_prompt.py`, `tests/test_subagent_prompt_security.py`, and `tests/test_lead_agent_prompt.py`.
|
||||
**User-scoped Skills**: Subagents resolve their configured skills through `get_or_new_user_skill_storage(user_id)` using the parent runtime identity, with `DEFAULT_USER_ID` only when no identity is available. This keeps custom-skill shadowing and visibility aligned with the lead agent instead of reading the global-only catalog.
|
||||
**Registry and managed definitions**: Resolve built-ins → `config.yaml custom_agents` → enabled admin-managed definitions → explicit `subagents.agents.<name>` overrides. Managed definitions are deployment-wide, use Custom Agents' `agent_storage.backend`, and remain stored but inactive on built-in/config name collisions. Default Lead Agent sees the full enabled catalog. Snapshot a Custom Agent's `allowed_subagents` into run metadata (`None` = all, `[]` = deny, list = allowlist); enforce it in prompt discovery and `task`, never reload mutable caller policy inside the tool.
|
||||
**Benefit-based routing policy**: Default to direct execution. Permit `task` only when parallel latency, specialist skill or context isolation clearly outweigh startup, duplicate discovery, synthesis, state conflicts and side effects. Parallel scopes need independent outputs and disjoint mutable state; use the fewest useful agents, reassessing each batch. Bounded sequential delegation requires strong specialist/context benefit. A per-response limit of 1 permits only material specialist/context benefit: omit parallel/multi-batch benefits. Align `lead_agent/prompt.py`, the `task` description and built-in roles. Tests: `test_subagent_routing_prompt.py`, `test_subagent_prompt_security.py`, `test_lead_agent_prompt.py`.
|
||||
**User-scoped Skills**: Use `get_or_new_user_skill_storage(user_id)` with parent runtime identity (`DEFAULT_USER_ID` only if absent), preserving lead-agent custom-skill shadowing/visibility; never use the global-only catalog.
|
||||
**Upload-state boundary**: Ordinary `task` delegation snapshots a valid parent `ThreadState.uploaded_files` list at dispatch, deep-copies it across the isolated-loop boundary, seeds it into the child's fresh state, and only then makes `list_uploaded_files` eligible for normal tool-policy filtering. An explicit empty list is valid and must be preserved because it means every upload in the thread is historical for this run. Missing or malformed state fails closed with the tool disabled. Durable `batch_task` execution intentionally keeps the tool disabled: delayed and recovered items have no valid parent-run upload boundary, and supporting that case requires a separate persisted-state contract.
|
||||
**Knowledge-scope boundary**: `task` and `batch_task` inherit the parent's canonical execution scope, never display data. Batch specs persist it across lease recovery; malformed scope fails. Model input cannot set or broaden it, and shared middleware/retrieval enforce `disabled`.
|
||||
**Date context (#4781)**: Every built-in subagent execution registers `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` hook adds a hidden framework-owned `SystemMessage` containing only `<current_date>` before the first model call; it does not read `AppConfig.memory`, call the memory manager, rewrite the task `HumanMessage`, or inherit the lead agent's frozen-conversation/midnight lifecycle. The coalescer merges that reminder with the subagent's static prompt so strict providers still receive exactly one leading `SystemMessage`. The lead-only `DynamicContextMiddleware` registration and its date, optional-memory, and midnight-update behavior remain unchanged.
|
||||
**Date context (#4781)**: Built-in executions register `SubagentDateContextMiddleware` immediately before `SystemMessageCoalescingMiddleware`. Its one-time `before_agent` injects only a hidden framework `SystemMessage` containing `<current_date>` before the first model call: no `AppConfig.memory`, memory-manager call, task `HumanMessage` rewrite or inherited lead frozen-conversation/midnight lifecycle. Coalescing with the static prompt yields one leading `SystemMessage` for strict providers. Lead-only `DynamicContextMiddleware` retains its date, optional-memory and midnight-update behavior.
|
||||
**Execution**: Ordinary and durable-batch native subagents share one persistent isolated event loop. Retain loop/thread ownership until the worker stops; bounded joins do not prove termination. Fence submissions during cleanup. After exit, the next submission reaps/closes the old loop under the shutdown lifecycle lock before replacement. Gateway/embedded startup installs one process-wide async FIFO admission controller (default 3 running, bounded queue). Direct `create_deerflow_agent` callers may pass a caller-owned `SubagentRuntime`; reuse it across graphs so bound `task`, optional batch tools/service, middleware limits and `SubagentExecutor` share one controller without global YAML. Start owned batch services before graph construction; stop them at application shutdown. Waiters hold no scheduler thread; cancellation/timeout release queue/slot ownership.
|
||||
**Shared sandbox execution lifecycle** (#5128): every admitted subagent run carries a stable task-derived `sandbox_lease_owner_id` and matching `sandbox_command_scope_id` in its runtime context. Sandbox middleware retains that execution against the lead thread's active provider client, so one child finishing cannot close the sandbox while siblings still run; the final holder performs any pending provider release. A rollback/fork-restored child reusing the parent's live client binds a non-releasing holder: it fences parent cleanup and owns its command scope without requesting a park itself; a parent's earlier park request waits for the child, while a missing inherited client falls through to a normal fresh acquire. On AIO, the command scope selects one explicit persistent shell session per subagent, allowing independent scopes to run concurrently while preserving in-order shell state within one child. Sync sandbox tool bodies offloaded with `asyncio.to_thread` are shielded and drained across repeated cancellation before the outer execution can clean its holder; a cancelled worker can therefore neither re-admit an already-released owner nor run after subagent terminalization. `SubagentExecutor` invokes active LangGraph stream cleanup before releasing its sandbox lease or notifying task stop. Slow cooperative cleanup keeps the result non-terminal and retains its sandbox lease and capacity slot; its warning is diagnostic, not a safe hard-timeout boundary. Middleware performs the normal release, and `SubagentExecutor` repeats it idempotently in `finally` so exceptions, cooperative cancellation, and timeout unwind paths cannot leak a lease or scoped session.
|
||||
**Concurrency and total delegation cap**: Ordinary `task` concurrency is resolved once as the minimum of the per-run request, the startup-frozen `subagent_runtime.max_running`, and the schema safety ceiling (1-64), then shared by the lead prompt and `SubagentLimitMiddleware`. Hot reloads must not make either layer advertise more capacity than the already-created process controller; a changed startup-only value takes effect only after restart. The same middleware separately enforces `subagents.max_total_per_run` (default 6, config schema 1-50, runtime override `max_total_subagents` clamped to the same range) against current-run entries in the durable delegation ledger, so a long lead-agent run cannot bypass concurrency limits by launching repeated legal-sized batches at each planning checkpoint, but historical delegations from previous runs in the same thread do not consume the new run's budget. Explicit `batch_task` work does not consume or relax that ordinary-run ledger: its persisted total/live/running limits live under `subagent_batches`. Gateway `run_agent()` and embedded `DeerFlowClient.stream()` both provide a per-invocation `run_id` in runtime context; `DeerFlowClient.stream()` also tags its input `HumanMessage` with that same id so durable-context capture can identify the current request boundary. Gateway resume paths may not append a new `HumanMessage`, so the worker also exposes the pre-run checkpoint's message ids in runtime context; durable-context capture uses that as the current-run boundary and never re-tags older task calls as the resumed run. When no delegation slots remain, task calls are stripped, provider raw tool-call metadata is synced, `finish_reason` is forced to `stop`, and a visible "subagent delegation limit" note is appended so the agent can synthesize already-collected results. Default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150`.
|
||||
|
||||
@ -10,6 +10,10 @@ ownership and capture cleanup. Retain the real XLSX conversion coverage.
|
||||
|
||||
Backend tests must preserve the runtime invariants they exercise without changing production execution topology.
|
||||
|
||||
Runtime-callback cleanup tests keep DEBUG capture enabled and implement graph
|
||||
checkpoint reads. Fixture-only exception tracebacks can retain the journal;
|
||||
preserve the strict weak-reference assertion rather than clearing captured logs.
|
||||
|
||||
Config-upgrade checkout fixtures clear inherited `DEER_FLOW_CONFIG_PATH`,
|
||||
`DEER_FLOW_PROJECT_ROOT` and `DEER_FLOW_ENV_FILE` before applying explicit test
|
||||
overrides. Exercise selector leaks through the real script with temporary
|
||||
|
||||
@ -377,12 +377,13 @@ async def test_async_summarization_observes_each_provider_attempt_and_live_store
|
||||
|
||||
middleware = DeerFlowSummarizationMiddleware.__new__(DeerFlowSummarizationMiddleware)
|
||||
middleware._extensions = extensions
|
||||
middleware._prepare_summary_prompt = lambda messages, previous_summary=None: "prompt"
|
||||
middleware._app_config = None
|
||||
middleware.trim_tokens_to_summarize = None
|
||||
middleware._generation_candidate_names = lambda: ["first", "second"]
|
||||
models = {"first": _Failing(), "second": _Working()}
|
||||
middleware._model_for = lambda name: models[name]
|
||||
|
||||
result = await middleware._asummarize_with(["message"], task_store=store)
|
||||
result = await middleware._asummarize_prompt("prompt", task_store=store)
|
||||
|
||||
assert result == "compact summary"
|
||||
assert [call[2].error is None for call in observer.calls] == [False, True]
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
"""Standalone package through the real extension isolation/graph/checkpoint path."""
|
||||
|
||||
import asyncio
|
||||
import html
|
||||
import json
|
||||
import logging
|
||||
import runpy
|
||||
@ -14,7 +15,7 @@ from deerflow_extension_api.auth import ExtensionPrincipal
|
||||
from deerflow_extension_api.plugins import ActionContext
|
||||
from langchain.agents import create_agent
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage, get_buffer_string
|
||||
from langchain_core.messages.utils import count_tokens_approximately
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
@ -134,7 +135,13 @@ def test_native_summary_still_handles_disabled_or_failed_pruning(jev, monkeypatc
|
||||
requests = transport(monkeypatch, error=httpx.ReadTimeout("private upstream details") if failure else None)
|
||||
graph, _, summary, _ = build_graph(jev, enabled=enabled)
|
||||
graph.invoke({"messages": history()}, {"configurable": {"thread_id": "fallback"}})
|
||||
assert len(summary._seen) == 1
|
||||
assert summary._seen
|
||||
source = ""
|
||||
for messages in summary._seen:
|
||||
prompt = "\n".join(str(message.content) for message in messages)
|
||||
source += html.unescape(prompt.split("<new_messages>\n", 1)[1].split("\n</new_messages>", 1)[0])
|
||||
assert source == get_buffer_string(history()[:3])
|
||||
assert graph.get_state({"configurable": {"thread_id": "fallback"}}).values["summary_text"] == "ORCHID"
|
||||
assert len(requests) == int(enabled)
|
||||
|
||||
|
||||
|
||||
@ -1449,9 +1449,10 @@ async def test_terminal_cycle_collection_is_coalesced_contextless_and_off_loop(m
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_agent_releases_terminal_runtime_callbacks():
|
||||
async def test_run_agent_releases_terminal_runtime_callbacks(caplog):
|
||||
from deerflow.runtime.journal import RunJournal
|
||||
|
||||
caplog.set_level(logging.DEBUG, logger="deerflow.runtime.runs.worker")
|
||||
run_manager = RunManager()
|
||||
record = await run_manager.create("thread-runtime-release")
|
||||
bridge = SimpleNamespace(
|
||||
@ -1469,6 +1470,11 @@ async def test_run_agent_releases_terminal_runtime_callbacks():
|
||||
captured["journal"] = next(callback for callback in callbacks if isinstance(callback, RunJournal))
|
||||
yield {"messages": []}
|
||||
|
||||
async def aget_state(self, _config):
|
||||
# Finalization reads checkpoint messages. A missing-method traceback
|
||||
# retained by DEBUG log capture would itself keep the journal alive.
|
||||
return SimpleNamespace(values={"messages": []})
|
||||
|
||||
await run_agent(
|
||||
bridge,
|
||||
run_manager,
|
||||
@ -1484,6 +1490,7 @@ async def test_run_agent_releases_terminal_runtime_callbacks():
|
||||
|
||||
stream_config = captured["config"]
|
||||
journal = captured["journal"]
|
||||
assert "Failed to inspect trailing human_input_request" not in caplog.text
|
||||
assert "__pregel_runtime" not in stream_config["configurable"]
|
||||
assert "__run_journal" not in stream_config["context"]
|
||||
assert journal not in (stream_config.get("callbacks") or [])
|
||||
|
||||
142
backend/tests/test_summarization_chunk_cooldown.py
Normal file
142
backend/tests/test_summarization_chunk_cooldown.py
Normal file
@ -0,0 +1,142 @@
|
||||
"""A breached compaction budget must not be re-spent on every turn."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from collections import OrderedDict
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from test_summarization_cost_recovery import FreeModel, _middleware, _state, _text
|
||||
|
||||
from deerflow.agents.middlewares import summarization_middleware as module
|
||||
from deerflow.agents.middlewares.summarization_middleware import SummaryGenerationError
|
||||
|
||||
|
||||
class ProgressModel(FreeModel):
|
||||
"""Produce a new carry each time so the no-op cache cannot hide spend."""
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
self.reply = f"summary-{len(self.prompts) + 1}"
|
||||
return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_cooldowns(monkeypatch):
|
||||
monkeypatch.setattr(module, "_SUMMARY_CHUNK_RETRY_CACHE", OrderedDict(), raising=False)
|
||||
|
||||
|
||||
async def compact(middleware, state, runtime, asynchronous, manual=False):
|
||||
method = middleware.acompact_state if asynchronous else middleware.compact_state
|
||||
if manual:
|
||||
with pytest.raises(SummaryGenerationError):
|
||||
if asynchronous:
|
||||
await method(state, runtime, force=True, raise_on_failure=True)
|
||||
else:
|
||||
method(state, runtime, force=True, raise_on_failure=True)
|
||||
return None
|
||||
return await middleware.abefore_model(state, runtime) if asynchronous else middleware.before_model(state, runtime)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("manual", [False, True], ids=["automatic", "manual"])
|
||||
@pytest.mark.parametrize("rebuild", [False, True], ids=["same-instance", "new-instance"])
|
||||
async def test_breach_cooldown_survives_new_turns_and_middleware_rebuilds(monkeypatch, caplog, asynchronous, manual, rebuild):
|
||||
runtime = SimpleNamespace(context={"thread_id": "over-limit", "user_id": "alice"})
|
||||
state, _ = _state(_text(128 * 80), "S")
|
||||
original = copy.deepcopy(state)
|
||||
model = ProgressModel(hard_limit=140)
|
||||
hook, observed = Mock(), Mock()
|
||||
monkeypatch.setattr(module, "notify_context_compacted", observed)
|
||||
middleware = _middleware(model, budget=128, before_summarization=[hook])
|
||||
assert await compact(middleware, state, runtime, asynchronous, manual) is None
|
||||
assert len(model.prompts) == 64 and model.guard_hits == 0
|
||||
assert state == original
|
||||
# A full-source hash would miss the next trigger as history grows.
|
||||
state["messages"].extend([HumanMessage(content="next request", id="next-user"), AIMessage(content="next answer", id="next-ai")])
|
||||
original = copy.deepcopy(state)
|
||||
if rebuild:
|
||||
middleware = _middleware(model, budget=128, before_summarization=[hook])
|
||||
assert await compact(middleware, state, runtime, asynchronous, manual) is None
|
||||
assert len(model.prompts) == 64, "the blocked retry must make no additional model calls"
|
||||
assert state == original
|
||||
assert middleware.summary_call_count == middleware.summary_noop_count == middleware.summary_skip_count == 0
|
||||
hook.assert_not_called()
|
||||
observed.assert_not_called()
|
||||
assert "cooldown" in caplog.text
|
||||
assert "trim_tokens_to_summarize" in caplog.text
|
||||
assert "evidence_" not in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_cooldown_expires_without_sleeping(monkeypatch, asynchronous):
|
||||
clock = [1000.0]
|
||||
monkeypatch.setattr(module, "monotonic", lambda: clock[0], raising=False)
|
||||
runtime = SimpleNamespace(context={"thread_id": "expiry", "user_id": "alice"})
|
||||
state, _ = _state(_text(128 * 80), "S")
|
||||
model = ProgressModel(hard_limit=140)
|
||||
middleware = _middleware(model, budget=128)
|
||||
assert await compact(middleware, state, runtime, asynchronous) is None
|
||||
assert len(model.prompts) == 64
|
||||
clock[0] += 299
|
||||
assert await compact(middleware, state, runtime, asynchronous) is None
|
||||
assert len(model.prompts) == 64
|
||||
clock[0] += 1
|
||||
assert await compact(middleware, state, runtime, asynchronous) is None
|
||||
assert len(model.prompts) == 128 and model.guard_hits == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("remedy", ["larger-budget", "smaller-source"])
|
||||
async def test_remediation_can_retry_immediately(asynchronous, remedy):
|
||||
runtime = SimpleNamespace(context={"thread_id": "remedy", "user_id": "alice"})
|
||||
state, _ = _state(_text(128 * 80), "S")
|
||||
model = ProgressModel(hard_limit=70)
|
||||
middleware = _middleware(model, budget=128)
|
||||
assert await compact(middleware, state, runtime, asynchronous) is None
|
||||
assert len(model.prompts) == 64
|
||||
if remedy == "larger-budget":
|
||||
middleware = _middleware(model, budget=20000)
|
||||
else:
|
||||
state["messages"][0] = state["messages"][0].model_copy(update={"content": "shortened input"})
|
||||
assert await compact(middleware, state, runtime, asynchronous) is not None
|
||||
assert 64 < len(model.prompts) <= 66 and model.guard_hits == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scope", ["user", "thread", "agent", "window", "subagent"])
|
||||
async def test_breach_does_not_suppress_another_execution_scope(scope):
|
||||
runtime = SimpleNamespace(context={"thread_id": "shared-id", "user_id": "alice", "agent_name": "researcher"})
|
||||
if scope == "subagent":
|
||||
runtime.context.update(is_subagent=True, agent_id="child", sandbox_lease_owner_id="execution-1")
|
||||
state, _ = _state(_text(128 * 80), "S")
|
||||
first = ProgressModel()
|
||||
assert await compact(_middleware(first, budget=128), state, runtime, True) is None
|
||||
assert len(first.prompts) == 64
|
||||
if scope == "subagent":
|
||||
runtime.context["sandbox_lease_owner_id"] = "execution-2"
|
||||
elif scope == "window":
|
||||
state["messages"][0] = state["messages"][0].model_copy(update={"id": "different-window"})
|
||||
else:
|
||||
key = {"user": "user_id", "thread": "thread_id", "agent": "agent_name"}[scope]
|
||||
runtime.context[key] = "another-value"
|
||||
other = ProgressModel()
|
||||
assert await compact(_middleware(other, budget=128), state, runtime, True) is None
|
||||
assert len(other.prompts) == 64 and other.guard_hits == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ordinary_generation_failure_does_not_start_chunk_cooldown():
|
||||
runtime = SimpleNamespace(context={"thread_id": "transient", "user_id": "alice"})
|
||||
state, _ = _state("short source", "S")
|
||||
model = FreeModel(reply="")
|
||||
middleware = _middleware(model)
|
||||
assert await compact(middleware, state, runtime, True) is None
|
||||
model.reply = "Recovered summary"
|
||||
assert await compact(middleware, state, runtime, True) is not None
|
||||
assert len(model.prompts) == 2
|
||||
359
backend/tests/test_summarization_cost_recovery.py
Normal file
359
backend/tests/test_summarization_cost_recovery.py
Normal file
@ -0,0 +1,359 @@
|
||||
"""Cost guards and recoverable carries must never truncate selected source."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import html
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, get_buffer_string
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from pydantic import Field
|
||||
|
||||
from deerflow.agents.middlewares.pii_redaction_middleware import redact_text
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummaryGenerationError
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.config.task_continuity_config import TaskContinuityConfig
|
||||
|
||||
|
||||
def _text(length):
|
||||
return " ".join(f"evidence_{index:05d} remains relevant" for index in range(length // 20 + 2))[:length]
|
||||
|
||||
|
||||
class FreeModel(BaseChatModel):
|
||||
reply: str = "S"
|
||||
hard_limit: int = 70
|
||||
prompts: list[str] = Field(default_factory=list)
|
||||
guard_hits: int = 0
|
||||
|
||||
@property
|
||||
def _llm_type(self):
|
||||
return "free-cost-recovery-regression"
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
if len(self.prompts) >= self.hard_limit:
|
||||
self.guard_hits += 1
|
||||
raise RuntimeError("free regression model call limit reached")
|
||||
self.prompts.append("\n".join(str(message.content) for message in messages))
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=self.reply))])
|
||||
|
||||
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
|
||||
class CountingChars:
|
||||
def __init__(self):
|
||||
self.input_characters = 0
|
||||
|
||||
def __call__(self, messages):
|
||||
length = sum(len(message.content) for message in messages)
|
||||
self.input_characters += length
|
||||
return length
|
||||
|
||||
|
||||
def _middleware(model, *, budget=512, default_counter=False, **kwargs):
|
||||
options = {
|
||||
"model": model,
|
||||
"trigger": ("messages", 4),
|
||||
"keep": ("messages", 2),
|
||||
"trim_tokens_to_summarize": budget,
|
||||
"summary_prompt": "{messages}",
|
||||
"anchor_model_name": "offline",
|
||||
"run_model_name": "offline",
|
||||
}
|
||||
if not default_counter:
|
||||
options["token_counter"] = CountingChars()
|
||||
options.update(kwargs)
|
||||
return DeerFlowSummarizationMiddleware(**options)
|
||||
|
||||
|
||||
def _state(text, previous=None):
|
||||
selected = [HumanMessage(content=text, id="old-user"), AIMessage(content="old answer", id="old-ai")]
|
||||
state = {
|
||||
"messages": [*selected, HumanMessage(content="current request", id="current-user"), AIMessage(content="current answer", id="current-ai")],
|
||||
"summary_text": previous,
|
||||
}
|
||||
return state, selected
|
||||
|
||||
|
||||
def _sections(prompt):
|
||||
return {name: html.unescape(body) for name, body in re.findall(r"<(existing_summary|new_messages)>\n(.*?)\n</\1>", prompt, re.DOTALL)}
|
||||
|
||||
|
||||
def _assert_full_source_and_budget(model, selected, counter, budget, pii=None):
|
||||
sections = [_sections(prompt) for prompt in model.prompts]
|
||||
source = "".join(part["new_messages"] for part in sections)
|
||||
assert source == redact_text(get_buffer_string(selected), pii)
|
||||
for part in sections:
|
||||
assert sum(counter([HumanMessage(content=text)]) for text in part.values()) <= budget
|
||||
|
||||
|
||||
def _assert_shortened(text, start, end):
|
||||
assert text.startswith(start) and text.endswith(end)
|
||||
assert re.search(r"shorten|truncat|omitt|缩短|省略|截断", text, re.IGNORECASE), "summary shortening needs an explicit marker"
|
||||
|
||||
|
||||
async def _compact(middleware, state, asynchronous):
|
||||
runtime = SimpleNamespace(context={})
|
||||
return await middleware.acompact_state(state, runtime, force=True) if asynchronous else middleware.compact_state(state, runtime, force=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_near_budget_carry_recovers_with_default_counter_and_complete_source(asynchronous):
|
||||
model = FreeModel(hard_limit=8)
|
||||
middleware = _middleware(model, budget=4000, default_counter=True)
|
||||
# Default approximate counter, not a custom tiny-budget token model.
|
||||
low, high, previous = 0, 32000, ""
|
||||
while low <= high:
|
||||
middle = (low + high) // 2
|
||||
candidate = _text(middle)
|
||||
if middleware.token_counter([HumanMessage(content=candidate)]) <= 3990:
|
||||
previous, low = candidate, middle + 1
|
||||
else:
|
||||
high = middle - 1
|
||||
assert middleware.token_counter([HumanMessage(content=previous)]) == 3990
|
||||
model.reply = previous
|
||||
state, selected = _state(_text(780), previous)
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None, "a verbose carry must not consume all useful source progress"
|
||||
assert model.guard_hits == 0 and len(model.prompts) <= 4
|
||||
assert middleware.token_counter([HumanMessage(content=result.summary_text)]) <= 2000
|
||||
_assert_full_source_and_budget(model, selected, middleware.token_counter, 4000)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("manual", [False, True], ids=["automatic", "manual"])
|
||||
async def test_64_chunk_policy_aborts_without_committing_any_partial_result(monkeypatch, asynchronous, manual):
|
||||
from deerflow.agents.middlewares import summarization_middleware as module
|
||||
from deerflow.agents.task_continuity import archive
|
||||
|
||||
budget = 128
|
||||
model = FreeModel(hard_limit=70)
|
||||
hook = Mock()
|
||||
observed = Mock()
|
||||
capture, acapture = Mock(), AsyncMock()
|
||||
monkeypatch.setattr(module, "notify_context_compacted", observed)
|
||||
monkeypatch.setattr(archive, "capture", capture)
|
||||
monkeypatch.setattr(archive, "acapture", acapture)
|
||||
middleware = _middleware(model, budget=budget, before_summarization=[hook], task_continuity_config=TaskContinuityConfig(enabled=True))
|
||||
state, selected = _state(_text(budget * 80), "S")
|
||||
state["task_history"] = {"already": "captured"}
|
||||
original = copy.deepcopy(state)
|
||||
assert len(get_buffer_string(selected)) > 64 * budget
|
||||
runtime = SimpleNamespace(context={})
|
||||
|
||||
if manual:
|
||||
with pytest.raises(SummaryGenerationError):
|
||||
if asynchronous:
|
||||
await middleware.acompact_state(state, runtime, force=True, raise_on_failure=True)
|
||||
else:
|
||||
middleware.compact_state(state, runtime, force=True, raise_on_failure=True)
|
||||
else:
|
||||
result = await middleware.abefore_model(state, runtime) if asynchronous else middleware.before_model(state, runtime)
|
||||
assert result is None
|
||||
|
||||
# The fake has its own independent 70-call stop to keep a broken product
|
||||
# guard free and bounded; hitting that fake stop is NOT passing this test.
|
||||
assert 0 < len(model.prompts) <= 64
|
||||
assert model.guard_hits == 0
|
||||
assert state == original
|
||||
assert middleware.summary_call_count == middleware.summary_noop_count == middleware.summary_skip_count == 0
|
||||
hook.assert_not_called()
|
||||
observed.assert_not_called()
|
||||
capture.assert_not_called()
|
||||
acapture.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_final_generated_summary_is_shortened_and_next_turn_can_compact(asynchronous):
|
||||
reply = "GENERATED_START " + _text(1400) + " GENERATED_END"
|
||||
model = FreeModel(reply=reply)
|
||||
middleware = _middleware(model)
|
||||
state, selected = _state("a short older request")
|
||||
|
||||
first = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert first is not None
|
||||
assert middleware.token_counter([HumanMessage(content=first.summary_text)]) <= 256
|
||||
_assert_shortened(first.summary_text, "GENERATED_START", "GENERATED_END")
|
||||
_assert_full_source_and_budget(model, selected, middleware.token_counter, 512)
|
||||
followup = {
|
||||
"messages": [*first.preserved_messages, HumanMessage(content="next request", id="next-user"), AIMessage(content="next answer", id="next-ai")],
|
||||
"summary_text": first.summary_text,
|
||||
}
|
||||
model.reply = "S"
|
||||
calls_before = len(model.prompts)
|
||||
second = await _compact(middleware, followup, asynchronous)
|
||||
assert second is not None and len(model.prompts) > calls_before
|
||||
assert middleware.summary_call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_oversized_legacy_summary_recovers_without_clipping_new_source(asynchronous):
|
||||
previous = "LEGACY_START " + _text(1500) + " LEGACY_END"
|
||||
model = FreeModel(reply="A compact carry")
|
||||
middleware = _middleware(model)
|
||||
state, selected = _state(_text(2400), previous)
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None and model.guard_hits == 0
|
||||
first_carry = _sections(model.prompts[0])["existing_summary"]
|
||||
assert middleware.token_counter([HumanMessage(content=first_carry)]) <= 256
|
||||
_assert_shortened(first_carry, "LEGACY_START", "LEGACY_END")
|
||||
_assert_full_source_and_budget(model, selected, middleware.token_counter, 512)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_shortening_preserves_unicode_ends_and_redacts_pii_before_slicing(asynchronous):
|
||||
pii = PiiRedactionConfig(enabled=True, token_secret="carry-shortening-test-secret-0123456789")
|
||||
previous = "LEGACY_START中🦉e\u0301 alice@example.com " + "中文<字段>& " * 180 + " bob@example.org LEGACY_END火🦊"
|
||||
reply = "GENERATED_START中🦉e\u0301 " + "摘要<内容>& " * 180 + " GENERATED_END火🦊"
|
||||
model = FreeModel(reply=reply)
|
||||
middleware = _middleware(model, app_config=SimpleNamespace(pii_redaction=pii))
|
||||
state, selected = _state("source contact source.person@example.net " + "完整事实🧭<项目>& " * 120, previous)
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None and model.guard_hits == 0
|
||||
result.summary_text.encode("utf-8")
|
||||
_assert_shortened(result.summary_text, "GENERATED_START中🦉e\u0301", "GENERATED_END火🦊")
|
||||
_assert_shortened(_sections(model.prompts[0])["existing_summary"], "LEGACY_START中🦉e\u0301", "LEGACY_END火🦊")
|
||||
assert all(email not in prompt for email in ("alice@example.com", "bob@example.org", "source.person@example.net") for prompt in model.prompts)
|
||||
_assert_full_source_and_budget(model, selected, middleware.token_counter, 512, pii)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unbounded_budget_keeps_full_carry_and_uses_one_step():
|
||||
previous = "LEGACY_START " + _text(1500) + " LEGACY_END"
|
||||
reply = "GENERATED_START " + _text(2000) + " GENERATED_END"
|
||||
model = FreeModel(reply=reply)
|
||||
middleware = _middleware(model, budget=None)
|
||||
state, selected = _state(_text(2400), previous)
|
||||
|
||||
result = await _compact(middleware, state, True)
|
||||
|
||||
assert result is not None and result.summary_text == reply
|
||||
assert len(model.prompts) == 1
|
||||
assert _sections(model.prompts[0])["existing_summary"] == previous
|
||||
assert _sections(model.prompts[0])["new_messages"] == get_buffer_string(selected)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_counter_input_work_grows_with_source_without_full_suffix_rescans():
|
||||
work = []
|
||||
for length in (2048, 8192):
|
||||
counter = CountingChars()
|
||||
model = FreeModel(reply="short carry")
|
||||
middleware = _middleware(model, token_counter=counter)
|
||||
state, selected = _state(_text(length), "S")
|
||||
result = await _compact(middleware, state, True)
|
||||
assert result is not None and model.guard_hits == 0
|
||||
input_characters = counter.input_characters
|
||||
_assert_full_source_and_budget(model, selected, CountingChars(), 512)
|
||||
work.append(input_characters)
|
||||
# Four times the source may have fixed overhead and chunk-boundary effects;
|
||||
# allow generous slack, while rejecting the old full-suffix quadratic term.
|
||||
assert work[1] <= work[0] * 5.5, f"counter input work grew superlinearly: {work}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_exactly_64_chunks_can_finish_with_an_extra_fallback_attempt(asynchronous):
|
||||
class FailFirstModel(FreeModel):
|
||||
failed: bool = False
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
if not self.failed:
|
||||
self.failed = True
|
||||
self.prompts.append("\n".join(str(message.content) for message in messages))
|
||||
raise RuntimeError("one primary attempt needs the configured fallback")
|
||||
return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
primary, fallback = FailFirstModel(), FreeModel()
|
||||
budget = 64
|
||||
middleware = _middleware(primary, budget=budget, anchor_model_name="primary", configured_model_name="primary", run_model_name="fallback", prebuilt_models={"fallback": fallback})
|
||||
# A full one-character carry leaves 63 source characters per bounded call.
|
||||
# A single serialized line avoids optional newline boundary adjustments.
|
||||
text = _text((budget - 1) * 64 - len("Human: "))
|
||||
selected = [HumanMessage(content=text, id="old-user")]
|
||||
state = {"messages": [*selected, HumanMessage(content="current", id="current-user"), AIMessage(content="recent", id="current-ai")], "summary_text": "S"}
|
||||
assert len(get_buffer_string(selected)) == (budget - 1) * 64
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None, "the 64th completed chunk is permitted; the guard applies before a 65th chunk"
|
||||
assert len(primary.prompts) == 64
|
||||
assert len(fallback.prompts) == 1
|
||||
assert primary.guard_hits == fallback.guard_hits == 0
|
||||
assert fallback.prompts[0] == primary.prompts[0]
|
||||
_assert_full_source_and_budget(primary, selected, middleware.token_counter, budget)
|
||||
assert middleware.summary_call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_clamped_noop_cache_skips_calls_without_claiming_original_summary_unchanged(asynchronous):
|
||||
previous = "LEGACY_START " + _text(1400) + " LEGACY_END"
|
||||
model = FreeModel(reply=previous)
|
||||
middleware = _middleware(model)
|
||||
state, _ = _state("selected source remains present until a successful commit", previous)
|
||||
original = copy.deepcopy(state)
|
||||
|
||||
first = await _compact(middleware, state, asynchronous)
|
||||
assert first is not None and first.summary_text != previous
|
||||
assert len(model.prompts) == 1
|
||||
assert middleware.summary_noop_count == middleware.summary_skip_count == 0
|
||||
|
||||
second = await _compact(middleware, state, asynchronous)
|
||||
assert second is not None and second.summary_text == first.summary_text
|
||||
assert len(model.prompts) == 1
|
||||
assert middleware.summary_skip_count == 1
|
||||
assert middleware.summary_noop_count == 0, "cache reuse did not preserve the original oversized summary byte-for-byte"
|
||||
assert state == original
|
||||
|
||||
changed = copy.deepcopy(state)
|
||||
changed["messages"][0] = changed["messages"][0].model_copy(update={"content": "a newly selected source"})
|
||||
third = await _compact(middleware, changed, asynchronous)
|
||||
assert third is not None
|
||||
assert len(model.prompts) == 2
|
||||
assert _sections(model.prompts[-1])["new_messages"] == get_buffer_string(changed["messages"][:2])
|
||||
assert middleware.summary_call_count == 3
|
||||
assert middleware.summary_skip_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_weighted_unicode_counter_can_exceed_character_count_without_losing_source(asynchronous):
|
||||
def weighted_tokens(messages):
|
||||
return sum(1 if character.isascii() else 5 for message in messages for character in message.content)
|
||||
|
||||
# Unique numeric labels ensure every source fragment has a distinct prompt,
|
||||
# so no-op cache reuse cannot hide repeated source from the model recorder.
|
||||
text = " ".join(f"{index:03d}中🦉e\u0301火<{index:03d}>&" for index in range(40))
|
||||
state, selected = _state(text, "S")
|
||||
serialized = get_buffer_string(selected)
|
||||
assert weighted_tokens([HumanMessage(content=serialized)]) > len(serialized)
|
||||
assert weighted_tokens([HumanMessage(content="中🦉e\u0301火")]) > len("中🦉e\u0301火")
|
||||
model = FreeModel()
|
||||
budget = 256
|
||||
middleware = _middleware(model, budget=budget, token_counter=weighted_tokens)
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None and model.guard_hits == 0
|
||||
assert len(model.prompts) > 1
|
||||
assert len(set(model.prompts)) == len(model.prompts)
|
||||
_assert_full_source_and_budget(model, selected, weighted_tokens, budget)
|
||||
assert weighted_tokens([HumanMessage(content=result.summary_text)]) <= budget // 2
|
||||
230
backend/tests/test_summarization_rolling_contract.py
Normal file
230
backend/tests/test_summarization_rolling_contract.py
Normal file
@ -0,0 +1,230 @@
|
||||
"""Rolling compaction preserves candidate, cache, observation and PII contracts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import html
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from deerflow_extension_api import SystemOperationKind
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, get_buffer_string
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from pydantic import Field
|
||||
|
||||
from deerflow.agents.middlewares.pii_redaction_middleware import redact_text
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
|
||||
from deerflow.config.pii_redaction_config import PiiRedactionConfig
|
||||
from deerflow.extensions import ExtensionRegistry
|
||||
|
||||
|
||||
class RecordingModel(BaseChatModel):
|
||||
answers: list[Any] = Field(default_factory=list)
|
||||
prompts: list[str] = Field(default_factory=list)
|
||||
default_answer: str = "S"
|
||||
|
||||
@property
|
||||
def _llm_type(self) -> str:
|
||||
return "rolling-contract-offline"
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
self.prompts.append("\n".join(str(message.content) for message in messages))
|
||||
answer = self.answers.pop(0) if self.answers else self.default_answer
|
||||
if isinstance(answer, Exception):
|
||||
raise answer
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=answer))])
|
||||
|
||||
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
|
||||
def _char_count(messages):
|
||||
return sum(len(message.content) for message in messages)
|
||||
|
||||
|
||||
def _sections(prompt):
|
||||
return {name: html.unescape(body) for name, body in re.findall(r"<(existing_summary|new_messages)>\n(.*?)\n</\1>", prompt, re.DOTALL)}
|
||||
|
||||
|
||||
def _source(prompts):
|
||||
return "".join(_sections(prompt)["new_messages"] for prompt in prompts)
|
||||
|
||||
|
||||
def _state(text=None):
|
||||
if text is None:
|
||||
text = " ".join(f"fact_{index:03d}: distinct evidence" for index in range(24))
|
||||
selected = [HumanMessage(content=text, id="old-user"), AIMessage(content="old answer", id="old-ai")]
|
||||
return {
|
||||
"messages": [*selected, HumanMessage(content="current request", id="current-user"), AIMessage(content="current answer", id="current-ai")],
|
||||
"summary_text": "S",
|
||||
}, selected
|
||||
|
||||
|
||||
def _middleware(model, *, default_counter=False, **kwargs):
|
||||
options = {
|
||||
"trigger": ("messages", 4),
|
||||
"keep": ("messages", 2),
|
||||
"trim_tokens_to_summarize": 96,
|
||||
"summary_prompt": "{messages}",
|
||||
"anchor_model_name": "primary",
|
||||
"run_model_name": "primary",
|
||||
}
|
||||
if not default_counter:
|
||||
options["token_counter"] = _char_count
|
||||
options.update(kwargs)
|
||||
return DeerFlowSummarizationMiddleware(model=model, **options)
|
||||
|
||||
|
||||
async def _compact(middleware, state, asynchronous, runtime=None):
|
||||
runtime = runtime or SimpleNamespace(context={})
|
||||
if asynchronous:
|
||||
return await middleware.acompact_state(state, runtime, force=True)
|
||||
return middleware.compact_state(state, runtime, force=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_middle_chunk_fallback_becomes_the_next_full_carry(asynchronous):
|
||||
primary = RecordingModel(answers=["PRIMARY_ONE", RuntimeError("primary unavailable")], default_answer="PRIMARY_LATER")
|
||||
fallback = RecordingModel(default_answer="FALLBACK_TWO")
|
||||
middleware = _middleware(primary, configured_model_name="primary", run_model_name="fallback", prebuilt_models={"fallback": fallback})
|
||||
state, selected = _state()
|
||||
original = copy.deepcopy(state)
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None
|
||||
assert len(primary.prompts) > 2
|
||||
assert len(fallback.prompts) == 1
|
||||
assert fallback.prompts[0] == primary.prompts[1]
|
||||
assert _sections(primary.prompts[1])["existing_summary"] == "PRIMARY_ONE"
|
||||
assert _sections(primary.prompts[2])["existing_summary"] == "FALLBACK_TWO"
|
||||
assert _source(primary.prompts) == get_buffer_string(selected)
|
||||
assert middleware.summary_call_count == 1
|
||||
assert middleware.summary_skip_count == 0
|
||||
assert state == original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_only_a_fully_cached_rolling_pass_counts_as_skipped(asynchronous):
|
||||
model = RecordingModel()
|
||||
middleware = _middleware(model)
|
||||
state, _ = _state()
|
||||
|
||||
first = await _compact(middleware, state, asynchronous)
|
||||
first_calls = len(model.prompts)
|
||||
assert first is not None and first_calls > 2
|
||||
second = await _compact(middleware, state, asynchronous)
|
||||
assert second is not None
|
||||
assert len(model.prompts) == first_calls
|
||||
assert middleware.summary_skip_count == 1
|
||||
|
||||
changed = copy.deepcopy(state)
|
||||
changed["messages"][0] = changed["messages"][0].model_copy(update={"content": state["messages"][0].content.replace("fact_012", "FACT_012")})
|
||||
third = await _compact(middleware, changed, asynchronous)
|
||||
assert third is not None
|
||||
fresh_calls = len(model.prompts) - first_calls
|
||||
assert 0 < fresh_calls < first_calls, "cached prefix and suffix may not hide a model call for the changed middle"
|
||||
assert middleware.summary_call_count == 3
|
||||
assert middleware.summary_noop_count == 3
|
||||
assert middleware.summary_skip_count == 1, "a partly reused pass still made an LLM call"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_observer_counts_actual_attempts_and_telemetry_counts_whole_passes():
|
||||
observations = []
|
||||
|
||||
class Observer:
|
||||
async def on_system_model_call(self, app_store, task_store, kind, request, result):
|
||||
observations.append((kind, request, result))
|
||||
|
||||
registry = ExtensionRegistry()
|
||||
with registry.attributed_to("rolling-contract"):
|
||||
registry.system_model_observer(Observer())
|
||||
primary = RecordingModel(answers=["S", RuntimeError("middle primary failure")])
|
||||
fallback = RecordingModel()
|
||||
middleware = _middleware(primary, configured_model_name="primary", run_model_name="fallback", prebuilt_models={"fallback": fallback}, extensions=registry.build())
|
||||
state, _ = _state()
|
||||
journal = SimpleNamespace(record_middleware=Mock())
|
||||
runtime = SimpleNamespace(context={"__run_journal": journal})
|
||||
|
||||
first = await _compact(middleware, state, True, runtime)
|
||||
|
||||
assert first is not None
|
||||
attempts = len(primary.prompts) + len(fallback.prompts)
|
||||
assert attempts > 3 and len(observations) == attempts
|
||||
assert all(kind == SystemOperationKind.SUMMARIZATION for kind, _, _ in observations)
|
||||
assert sum(result.error is not None for _, _, result in observations) == 1
|
||||
assert [request.model_name for _, request, result in observations if result.error is not None] == ["primary"]
|
||||
assert [request.model_name for _, request, _ in observations].count("fallback") == 1
|
||||
assert all(len(request.messages) == 1 for _, request, _ in observations)
|
||||
assert journal.record_middleware.call_count == 1
|
||||
assert journal.record_middleware.call_args.kwargs["changes"]["call_count"] == 1
|
||||
assert journal.record_middleware.call_args.kwargs["changes"]["summarized_message_count"] == 2
|
||||
|
||||
# The failed primary's fallback no-op was not cached. Prefix and suffix
|
||||
# cache hits must still allow that primary candidate to recover next time.
|
||||
second = await _compact(middleware, state, True, runtime)
|
||||
assert second is not None
|
||||
assert len(observations) == attempts + 1
|
||||
assert middleware.summary_skip_count == 0
|
||||
third = await _compact(middleware, state, True, runtime)
|
||||
assert third is not None
|
||||
assert len(observations) == attempts + 1
|
||||
assert middleware.summary_call_count == journal.record_middleware.call_count == 3
|
||||
assert middleware.summary_skip_count == 1
|
||||
assert journal.record_middleware.call_args.kwargs["changes"]["llm_call_skipped"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_pii_spanning_a_source_chunk_boundary_is_redacted_before_splitting(asynchronous):
|
||||
email = "boundary.user@example.com"
|
||||
state, selected = _state("x" * 60 + " " + email + " trailing evidence " * 10)
|
||||
source = get_buffer_string(selected)
|
||||
control_model = RecordingModel()
|
||||
control = _middleware(control_model, trim_tokens_to_summarize=80)
|
||||
assert await _compact(control, state, asynchronous) is not None
|
||||
first_edge = len(_sections(control_model.prompts[0])["new_messages"])
|
||||
assert source.index(email) < first_edge < source.index(email) + len(email), "the raw identifier must straddle a real model-input boundary"
|
||||
|
||||
config = PiiRedactionConfig(enabled=True, token_secret="rolling-contract-test-secret-0123456789")
|
||||
model = RecordingModel()
|
||||
middleware = _middleware(model, trim_tokens_to_summarize=80, app_config=SimpleNamespace(pii_redaction=config))
|
||||
|
||||
assert await _compact(middleware, state, asynchronous) is not None
|
||||
|
||||
captured = _source(model.prompts)
|
||||
assert captured == redact_text(source, config)
|
||||
assert email not in captured
|
||||
assert "boundary.user" not in captured and "example.com" not in captured
|
||||
assert all(email not in prompt for prompt in model.prompts)
|
||||
replacement = redact_text(email, config)
|
||||
first_redacted_edge = len(_sections(model.prompts[0])["new_messages"])
|
||||
assert captured.index(replacement) < first_redacted_edge < captured.index(replacement) + len(replacement)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
async def test_default_approximate_counter_keeps_unicode_source_and_each_raw_budget(asynchronous):
|
||||
text = " ".join(f"条例{index:03d}:中文🦉e\u0301 <字段>& 完整记录" for index in range(100))
|
||||
state, selected = _state(text)
|
||||
model = RecordingModel()
|
||||
budget = 100
|
||||
middleware = _middleware(model, default_counter=True, trim_tokens_to_summarize=budget)
|
||||
|
||||
result = await _compact(middleware, state, asynchronous)
|
||||
|
||||
assert result is not None and len(model.prompts) > 1
|
||||
captured = _source(model.prompts)
|
||||
assert captured.encode("utf-8") == get_buffer_string(selected).encode("utf-8")
|
||||
assert "🦉e\u0301" in captured
|
||||
for prompt in model.prompts:
|
||||
raw = _sections(prompt)
|
||||
assert sum(middleware.token_counter([HumanMessage(content=section)]) for section in raw.values()) <= budget
|
||||
assert middleware.summary_call_count == 1
|
||||
304
backend/tests/test_summarization_source_coverage.py
Normal file
304
backend/tests/test_summarization_source_coverage.py
Normal file
@ -0,0 +1,304 @@
|
||||
"""A successful compaction must actually send every selected source to its model."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import html
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage, get_buffer_string
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from pydantic import Field
|
||||
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummaryGenerationError
|
||||
from deerflow.config.task_continuity_config import TaskContinuityConfig
|
||||
from deerflow.extensions import ExtensionRegistry
|
||||
|
||||
_RAW_BUDGET = 1200
|
||||
_OLD_MEMORY = "PREVIOUS_DURABLE_MEMORY: preserve the user's release requirements. " + "Already verified context. " * 8
|
||||
|
||||
|
||||
def _char_count(messages) -> int:
|
||||
"""A deterministic test budget in raw text characters, without a tokenizer."""
|
||||
return sum(len(str(message.content)) for message in messages)
|
||||
|
||||
|
||||
class RecordingSummaryModel(BaseChatModel):
|
||||
prompts: list[str] = Field(default_factory=list)
|
||||
fail_after: int | None = None
|
||||
output_text: str | None = None
|
||||
|
||||
@property
|
||||
def _llm_type(self) -> str:
|
||||
return "source-coverage-offline-summary"
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
self.prompts.append("\n".join(str(message.content) for message in messages))
|
||||
if self.fail_after is not None and len(self.prompts) > self.fail_after:
|
||||
raise RuntimeError("offline provider failed after a successful summary step")
|
||||
# Deliberately short: the configured budget leaves room for this rolling
|
||||
# summary and progress instructions, rather than creating a no-progress case.
|
||||
answer = self.output_text if self.output_text is not None else f"DIGEST_STEP_{len(self.prompts):03d}: remembered prior evidence."
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=answer))])
|
||||
|
||||
async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
return self._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
|
||||
def _payload(label: str, padding: int = 500) -> tuple[str, tuple[str, ...]]:
|
||||
markers = (f"{label}_BEGIN", f"{label}_MIDDLE", f"{label}_END")
|
||||
# HTML-sensitive text exercises the raw-before-escaping budget as well.
|
||||
filler = " factual <value> & " * (padding // 18 + 1)
|
||||
return f"{markers[0]} {filler} {markers[1]} {filler} {markers[2]}", markers
|
||||
|
||||
|
||||
def _history(case: str):
|
||||
selected = []
|
||||
markers = []
|
||||
current = HumanMessage(content="CURRENT_REQUEST_STAYS_RAW", id="current-user")
|
||||
tail = AIMessage(content="RECENT_ANSWER_STAYS_RAW", id="recent-answer")
|
||||
|
||||
if case == "many-turns":
|
||||
for index in range(4):
|
||||
for role, message_class in (("USER", HumanMessage), ("ANSWER", AIMessage)):
|
||||
text, labels = _payload(f"TURN_{index}_{role}")
|
||||
selected.append(message_class(content=text, id=f"old-{index}-{role}"))
|
||||
markers.extend(labels)
|
||||
messages = [*selected, current, tail]
|
||||
elif case == "first-turn":
|
||||
for index in range(3):
|
||||
plan, plan_labels = _payload(f"FIRST_PLAN_{index}")
|
||||
result, result_labels = _payload(f"FIRST_RESULT_{index}")
|
||||
selected.extend(
|
||||
[
|
||||
AIMessage(content=plan, id=f"plan-{index}", tool_calls=[{"name": "read_file", "args": {}, "id": f"call-{index}"}]),
|
||||
ToolMessage(content=result, id=f"result-{index}", tool_call_id=f"call-{index}"),
|
||||
]
|
||||
)
|
||||
markers.extend((*plan_labels, *result_labels))
|
||||
# Preserve the final complete AI/Tool exchange. The first real user
|
||||
# request is rescued independently of the selected older exchanges.
|
||||
recent = [
|
||||
AIMessage(content="RECENT_PLAN_STAYS_RAW", id="recent-plan", tool_calls=[{"name": "read_file", "args": {}, "id": "recent-call"}]),
|
||||
ToolMessage(content="RECENT_RESULT_STAYS_RAW", id="recent-result", tool_call_id="recent-call"),
|
||||
]
|
||||
messages = [current, *selected, *recent]
|
||||
elif case == "oversized-message":
|
||||
text, labels = _payload("ONE_OVERSIZED_MESSAGE", padding=6500)
|
||||
selected = [HumanMessage(content=text, id="oversized-user"), AIMessage(content="old response", id="old-response")]
|
||||
markers.extend(labels)
|
||||
messages = [*selected, current, tail]
|
||||
else:
|
||||
assert case == "before-first-human"
|
||||
plan, plan_labels = _payload("BEFORE_HUMAN_PLAN")
|
||||
result, result_labels = _payload("BEFORE_HUMAN_TOOL")
|
||||
selected = [
|
||||
AIMessage(content=plan, id="leading-ai", tool_calls=[{"name": "read_file", "args": {}, "id": "leading-call"}]),
|
||||
ToolMessage(content=result, id="leading-tool", tool_call_id="leading-call"),
|
||||
]
|
||||
markers.extend((*plan_labels, *result_labels))
|
||||
messages = [*selected, current, tail]
|
||||
return {"messages": messages}, selected, markers
|
||||
|
||||
|
||||
def _raw_sections(prompt: str) -> list[tuple[str, str]]:
|
||||
# These are the documented model-facing source sections, not an internal
|
||||
# chunk planner. Decode after extraction so escaped source text cannot forge
|
||||
# a section boundary. No assertion depends on how many chunks are produced.
|
||||
matches = re.findall(r"<(existing_summary|new_messages)>\n(.*?)\n</\1>", prompt, re.DOTALL)
|
||||
assert matches, "the captured model input must contain the actual source text"
|
||||
return [(name, html.unescape(body)) for name, body in matches]
|
||||
|
||||
|
||||
def _middleware(model, *, raw_budget=_RAW_BUDGET, **kwargs):
|
||||
return DeerFlowSummarizationMiddleware(
|
||||
model=model,
|
||||
trigger=("messages", 4),
|
||||
keep=("messages", 2),
|
||||
token_counter=_char_count,
|
||||
trim_tokens_to_summarize=raw_budget,
|
||||
summary_prompt="{messages}",
|
||||
run_model_name="offline-summary",
|
||||
anchor_model_name="offline-summary",
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("previous_summary", [None, _OLD_MEMORY], ids=["first-summary", "existing-summary"])
|
||||
@pytest.mark.parametrize("case", ["many-turns", "first-turn", "oversized-message", "before-first-human"])
|
||||
async def test_every_selected_source_reaches_bounded_summary_input(case, previous_summary, asynchronous):
|
||||
state, selected, markers = _history(case)
|
||||
state["summary_text"] = previous_summary
|
||||
original = copy.deepcopy(state)
|
||||
model = RecordingSummaryModel()
|
||||
hooks = []
|
||||
middleware = _middleware(model, before_summarization=[hooks.append])
|
||||
runtime = SimpleNamespace(context={"thread_id": "source-coverage"})
|
||||
|
||||
result = await middleware.acompact_state(state, runtime) if asynchronous else middleware.compact_state(state, runtime)
|
||||
|
||||
assert result is not None
|
||||
assert {message.id for message in result.messages_to_summarize} == {message.id for message in selected}
|
||||
assert "current-user" in {message.id for message in result.preserved_messages}
|
||||
assert model.prompts
|
||||
sections = [_raw_sections(prompt) for prompt in model.prompts]
|
||||
# Only new source participates in this stream: inserting the rolling
|
||||
# existing_summary between chunks would break a marker at a valid edge.
|
||||
captured_source = "".join(body for prompt_sections in sections for name, body in prompt_sections if name == "new_messages")
|
||||
missing = [marker for marker in markers if marker not in captured_source]
|
||||
assert not missing, f"sources were removed without reaching the summary model: {missing}"
|
||||
assert captured_source == get_buffer_string(selected), "every selected raw source character must reach the model without omissions or duplication"
|
||||
for raw in sections:
|
||||
assert sum(len(body) for _, body in raw) <= _RAW_BUDGET
|
||||
if previous_summary:
|
||||
assert dict(sections[0])["existing_summary"] == previous_summary
|
||||
assert len(hooks) == 1
|
||||
assert state == original
|
||||
|
||||
|
||||
async def _assert_failed_compaction_is_atomic(monkeypatch, model, *, asynchronous, manual, previous_summary=_OLD_MEMORY, raw_budget=_RAW_BUDGET):
|
||||
from deerflow.agents.task_continuity import archive
|
||||
|
||||
state, _, _ = _history("many-turns")
|
||||
state.update(summary_text=previous_summary, task_history={"existing": "untouched"})
|
||||
original = copy.deepcopy(state)
|
||||
hooks = []
|
||||
observed = []
|
||||
registry = ExtensionRegistry()
|
||||
with registry.attributed_to("test-source-coverage"):
|
||||
registry.context_compaction_observer(lambda event, context=None: observed.append(event))
|
||||
capture = Mock(return_value={"unexpected": "archive commit"})
|
||||
acapture = AsyncMock(return_value={"unexpected": "archive commit"})
|
||||
monkeypatch.setattr(archive, "capture", capture)
|
||||
monkeypatch.setattr(archive, "acapture", acapture)
|
||||
middleware = _middleware(
|
||||
model,
|
||||
raw_budget=raw_budget,
|
||||
before_summarization=[hooks.append],
|
||||
task_continuity_config=TaskContinuityConfig(enabled=True),
|
||||
extensions=registry.build(),
|
||||
)
|
||||
runtime = SimpleNamespace(context={"thread_id": "source-coverage"})
|
||||
|
||||
if manual:
|
||||
with pytest.raises(SummaryGenerationError):
|
||||
if asynchronous:
|
||||
await middleware.acompact_state(state, runtime, force=True, raise_on_failure=True)
|
||||
else:
|
||||
middleware.compact_state(state, runtime, force=True, raise_on_failure=True)
|
||||
else:
|
||||
result = await middleware.abefore_model(state, runtime) if asynchronous else middleware.before_model(state, runtime)
|
||||
assert result is None, "automatic failure may not return a removal or partial summary update"
|
||||
|
||||
assert state == original
|
||||
assert hooks == []
|
||||
assert observed == []
|
||||
capture.assert_not_called()
|
||||
acapture.assert_not_awaited()
|
||||
assert middleware.summary_call_count == 0
|
||||
for prompt in model.prompts:
|
||||
assert sum(len(body) for _, body in _raw_sections(prompt)) <= raw_budget
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("manual", [False, True], ids=["automatic", "manual"])
|
||||
async def test_midway_provider_failure_cannot_commit_partial_compaction(monkeypatch, asynchronous, manual):
|
||||
model = RecordingSummaryModel(fail_after=1)
|
||||
await _assert_failed_compaction_is_atomic(monkeypatch, model, asynchronous=asynchronous, manual=manual)
|
||||
assert len(model.prompts) > 1, "the fault must occur after a successful provider step"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("manual", [False, True], ids=["automatic", "manual"])
|
||||
async def test_tiny_carry_budget_cannot_commit_partial_compaction(monkeypatch, asynchronous, manual):
|
||||
model = RecordingSummaryModel()
|
||||
await _assert_failed_compaction_is_atomic(monkeypatch, model, asynchronous=asynchronous, manual=manual, previous_summary=None, raw_budget=1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("manual", [False, True], ids=["automatic", "manual"])
|
||||
@pytest.mark.parametrize("boundary", ["oversized-previous", "previous-fills-budget", "intermediate-model-output", "final-model-output"])
|
||||
async def test_generated_carry_is_bounded_without_clipping_selected_sources(asynchronous, manual, boundary):
|
||||
state, selected, _ = _history("many-turns")
|
||||
model = RecordingSummaryModel()
|
||||
previous = _OLD_MEMORY
|
||||
carry_begin, carry_end = "LEGACY_CARRY_BEGIN|", "|LEGACY_CARRY_END"
|
||||
generated = "MODEL_OUTPUT_BEGIN|" + "expanded summary " * 150 + "|MODEL_OUTPUT_END"
|
||||
if boundary == "oversized-previous":
|
||||
previous = carry_begin + "P" * (_RAW_BUDGET + 1) + carry_end
|
||||
elif boundary == "previous-fills-budget":
|
||||
previous = carry_begin + "P" * (_RAW_BUDGET - len(carry_begin) - len(carry_end)) + carry_end
|
||||
else:
|
||||
model.output_text = generated
|
||||
if boundary == "final-model-output":
|
||||
selected = [HumanMessage(content="old question", id="old-user"), AIMessage(content="old answer", id="old-ai")]
|
||||
state["messages"] = [*selected, HumanMessage(content="current request", id="current-user"), AIMessage(content="recent answer", id="recent-answer")]
|
||||
state["summary_text"] = previous
|
||||
original = copy.deepcopy(state)
|
||||
hooks = []
|
||||
middleware = _middleware(model, before_summarization=[hooks.append])
|
||||
runtime = SimpleNamespace(context={"thread_id": "source-coverage"})
|
||||
|
||||
if manual:
|
||||
result = await middleware.acompact_state(state, runtime, force=True, raise_on_failure=True) if asynchronous else middleware.compact_state(state, runtime, force=True, raise_on_failure=True)
|
||||
assert result is not None
|
||||
summary = result.summary_text
|
||||
else:
|
||||
update = await middleware.abefore_model(state, runtime) if asynchronous else middleware.before_model(state, runtime)
|
||||
assert update is not None
|
||||
summary = update["summary_text"]
|
||||
|
||||
assert state == original
|
||||
assert len(hooks) == 1
|
||||
assert list(hooks[0].messages_to_summarize) == selected
|
||||
assert "current-user" in {message.id for message in hooks[0].preserved_messages}
|
||||
sections = [_raw_sections(prompt) for prompt in model.prompts]
|
||||
captured = "".join(body for step in sections for name, body in step if name == "new_messages")
|
||||
assert captured == get_buffer_string(selected)
|
||||
for step in sections:
|
||||
assert sum(len(body) for _, body in step) <= _RAW_BUDGET
|
||||
for name, body in step:
|
||||
if name == "existing_summary":
|
||||
assert _char_count([HumanMessage(content=body)]) <= _RAW_BUDGET // 2
|
||||
assert _char_count([HumanMessage(content=summary)]) <= _RAW_BUDGET // 2
|
||||
if boundary in {"oversized-previous", "previous-fills-budget"}:
|
||||
first_carry = dict(sections[0])["existing_summary"]
|
||||
assert first_carry.startswith(carry_begin) and first_carry.endswith(carry_end)
|
||||
assert "[summary shortened]" in first_carry
|
||||
assert summary == f"DIGEST_STEP_{len(model.prompts):03d}: remembered prior evidence."
|
||||
else:
|
||||
assert summary.startswith("MODEL_OUTPUT_BEGIN|") and summary.endswith("|MODEL_OUTPUT_END")
|
||||
assert "[summary shortened]" in summary
|
||||
if boundary == "intermediate-model-output":
|
||||
assert any("[summary shortened]" in dict(step).get("existing_summary", "") for step in sections[1:])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.parametrize("large_output", [False, True], ids=["normal-summary", "large-summary"])
|
||||
async def test_unbounded_budget_preserves_generated_summary_bytes(asynchronous, large_output):
|
||||
state, selected, _ = _history("many-turns")
|
||||
previous = "PREVIOUS_UNBOUNDED_BEGIN|" + "P" * 2000 + "|PREVIOUS_UNBOUNDED_END"
|
||||
state["summary_text"] = previous
|
||||
expected = "UNBOUNDED_OUTPUT_BEGIN|" + "V" * 2000 + "|UNBOUNDED_OUTPUT_END" if large_output else "normal generated summary"
|
||||
model = RecordingSummaryModel(output_text=expected)
|
||||
middleware = _middleware(model, raw_budget=None)
|
||||
runtime = SimpleNamespace(context={})
|
||||
|
||||
result = await middleware.acompact_state(state, runtime, force=True, raise_on_failure=True) if asynchronous else middleware.compact_state(state, runtime, force=True, raise_on_failure=True)
|
||||
|
||||
assert result is not None
|
||||
assert result.summary_text == expected
|
||||
sections = [_raw_sections(prompt) for prompt in model.prompts]
|
||||
assert dict(sections[0])["existing_summary"] == previous
|
||||
captured = "".join(body for step in sections for name, body in step if name == "new_messages")
|
||||
assert captured == get_buffer_string(selected)
|
||||
@ -1,11 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage, get_buffer_string
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
from pydantic import Field
|
||||
|
||||
@ -55,9 +56,14 @@ class _StaticChatModel(BaseChatModel):
|
||||
|
||||
class _RecordingSummaryModel(_StaticChatModel):
|
||||
prompts: list[str] = Field(default_factory=list)
|
||||
vary_responses: bool = False
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
self.prompts.append("\n".join(str(getattr(message, "content", message)) for message in messages))
|
||||
if self.vary_responses:
|
||||
# Coverage tests must observe each input rather than exercising the
|
||||
# separately tested exact-prompt/no-op cache with repetitive padding.
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=f"summary step {len(self.prompts)}"))])
|
||||
return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs)
|
||||
|
||||
|
||||
@ -69,6 +75,21 @@ def _big_history(n: int = 12) -> list:
|
||||
return messages
|
||||
|
||||
|
||||
def _assert_complete_inputs(model, selected, budget=None) -> str:
|
||||
assert model.prompts
|
||||
captured = []
|
||||
for prompt in model.prompts:
|
||||
sections = re.findall(r"<(existing_summary|new_messages)>\n(.*?)\n</\1>", prompt, re.DOTALL)
|
||||
assert sections
|
||||
raw = [(name, html.unescape(body)) for name, body in sections]
|
||||
captured.extend(body for name, body in raw if name == "new_messages")
|
||||
if budget is not None:
|
||||
assert sum(len(body) for _, body in raw) <= budget
|
||||
new_text = "".join(captured)
|
||||
assert new_text == get_buffer_string(selected)
|
||||
return new_text
|
||||
|
||||
|
||||
class TestSummaryFailureSafety:
|
||||
def test_summary_model_failure_does_not_destroy_history(self):
|
||||
middleware = DeerFlowSummarizationMiddleware(
|
||||
@ -113,14 +134,14 @@ class TestSummaryWritesChannel:
|
||||
assert result is not None
|
||||
assert user in result.preserved_messages
|
||||
assert list(result.messages_to_summarize) == history[1:5]
|
||||
assert len(model.prompts) == 1
|
||||
new_text = _assert_complete_inputs(model, result.messages_to_summarize, 4000)
|
||||
for sentinel in ("PLAN_0", "RESULT_0", "PLAN_1", "RESULT_1"):
|
||||
assert sentinel in model.prompts[0]
|
||||
assert "CURRENT_REQUEST" not in model.prompts[0]
|
||||
assert "RESULT_2" not in model.prompts[0]
|
||||
assert sentinel in new_text
|
||||
assert "CURRENT_REQUEST" not in new_text
|
||||
assert "RESULT_2" not in new_text
|
||||
|
||||
@pytest.mark.parametrize("async_mode", [False, True])
|
||||
@pytest.mark.parametrize("previous_summary", [None, "O" * 1000 + " OLD_END"], ids=["without-summary", "with-summary"])
|
||||
@pytest.mark.parametrize("previous_summary", [None, "PRIOR_SUMMARY OLD_END"], ids=["without-summary", "with-summary"])
|
||||
@pytest.mark.parametrize("trim_limit", [120, None], ids=["bounded", "untrimmed"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_rescued_user_window_keeps_recent_exchanges(self, async_mode, previous_summary, trim_limit):
|
||||
@ -151,18 +172,16 @@ class TestSummaryWritesChannel:
|
||||
assert result is not None
|
||||
assert user in result.preserved_messages
|
||||
assert list(result.messages_to_summarize) == history[1:13]
|
||||
assert len(model.prompts) == 1
|
||||
new_text = model.prompts[0].split("<new_messages>\n", 1)[1].split("\n</new_messages>", 1)[0]
|
||||
new_text = _assert_complete_inputs(model, result.messages_to_summarize, trim_limit)
|
||||
assert "RESULT_5" in new_text
|
||||
for i in range(6):
|
||||
assert f"PLAN_{i}" in new_text
|
||||
assert f"RESULT_{i}" in new_text
|
||||
if trim_limit is None:
|
||||
for i in range(6):
|
||||
assert f"PLAN_{i}" in new_text
|
||||
assert f"RESULT_{i}" in new_text
|
||||
assert len(model.prompts) == 1
|
||||
assert len(new_text) > 120
|
||||
if previous_summary:
|
||||
assert previous_summary in model.prompts[0]
|
||||
else:
|
||||
assert "RESULT_0" not in new_text
|
||||
assert "CURRENT_REQUEST" not in new_text
|
||||
assert "RESULT_6" not in new_text
|
||||
if previous_summary:
|
||||
@ -172,8 +191,8 @@ class TestSummaryWritesChannel:
|
||||
|
||||
@pytest.mark.parametrize("async_mode", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_history_empty_trim_preserves_recent_tool_result(self, async_mode):
|
||||
model = _RecordingSummaryModel()
|
||||
async def test_mixed_history_covers_old_request_and_tool_results(self, async_mode):
|
||||
model = _RecordingSummaryModel(vary_responses=True)
|
||||
middleware = DeerFlowSummarizationMiddleware(
|
||||
model=model,
|
||||
trigger=("messages", 4),
|
||||
@ -199,37 +218,40 @@ class TestSummaryWritesChannel:
|
||||
assert result is not None
|
||||
assert list(result.messages_to_summarize) == [history[0], *history[2:12]]
|
||||
assert list(result.preserved_messages) == [user, *history[12:]]
|
||||
assert len(model.prompts) == 1
|
||||
new_text = model.prompts[0].split("<new_messages>\n", 1)[1].split("\n</new_messages>", 1)[0]
|
||||
assert new_text == "Tool: RESULT_4"
|
||||
assert "CURRENT_REQUEST" not in model.prompts[0]
|
||||
assert "RESULT_5" not in model.prompts[0]
|
||||
new_text = _assert_complete_inputs(model, result.messages_to_summarize, 120)
|
||||
assert "OLD_REQUEST " + "x" * 1000 in new_text
|
||||
for i in range(5):
|
||||
assert f"PLAN_{i}" in new_text
|
||||
assert f"RESULT_{i}" in new_text
|
||||
assert "CURRENT_REQUEST" not in new_text
|
||||
assert "RESULT_5" not in new_text
|
||||
|
||||
def test_tool_only_fallback_applies_budget_before_escaping(self):
|
||||
def test_tool_source_coverage_applies_budget_before_escaping(self):
|
||||
model = _RecordingSummaryModel(vary_responses=True)
|
||||
middleware = DeerFlowSummarizationMiddleware(
|
||||
model=_StaticChatModel(),
|
||||
model=model,
|
||||
trigger=("messages", 4),
|
||||
keep=("messages", 2),
|
||||
token_counter=_char_count,
|
||||
trim_tokens_to_summarize=80,
|
||||
)
|
||||
prompt = middleware._build_summary_prompt(
|
||||
[ToolMessage(content="<" * 1000 + " TOOL_END", tool_call_id="call")],
|
||||
previous_summary="&" * 1000 + " OLD_END",
|
||||
)
|
||||
selected = [
|
||||
AIMessage(content="read", id="plan", tool_calls=[{"name": "read_file", "args": {}, "id": "call"}]),
|
||||
ToolMessage(content="<" * 1000 + " TOOL_END", tool_call_id="call", id="tool"),
|
||||
]
|
||||
tail = [HumanMessage(content="current", id="current"), AIMessage(content="recent", id="recent")]
|
||||
previous = "&" * 16 + " OLD_END"
|
||||
result = middleware.compact_state({"messages": [*selected, *tail], "summary_text": previous}, SimpleNamespace(context={}), force=True)
|
||||
|
||||
assert prompt is not None
|
||||
new_text = prompt.split("<new_messages>\n", 1)[1].split("\n</new_messages>", 1)[0]
|
||||
old_text = prompt.split("<existing_summary>\n", 1)[1].split("\n</existing_summary>", 1)[0]
|
||||
# The raw input budget excludes escaping and the surrounding prompt.
|
||||
assert len(html.unescape(new_text)) + len(html.unescape(old_text)) <= 80
|
||||
assert len(new_text) + len(old_text) > 80
|
||||
assert "<" in new_text
|
||||
assert "&" in old_text
|
||||
assert "<" not in new_text
|
||||
assert "<" not in old_text
|
||||
assert result is not None
|
||||
assert list(result.messages_to_summarize) == selected
|
||||
assert list(result.preserved_messages) == tail
|
||||
new_text = _assert_complete_inputs(model, selected, 80)
|
||||
assert "TOOL_END" in new_text
|
||||
assert "OLD_END" in old_text
|
||||
assert html.escape(previous, quote=False) in model.prompts[0]
|
||||
assert any("<" in prompt for prompt in model.prompts)
|
||||
encoded_sections = [re.findall(r"<(existing_summary|new_messages)>\n(.*?)\n</\1>", prompt, re.DOTALL) for prompt in model.prompts]
|
||||
assert any(sum(len(body) for _, body in sections) > 80 for sections in encoded_sections)
|
||||
|
||||
def _middleware(self) -> DeerFlowSummarizationMiddleware:
|
||||
return DeerFlowSummarizationMiddleware(
|
||||
@ -332,7 +354,7 @@ class TestSummaryWritesChannel:
|
||||
assert len(result.preserved_messages) == 2
|
||||
assert len(result.messages_to_summarize) > 0
|
||||
|
||||
def test_previous_summary_is_trimmed_with_summary_prompt_input(self):
|
||||
def test_single_prompt_builder_rejects_previous_summary_requiring_clipping(self):
|
||||
middleware = DeerFlowSummarizationMiddleware(
|
||||
model=_StaticChatModel(text="UPDATED_SUMMARY"),
|
||||
trigger=("messages", 4),
|
||||
@ -347,11 +369,9 @@ class TestSummaryWritesChannel:
|
||||
previous_summary=previous_summary,
|
||||
)
|
||||
|
||||
assert prompt is not None
|
||||
assert previous_summary not in prompt
|
||||
assert "NEW_MESSAGE_SENTINEL" in prompt
|
||||
assert prompt is None
|
||||
|
||||
def test_new_message_summary_prompt_trim_uses_token_counter_budget(self):
|
||||
def test_single_prompt_builder_rejects_source_exceeding_budget(self):
|
||||
middleware = DeerFlowSummarizationMiddleware(
|
||||
model=_StaticChatModel(text="UPDATED_SUMMARY"),
|
||||
trigger=("messages", 4),
|
||||
@ -362,33 +382,30 @@ class TestSummaryWritesChannel:
|
||||
|
||||
body = middleware._build_summary_input_text("Human: NEW_MESSAGE_SENTINEL " + ("N" * 200))
|
||||
|
||||
assert body is not None
|
||||
new_messages = body.split("<new_messages>\n", 1)[1].split("\n</new_messages>", 1)[0]
|
||||
assert len(new_messages) <= 40
|
||||
assert "NEW_MESSAGE_SENTINEL" in new_messages
|
||||
assert body is None
|
||||
|
||||
@pytest.mark.parametrize(("strategy", "expected"), [("first", "ab"), ("last", "ef")])
|
||||
def test_summary_prompt_fallback_bound_respects_small_budget(self, strategy, expected):
|
||||
@pytest.mark.parametrize("async_mode", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_summary_input_counter_failure_preserves_complete_sources(self, async_mode):
|
||||
def counter(messages):
|
||||
if len(messages) == 1 and isinstance(messages[0], HumanMessage):
|
||||
return _raising_count(messages)
|
||||
return _char_count(messages)
|
||||
|
||||
model = _RecordingSummaryModel(vary_responses=True)
|
||||
middleware = DeerFlowSummarizationMiddleware(
|
||||
model=_StaticChatModel(text="UPDATED_SUMMARY"),
|
||||
model=model,
|
||||
trigger=("messages", 4),
|
||||
keep=("messages", 2),
|
||||
token_counter=_raising_count,
|
||||
trim_tokens_to_summarize=2,
|
||||
token_counter=counter,
|
||||
trim_tokens_to_summarize=90,
|
||||
)
|
||||
history = _big_history(4)
|
||||
state = {"messages": history, "summary_text": "PRIOR_SUMMARY"}
|
||||
runtime = SimpleNamespace(context={})
|
||||
result = await middleware.acompact_state(state, runtime, force=True) if async_mode else middleware.compact_state(state, runtime, force=True)
|
||||
|
||||
text = middleware._trim_summary_section_text("abcdef", 2, strategy=strategy)
|
||||
|
||||
assert text == expected
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cap", "expected"),
|
||||
[(1, "i"), (2, "hi"), (5, "efghi"), (6, "\n...\ni"), (8, "\n...\nghi"), (9, "abcdefghi"), (20, "abcdefghi")],
|
||||
)
|
||||
def test_tail_fallback_marks_omitted_text_within_budget(self, cap, expected):
|
||||
middleware = DeerFlowSummarizationMiddleware(model=_StaticChatModel(), token_counter=_raising_count)
|
||||
|
||||
text = middleware._trim_summary_section_text("abcdefghi", cap, strategy="last")
|
||||
|
||||
assert text == expected
|
||||
assert len(text) <= cap
|
||||
assert result is not None
|
||||
assert list(result.messages_to_summarize) == history[:-2]
|
||||
assert list(result.preserved_messages) == history[-2:]
|
||||
_assert_complete_inputs(model, result.messages_to_summarize, 90)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user