fix(middleware): prevent title middleware from streaming tokens (#3566)

This commit is contained in:
AnoobFeng 2026-06-21 22:30:35 +08:00 committed by GitHub
parent 78fff5a5e2
commit c495736f0a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 2 deletions

View File

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, NotRequired, override
from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langgraph.config import get_config
from langgraph.constants import TAG_NOSTREAM
from langgraph.runtime import Runtime
from deerflow.agents.middlewares.dynamic_context_middleware import is_dynamic_context_reminder
@ -140,7 +141,11 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]):
parent = {}
config = {**parent}
config["run_name"] = "title_agent"
config["tags"] = [*(config.get("tags") or []), "middleware:title"]
config["tags"] = [
*(config.get("tags") or []),
"middleware:title",
TAG_NOSTREAM,
]
return config
def _generate_title_result(self, state: TitleMiddlewareState) -> dict | None:

View File

@ -5,6 +5,7 @@ from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.constants import TAG_NOSTREAM
from deerflow.agents.middlewares import title_middleware as title_middleware_module
from deerflow.agents.middlewares.dynamic_context_middleware import _DYNAMIC_CONTEXT_REMINDER_KEY
@ -113,9 +114,22 @@ class TestTitleMiddlewareCoreLogic:
model.ainvoke.assert_awaited_once()
assert model.ainvoke.await_args.kwargs["config"] == {
"run_name": "title_agent",
"tags": ["middleware:title"],
"tags": ["middleware:title", TAG_NOSTREAM],
}
def test_title_model_config_preserves_parent_tags_and_adds_nostream(self, monkeypatch):
middleware = TitleMiddleware()
monkeypatch.setattr(
title_middleware_module,
"get_config",
MagicMock(return_value={"tags": ["parent"]}),
)
config = middleware._get_runnable_config()
assert config["run_name"] == "title_agent"
assert config["tags"] == ["parent", "middleware:title", TAG_NOSTREAM]
def test_generate_title_uses_explicit_app_config_without_global_config(self, monkeypatch):
title_config = TitleConfig(enabled=True, model_name="title-model", max_chars=20)
app_config = SimpleNamespace(title=title_config)