diff --git a/CHANGELOG.md b/CHANGELOG.md index 3afabcf5c7..f1e7715009 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,6 +47,10 @@ This section accumulates work toward the **2.3.0** milestone ### Fixed +- **summarization:** vLLM compaction counts reasoning tokens replayed to the + model when checking the trigger and estimating how much history to keep. + ([#8077](https://github.com/bytedance/deer-flow/pull/8077)) + - **gateway:** Posting feedback twice for the same run now returns a 409 instead of a 500. `POST /api/threads/{id}/runs/{run_id}/feedback` allows one feedback row per user and run; a retry, or feedback on another message of the diff --git a/README.md b/README.md index d76616f45e..aefea9965d 100644 --- a/README.md +++ b/README.md @@ -2737,7 +2737,8 @@ A pass is limited to 64 chunks and leaves history unchanged on failure or limit exhaustion. A limit breach starts a five-minute per-process cooldown for the same thread/history window and policy, including rebuilt middleware. A larger input budget or smaller source window permits an immediate retry. Model fallback -and SDK retries can add calls. See the +and SDK retries can add calls. For vLLM run models, history and retention +estimates include assistant reasoning replayed to the provider. See the [summarization guide](backend/docs/summarization.md). With `task_continuity.enabled`, `history_search` searches the current task's active diff --git a/backend/docs/summarization.md b/backend/docs/summarization.md index a5d937e392..916ea7760d 100644 --- a/backend/docs/summarization.md +++ b/backend/docs/summarization.md @@ -282,6 +282,11 @@ The default LangChain prompt instructs the model to: - Uses approximate token counting based on character count - For Anthropic models: ~3.3 characters per token - For other models: Uses LangChain's default estimation +- For `VllmChatModel` using Chat Completions: History and token-based retention + estimates include the assistant reasoning replayed to vLLM. A non-null `reasoning` field takes + precedence over `reasoning_content`, including when it is empty. This applies + to both automatic and manual compaction. The estimate is approximate and + leaves the summary prompt and its trimming budget unchanged. - Can be customized with a custom `token_counter` function ### Message Preservation diff --git a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md index c0605afa37..70e433a0bc 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md +++ b/backend/packages/harness/deerflow/agents/middlewares/AGENTS.md @@ -25,6 +25,8 @@ PII before splitting, budgets before escaping; hooks/history/removal on success. Chunk/carry/cooldown invariants and tests: [contract](../../../../../docs/summarization.md). +vLLM compaction counts replayed reasoning. Keep it out of summary input. + Todo compaction: [contract](../../../../../docs/summarization.md#todo-reminders). Delegation verdicts are untrusted: revalidate persisted values, ignore malformed diff --git a/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py index 34e5ef62a5..bef9ccba0f 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py @@ -10,6 +10,7 @@ import threading from collections import OrderedDict from collections.abc import Mapping from dataclasses import dataclass +from functools import partial from time import monotonic from typing import Any, Protocol, override, runtime_checkable @@ -17,6 +18,7 @@ from deerflow_extension_api import CompactionEvent, ExtensionData, canonical_has from langchain.agents import AgentState from langchain.agents.middleware import SummarizationMiddleware from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, SystemMessage, get_buffer_string +from langchain_core.runnables import RunnableBinding from langgraph.config import get_config from langgraph.constants import TAG_NOSTREAM from langgraph.graph.message import REMOVE_ALL_MESSAGES @@ -32,6 +34,7 @@ from deerflow.config.task_continuity_config import TaskContinuityConfig from deerflow.extensions.notify import notify_context_compacted from deerflow.models import create_chat_model from deerflow.models.reasoning import resolve_reasoning_contract +from deerflow.models.vllm_provider import VllmChatModel, count_tokens_with_replayed_reasoning from deerflow.runtime.events.catalog import MIDDLEWARE_SUMMARIZE_TAG from deerflow.utils.assembly_io import run_assembly @@ -177,6 +180,15 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): **kwargs, ) -> None: super().__init__(*args, **kwargs) + run_model = self.model + seen_bindings: set[int] = set() + while isinstance(run_model, RunnableBinding) and id(run_model) not in seen_bindings: + seen_bindings.add(id(run_model)) + run_model = run_model.bound + if "token_counter" not in kwargs and isinstance(run_model, VllmChatModel) and run_model.replays_reasoning(): + self.token_counter = partial(count_tokens_with_replayed_reasoning, token_counter=self.token_counter) + self._partial_token_counter = partial(count_tokens_with_replayed_reasoning, token_counter=self._partial_token_counter) + # Preserve LangChain's unscaled suffix estimate. self._task_continuity_config = task_continuity_config if task_continuity_config is not None and task_continuity_config.enabled is True else None self._before_summarization_hooks = before_summarization or [] # Model-ownership state. The model that actually executes the run is selected @@ -1350,8 +1362,6 @@ def create_summarization_middleware( hooks: list[BeforeSummarizationHook] = [] if resolved_app_config.memory.enabled and not skip_memory_flush: - from functools import partial - from deerflow.agents.memory.summarization_hook import memory_flush_hook hooks.append(partial(memory_flush_hook, pii_redaction_config=getattr(resolved_app_config, "pii_redaction", None))) diff --git a/backend/packages/harness/deerflow/models/vllm_provider.py b/backend/packages/harness/deerflow/models/vllm_provider.py index 0a3ebf094f..a8db33fb81 100644 --- a/backend/packages/harness/deerflow/models/vllm_provider.py +++ b/backend/packages/harness/deerflow/models/vllm_provider.py @@ -23,7 +23,7 @@ import json import threading import time from collections import OrderedDict -from collections.abc import Mapping +from collections.abc import Callable, Iterable, Mapping from typing import Any, cast import openai @@ -31,12 +31,15 @@ from langchain_core.language_models import LanguageModelInput from langchain_core.messages import ( AIMessage, AIMessageChunk, + BaseMessage, BaseMessageChunk, ChatMessageChunk, FunctionMessageChunk, HumanMessageChunk, + MessageLikeRepresentation, SystemMessageChunk, ToolMessageChunk, + convert_to_messages, ) from langchain_core.messages.ai import UsageMetadata, subtract_usage from langchain_core.messages.tool import tool_call_chunk @@ -112,6 +115,24 @@ def _pick_reasoning(fields: Mapping[str, Any]) -> Any: return fields.get("reasoning_content") if reasoning is None else reasoning +def count_tokens_with_replayed_reasoning( + messages: Iterable[MessageLikeRepresentation], + *, + token_counter: Callable[[Iterable[MessageLikeRepresentation]], int], +) -> int: + counted: list[BaseMessage] = [] + for message in convert_to_messages(messages): + if isinstance(message, AIMessage): + reasoning = _pick_reasoning(message.additional_kwargs) + text = _reasoning_to_text(reasoning) if reasoning is not None else "" + if text: + content = message.content + text if isinstance(message.content, str) else [*message.content, text] + message = message.model_copy(update={"content": content}) + # Count replayed reasoning on a copy. + counted.append(message) + return token_counter(counted) + + def _convert_delta_to_message_chunk_with_reasoning(_dict: Mapping[str, Any], default_class: type[BaseMessageChunk]) -> BaseMessageChunk: """Convert a streaming delta to a LangChain message chunk while preserving reasoning.""" id_ = _dict.get("id") @@ -204,6 +225,9 @@ class VllmChatModel(GuardedChatOpenAI): def _llm_type(self) -> str: return "vllm-openai-compatible" + def replays_reasoning(self) -> bool: + return not self._use_responses_api(self._default_params) # Chat Completions only. + def _usage_delta(self, completion_id: str, usage: UsageMetadata, *, terminal: bool) -> UsageMetadata: """Convert a completion's cumulative usage snapshot into a delta.""" with self._cumulative_usage_lock: diff --git a/backend/tests/test_summarization_middleware.py b/backend/tests/test_summarization_middleware.py index ead7a9bb5a..242eeb3099 100644 --- a/backend/tests/test_summarization_middleware.py +++ b/backend/tests/test_summarization_middleware.py @@ -10,7 +10,9 @@ import pytest from langchain.agents import create_agent from langchain_core.language_models import BaseChatModel from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage, SystemMessage, ToolMessage +from langchain_core.messages.utils import count_tokens_approximately from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_core.runnables import RunnableBinding from langgraph.constants import TAG_NOSTREAM from pydantic import ValidationError @@ -23,6 +25,7 @@ from deerflow.config.memory_config import MemoryConfig from deerflow.config.model_config import ModelConfig, ReasoningCapabilities from deerflow.config.sandbox_config import SandboxConfig from deerflow.config.summarization_config import ContextSize, SummarizationConfig +from deerflow.models.vllm_provider import VllmChatModel def _messages() -> list: @@ -1207,6 +1210,107 @@ def _factory_app_config(model_names, *, summary_model_name=None, summarization_k ) +@pytest.fixture +def vllm_compaction(monkeypatch): + models = { + "vllm": VllmChatModel(model="test-model", api_key="dummy", base_url="http://localhost:8000/v1"), + "summary": _StaticChatModel(text="compressed summary"), + } + monkeypatch.setattr( + "deerflow.agents.middlewares.summarization_middleware.create_chat_model", + lambda name, **_: models[name], + ) + config = _factory_app_config( + ("vllm", "summary"), + summary_model_name="summary", + summarization_kwargs={"trigger": ContextSize(type="tokens", value=150)}, + ) + middleware = create_summarization_middleware(app_config=config, run_model_name="vllm", keep=("tokens", 100)) + assert middleware is not None + return middleware + + +_REPLAYED_REASONING = "Review the tool results. " * 100 + + +@pytest.mark.parametrize("nested", [False, True], ids=["single-binding", "double-binding"]) +def test_vllm_compaction_counts_reasoning_through_bindings(vllm_compaction, nested): + model = vllm_compaction.model + if nested: + model = RunnableBinding(bound=model) + middleware = DeerFlowSummarizationMiddleware(model=model) + message = AIMessage(content="Checked.", additional_kwargs={"reasoning": _REPLAYED_REASONING}) + counted = AIMessage(content="Checked." + _REPLAYED_REASONING) + + assert middleware.token_counter([message]) == count_tokens_approximately([counted]) + + +@pytest.mark.parametrize( + ("fields", "replayed"), + [ + pytest.param({}, None, id="ordinary-message"), + pytest.param({"reasoning": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="reasoning-only"), + pytest.param({"reasoning_content": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="legacy-only"), + pytest.param({"reasoning": None, "reasoning_content": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="null-canonical"), + pytest.param({"reasoning": _REPLAYED_REASONING, "reasoning_content": _REPLAYED_REASONING}, _REPLAYED_REASONING, id="duplicate-alias"), + pytest.param({"reasoning": _REPLAYED_REASONING, "reasoning_content": "different text" * 500}, _REPLAYED_REASONING, id="canonical-wins"), + pytest.param({"reasoning": "", "reasoning_content": _REPLAYED_REASONING}, "", id="empty-canonical"), + ], +) +def test_vllm_compaction_counts_replayed_reasoning_once(vllm_compaction, fields, replayed): + message = AIMessage( + content="Checked.", + additional_kwargs=fields, + tool_calls=[{"name": "check", "args": {}, "id": "call-1"}], + ) + payload = vllm_compaction.model._get_request_payload([message]) + assert payload["messages"][0].get("reasoning") == replayed + assert "reasoning_content" not in payload["messages"][0] + + counted = message.model_copy(update={"content": message.content + (replayed or "")}) + assert vllm_compaction.token_counter([message]) == count_tokens_approximately([counted]) + assert message.content == "Checked." + assert message.additional_kwargs == fields + + +@pytest.mark.asyncio +@pytest.mark.parametrize("manual", [False, True], ids=["auto", "manual"]) +async def test_vllm_compaction_uses_replayed_reasoning_for_trigger_and_retention(vllm_compaction, manual): + messages = _messages() + plain_state = {"messages": messages} + assert await vllm_compaction.abefore_model(plain_state, _runtime()) is None + messages[1] = messages[1].model_copy(update={"additional_kwargs": {"reasoning": _REPLAYED_REASONING}}) + state = {"messages": messages} + + if manual: + result = await vllm_compaction.acompact_state(state, _runtime(), force=True) + assert result is not None + assert result.total_tokens >= 150 + assert list(result.preserved_messages) == messages[2:] + assert result.summary_text == "compressed summary" + else: + update = await vllm_compaction.abefore_model(state, _runtime()) + assert update is not None + assert update["messages"][1:] == messages[2:] + assert update["summary_text"] == "compressed summary" + + +def test_non_vllm_compaction_keeps_existing_token_estimate(): + middleware = DeerFlowSummarizationMiddleware(model=_StaticChatModel()) + message = AIMessage(content="Checked.", additional_kwargs={"reasoning": _REPLAYED_REASONING}) + assert middleware.token_counter([message]) == count_tokens_approximately([message]) + + +def test_vllm_compaction_ignores_reasoning_when_using_responses_api(vllm_compaction): + model = vllm_compaction.model.bound.model_copy(update={"use_responses_api": True}) + middleware = DeerFlowSummarizationMiddleware(model=model) + message = AIMessage(content="Checked.", additional_kwargs={"reasoning": _REPLAYED_REASONING}) + payload = model._get_request_payload([message]) + assert "input" in payload + assert _REPLAYED_REASONING not in repr(payload["input"]) + assert middleware.token_counter([message]) == count_tokens_approximately([message]) + + @pytest.mark.parametrize("trim_limit", [None, 80, 4000]) def test_factory_preserves_explicit_summary_input_limit(monkeypatch, trim_limit): monkeypatch.setattr("deerflow.agents.middlewares.summarization_middleware.create_chat_model", _tracking_create_chat_model([])) diff --git a/backend/tests/test_vllm_provider.py b/backend/tests/test_vllm_provider.py index b1002c0beb..539e69c394 100644 --- a/backend/tests/test_vllm_provider.py +++ b/backend/tests/test_vllm_provider.py @@ -6,8 +6,9 @@ import json import httpx import pytest from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage +from langchain_core.messages.utils import count_tokens_approximately -from deerflow.models.vllm_provider import VllmChatModel +from deerflow.models.vllm_provider import VllmChatModel, count_tokens_with_replayed_reasoning def _make_model(*, cumulative_stream_usage: bool = False) -> VllmChatModel: @@ -62,6 +63,23 @@ def _assert_usage(message, *, input_tokens: int, output_tokens: int, total_token assert message.usage_metadata["total_tokens"] == total_tokens +def test_count_tokens_with_replayed_reasoning_preserves_content_blocks(): + blocks = [ + {"type": "text", "text": "Checked the source."}, + {"type": "text", "text": "The output matches."}, + ] + reasoning = "Compare each output field with its source." + message = AIMessage(content=blocks, additional_kwargs={"reasoning": reasoning}) + original = message.model_copy(deep=True) + + count = count_tokens_with_replayed_reasoning([message], token_counter=count_tokens_approximately) + + assert count == count_tokens_approximately([AIMessage(content=[*blocks, reasoning])]) + assert count > count_tokens_approximately([original]) + assert count > count_tokens_approximately([AIMessage(content=[reasoning])]) + assert message == original + + def test_vllm_provider_restores_reasoning_in_request_payload(): model = _make_model() payload = model._get_request_payload(