mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-18 18:46:17 +00:00
* feat(extensions): observe task lifecycle and system model calls PR 1 (#4636) gave extensions a middleware chain, and a middleware only sees what passes through the agent graph. Two runtime surfaces stay invisible to it: when a lead run or a subagent begins and ends, and the DeerFlow-owned model calls made outside the graph. This slice adds both, with no new Gateway surface -- routers, services, and the reference extension stay in PR 3. Contract (deerflow-extension-api 0.1.1) --------------------------------------- Two contribution kinds join `middlewares` on the registry: `task_lifecycle` (`on_task_start` / `on_task_stop`, receiving a `TaskInfo` and a conservative `TaskOutcome` of completed / aborted / failed) and `system_model_observer` (`on_system_model_call`, receiving a `SystemOperationKind`, a `SystemModelRequest` snapshot, and a `SystemModelResult` carrying either the response or the provider exception plus a duration). `SystemModelRequest.messages` normalizes to a tuple at construction. Goal evaluation and memory extraction pass a message list while title generation and summarization pass one prompt string, and a bare `str` already satisfies `Sequence` -- without normalization an observer iterating `request.messages` would silently walk characters. Copying also makes the frozen snapshot immutable in fact rather than only by declaration, since observations may run after the call site returns and keeps mutating its own list. Registry marks and rollbacks become per-bucket and positional, so an `install()` that fails after registering two different kinds cannot leave one of them behind. `needs_task_store` now covers all three kinds: a deployment that registers only lifecycle hooks still gets a task store. Task lifecycle -------------- The lead worker notifies start after the run has started and stop after completion persistence and the completion hook, but before clearing the finalizing barrier and publishing the stream end -- holding the barrier across stop is what keeps a same-thread replacement run from overlapping this task's lifecycle. Cancellation raised out of the stop notification is deferred, not propagated in place, so a cancelled run still clears the barrier and emits its end frame. A subagent with a parent `run_id` wraps its execution in the same pair inside `finally`, reporting `parent_task_id` so a delegation tree is reconstructable; a subagent without a `run_id` (embedded client, standalone LangGraph Server) logs and skips rather than inventing a parent. Contributors run in registration order inside one shared 3s budget and every failure is logged and failed open. System model calls ------------------ Four kinds cover the model calls the middleware chain cannot see: goal evaluation, memory extraction, title generation, and summarization. Each site reports both terminal paths without changing the provider exception the host observes, short-circuits on `has_system_model_observers`, and passes the live task store when the runtime has one (detached work gets an isolated store). The sync summarization half stays unobserved on purpose -- it and its only host caller are the sync side of an async-only runtime, so notifying there would block a thread on a call site the host never reaches; the reason is recorded at the call site. The DeerMem backend must stay vendorable and cannot import the extension API, so it reports through a new `MemoryCallbacks.on_memory_llm_result` host hook that the DeerFlow-side callbacks translate into an observation. Notification loop ----------------- Extension resources must be touched on the loop that created them, but subagents can execute on isolated loops and DeerMem runs on a worker thread. The Gateway registers its serving loop before any runtime dependency starts and resets it last through the exit stack, so every startup-failure and cancellation path is covered. Awaited hooks raised on another loop are dispatched across with `run_coroutine_threadsafe` and awaited under the same budget; synchronous sites submit fire-and-forget work. Shutdown stops accepting detached observations before the memory flush -- that flush runs on a worker thread and can emit memory observations -- while keeping the loop alive for awaited task hooks until run and subagent drain completes. Tests ----- `test_extension_task_lifecycle.py`, `test_extension_subagent_lifecycle.py`, and `test_extension_system_model_calls.py` cover ordering, fail-open, budget exhaustion, snapshot binding under a concurrent singleton replacement, the loop-dispatch and shutdown-suspension paths, and both terminal paths at every call site. `test_gateway_run_drain_shutdown.py` pins the stop-before-barrier and drain ordering. * fix(extensions): decide notification fail-open by origin, observe cancellation `_notify_each` only guarded `Exception`, so a contributor letting a `CancelledError` escape — an extension implementing an internal timeout with cancellation, say — skipped its successors and reached the worker's deferred-interrupt path, ending an otherwise successful run as cancelled. Fail-open is about where a failure came from, not its base class: only a genuine cancellation of the host task increments `Task.cancelling()`, so propagate on that and contain everything else. `KeyboardInterrupt` / `SystemExit` still propagate. `observe_system_model_call` skipped observers on cancellation for the same base-class reason, leaving goal / title / summarization silent on a terminal path that is routine — interrupt/rollback admission and shutdown both cancel the run task, with the provider tokens already spent. Awaiting observers there is unreliable (a repeated cancel interrupts that await before any of them runs), so report through the same non-blocking submission the synchronous memory bridge uses, then propagate the cancellation untouched. DeerMem keeps `BaseException` around its provider call, now with the reason recorded: that path runs on a worker thread, where cancelling the awaiting side never interrupts the running thread, so `CancelledError` cannot arrive at all. Its host-hook wrapper narrows to `Exception` — only the hook's own failures are non-fatal, and an observability path must not swallow a process teardown signal. * fix(extensions): warn on budget exhaustion, scope observer logs by task, propagate teardown Review response on #4684: - The memory observation bridge caught BaseException, which would swallow a teardown signal raised while dispatching; it now catches Exception, matching the boundary the DeerMem-side call site documents and tests. - A notification-budget timeout raised mid-hook fell into the generic hook-failure path and logged an asyncio-internal traceback; it now logs a warning like the pre-hook budget skip, while a TimeoutError a contributor raises on its own stays classified as a hook failure. - System model observer logs passed the operation kind as the task id, so log lines said "task goal/title/..."; they now carry the task scope id alongside the kind.
815 lines
36 KiB
Python
815 lines
36 KiB
Python
"""Summarization middleware extensions for DeerFlow."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import html
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import Any, Protocol, override, runtime_checkable
|
|
|
|
from langchain.agents import AgentState
|
|
from langchain.agents.middleware import SummarizationMiddleware
|
|
from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, get_buffer_string, trim_messages
|
|
from langgraph.config import get_config
|
|
from langgraph.constants import TAG_NOSTREAM
|
|
from langgraph.graph.message import REMOVE_ALL_MESSAGES
|
|
from langgraph.runtime import Runtime
|
|
|
|
from deerflow.agents.middlewares.dynamic_context_middleware import is_dynamic_context_reminder
|
|
from deerflow.config.app_config import get_app_config
|
|
from deerflow.models import create_chat_model
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_SUMMARY_TRIGGER_MESSAGE_NAME = "summary"
|
|
_UNSET = object()
|
|
# 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.
|
|
|
|
Raised only when a caller opts in via ``raise_on_failure`` (the manual
|
|
``/compact`` path) so a real failure is reported distinctly from "nothing to
|
|
compact". The automatic path leaves ``raise_on_failure`` False and swallows the
|
|
failure, leaving compaction state unchanged for the turn.
|
|
"""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SummarizationEvent:
|
|
"""Context emitted before conversation history is summarized away."""
|
|
|
|
messages_to_summarize: tuple[AnyMessage, ...]
|
|
preserved_messages: tuple[AnyMessage, ...]
|
|
thread_id: str | None
|
|
agent_name: str | None
|
|
runtime: Runtime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ContextCompactionResult:
|
|
"""Result of summarizing old context and retaining the active tail."""
|
|
|
|
summary_text: str
|
|
messages_to_summarize: tuple[AnyMessage, ...]
|
|
preserved_messages: tuple[AnyMessage, ...]
|
|
total_tokens: int
|
|
|
|
|
|
@runtime_checkable
|
|
class BeforeSummarizationHook(Protocol):
|
|
"""Hook invoked before summarization removes messages from state."""
|
|
|
|
def __call__(self, event: SummarizationEvent) -> None: ...
|
|
|
|
|
|
def _resolve_thread_id(runtime: Runtime) -> str | None:
|
|
"""Resolve the current thread ID from runtime context or LangGraph config."""
|
|
thread_id = runtime.context.get("thread_id") if runtime.context else None
|
|
if thread_id is None:
|
|
try:
|
|
config_data = get_config()
|
|
except RuntimeError:
|
|
return None
|
|
thread_id = config_data.get("configurable", {}).get("thread_id")
|
|
return thread_id
|
|
|
|
|
|
def _resolve_agent_name(runtime: Runtime) -> str | None:
|
|
"""Resolve the current agent name from runtime context or LangGraph config."""
|
|
agent_name = runtime.context.get("agent_name") if runtime.context else None
|
|
if agent_name is None:
|
|
try:
|
|
config_data = get_config()
|
|
except RuntimeError:
|
|
return None
|
|
agent_name = config_data.get("configurable", {}).get("agent_name")
|
|
return agent_name
|
|
|
|
|
|
class DeerFlowSummarizationMiddleware(SummarizationMiddleware):
|
|
"""Summarization middleware with pre-compression hook dispatch."""
|
|
|
|
def __init__(
|
|
self,
|
|
*args,
|
|
before_summarization: list[BeforeSummarizationHook] | None = None,
|
|
app_config: Any | None = None,
|
|
configured_model_name: str | None = None,
|
|
run_model_name: str | None = None,
|
|
anchor_model_name: str | None = _UNSET, # type: ignore[assignment]
|
|
extensions=None,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self._before_summarization_hooks = before_summarization or []
|
|
# Model-ownership state. The model that actually executes the run is selected
|
|
# per run and is the authoritative source of truth, so the caller (lead /
|
|
# subagent / manual builders) supplies it directly as ``run_model_name``
|
|
# instead of the middleware re-deriving it from ``runtime.context`` /
|
|
# ``get_config()`` — those fields do not carry a custom agent's or a subagent's
|
|
# resolved model.
|
|
#
|
|
# ``configured_model_name`` is the explicitly configured summary model
|
|
# (``None`` => summarize with the run's own model). ``run_model_name`` is the
|
|
# model the run executes with; when they differ and the summary provider is
|
|
# broken (expired key, quota, outage) the run's own working model can still
|
|
# compact.
|
|
self._app_config = app_config
|
|
self._configured_summary_model_name = configured_model_name
|
|
self._run_model_name = run_model_name
|
|
# The summary LLM call runs inside a LangGraph middleware hook, so its token
|
|
# stream would otherwise be captured by the messages-tuple stream callback and
|
|
# broadcast to the frontend as a phantom AI message. Tag a dedicated model copy
|
|
# with TAG_NOSTREAM so the streaming handler skips it.
|
|
# Keep self.model untagged so the parent's profile / ls_params inspection still works.
|
|
self._summary_model = self._tag_nostream(self.model)
|
|
# ``self.model`` is the pre-built *anchor* model: it drives the parent's token
|
|
# counter / profile inspection and is reused verbatim by generation when a
|
|
# candidate matches its name. The factory builds it guarded and passes its name
|
|
# explicitly; direct construction (tests) mirrors the old factory choice
|
|
# (configured model, else default) so the passed ``model`` is the primary.
|
|
if anchor_model_name is _UNSET:
|
|
self._anchor_model_name = configured_model_name or self._default_model_name()
|
|
else:
|
|
self._anchor_model_name = anchor_model_name
|
|
if extensions is None:
|
|
from deerflow.extensions import get_agent_build_extensions
|
|
|
|
extensions = get_agent_build_extensions()
|
|
self._extensions = extensions
|
|
# Nostream generation models built lazily by name and cached (None = a build
|
|
# that failed, so a broken candidate config is not retried every turn and does
|
|
# not escape the fail-open boundary).
|
|
self._model_cache: dict[str | None, Any] = {}
|
|
|
|
def _tag_nostream(self, model: Any) -> Any:
|
|
"""Return a copy of ``model`` carrying TAG_NOSTREAM without clobbering tags.
|
|
|
|
lead_agent/agent.py binds "middleware:summarize" for RunJournal attribution;
|
|
RunnableBinding.with_config shallow-merges config, so existing tags must be
|
|
preserved explicitly instead of being overwritten with just [TAG_NOSTREAM].
|
|
"""
|
|
existing_tags = list((getattr(model, "config", None) or {}).get("tags") or [])
|
|
merged_tags = [*existing_tags, TAG_NOSTREAM] if TAG_NOSTREAM not in existing_tags else existing_tags
|
|
return model.with_config(tags=merged_tags)
|
|
|
|
def _default_model_name(self) -> str | None:
|
|
if self._app_config is None:
|
|
return None
|
|
models = getattr(self._app_config, "models", None)
|
|
return models[0].name if models else None
|
|
|
|
def _generation_candidate_names(self) -> list[str | None]:
|
|
"""Ordered summary-generation candidates by name (deduplicated).
|
|
|
|
Explicit summary model: the configured model first, then the run's own model
|
|
as a distinct fallback. ``model_name: null``: the run's own model only — its
|
|
construction is the primary, so there is no eager dependency on
|
|
``config.models[0]`` (a bare default is used only when no run model was
|
|
resolved). A ``None`` entry means "let ``create_chat_model`` pick the default",
|
|
which only occurs when nothing resolves a name.
|
|
"""
|
|
default = self._default_model_name()
|
|
if self._configured_summary_model_name is not None:
|
|
names = [self._configured_summary_model_name, self._run_model_name or default]
|
|
else:
|
|
names = [self._run_model_name or default]
|
|
deduped: list[str | None] = []
|
|
seen: set[str | None] = set()
|
|
for name in names:
|
|
if name in seen:
|
|
continue
|
|
seen.add(name)
|
|
deduped.append(name)
|
|
return deduped
|
|
|
|
def _model_for(self, name: str | None) -> Any | None:
|
|
"""The nostream summary model for ``name``, built lazily and guarded.
|
|
|
|
Returns the pre-built anchor when ``name`` matches it (no rebuild), otherwise
|
|
constructs and caches. A construction failure is caught and cached as ``None``
|
|
so a broken candidate config never escapes the fail-open boundary, is never
|
|
retried this turn, and still lets the next candidate run.
|
|
"""
|
|
if name == self._anchor_model_name:
|
|
return self._summary_model
|
|
if name in self._model_cache:
|
|
return self._model_cache[name]
|
|
try:
|
|
model = create_chat_model(
|
|
name=name,
|
|
thinking_enabled=False,
|
|
app_config=self._app_config,
|
|
attach_tracing=False,
|
|
)
|
|
built = self._tag_nostream(model.with_config(tags=["middleware:summarize"]))
|
|
except Exception:
|
|
logger.exception("Failed to build summary model %r; trying the next candidate", name)
|
|
built = None
|
|
self._model_cache[name] = built
|
|
return built
|
|
|
|
@override
|
|
def _create_summary(self, messages_to_summarize: list[AnyMessage]) -> str | None:
|
|
return self._summarize_with(messages_to_summarize)
|
|
|
|
@override
|
|
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.
|
|
"""
|
|
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.
|
|
|
|
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``).
|
|
|
|
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.
|
|
"""
|
|
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.
|
|
names = self._generation_candidate_names()
|
|
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:
|
|
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
|
|
names = self._generation_candidate_names()
|
|
for index, name in enumerate(names):
|
|
text = await self._ainvoke_summary(
|
|
self._model_for(name),
|
|
prompt,
|
|
last=index == len(names) - 1,
|
|
model_name=name,
|
|
task_store=task_store,
|
|
)
|
|
if text is not None:
|
|
return text
|
|
return None
|
|
|
|
def _invoke_summary(self, model: Any | None, prompt: str, *, last: bool = False) -> str | None:
|
|
"""Invoke ``model`` for a summary; ``None`` on error or a blank response.
|
|
|
|
Text extraction / non-empty validation runs *inside* the try: reading the
|
|
response's ``.text`` is part of consuming the provider result, so a failing
|
|
accessor must convert to a candidate failure (fall through) rather than escape
|
|
the fail-open boundary.
|
|
|
|
Deliberately unobserved by system-model-call extensions, unlike
|
|
:meth:`_ainvoke_summary`. Both this method and its only host caller
|
|
(``compact_state``) are the sync half of an async-only runtime: the agent runs
|
|
through ``abefore_model``, and ``runtime/context_compaction.py`` calls
|
|
``acompact_state``. Notifying from here would have to block the calling thread
|
|
on the extension loop for a call site the host never reaches.
|
|
"""
|
|
if model is None:
|
|
return None
|
|
try:
|
|
response = model.invoke(prompt, config={"metadata": {"lc_source": "summarization"}})
|
|
return self._checked_summary(response, last)
|
|
except Exception:
|
|
self._log_summary_error(last)
|
|
return None
|
|
|
|
async def _ainvoke_summary(
|
|
self,
|
|
model: Any | None,
|
|
prompt: str,
|
|
*,
|
|
last: bool = False,
|
|
model_name: str | None = None,
|
|
task_store=None,
|
|
) -> str | None:
|
|
"""Async counterpart of :meth:`_invoke_summary`."""
|
|
if model is None:
|
|
return None
|
|
try:
|
|
invoke_config = {"metadata": {"lc_source": "summarization"}}
|
|
extensions = getattr(self, "_extensions", None)
|
|
if extensions is None:
|
|
response = await model.ainvoke(prompt, config=invoke_config)
|
|
else:
|
|
from deerflow_extension_api import SystemOperationKind
|
|
|
|
from deerflow.extensions.notify import observe_system_model_call
|
|
|
|
response = await observe_system_model_call(
|
|
extensions,
|
|
SystemOperationKind.SUMMARIZATION,
|
|
messages=prompt,
|
|
model_name=model_name,
|
|
invoke_config=invoke_config,
|
|
invoke=lambda: model.ainvoke(prompt, config=invoke_config),
|
|
task_store=task_store,
|
|
)
|
|
return self._checked_summary(response, last)
|
|
except Exception:
|
|
self._log_summary_error(last)
|
|
return None
|
|
|
|
def _checked_summary(self, response: Any, last: bool) -> str | None:
|
|
summary = self._nonempty_summary(getattr(response, "text", None))
|
|
if summary is None:
|
|
self._log_summary_empty(last)
|
|
return summary
|
|
|
|
@staticmethod
|
|
def _log_summary_error(last: bool) -> None:
|
|
if last:
|
|
logger.exception("Summary generation failed; skipping compaction this turn")
|
|
else:
|
|
logger.warning("Summary generation failed; falling back to the run model", exc_info=True)
|
|
|
|
@staticmethod
|
|
def _log_summary_empty(last: bool) -> None:
|
|
if last:
|
|
logger.warning("Summary model returned empty text; skipping compaction this turn")
|
|
else:
|
|
logger.warning("Summary model returned empty text; falling back to the run model")
|
|
|
|
@staticmethod
|
|
def _summary_count_message(summary_text: str) -> HumanMessage:
|
|
return HumanMessage(content=summary_text, name=_SUMMARY_TRIGGER_MESSAGE_NAME)
|
|
|
|
def _messages_for_trigger_count(self, messages: list[AnyMessage], summary_text: str | None) -> list[AnyMessage]:
|
|
if not summary_text:
|
|
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)
|
|
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
|
|
except Exception:
|
|
logger.debug("Failed to trim summary prompt section with token counter; falling back to deterministic text cap", exc_info=True)
|
|
return self._bound_text(text, max_tokens)
|
|
|
|
def _build_summary_input_text(self, formatted_messages: str, previous_summary: str | None = None) -> str | 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="first",
|
|
)
|
|
else:
|
|
trimmed_previous_summary = ""
|
|
trimmed_new_messages = self._trim_summary_section_text(
|
|
formatted_messages,
|
|
max_tokens,
|
|
strategy="first",
|
|
)
|
|
|
|
# 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:
|
|
return None
|
|
return "\n".join(parts)
|
|
|
|
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)
|
|
if not trimmed_messages:
|
|
trimmed_messages = messages_to_summarize[-1:]
|
|
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)
|
|
if not formatted_messages:
|
|
return None
|
|
return self.summary_prompt.format(messages=formatted_messages).rstrip()
|
|
|
|
def before_model(self, state: AgentState, runtime: Runtime) -> dict | None:
|
|
return self._maybe_summarize(state, runtime)
|
|
|
|
async def abefore_model(self, state: AgentState, runtime: Runtime) -> dict | None:
|
|
return await self._amaybe_summarize(state, runtime)
|
|
|
|
def _prepare_compaction(
|
|
self,
|
|
state: AgentState,
|
|
*,
|
|
force: bool = False,
|
|
) -> tuple[list[AnyMessage], list[AnyMessage], str | None, int] | None:
|
|
messages = state["messages"]
|
|
self._ensure_message_ids(messages)
|
|
|
|
previous_summary = state.get("summary_text") if isinstance(state.get("summary_text"), str) else None
|
|
trigger_messages = self._messages_for_trigger_count(messages, previous_summary)
|
|
total_tokens = self.token_counter(trigger_messages)
|
|
if not force and not self._should_summarize(trigger_messages, total_tokens):
|
|
return None
|
|
|
|
cutoff_index = self._determine_cutoff_index(messages)
|
|
if cutoff_index <= 0:
|
|
return None
|
|
|
|
messages_to_summarize, preserved_messages = self._partition_messages(messages, cutoff_index)
|
|
messages_to_summarize, preserved_messages = self._preserve_dynamic_context_reminders(messages_to_summarize, preserved_messages)
|
|
if not messages_to_summarize:
|
|
return None
|
|
return messages_to_summarize, preserved_messages, previous_summary, total_tokens
|
|
|
|
def compact_state(
|
|
self,
|
|
state: AgentState,
|
|
runtime: Runtime,
|
|
*,
|
|
force: bool = False,
|
|
raise_on_failure: bool = False,
|
|
) -> ContextCompactionResult | None:
|
|
"""Summarize old context and retain the active tail.
|
|
|
|
``force`` bypasses the automatic trigger threshold (a manual caller always
|
|
wants to compact). ``raise_on_failure`` is a *separate* concern: when set (the
|
|
manual ``/compact`` path), a generation failure raises ``SummaryGenerationError``
|
|
so it can be reported distinctly from "nothing to compact"; the automatic path
|
|
leaves it False and swallows the failure, retrying on a later triggered turn.
|
|
"""
|
|
prepared = self._prepare_compaction(state, force=force)
|
|
if prepared is None:
|
|
return None
|
|
messages_to_summarize, preserved_messages, previous_summary, total_tokens = prepared
|
|
summary = self._summarize_with(messages_to_summarize, previous_summary=previous_summary)
|
|
if summary is None:
|
|
if raise_on_failure:
|
|
raise SummaryGenerationError("summary generation failed")
|
|
return None
|
|
# Fire hooks only once a replacement summary exists — flushing pre-compaction
|
|
# messages into durable memory for a summary that never materializes would
|
|
# duplicate that work on the next attempt. Messages are still removed after
|
|
# this returns (in _maybe_summarize), so hooks run before they are gone.
|
|
self._fire_hooks(messages_to_summarize, preserved_messages, runtime)
|
|
return ContextCompactionResult(
|
|
summary_text=summary,
|
|
messages_to_summarize=tuple(messages_to_summarize),
|
|
preserved_messages=tuple(preserved_messages),
|
|
total_tokens=total_tokens,
|
|
)
|
|
|
|
async def acompact_state(
|
|
self,
|
|
state: AgentState,
|
|
runtime: Runtime,
|
|
*,
|
|
force: bool = False,
|
|
raise_on_failure: bool = False,
|
|
) -> ContextCompactionResult | None:
|
|
"""Async counterpart of :meth:`compact_state` (see it for ``raise_on_failure``)."""
|
|
prepared = self._prepare_compaction(state, force=force)
|
|
if prepared is None:
|
|
return None
|
|
messages_to_summarize, preserved_messages, previous_summary, total_tokens = prepared
|
|
from deerflow_extension_api import task_store_from_runtime
|
|
|
|
summary = await self._asummarize_with(
|
|
messages_to_summarize,
|
|
previous_summary=previous_summary,
|
|
task_store=task_store_from_runtime(runtime),
|
|
)
|
|
if summary is None:
|
|
if raise_on_failure:
|
|
raise SummaryGenerationError("summary generation failed")
|
|
return None
|
|
# Fire hooks only once a replacement summary exists (see compact_state).
|
|
self._fire_hooks(messages_to_summarize, preserved_messages, runtime)
|
|
return ContextCompactionResult(
|
|
summary_text=summary,
|
|
messages_to_summarize=tuple(messages_to_summarize),
|
|
preserved_messages=tuple(preserved_messages),
|
|
total_tokens=total_tokens,
|
|
)
|
|
|
|
def _maybe_summarize(self, state: AgentState, runtime: Runtime) -> dict | None:
|
|
result = self.compact_state(state, runtime, force=False)
|
|
if result is None:
|
|
return None
|
|
return {
|
|
"messages": [
|
|
RemoveMessage(id=REMOVE_ALL_MESSAGES),
|
|
*result.preserved_messages,
|
|
],
|
|
"summary_text": result.summary_text,
|
|
}
|
|
|
|
async def _amaybe_summarize(self, state: AgentState, runtime: Runtime) -> dict | None:
|
|
result = await self.acompact_state(state, runtime, force=False)
|
|
if result is None:
|
|
return None
|
|
return {
|
|
"messages": [
|
|
RemoveMessage(id=REMOVE_ALL_MESSAGES),
|
|
*result.preserved_messages,
|
|
],
|
|
"summary_text": result.summary_text,
|
|
}
|
|
|
|
def _preserve_dynamic_context_reminders(
|
|
self,
|
|
messages_to_summarize: list[AnyMessage],
|
|
preserved_messages: list[AnyMessage],
|
|
) -> tuple[list[AnyMessage], list[AnyMessage]]:
|
|
"""Keep hidden dynamic-context reminders and their ID-swap peers out of summary compression.
|
|
|
|
These reminders carry the current date and optional memory. If summarization
|
|
removes them, DynamicContextMiddleware can lose the already-injected reminder
|
|
and inject a replacement into the wrong point of the conversation.
|
|
|
|
The ID-swap triplet produced by ``_make_reminder_and_user_messages`` contains
|
|
three messages: ``SystemMessage(id=X)`` and ``HumanMessage(id=X__memory)`` are
|
|
both tagged with ``dynamic_context_reminder=True``, but ``HumanMessage(id=X__user)``
|
|
carries the original user content and is **not** tagged. Without peer rescue,
|
|
``__user`` would stay in ``to_summarize`` and be compressed into prose — orphaning
|
|
the tagged messages and losing the user question from the model's direct context.
|
|
|
|
This method rescues tagged reminders and also rescues any untagged messages whose
|
|
``id`` shares the same ``stable_id`` prefix (i.e. ``X__user``, ``X__memory``).
|
|
"""
|
|
reminders = [msg for msg in messages_to_summarize if is_dynamic_context_reminder(msg)]
|
|
if not reminders:
|
|
return messages_to_summarize, preserved_messages
|
|
|
|
# Collect the base IDs (the stable_id prefix) from tagged reminders.
|
|
# For a reminder with id="ctx-001__memory", the base is "ctx-001".
|
|
# For a reminder with id="ctx-001" (SystemMessage), the base is "ctx-001".
|
|
# removesuffix is suffix-only — it won't strip a "__" that sits in the
|
|
# middle of a stable_id (e.g. "ctx__001" stays intact, unlike rsplit
|
|
# which would mis-derive "ctx"). Only known ID-swap suffixes (__memory,
|
|
# __user) are stripped; __user is not tagged so won't appear in reminders,
|
|
# but is included defensively.
|
|
reminder_base_ids: set[str] = set()
|
|
for msg in reminders:
|
|
if msg.id:
|
|
base = msg.id.removesuffix("__memory").removesuffix("__user")
|
|
reminder_base_ids.add(base)
|
|
|
|
# Single-pass partition: walk messages_to_summarize in chronological order
|
|
# and rescue both tagged reminders and untagged ID-swap peers (whose id
|
|
# starts with a known base + "__"). This preserves the original message
|
|
# order within rescued — critical when multiple triplets land in one
|
|
# summarization window — and eliminates the need for id(m)-based dedup
|
|
# that the previous reminders+peers concatenation required.
|
|
rescued: list[AnyMessage] = []
|
|
remaining: list[AnyMessage] = []
|
|
for msg in messages_to_summarize:
|
|
if is_dynamic_context_reminder(msg) or (msg.id and any(msg.id.startswith(b + "__") for b in reminder_base_ids)):
|
|
rescued.append(msg)
|
|
else:
|
|
remaining.append(msg)
|
|
return remaining, rescued + preserved_messages
|
|
|
|
def _fire_hooks(
|
|
self,
|
|
messages_to_summarize: list[AnyMessage],
|
|
preserved_messages: list[AnyMessage],
|
|
runtime: Runtime,
|
|
) -> None:
|
|
if not self._before_summarization_hooks:
|
|
return
|
|
|
|
event = SummarizationEvent(
|
|
messages_to_summarize=tuple(messages_to_summarize),
|
|
preserved_messages=tuple(preserved_messages),
|
|
thread_id=_resolve_thread_id(runtime),
|
|
agent_name=_resolve_agent_name(runtime),
|
|
runtime=runtime,
|
|
)
|
|
|
|
for hook in self._before_summarization_hooks:
|
|
try:
|
|
hook(event)
|
|
except Exception:
|
|
hook_name = getattr(hook, "__name__", None) or type(hook).__name__
|
|
logger.exception("before_summarization hook %s failed", hook_name)
|
|
|
|
|
|
def _build_summary_anchor(candidate_names: list[str | None], app_config: Any) -> tuple[Any | None, str | None]:
|
|
"""Build the first constructible model among ``candidate_names`` (guarded).
|
|
|
|
The returned model is tagged for RunJournal attribution but *not* TAG_NOSTREAM (the
|
|
middleware wraps a nostream copy). It becomes the parent's token-counter / profile
|
|
anchor and is reused for generation when a candidate matches its name. A per-name
|
|
construction failure is swallowed and the next candidate tried, so a broken primary
|
|
constructor neither breaks agent construction nor skips the healthy run model; a
|
|
trailing ``None`` name asks ``create_chat_model`` for its own default. Returns
|
|
``(None, None)`` when nothing can be constructed.
|
|
"""
|
|
tried: set[str | None] = set()
|
|
for name in candidate_names:
|
|
if name in tried:
|
|
continue
|
|
tried.add(name)
|
|
try:
|
|
model = create_chat_model(name=name, thinking_enabled=False, app_config=app_config, attach_tracing=False)
|
|
except Exception:
|
|
logger.exception("Failed to build summary anchor model %r; trying the next candidate", name)
|
|
continue
|
|
return model.with_config(tags=["middleware:summarize"]), name
|
|
return None, None
|
|
|
|
|
|
def create_summarization_middleware(
|
|
*,
|
|
app_config: Any | None = None,
|
|
keep: tuple[str, int | float] | None = None,
|
|
skip_memory_flush: bool = False,
|
|
run_model_name: str | None = None,
|
|
extensions=None,
|
|
) -> 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.
|
|
|
|
``run_model_name`` is the model the run actually executes with, resolved by the
|
|
caller (the lead / subagent / manual builders each already resolve it) and passed
|
|
in as the authoritative source of truth for ``model_name: null`` summarization and
|
|
the explicit-summary-model fallback. The middleware does not re-derive it from
|
|
``runtime.context`` / ``get_config()``, which do not carry a custom agent's or a
|
|
subagent's resolved model.
|
|
|
|
``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
|
|
|
|
if not config.enabled:
|
|
return None
|
|
|
|
trigger = None
|
|
if config.trigger is not None:
|
|
if isinstance(config.trigger, list):
|
|
trigger = [item.to_tuple() for item in config.trigger]
|
|
else:
|
|
trigger = config.trigger.to_tuple()
|
|
|
|
default_name = resolved_app_config.models[0].name if getattr(resolved_app_config, "models", None) else None
|
|
# Build the anchor (token-counter / profile model, reused for generation) guarded,
|
|
# rather than eagerly building the configured/default model and letting a broken
|
|
# constructor escape. Candidates in order: the primary generation model (configured
|
|
# summary model, else the run's own model), then the run model, then the default,
|
|
# then ``None`` (create_chat_model's default) as a last resort. So the null case
|
|
# builds from ``run_model_name`` — not ``config.models[0]`` — and a broken primary
|
|
# falls through to the healthy run model instead of failing agent construction.
|
|
primary_name = config.model_name or run_model_name or default_name
|
|
anchor_model, anchor_name = _build_summary_anchor(
|
|
[primary_name, run_model_name or default_name, default_name, None],
|
|
resolved_app_config,
|
|
)
|
|
if anchor_model is None:
|
|
logger.warning("Summarization is enabled but no summary model could be constructed; compaction is unavailable for this build")
|
|
return None
|
|
|
|
kwargs: dict[str, Any] = {
|
|
"model": anchor_model,
|
|
"trigger": trigger,
|
|
"keep": keep or config.keep.to_tuple(),
|
|
}
|
|
if config.trim_tokens_to_summarize is not None:
|
|
kwargs["trim_tokens_to_summarize"] = config.trim_tokens_to_summarize
|
|
if config.summary_prompt is not None:
|
|
kwargs["summary_prompt"] = config.summary_prompt
|
|
|
|
hooks: list[BeforeSummarizationHook] = []
|
|
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(
|
|
**kwargs,
|
|
before_summarization=hooks,
|
|
app_config=resolved_app_config,
|
|
configured_model_name=config.model_name,
|
|
run_model_name=run_model_name,
|
|
anchor_model_name=anchor_name,
|
|
extensions=extensions,
|
|
)
|