diff --git a/app/agent/__init__.py b/app/agent/__init__.py index bfc13eae5..3305f5df8 100644 --- a/app/agent/__init__.py +++ b/app/agent/__init__.py @@ -11,9 +11,6 @@ from typing import Any, Callable, Dict, List, Optional from fastapi.concurrency import run_in_threadpool from langchain.agents import create_agent -from langchain.agents.middleware import ( - SummarizationMiddleware, -) from langchain_core.messages import ( # noqa: F401 HumanMessage, BaseMessage, @@ -43,6 +40,9 @@ from app.agent.middleware.patch_tool_calls import PatchToolCallsMiddleware from app.agent.middleware.policy import AgentPolicyMiddleware from app.agent.middleware.runtime_config import RuntimeConfigMiddleware from app.agent.middleware.skills import SKILL_TOOL_NAME, SkillsMiddleware +from app.agent.middleware.summarization import ( + ContextPreservingSummarizationMiddleware as SummarizationMiddleware, +) from app.agent.middleware.subagents import ( SUBAGENT_CONTROL_TOOL_NAME, SUBAGENT_TASK_TOOL_NAME, diff --git a/app/agent/middleware/summarization.py b/app/agent/middleware/summarization.py new file mode 100644 index 000000000..a3d2cd70d --- /dev/null +++ b/app/agent/middleware/summarization.py @@ -0,0 +1,59 @@ +"""Agent 会话上下文压缩中间件。""" + +from langchain.agents.middleware import SummarizationMiddleware +from langchain_core.messages import AnyMessage +from langchain_core.messages.utils import get_buffer_string + + +class ContextSummarizationError(RuntimeError): + """摘要不可用且原有会话上下文未被替换。""" + + +class ContextPreservingSummarizationMiddleware(SummarizationMiddleware): + """摘要失败时中止状态更新，避免永久丢失既有会话上下文。""" + + _ERROR_MESSAGE = "会话上下文压缩失败，原有上下文已保留，请稍后重试" + _UNSUMMARIZABLE_MESSAGE = ( + "会话历史中存在无法压缩的超长内容，原有上下文已保留，" + "请新建或清空会话后继续" + ) + + @classmethod + def _require_valid_summary(cls, summary: str) -> str: + """拒绝无法继续承载会话上下文的空摘要。""" + if not summary: + raise ContextSummarizationError(cls._ERROR_MESSAGE) + return summary + + def _prepare_summary_input( + self, messages_to_summarize: list[AnyMessage] + ) -> str: + """复用 LangChain 裁剪策略生成摘要模型输入。""" + trimmed_messages = self._trim_messages_for_summary(messages_to_summarize) + if not trimmed_messages: + raise ContextSummarizationError(self._UNSUMMARIZABLE_MESSAGE) + return get_buffer_string(trimmed_messages, format="xml") + + def _create_summary(self, messages_to_summarize: list[AnyMessage]) -> str: + """同步摘要失败时保持原图状态。""" + formatted_messages = self._prepare_summary_input(messages_to_summarize) + try: + response = self.model.invoke( + self.summary_prompt.format(messages=formatted_messages).rstrip(), + config={"metadata": {"lc_source": "summarization"}}, + ) + except Exception as err: + raise ContextSummarizationError(self._ERROR_MESSAGE) from err + return self._require_valid_summary(response.text.strip()) + + async def _acreate_summary(self, messages_to_summarize: list[AnyMessage]) -> str: + """异步摘要失败时保持原图状态。""" + formatted_messages = self._prepare_summary_input(messages_to_summarize) + try: + response = await self.model.ainvoke( + self.summary_prompt.format(messages=formatted_messages).rstrip(), + config={"metadata": {"lc_source": "summarization"}}, + ) + except Exception as err: + raise ContextSummarizationError(self._ERROR_MESSAGE) from err + return self._require_valid_summary(response.text.strip()) diff --git a/tests/test_agent_summarization_streaming.py b/tests/test_agent_summarization_streaming.py index a7a4afacb..c38c3e46f 100644 --- a/tests/test_agent_summarization_streaming.py +++ b/tests/test_agent_summarization_streaming.py @@ -1,10 +1,18 @@ import asyncio -from unittest.mock import patch +import uuid +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch -from langchain.agents.middleware import SummarizationMiddleware +import pytest +from langchain_core.messages import AIMessage, HumanMessage import app.agent as agent_module +from app.agent.memory import memory_manager from app.agent.middleware.runtime_config import RuntimeConfigMiddleware +from app.agent.middleware.summarization import ( + ContextSummarizationError, + ContextPreservingSummarizationMiddleware, +) class _FakeLLM: @@ -15,6 +23,40 @@ class _FakeLLM: self.profile = {"max_input_tokens": 64000} +class _FailingSummaryLLM(_FakeLLM): + """模拟摘要模型暂时不可用。""" + + async def ainvoke(self, *_args, **_kwargs): + """摘要请求始终超时。""" + raise TimeoutError("summary provider unavailable") + + def invoke(self, *_args, **_kwargs): + """同步摘要请求始终超时。""" + raise TimeoutError("summary provider unavailable") + + +class _SuccessfulSummaryLLM(_FakeLLM): + """提供稳定摘要结果。""" + + async def ainvoke(self, *_args, **_kwargs): + """返回异步摘要。""" + return AIMessage(content="保留旧事实的摘要") + + def invoke(self, *_args, **_kwargs): + """返回同步摘要。""" + return AIMessage(content="保留旧事实的摘要") + + +class _FailingGraph: + """模拟上下文压缩阶段失败的 Agent 图。""" + + async def ainvoke(self, _payload, config=None): + """在提交图状态前终止执行。""" + raise ContextSummarizationError( + "会话上下文压缩失败，原有上下文已保留，请稍后重试" + ) + + def test_streaming_agent_uses_non_streaming_llm_for_summary(): """流式 Agent 的摘要中间件应使用非流式 LLM。""" agent = agent_module.MoviePilotAgent(session_id="session-1", user_id="10001") @@ -46,7 +88,7 @@ def test_streaming_agent_uses_non_streaming_llm_for_summary(): summary_middleware = next( middleware for middleware in captured["middleware"] - if isinstance(middleware, SummarizationMiddleware) + if isinstance(middleware, ContextPreservingSummarizationMiddleware) ) assert captured["model"] is main_llm @@ -169,13 +211,137 @@ def test_non_streaming_agent_reuses_main_llm_for_summary(): summary_middleware = next( middleware for middleware in captured["middleware"] - if isinstance(middleware, SummarizationMiddleware) + if isinstance(middleware, ContextPreservingSummarizationMiddleware) ) assert captured["model"] is main_llm assert summary_middleware.model is main_llm +def test_summary_failure_does_not_replace_existing_context(): + """摘要模型失败时应中止压缩，避免错误文本替换既有上下文。""" + agent = agent_module.MoviePilotAgent(session_id="session-1", user_id="10001") + failing_llm = _FailingSummaryLLM("summary") + captured: dict = {} + + def _fake_create_agent(**kwargs): + """捕获 create_agent 参数。""" + captured.update(kwargs) + return object() + + with ( + patch.object(agent, "_initialize_llm", return_value=failing_llm), + patch.object(agent, "_initialize_tools", return_value=[]), + patch.object( + agent_module.prompt_manager, "get_agent_prompt", return_value="prompt" + ), + patch.object( + agent_module, "create_subagent_middlewares", return_value=([], []) + ), + patch.object(agent_module, "create_agent", side_effect=_fake_create_agent), + patch.object(agent_module.settings, "LLM_MAX_TOOLS", 0), + ): + asyncio.run(agent._create_agent(streaming=False)) + + summary_middleware = next( + middleware + for middleware in captured["middleware"] + if isinstance(middleware, ContextPreservingSummarizationMiddleware) + ) + messages = [ + HumanMessage(content=f"必须保留的旧上下文 {index} " * 200) + for index in range(160) + ] + assert summary_middleware.token_counter(messages) >= 64000 * 0.85 + cutoff_index = summary_middleware._determine_cutoff_index(messages) + assert summary_middleware._trim_messages_for_summary(messages[:cutoff_index]) + + with pytest.raises(RuntimeError, match="会话上下文压缩失败"): + asyncio.run(summary_middleware.abefore_model({"messages": messages}, None)) + + with pytest.raises(RuntimeError, match="会话上下文压缩失败"): + summary_middleware.before_model({"messages": messages}, None) + + +def test_summary_success_still_replaces_old_context(): + """摘要成功时仍应保留上游的压缩行为。""" + middleware = ContextPreservingSummarizationMiddleware( + model=_SuccessfulSummaryLLM("summary"), + trigger=("messages", 25), + ) + messages = [HumanMessage(content=f"消息 {index}") for index in range(25)] + + update = asyncio.run(middleware.abefore_model({"messages": messages}, None)) + + assert update is not None + assert "保留旧事实的摘要" in update["messages"][1].content + assert len(update["messages"]) < len(messages) + + +def test_unsummarizable_message_requires_new_context_instead_of_retry(): + """确定性不可裁剪的消息应给出可前进路径，而非建议无效重试。""" + middleware = ContextPreservingSummarizationMiddleware( + model=_SuccessfulSummaryLLM("summary"), + trigger=("messages", 21), + keep=("messages", 20), + trim_tokens_to_summarize=1000, + ) + messages = [ + HumanMessage(content="无法裁剪的单条超长消息" * 4000), + *[HumanMessage(content=f"后续消息 {index}") for index in range(20)], + ] + assert not middleware._trim_messages_for_summary(messages[:1]) + + errors = [] + for _ in range(2): + with pytest.raises(ContextSummarizationError) as error: + middleware.before_model({"messages": messages}, None) + errors.append(str(error.value)) + + assert errors == [ + "会话历史中存在无法压缩的超长内容，原有上下文已保留，请新建或清空会话后继续" + ] * 2 + assert all("稍后重试" not in error for error in errors) + + +def test_summary_failure_preserves_database_history(): + """上下文压缩失败时不得覆盖数据库中的上一轮消息。""" + session_id = f"summary-failure-{uuid.uuid4().hex}" + user_id = "10001" + memory_manager.save_agent_messages( + session_id=session_id, + user_id=user_id, + messages=[HumanMessage(content="数据库中的旧事实")], + ) + memory_manager.clear_memory(session_id, user_id) + restored_messages = memory_manager.get_agent_messages(session_id, user_id) + agent = agent_module.MoviePilotAgent(session_id=session_id, user_id=user_id) + agent._compiled_agent_bundle = object() + agent._should_stream = lambda: False + agent._create_agent = AsyncMock(return_value=_FailingGraph()) + agent.stream_handler = SimpleNamespace( + stop_streaming=AsyncMock(return_value=(False, "")) + ) + agent.send_agent_message = AsyncMock() + + with ( + patch("app.agent.eventmanager.send_event") as send_usage_event, + ): + result, _ = asyncio.run( + agent._execute_agent( + [*restored_messages, HumanMessage(content="继续原来的任务")] + ) + ) + + memory_manager.clear_memory(session_id, user_id) + recovered_messages = memory_manager.get_agent_messages(session_id, user_id) + assert result == "智能助手执行失败: 会话上下文压缩失败，原有上下文已保留，请稍后重试" + assert agent._compiled_agent_bundle is None + assert [message.content for message in recovered_messages] == ["数据库中的旧事实"] + send_usage_event.assert_called_once() + assert not send_usage_event.call_args.args[1].success + + def test_agent_uses_runtime_config_middleware_instead_of_hooks(): """Agent 应使用运行时配置中间件而不是旧 hooks。""" agent = agent_module.MoviePilotAgent(session_id="session-1", user_id="10001")