mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-26 07:57:57 +00:00
Phase 3 of #3875 — subagents previously inherited none of the lead's context-compaction, so a deep-research subagent (max_turns up to 150) could accumulate >1M cumulative input before max_turns/timeout/token_budget engaged, even after Phase 2's budget capped the pathological tail. - Gate the subagent runtime chain on the SAME ``app_config.summarization.enabled`` switch the lead reads (per maintainer guidance in #3875), via the shared ``create_summarization_middleware`` factory. One config covers both chains; no separate ``subagents.summarization`` field. No-op when summarization is off (factory returns None). - ``skip_memory_flush=True`` on the subagent path: the factory otherwise attaches ``memory_flush_hook`` (when memory.enabled), which flushes pre-compaction messages into durable memory keyed by thread_id. Subagents share the parent's thread_id, so without skipping the hook a subagent's internal turns would pollute the PARENT thread's durable memory (#3875 Phase 3 review point). - Harden ``capture_new_step_messages`` to tolerate history contraction: summarization rewrites the messages channel via ``RemoveMessage(id=REMOVE_ALL_MESSAGES)``, shrinking len(messages) below the step-capture cursor. Without a reset, every step appended after the compaction point was dropped until length overtook the stale cursor (#3845 interaction, maintainer validation point (a)). Cursor now resets to the new tail; id/content dedup prevents re-emitting pre-compaction steps. - Couple the DEFAULT token-budget ceiling to ``summarization.enabled`` (#3875 Phase 3 review point): 1M when compaction is on, 2M when off (preserves Phase 2's deliberate headroom for summarization-off deep-research runs that can exceed 1M). A user-set budget (global or per-agent) always wins regardless of the switch. Flagged tunable. The summarization middleware does not implement ``consume_stop_reason``, so the Phase 2 guard-cap stop-reason channel is unaffected. Refs: https://github.com/bytedance/deer-flow/issues/3875
This commit is contained in:
hataa
GitHub
parent
eb8eec21ce
commit
266883b3dd
@ -379,8 +379,9 @@ Proxied through nginx: `/api/langgraph/*` → Gateway LangGraph-compatible runti
|
||||
**Concurrency**: `MAX_CONCURRENT_SUBAGENTS = 3` enforced by `SubagentLimitMiddleware` (truncates excess tool calls in `after_model`); default subagent timeout `subagents.timeout_seconds=1800` (30 min) and built-in `general-purpose` `max_turns=150` (raised from 100/15-min so deep-research subtasks stop hitting `GraphRecursionError` out of the box)
|
||||
**Flow**: `task()` tool → `SubagentExecutor` → background thread → poll 5s → SSE events → result
|
||||
**Events**: `task_started`, `task_running`, `task_completed`/`task_failed`/`task_timed_out`
|
||||
**Guardrail caps & `stop_reason` (#3875 Phase 2)**: three independent axes can end a subagent run early, and all now surface *why* through one additive field rather than a new status enum. **Turn axis**: `recursion_limit` on the subagent `run_config` equals `max_turns`, so exhausting the turn budget raises `GraphRecursionError` from `agent.astream`; `executor.py::_aexecute` catches it specifically (before the generic `except Exception`). **Token axis**: `TokenBudgetMiddleware` is attached per-agent via `build_subagent_runtime_middlewares` from `subagents.token_budget` (default 2,000,000 tokens, warn at 0.7, hard-stop at 1.0 — a backstop against a subagent that burns tokens on trivial work). It does *not* raise: at the hard-stop threshold it strips the in-flight turn's tool calls, forces `finish_reason="stop"`, and lets the run complete naturally with a final answer. **Loop axis**: `LoopDetectionMiddleware` (attached at the same point) catches repeated identical tool-call sets — or one tool *type* called many times with varying args — and its hard-stop likewise strips `tool_calls` and forces a final answer without raising, recording `loop_capped`. Each guard exposes its cap on a per-`run_id` `consume_stop_reason(run_id)` accessor; `_aexecute` collects **every** middleware with that method (duck-typed via `hasattr`, so the executor has no import coupling to the guard classes) and surfaces the first non-`None` reason — adding a future guard needs no executor change. **Surfacing**: whichever axis fired, `_aexecute` stamps a normal status plus an additive reason — `completed` + `stop_reason=token_capped|turn_capped|loop_capped` when a usable final answer (or partial recovered from the last streamed chunk via `_extract_final_result` → `utils/messages.py::message_content_to_text`, returning a `"No response Generated"` sentinel when no text survived) was produced; `failed` + `stop_reason=turn_capped` when nothing usable survived. `SubagentResult.stop_reason` flows through `task_tool.py::_task_result_command` → `format_subagent_result_message` (renders `Task Succeeded (capped: ...)` / `Task failed (capped: ...)`) and `make_subagent_additional_kwargs`, which stamps the additive `subagent_stop_reason` key alongside the normal `subagent_status`. **Why additive, not an enum**: a new status value would break v1 consumers; an optional field is ignored by older frontends and ledger readers, so the cross-language contract (`contracts/subagent_status_contract.json` v2 + `subagents/status_contract.py` + `frontend/.../subtask-result.ts`, pinned by `test_status_values_match_contract` / `test_stop_reason_values_match_contract`) stays backward-compatible. The durable delegation ledger captures `stop_reason` onto the entry and renders model-facing guidance ("hit a guardrail cap with a partial result; reuse it, retry tighter, or raise the per-agent budget (`max_turns` / `token_budget`)") so the lead reuses a capped completion knowingly instead of mistaking it for a clean one. (Phase 1 shipped this surfacing as a `MAX_TURNS_REACHED` status enum in #3949; Phase 2 replaced that enum with the additive `stop_reason` field per the agreed design — the `max_turns_reached` status value and `SubagentStatus.MAX_TURNS_REACHED` are gone.)
|
||||
**Step capture & persistence (#3779)**: `executor.py` captures both assistant turns (`AIMessage`) **and** tool outputs (`ToolMessage`) via `subagents/step_events.py::capture_new_step_messages`, which walks the *newly-appended tail* of each `stream_mode="values"` chunk (not just `messages[-1]`) so a multi-tool-call turn — where LangGraph's `ToolNode` appends several `ToolMessage`s in one super-step — keeps every tool output instead of dropping all but the last. `runtime/runs/worker.py::_SubagentEventBuffer` additionally persists these `task_*` custom events to the `RunEventStore` as `subagent.start`/`subagent.step`/`subagent.end` (`category="subagent"`, `task_id` in `metadata`). It **batches** writes via `put_batch` (flushing on a terminal `subagent.end`, at `FLUSH_THRESHOLD` events, and in the worker's `finally`) rather than one `put()` per step, since `put()` is a documented low-frequency path (per-thread advisory lock per call) and a deep subagent (`max_turns=150`) emits hundreds of steps on the hot stream loop. `build_subagent_step` caps both the per-step `text` and each tool call's serialized `args` at `SUBAGENT_STEP_MAX_CHARS` (flagged `truncated` / `args_truncated`) so a large `write_file`/`bash` payload can't produce an unbounded row. The dedicated category keeps them out of `list_messages` (the thread feed) while `list_events` returns them for the frontend's fetch-on-expand backfill. `list_events` accepts `task_id` (filters on `metadata["task_id"]` — SQL-side in `DbRunEventStore` via `event_metadata["task_id"].as_string()`, in-memory in the JSONL/memory stores) plus an `after_seq` forward cursor, so the card pages through one subagent's steps without the run-wide `limit` truncating the tail (no schema migration: the filter rides the existing run-scoped index). `step_events.py` is a pure, unit-tested layer (`build_subagent_step` / `subagent_run_event`).
|
||||
**Guardrail caps & `stop_reason` (#3875 Phase 2)**: three independent axes can end a subagent run early, and all now surface *why* through one additive field rather than a new status enum. **Turn axis**: `recursion_limit` on the subagent `run_config` equals `max_turns`, so exhausting the turn budget raises `GraphRecursionError` from `agent.astream`; `executor.py::_aexecute` catches it specifically (before the generic `except Exception`). **Token axis**: `TokenBudgetMiddleware` is attached per-agent via `build_subagent_runtime_middlewares` from `subagents.token_budget` (default `max_tokens` **coupled to `summarization.enabled`** — 1,000,000 when subagent summarization is on, 2,000,000 when off, warn at 0.7, hard-stop at 1.0; a user-set budget always wins regardless of the switch — #3875 Phase 3; a backstop against a subagent that burns tokens on trivial work). It does *not* raise: at the hard-stop threshold it strips the in-flight turn's tool calls, forces `finish_reason="stop"`, and lets the run complete naturally with a final answer. **Loop axis**: `LoopDetectionMiddleware` (attached at the same point) catches repeated identical tool-call sets — or one tool *type* called many times with varying args — and its hard-stop likewise strips `tool_calls` and forces a final answer without raising, recording `loop_capped`. Each guard exposes its cap on a per-`run_id` `consume_stop_reason(run_id)` accessor; `_aexecute` collects **every** middleware with that method (duck-typed via `hasattr`, so the executor has no import coupling to the guard classes) and surfaces the first non-`None` reason — adding a future guard needs no executor change. **Surfacing**: whichever axis fired, `_aexecute` stamps a normal status plus an additive reason — `completed` + `stop_reason=token_capped|turn_capped|loop_capped` when a usable final answer (or partial recovered from the last streamed chunk via `_extract_final_result` → `utils/messages.py::message_content_to_text`, returning a `"No response Generated"` sentinel when no text survived) was produced; `failed` + `stop_reason=turn_capped` when nothing usable survived. `SubagentResult.stop_reason` flows through `task_tool.py::_task_result_command` → `format_subagent_result_message` (renders `Task Succeeded (capped: ...)` / `Task failed (capped: ...)`) and `make_subagent_additional_kwargs`, which stamps the additive `subagent_stop_reason` key alongside the normal `subagent_status`. **Why additive, not an enum**: a new status value would break v1 consumers; an optional field is ignored by older frontends and ledger readers, so the cross-language contract (`contracts/subagent_status_contract.json` v2 + `subagents/status_contract.py` + `frontend/.../subtask-result.ts`, pinned by `test_status_values_match_contract` / `test_stop_reason_values_match_contract`) stays backward-compatible. The durable delegation ledger captures `stop_reason` onto the entry and renders model-facing guidance ("hit a guardrail cap with a partial result; reuse it, retry tighter, or raise the per-agent budget (`max_turns` / `token_budget`)") so the lead reuses a capped completion knowingly instead of mistaking it for a clean one. (Phase 1 shipped this surfacing as a `MAX_TURNS_REACHED` status enum in #3949; Phase 2 replaced that enum with the additive `stop_reason` field per the agreed design — the `max_turns_reached` status value and `SubagentStatus.MAX_TURNS_REACHED` are gone.)
|
||||
**Context compaction (#3875 Phase 3)**: subagents inherit `DeerFlowSummarizationMiddleware` via `build_subagent_runtime_middlewares`, gated on the **same** `summarization.enabled` switch the lead reads (one config covers both chains; trigger/keep/model/prompt come from the shared `summarization` config so they cannot drift). The factory is called with `skip_memory_flush=True` on the subagent path: the lead's `memory_flush_hook` (attached when `memory.enabled`) flushes pre-compaction messages into durable memory keyed by `thread_id`, and subagents share the parent's `thread_id`, so without skipping the hook a subagent's internal turns would pollute the **parent** thread's durable memory. Placement differs from the lead chain (lead appends summarization *before* the guard trio; subagent appends it *after*) — benign because the middleware implements only `before_model` (compaction) with no `after_model`/`consume_stop_reason`, so it cannot disturb the Phase 2 guard-cap stop-reason channel. Compaction rewrites the messages channel via `RemoveMessage(id=REMOVE_ALL_MESSAGES)`, which shrinks `len(messages)` below the step-capture cursor mid-run; `capture_new_step_messages` (see Step capture below) resets the cursor to the new tail on contraction so steps appended after the compaction point are not silently dropped.
|
||||
**Step capture & persistence (#3779)**: `executor.py` captures both assistant turns (`AIMessage`) **and** tool outputs (`ToolMessage`) via `subagents/step_events.py::capture_new_step_messages`, which walks the *newly-appended tail* of each `stream_mode="values"` chunk (not just `messages[-1]`) so a multi-tool-call turn — where LangGraph's `ToolNode` appends several `ToolMessage`s in one super-step — keeps every tool output instead of dropping all but the last. `runtime/runs/worker.py::_SubagentEventBuffer` additionally persists these `task_*` custom events to the `RunEventStore` as `subagent.start`/`subagent.step`/`subagent.end` (`category="subagent"`, `task_id` in `metadata`). It **batches** writes via `put_batch` (flushing on a terminal `subagent.end`, at `FLUSH_THRESHOLD` events, and in the worker's `finally`) rather than one `put()` per step, since `put()` is a documented low-frequency path (per-thread advisory lock per call) and a deep subagent (`max_turns=150`) emits hundreds of steps on the hot stream loop. `build_subagent_step` caps both the per-step `text` and each tool call's serialized `args` at `SUBAGENT_STEP_MAX_CHARS` (flagged `truncated` / `args_truncated`) so a large `write_file`/`bash` payload can't produce an unbounded row. The dedicated category keeps them out of `list_messages` (the thread feed) while `list_events` returns them for the frontend's fetch-on-expand backfill. `list_events` accepts `task_id` (filters on `metadata["task_id"]` — SQL-side in `DbRunEventStore` via `event_metadata["task_id"].as_string()`, in-memory in the JSONL/memory stores) plus an `after_seq` forward cursor, so the card pages through one subagent's steps without the run-wide `limit` truncating the tail (no schema migration: the filter rides the existing run-scoped index). `step_events.py` is a pure, unit-tested layer (`build_subagent_step` / `subagent_run_event`). **History contraction (#3875 Phase 3)**: `capture_new_step_messages` assumes append-only growth, but `DeerFlowSummarizationMiddleware` rewrites the messages channel via `RemoveMessage(id=REMOVE_ALL_MESSAGES)`, shrinking `len(messages)` below the cursor mid-run. On contraction (`total < processed_count`) the cursor resets to the new tail; `capture_step_message`'s id/content dedup prevents re-emitting pre-compaction steps, so steps appended after the compaction point are still captured instead of being dropped until `total` overtakes the stale cursor.
|
||||
**Deferred MCP tools** (if `tool_search.enabled`): `SubagentExecutor._build_initial_state` assembles deferral after policy filtering via the shared `assemble_deferred_tools` (fail-closed), appends the `tool_search` tool, injects the `<available-deferred-tools>` section into the subagent's `SystemMessage`, and threads the setup to `_create_agent`, which attaches `McpRoutingMiddleware` (when PR1 routing metadata matches deferred tools) before `DeferredToolFilterMiddleware` through `build_subagent_runtime_middlewares(...)`. Subagents thus withhold full MCP schemas until promotion, same as the lead agent; each task run gets a fresh `ThreadState` so promotion is isolated per run
|
||||
**Checkpointer isolation**: Subagent graphs are compiled with `checkpointer=False` to avoid inheriting the parent run's checkpointer, since subagents are one-shot and never resume.
|
||||
|
||||
|
||||
@ -437,15 +437,22 @@ def create_summarization_middleware(
|
||||
*,
|
||||
app_config: Any | None = None,
|
||||
keep: tuple[str, int | float] | None = None,
|
||||
skip_memory_flush: bool = False,
|
||||
) -> DeerFlowSummarizationMiddleware | None:
|
||||
"""Create the configured summarization middleware.
|
||||
|
||||
Both the lead-agent automatic path and the manual context-compaction path
|
||||
use this factory so model resolution, hooks, prompt config, and retention
|
||||
defaults cannot drift.
|
||||
"""
|
||||
from deerflow.agents.memory.summarization_hook import memory_flush_hook
|
||||
|
||||
``skip_memory_flush`` omits the ``memory_flush_hook`` that otherwise
|
||||
flushes pre-compaction messages into the durable memory queue. The lead
|
||||
chain keeps it (research should persist); the subagent chain sets it so a
|
||||
subagent's INTERNAL turns (the "Task" human message + intermediate AI/tool
|
||||
turns) are not written into the PARENT thread's durable memory — the hook
|
||||
is keyed by ``thread_id`` and subagents share the parent's ``thread_id``
|
||||
(#3875 Phase 3 review).
|
||||
"""
|
||||
resolved_app_config = app_config or get_app_config()
|
||||
config = resolved_app_config.summarization
|
||||
|
||||
@ -485,7 +492,9 @@ def create_summarization_middleware(
|
||||
kwargs["summary_prompt"] = config.summary_prompt
|
||||
|
||||
hooks: list[BeforeSummarizationHook] = []
|
||||
if resolved_app_config.memory.enabled:
|
||||
if resolved_app_config.memory.enabled and not skip_memory_flush:
|
||||
from deerflow.agents.memory.summarization_hook import memory_flush_hook
|
||||
|
||||
hooks.append(memory_flush_hook)
|
||||
|
||||
return DeerFlowSummarizationMiddleware(
|
||||
|
||||
@ -353,7 +353,17 @@ def build_subagent_runtime_middlewares(
|
||||
# builds a fresh middleware instance (see ``executor._create_agent``), so
|
||||
# parallel subagents cannot cross-contaminate even though they share the
|
||||
# parent thread_id/run_id in context.
|
||||
token_budget_config = app_config.subagents.get_token_budget_for(agent_name) if agent_name is not None else app_config.subagents.token_budget
|
||||
#
|
||||
# Default-ceiling coupling (#3875 Phase 3 review): the default ``max_tokens``
|
||||
# is re-coupled to ``summarization.enabled`` — 1M when compaction is on, 2M
|
||||
# when off. This ONLY applies to the default; a user-set budget (global or
|
||||
# per-agent) always wins, so a deployment that pinned a value is never
|
||||
# silently changed by flipping the summarization switch.
|
||||
summarization_enabled = app_config.summarization.enabled
|
||||
if agent_name is not None:
|
||||
token_budget_config = app_config.subagents.get_token_budget_for(agent_name, summarization_enabled=summarization_enabled)
|
||||
else:
|
||||
token_budget_config = app_config.subagents.token_budget
|
||||
if token_budget_config.enabled:
|
||||
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
|
||||
|
||||
@ -369,4 +379,45 @@ def build_subagent_runtime_middlewares(
|
||||
|
||||
middlewares.append(SafetyFinishReasonMiddleware.from_config(safety_config))
|
||||
|
||||
# DeerFlowSummarizationMiddleware — subagents inherit none of the lead's
|
||||
# context compaction today (#3875 Phase 3): a deep-research subagent
|
||||
# (``max_turns`` up to 150) can accumulate >1M cumulative input before
|
||||
# max_turns/timeout/token_budget engage, even though Phase 2's budget now
|
||||
# caps the pathological tail. Gated on the SAME
|
||||
# ``app_config.summarization.enabled`` switch the lead reads (per
|
||||
# maintainer guidance in #3875) so a single config covers both chains —
|
||||
# no separate ``subagents.summarization`` field. The shared factory
|
||||
# returns ``None`` when summarization is disabled, so this is a pure
|
||||
# no-op when the switch is off. Trigger/keep/model/prompt all come from
|
||||
# the same ``summarization`` config the lead reads, so the two chains
|
||||
# cannot drift.
|
||||
#
|
||||
# Placement differs from the lead chain: the lead appends summarization
|
||||
# BEFORE the guard trio (loop/token/safety), here it is appended AFTER.
|
||||
# This is benign — compaction runs in ``before_model`` regardless of
|
||||
# relative position, and the guard middlewares account in ``after_model``
|
||||
# — but noted because the relative order is not an exact mirror.
|
||||
#
|
||||
# ``skip_memory_flush=True``: the factory otherwise attaches
|
||||
# ``memory_flush_hook`` (when ``memory.enabled``), which flushes
|
||||
# pre-compaction messages into the durable memory queue keyed by
|
||||
# ``thread_id``. Subagents share the parent's ``thread_id`` in context, so
|
||||
# without skipping the hook a subagent's internal turns would be written
|
||||
# into the PARENT thread's durable memory (#3875 Phase 3 review).
|
||||
#
|
||||
# The middleware rewrites history via ``RemoveMessage(id=REMOVE_ALL_MESSAGES)``,
|
||||
# which shrinks the messages channel mid-run;
|
||||
# ``capture_new_step_messages`` must tolerate that contraction (see
|
||||
# ``step_events.py``) or it drops steps captured after the compaction
|
||||
# point. It does not implement ``consume_stop_reason``, so it does not
|
||||
# interfere with the Phase 2 guard-cap stop-reason channel.
|
||||
from deerflow.agents.middlewares.summarization_middleware import create_summarization_middleware
|
||||
|
||||
summarization_middleware = create_summarization_middleware(
|
||||
app_config=app_config,
|
||||
skip_memory_flush=True,
|
||||
)
|
||||
if summarization_middleware is not None:
|
||||
middlewares.append(summarization_middleware)
|
||||
|
||||
return middlewares
|
||||
|
||||
@ -9,18 +9,34 @@ from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def default_subagent_token_budget() -> TokenBudgetConfig:
|
||||
"""Default per-run token budget for subagents (#3875 Phase 2).
|
||||
def default_subagent_token_budget(*, summarization_enabled: bool = False) -> TokenBudgetConfig:
|
||||
"""Default per-run token budget for subagents (#3875 Phase 2 → Phase 3 coupling).
|
||||
|
||||
Enabled by default so the pathological-token-burn backstop actually
|
||||
engages (per umbrella #3857 point 4 — backstops must engage, not just
|
||||
exist). ``max_tokens`` is a deliberately loose ceiling: the reported 4.4M
|
||||
burn would have been cut roughly in half, while legitimate deep-research
|
||||
runs (``max_turns=150``, no summarization yet) can genuinely accumulate
|
||||
>1M cumulative input today. Tighten after Phase 3 lands subagent
|
||||
summarization. Flagged tunable in the PR description.
|
||||
exist). ``max_tokens`` is **coupled to whether subagent summarization is
|
||||
on** (#3875 Phase 3 review point):
|
||||
|
||||
- ``summarization_enabled=True`` (Phase 3 compacts the running context
|
||||
before it reaches pathological size): **1M** — tighter ceiling still
|
||||
covers legitimate deep research while catching degenerate runs earlier.
|
||||
- ``summarization_enabled=False``: **2M** — the Phase 2 ceiling. Phase 2's
|
||||
own docstring noted legitimate deep-research runs (``max_turns=150``,
|
||||
no summarization) "can genuinely accumulate >1M cumulative input," so a
|
||||
1M ceiling without compaction would prematurely cap them. Keeping 2M
|
||||
here preserves that headroom; the tighter 1M only applies when the
|
||||
compaction that justifies it is actually running.
|
||||
|
||||
The model-level ``default_factory`` (``SubagentsAppConfig.token_budget``)
|
||||
cannot read ``summarization.enabled`` (a sibling top-level field), so it
|
||||
falls back to the 2M no-compaction default; the builder
|
||||
(``build_subagent_runtime_middlewares``) recomputes via
|
||||
``get_token_budget_for(..., summarization_enabled=...)`` so the live value
|
||||
reflects the actual switch. A user-set ``token_budget`` (global or
|
||||
per-agent) always wins regardless of the switch. Flagged tunable.
|
||||
"""
|
||||
return TokenBudgetConfig(enabled=True, max_tokens=2_000_000, warn_threshold=0.7)
|
||||
max_tokens = 1_000_000 if summarization_enabled else 2_000_000
|
||||
return TokenBudgetConfig(enabled=True, max_tokens=max_tokens, warn_threshold=0.7)
|
||||
|
||||
|
||||
class SubagentOverrideConfig(BaseModel):
|
||||
@ -114,6 +130,20 @@ class SubagentsAppConfig(BaseModel):
|
||||
description="User-defined subagent types keyed by agent name",
|
||||
)
|
||||
|
||||
# True when ``token_budget`` was NOT explicitly provided by the user, i.e.
|
||||
# the field fell back to its default_factory. ``get_token_budget_for`` uses
|
||||
# this to decide whether the ceiling may be re-coupled to
|
||||
# ``summarization.enabled`` (#3875 Phase 3): a user-set budget is always
|
||||
# respected as-is. Set by ``__init__`` from ``model_fields_set`` and
|
||||
# preserved across the app-config reload path (which drops a default
|
||||
# ``token_budget`` before re-constructing — see
|
||||
# ``load_subagents_config_from_dict``).
|
||||
_token_budget_is_default: bool = True
|
||||
|
||||
def __init__(self, **data):
|
||||
super().__init__(**data)
|
||||
self._token_budget_is_default = "token_budget" not in self.model_fields_set
|
||||
|
||||
def get_timeout_for(self, agent_name: str) -> int:
|
||||
"""Get the effective timeout for a specific agent.
|
||||
|
||||
@ -165,7 +195,12 @@ class SubagentsAppConfig(BaseModel):
|
||||
return override.skills
|
||||
return None
|
||||
|
||||
def get_token_budget_for(self, agent_name: str) -> TokenBudgetConfig:
|
||||
def get_token_budget_for(
|
||||
self,
|
||||
agent_name: str,
|
||||
*,
|
||||
summarization_enabled: bool = False,
|
||||
) -> TokenBudgetConfig:
|
||||
"""Get the effective token-budget config for a specific agent.
|
||||
|
||||
Unlike ``max_turns``/``timeout_seconds`` (which keep a custom agent's
|
||||
@ -173,10 +208,21 @@ class SubagentsAppConfig(BaseModel):
|
||||
every subagent unless explicitly disabled — so the per-agent override
|
||||
wins when set, otherwise the global default applies to built-in AND
|
||||
custom agents alike (#3875 Phase 2 / umbrella #3857 point 4).
|
||||
|
||||
``summarization_enabled`` couples the DEFAULT ceiling to whether
|
||||
subagent summarization is on (#3875 Phase 3 review): 1M when
|
||||
compaction is running, 2M otherwise. It ONLY affects the default —
|
||||
any explicitly configured ``token_budget`` (global or per-agent)
|
||||
wins regardless, so a deployment that pinned a value is never
|
||||
silently changed by flipping the summarization switch.
|
||||
"""
|
||||
override = self.agents.get(agent_name)
|
||||
if override is not None and override.token_budget is not None:
|
||||
return override.token_budget
|
||||
# Only recompute when the caller is using the default (no explicit
|
||||
# global token_budget was set). A user-set global is respected as-is.
|
||||
if self._token_budget_is_default:
|
||||
return default_subagent_token_budget(summarization_enabled=summarization_enabled)
|
||||
return self.token_budget
|
||||
|
||||
|
||||
@ -191,6 +237,18 @@ def get_subagents_app_config() -> SubagentsAppConfig:
|
||||
def load_subagents_config_from_dict(config_dict: dict) -> None:
|
||||
"""Load subagents configuration from a dictionary."""
|
||||
global _subagents_config
|
||||
# The app-config reload path (app_config.py) round-trips via
|
||||
# ``config.subagents.model_dump()``, which serializes a default
|
||||
# ``token_budget`` into the dict. Re-constructing from that dict would make
|
||||
# ``model_fields_set`` contain ``token_budget`` and flip
|
||||
# ``_token_budget_is_default`` to False — breaking the
|
||||
# summarization-coupled recompute in ``get_token_budget_for`` (#3875 Phase
|
||||
# 3). Drop the key when its value still equals the no-compaction default so
|
||||
# the default_factory fires on reconstruction and the "user did not set
|
||||
# this" signal is preserved.
|
||||
tb = config_dict.get("token_budget")
|
||||
if tb is not None and tb == default_subagent_token_budget(summarization_enabled=False).model_dump():
|
||||
config_dict = {k: v for k, v in config_dict.items() if k != "token_budget"}
|
||||
_subagents_config = SubagentsAppConfig(**config_dict)
|
||||
|
||||
overrides_summary = {}
|
||||
|
||||
@ -95,8 +95,27 @@ def capture_new_step_messages(
|
||||
grow, re-examine only the trailing message so an id-less in-place replacement
|
||||
(same length, new content) is still captured — ``capture_step_message``'s
|
||||
dedup makes an unchanged re-yield a no-op. Returns the new cursor.
|
||||
|
||||
When the history *contracted* (``total < processed_count``) — which happens
|
||||
when ``DeerFlowSummarizationMiddleware`` rewrites the channel via
|
||||
``RemoveMessage(id=REMOVE_ALL_MESSAGES)`` (#3875 Phase 3) — reset the cursor
|
||||
to the new tail and let ``capture_step_message``'s id/content dedup prevent
|
||||
re-emitting steps captured before the compaction. Without this reset, every
|
||||
step appended after the compaction point is dropped until ``total`` overtakes
|
||||
the stale cursor.
|
||||
|
||||
INVARIANT: after the reset the no-growth branch only re-examines
|
||||
``messages[-1]``, so a genuinely new AIMessage/ToolMessage inserted at an
|
||||
index BELOW the reset cursor in a compacted list would be missed. This is
|
||||
not reachable today: the summarization middleware puts the summary into a
|
||||
separate ``summary_text`` state key, and the messages channel after
|
||||
compaction holds only already-seen preserved tail messages — compaction
|
||||
never inserts a NEW capturable message below the cursor. If a future
|
||||
middleware violates this invariant, the reset branch needs a full re-scan.
|
||||
"""
|
||||
total = len(messages)
|
||||
if total < processed_count:
|
||||
processed_count = total
|
||||
if total > processed_count:
|
||||
for message in messages[processed_count:total]:
|
||||
capture_step_message(message, captured, seen_ids)
|
||||
|
||||
@ -841,6 +841,60 @@ class TestAsyncExecutionPath:
|
||||
|
||||
assert [m["id"] for m in result.ai_messages] == ["ai-1", "tool-1", "tool-2", "tool-3", "ai-2"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aexecute_step_capture_survives_history_contraction(self, classes, base_config, mock_agent, msg):
|
||||
"""Regression for #3875 Phase 3: DeerFlowSummarizationMiddleware rewrites the
|
||||
messages channel mid-run via ``RemoveMessage(id=REMOVE_ALL_MESSAGES)``,
|
||||
so a later ``values`` snapshot hands the executor a SHORTER message list
|
||||
than the cursor it was tracking. Without the contraction reset in
|
||||
``capture_new_step_messages``, every step appended after the compaction
|
||||
is dropped until the list length overtakes the stale cursor.
|
||||
|
||||
Faithful to the real middleware: compaction puts the summary into a
|
||||
SEPARATE ``summary_text`` state key — the messages channel after
|
||||
compaction holds only the preserved recent tail (already-seen
|
||||
messages), NOT a synthetic summary AIMessage. So the contraction chunk
|
||||
is the already-seen tail (deduped, no new step); the real regression
|
||||
coverage is that POST-compaction growth is still captured."""
|
||||
SubagentExecutor = classes["SubagentExecutor"]
|
||||
|
||||
human = msg.human("Task")
|
||||
ai1 = msg.ai("turn one", "ai-1")
|
||||
tool1 = msg.tool("r1", "call_1", name="web_search", msg_id="tool-1")
|
||||
ai2 = msg.ai("turn two", "ai-2") # also the preserved tail after compaction
|
||||
tool2 = msg.tool("r2", "call_2", name="read_file", msg_id="tool-2")
|
||||
final = msg.ai("final answer", "ai-3")
|
||||
|
||||
chunks = [
|
||||
# Pre-compaction growth (cursor → 4).
|
||||
{"messages": [human, ai1]},
|
||||
{"messages": [human, ai1, tool1]},
|
||||
{"messages": [human, ai1, tool1, ai2]},
|
||||
# Compaction: channel rewrites to just the preserved tail (ai2) —
|
||||
# length drops from 4 to 1, below the cursor. ai2 is already seen
|
||||
# (deduped), so no new step is emitted. (The summary lives in
|
||||
# summary_text, out of channel.)
|
||||
{"messages": [ai2]},
|
||||
# Post-compaction growth — the bug: tool-2/final were dropped.
|
||||
{"messages": [ai2, tool2]},
|
||||
{"messages": [ai2, tool2, final]},
|
||||
]
|
||||
mock_agent.astream = lambda *args, **kwargs: async_iterator(chunks)
|
||||
|
||||
executor = SubagentExecutor(config=base_config, tools=[], thread_id="test-thread")
|
||||
with patch.object(executor, "_create_agent", return_value=mock_agent):
|
||||
result = await executor._aexecute("Task")
|
||||
|
||||
# Pre-compaction steps survive (ai2 not re-emitted — deduped), and
|
||||
# crucially the post-compaction tool + final answer are NOT dropped.
|
||||
assert [m["id"] for m in result.ai_messages] == [
|
||||
"ai-1",
|
||||
"tool-1",
|
||||
"ai-2",
|
||||
"tool-2",
|
||||
"ai-3",
|
||||
]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aexecute_handles_list_content(self, classes, base_config, mock_agent, msg):
|
||||
"""Test handling of list-type content in AIMessage."""
|
||||
|
||||
@ -243,6 +243,52 @@ def test_capture_new_step_messages_is_noop_on_values_reyield():
|
||||
assert len(captured) == 1
|
||||
|
||||
|
||||
def test_capture_new_step_messages_handles_history_contraction():
|
||||
# Regression for #3875 Phase 3: DeerFlowSummarizationMiddleware rewrites the
|
||||
# messages channel via RemoveMessage(id=REMOVE_ALL_MESSAGES), which shrinks
|
||||
# len(messages) below the cursor we were tracking. Without a contraction
|
||||
# reset, every step appended AFTER the compaction is dropped until total
|
||||
# overtakes the stale cursor.
|
||||
#
|
||||
# Faithful to the real middleware: compaction puts the summary into a
|
||||
# SEPARATE ``summary_text`` state key — the messages channel after
|
||||
# compaction holds only the preserved recent tail (already-seen messages),
|
||||
# NOT a synthetic summary AIMessage. So the contraction chunk is the
|
||||
# already-seen tail, deduped by id; the real regression coverage is that
|
||||
# POST-compaction growth is still captured.
|
||||
captured: list[dict] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
# Pre-compaction: a normal growing turn captures 3 steps (cursor → 4).
|
||||
ai1 = AIMessage(content="searching", id="ai-1")
|
||||
tool1 = ToolMessage(content="r1", tool_call_id="c1", name="web_search", id="tool-1")
|
||||
ai2 = AIMessage(content="done turn", id="ai-2")
|
||||
before = [HumanMessage(content="do research", id="h-1"), ai1, tool1, ai2]
|
||||
processed = capture_new_step_messages(before, captured, seen, 0)
|
||||
assert processed == 4
|
||||
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2"]
|
||||
|
||||
# Compaction rewrites the channel to just the preserved tail (ai2) —
|
||||
# length drops from 4 to 1, below the cursor. ai2 is already seen, so the
|
||||
# dedup makes it a no-op; no new step is emitted. (The summary itself lives
|
||||
# in summary_text and is never a capturable AIMessage — see step_events.py
|
||||
# INVARIANT.)
|
||||
compacted = [ai2]
|
||||
processed = capture_new_step_messages(compacted, captured, seen, processed)
|
||||
assert processed == 1
|
||||
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2"] # unchanged
|
||||
|
||||
# Post-compaction growth: a new turn appends after the preserved tail. This
|
||||
# is the bug the fix targets — without the reset, processed_count stays at
|
||||
# 4, total (3) never exceeds it, and tool-2/ai-3 are silently dropped.
|
||||
tool2 = ToolMessage(content="r2", tool_call_id="c2", name="read_file", id="tool-2")
|
||||
ai3 = AIMessage(content="final answer", id="ai-3")
|
||||
after = [ai2, tool2, ai3]
|
||||
processed = capture_new_step_messages(after, captured, seen, processed)
|
||||
assert processed == 3
|
||||
assert [c["id"] for c in captured] == ["ai-1", "tool-1", "ai-2", "tool-2", "ai-3"]
|
||||
|
||||
|
||||
def test_run_event_for_task_started():
|
||||
record = subagent_run_event({"type": "task_started", "task_id": "call_1", "description": "research X"})
|
||||
|
||||
|
||||
@ -14,6 +14,7 @@ import pytest
|
||||
from deerflow.config.subagents_config import (
|
||||
SubagentOverrideConfig,
|
||||
SubagentsAppConfig,
|
||||
default_subagent_token_budget,
|
||||
get_subagents_app_config,
|
||||
load_subagents_config_from_dict,
|
||||
)
|
||||
@ -124,6 +125,57 @@ class TestSubagentsAppConfigDefaults:
|
||||
with pytest.raises(ValueError):
|
||||
SubagentsAppConfig(max_turns=-60)
|
||||
|
||||
def test_default_token_budget_coupled_to_summarization_switch(self):
|
||||
"""The token-budget backstop engages by default (#3857 point 4). Its
|
||||
``max_tokens`` ceiling is coupled to whether subagent summarization is
|
||||
on (#3875 Phase 3 review): 1M when compaction runs, 2M otherwise —
|
||||
because Phase 2 acknowledged legitimate deep-research runs can exceed
|
||||
1M without compaction, so tightening to 1M unconditionally would
|
||||
prematurely cap summarization-off deployments."""
|
||||
# Default (no summarization): 2M — preserves Phase 2 headroom.
|
||||
budget_off = default_subagent_token_budget(summarization_enabled=False)
|
||||
assert budget_off.enabled is True
|
||||
assert budget_off.max_tokens == 2_000_000
|
||||
assert budget_off.warn_threshold == 0.7
|
||||
# Summarization on: 1M — compaction justifies the tighter ceiling.
|
||||
budget_on = default_subagent_token_budget(summarization_enabled=True)
|
||||
assert budget_on.max_tokens == 1_000_000
|
||||
# The AppConfig model-level default cannot read summarization.enabled,
|
||||
# so it falls back to the no-compaction 2M; the builder recomputes via
|
||||
# get_token_budget_for(summarization_enabled=...).
|
||||
config = SubagentsAppConfig()
|
||||
assert config.token_budget.max_tokens == 2_000_000
|
||||
assert config.token_budget.enabled is True
|
||||
|
||||
def test_get_token_budget_for_couples_default_to_summarization(self):
|
||||
"""``get_token_budget_for`` must re-couple the DEFAULT ceiling to the
|
||||
summarization switch, but a user-set budget (global or per-agent) must
|
||||
always win regardless of the switch (#3875 Phase 3 review)."""
|
||||
# Default global budget → re-coupled.
|
||||
config = SubagentsAppConfig()
|
||||
assert config.get_token_budget_for("general-purpose", summarization_enabled=True).max_tokens == 1_000_000
|
||||
assert config.get_token_budget_for("general-purpose", summarization_enabled=False).max_tokens == 2_000_000
|
||||
|
||||
def test_get_token_budget_for_respects_explicit_global(self):
|
||||
"""A user-set global ``token_budget`` is respected as-is — the
|
||||
summarization coupling only affects the default."""
|
||||
from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
|
||||
config = SubagentsAppConfig(token_budget=TokenBudgetConfig(enabled=True, max_tokens=500_000))
|
||||
# Explicit global wins for an agent with no per-agent override.
|
||||
assert config.get_token_budget_for("general-purpose", summarization_enabled=True).max_tokens == 500_000
|
||||
assert config.get_token_budget_for("general-purpose", summarization_enabled=False).max_tokens == 500_000
|
||||
|
||||
def test_get_token_budget_for_respects_per_agent_override(self):
|
||||
"""A per-agent ``token_budget`` override wins over both the default and
|
||||
the summarization coupling."""
|
||||
from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
|
||||
config = SubagentsAppConfig(
|
||||
agents={"bash": SubagentOverrideConfig(token_budget=TokenBudgetConfig(enabled=True, max_tokens=300_000))},
|
||||
)
|
||||
assert config.get_token_budget_for("bash", summarization_enabled=True).max_tokens == 300_000
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SubagentsAppConfig resolution helpers
|
||||
|
||||
@ -13,9 +13,10 @@ from langgraph.constants import TAG_NOSTREAM
|
||||
|
||||
from deerflow.agents.memory.summarization_hook import memory_flush_hook
|
||||
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY, DynamicContextMiddleware, is_dynamic_context_reminder
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummarizationEvent
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware, SummarizationEvent, create_summarization_middleware
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
from deerflow.config.memory_config import MemoryConfig
|
||||
from deerflow.config.summarization_config import SummarizationConfig
|
||||
|
||||
|
||||
def _messages() -> list:
|
||||
@ -629,3 +630,42 @@ def test_multiple_id_swap_triplets_preserve_chronological_order() -> None:
|
||||
f"{base2}__memory",
|
||||
f"{base2}__user",
|
||||
]
|
||||
|
||||
|
||||
def test_factory_attaches_memory_flush_hook_by_default(monkeypatch):
|
||||
"""The lead path keeps ``memory_flush_hook`` so pre-compaction messages
|
||||
persist into durable memory. Verified via the factory with memory enabled
|
||||
and the default ``skip_memory_flush=False``."""
|
||||
fake_model = MagicMock()
|
||||
fake_model.with_config.return_value = fake_model
|
||||
monkeypatch.setattr("deerflow.agents.middlewares.summarization_middleware.create_chat_model", lambda **kw: fake_model)
|
||||
|
||||
app_config = SimpleNamespace(
|
||||
summarization=SummarizationConfig(enabled=True),
|
||||
memory=MemoryConfig(enabled=True),
|
||||
)
|
||||
middleware = create_summarization_middleware(app_config=app_config)
|
||||
|
||||
assert middleware is not None
|
||||
assert memory_flush_hook in middleware._before_summarization_hooks
|
||||
|
||||
|
||||
def test_factory_skip_memory_flush_omits_hook(monkeypatch):
|
||||
"""``skip_memory_flush=True`` (the subagent path) must omit
|
||||
``memory_flush_hook``: subagents share the parent's ``thread_id``, so
|
||||
without skipping the hook a subagent's internal turns would flush into the
|
||||
PARENT thread's durable memory (#3875 Phase 3 review)."""
|
||||
fake_model = MagicMock()
|
||||
fake_model.with_config.return_value = fake_model
|
||||
monkeypatch.setattr("deerflow.agents.middlewares.summarization_middleware.create_chat_model", lambda **kw: fake_model)
|
||||
|
||||
app_config = SimpleNamespace(
|
||||
summarization=SummarizationConfig(enabled=True),
|
||||
memory=MemoryConfig(enabled=True),
|
||||
)
|
||||
middleware = create_summarization_middleware(app_config=app_config, skip_memory_flush=True)
|
||||
|
||||
assert middleware is not None
|
||||
# memory.enabled is True but the hook is skipped — the whole point.
|
||||
assert memory_flush_hook not in middleware._before_summarization_hooks
|
||||
assert middleware._before_summarization_hooks == []
|
||||
|
||||
@ -628,6 +628,57 @@ def test_subagent_runtime_middlewares_place_loop_detection_before_safety_finish(
|
||||
assert loop_idx < safety_idx
|
||||
|
||||
|
||||
def test_subagent_runtime_middlewares_attach_summarization_when_enabled(monkeypatch):
|
||||
"""Subagents must inherit the lead's DeerFlowSummarizationMiddleware so a
|
||||
long-running deep-research subagent compacts its context before it
|
||||
accumulates pathological input (#3875 Phase 3). Gated on the SAME
|
||||
``summarization.enabled`` switch the lead reads — one config covers both
|
||||
chains. The factory is patched at its source module so the test does not
|
||||
depend on a real chat model; only the gating + wiring is asserted."""
|
||||
from deerflow.agents.middlewares import summarization_middleware as sm
|
||||
|
||||
sentinel = object()
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_create_summarization_middleware(*, app_config=None, keep=None, skip_memory_flush=False):
|
||||
captured["app_config"] = app_config
|
||||
captured["keep"] = keep
|
||||
captured["skip_memory_flush"] = skip_memory_flush
|
||||
return sentinel
|
||||
|
||||
# summarization is enabled by default False; flip it on so the factory path
|
||||
# is taken (the factory early-returns None when disabled).
|
||||
from deerflow.config.summarization_config import SummarizationConfig
|
||||
|
||||
app_config = _make_app_config().model_copy(update={"summarization": SummarizationConfig(enabled=True)})
|
||||
monkeypatch.setattr(sm, "create_summarization_middleware", fake_create_summarization_middleware)
|
||||
_stub_runtime_middleware_imports(monkeypatch)
|
||||
|
||||
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name="test-model")
|
||||
|
||||
# The shared factory received the same app_config the builder did (no lead
|
||||
# wrapper, no config drift between the two chains).
|
||||
assert captured["app_config"] is app_config
|
||||
# skip_memory_flush=True so subagent-internal turns are not flushed into the
|
||||
# PARENT thread's durable memory (#3875 Phase 3 review).
|
||||
assert captured["skip_memory_flush"] is True
|
||||
assert sentinel in middlewares
|
||||
|
||||
|
||||
def test_subagent_runtime_middlewares_omit_summarization_when_factory_returns_none(monkeypatch):
|
||||
"""When ``summarization.enabled`` is False the shared factory returns None and
|
||||
the subagent chain must NOT carry a summarization middleware — the default
|
||||
state, since SummarizationConfig.enabled defaults to False."""
|
||||
from deerflow.agents.middlewares.summarization_middleware import DeerFlowSummarizationMiddleware
|
||||
|
||||
app_config = _make_app_config() # summarization.enabled defaults to False
|
||||
_stub_runtime_middleware_imports(monkeypatch)
|
||||
|
||||
middlewares = build_subagent_runtime_middlewares(app_config=app_config, model_name="test-model")
|
||||
|
||||
assert not any(isinstance(m, DeerFlowSummarizationMiddleware) for m in middlewares)
|
||||
|
||||
|
||||
def test_lead_runtime_chain_finds_historical_uploads_under_lazy_init_false(tmp_path, monkeypatch):
|
||||
"""Integration anchor for the ThreadData → Uploads ordering.
|
||||
|
||||
@ -680,3 +731,108 @@ def test_lead_runtime_chain_finds_historical_uploads_under_lazy_init_false(tmp_p
|
||||
assert "<uploaded_files>" in injected_content
|
||||
assert "prior-report.txt" in injected_content
|
||||
assert "previous messages" in injected_content # historical section header
|
||||
|
||||
|
||||
def test_subagent_summarization_fires_mid_run_and_produces_usable_result(monkeypatch):
|
||||
"""Integration coverage for #3875 Phase 3 review gap: drive the REAL
|
||||
``DeerFlowSummarizationMiddleware`` (the exact instance the subagent chain
|
||||
gets via ``create_summarization_middleware(skip_memory_flush=True)``) through
|
||||
a ``create_agent`` run, and assert that (a) compaction actually fires mid-run
|
||||
(messages channel contracts via ``RemoveMessage``) and (b) the run still
|
||||
completes with a usable final answer — not just wiring.
|
||||
|
||||
The builder-wiring test above proves the middleware lands on the chain; this
|
||||
proves the live middleware triggers and the run survives it. We bypass the
|
||||
full ``build_subagent_runtime_middlewares`` chain (whose sandbox/thread-data
|
||||
stubs aren't AgentMiddleware-compatible for a live run) and use the factory
|
||||
directly — the same instance the builder appends."""
|
||||
from langchain.agents import create_agent
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
from langchain_core.outputs import ChatGeneration, ChatResult
|
||||
|
||||
from deerflow.agents.middlewares.summarization_middleware import (
|
||||
DeerFlowSummarizationMiddleware,
|
||||
create_summarization_middleware,
|
||||
)
|
||||
from deerflow.agents.thread_state import ThreadState
|
||||
from deerflow.config.memory_config import MemoryConfig
|
||||
from deerflow.config.summarization_config import ContextSize, SummarizationConfig
|
||||
|
||||
# A model that always emits a plain AIMessage — no tools, so the run is a
|
||||
# single turn but the input already exceeds the trigger threshold, forcing
|
||||
# before_model compaction on the first (and only) model call.
|
||||
class _StaticModel(BaseChatModel):
|
||||
text: str = "final answer after compaction"
|
||||
|
||||
@property
|
||||
def _llm_type(self) -> str:
|
||||
return "static"
|
||||
|
||||
def bind_tools(self, tools, **kwargs):
|
||||
return self
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs):
|
||||
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=self.text))])
|
||||
|
||||
static_model = _StaticModel()
|
||||
# The factory resolves its summary model via create_chat_model; point it at
|
||||
# the same static model so no real provider is contacted.
|
||||
monkeypatch.setattr(
|
||||
"deerflow.agents.middlewares.summarization_middleware.create_chat_model",
|
||||
lambda **kwargs: static_model,
|
||||
)
|
||||
|
||||
app_config = SimpleNamespace(
|
||||
summarization=SummarizationConfig(
|
||||
enabled=True,
|
||||
trigger=ContextSize(type="messages", value=4),
|
||||
keep=ContextSize(type="messages", value=2),
|
||||
),
|
||||
# memory disabled + skip_memory_flush=True mirrors the subagent path:
|
||||
# no memory_flush_hook is attached.
|
||||
memory=MemoryConfig(enabled=False),
|
||||
)
|
||||
middleware = create_summarization_middleware(
|
||||
app_config=app_config,
|
||||
skip_memory_flush=True,
|
||||
)
|
||||
assert isinstance(middleware, DeerFlowSummarizationMiddleware), "the real middleware must be built"
|
||||
# Subagent invariant: skip_memory_flush means no durable-memory hook.
|
||||
assert not middleware._before_summarization_hooks
|
||||
|
||||
agent = create_agent(
|
||||
model=static_model,
|
||||
tools=[],
|
||||
middleware=[middleware],
|
||||
state_schema=ThreadState,
|
||||
)
|
||||
|
||||
# 6 messages > trigger(4) → compaction must fire in before_model.
|
||||
seed = [
|
||||
HumanMessage(content="q1", id="h1"),
|
||||
AIMessage(content="a1", id="a1"),
|
||||
HumanMessage(content="q2", id="h2"),
|
||||
AIMessage(content="a2", id="a2"),
|
||||
HumanMessage(content="q3", id="h3"),
|
||||
AIMessage(content="a3", id="a3"),
|
||||
]
|
||||
chunks = list(agent.stream({"messages": seed}, stream_mode="updates"))
|
||||
|
||||
# (a) Compaction fired: the middleware's before_model emitted a summary + RemoveMessage.
|
||||
before_model_chunks = [c for c in chunks if "DeerFlowSummarizationMiddleware.before_model" in c]
|
||||
assert before_model_chunks, "summarization before_model must fire when messages exceed the trigger"
|
||||
summary_update = before_model_chunks[0]["DeerFlowSummarizationMiddleware.before_model"]
|
||||
assert summary_update.get("summary_text"), "a summary must be produced"
|
||||
emitted = summary_update["messages"]
|
||||
assert isinstance(emitted[0], RemoveMessage), "compaction must lead with RemoveMessage"
|
||||
|
||||
# (b) The run completed with a usable final AIMessage despite compaction.
|
||||
# The model's output surfaces under the "model" node key in updates mode.
|
||||
final_messages: list = []
|
||||
for chunk in chunks:
|
||||
node_msg = chunk.get("model") or chunk.get("agent") or {}
|
||||
final_messages = node_msg.get("messages", final_messages)
|
||||
ai_finals = [m for m in final_messages if isinstance(m, AIMessage)]
|
||||
assert ai_finals, "the run must produce a final AIMessage after compaction"
|
||||
assert ai_finals[-1].content == "final answer after compaction"
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user