mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-23 06:28:34 +00:00
fix(middleware): prevent title middleware from streaming tokens (#3566)
This commit is contained in:
parent
78fff5a5e2
commit
c495736f0a
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user