"""Agent 会话上下文压缩中间件。""" from collections.abc import Awaitable, Callable from importlib import import_module from typing import Any from langchain.agents.middleware.summarization import ( DEFAULT_SUMMARY_PROMPT, ContextSize, SummarizationMiddleware, TokenCounter, TriggerClause, ) from langchain.agents.middleware.types import ( AgentMiddleware, ExtendedModelResponse, ModelRequest, ModelResponse, ) from langchain.chat_models import BaseChatModel from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, ToolMessage from langchain_core.messages.utils import count_tokens_approximately, get_buffer_string from langgraph.graph.message import REMOVE_ALL_MESSAGES from langgraph.types import Command from app.agent.middleware.usage import UsageMiddleware from app.runtime.log import logger try: _internal_call_metadata = import_module( "langchain.agents.middleware.internal_call_transformer" ).internal_call_metadata except ImportError: def _internal_call_metadata() -> dict[str, Any]: """旧版 LangChain 没有内部模型调用的流式过滤标记。""" return {} class ContextSummarizationError(RuntimeError): """摘要不可用且原有会话上下文未被替换。""" class ContextPreservingSummarizationMiddleware(SummarizationMiddleware): """摘要失败时中止状态更新，避免永久丢失既有会话上下文。""" _ERROR_MESSAGE = "会话上下文压缩失败，原有上下文已保留，请稍后重试" _UNSUMMARIZABLE_MESSAGE = ( "会话历史中存在无法压缩的超长内容，原有上下文已保留，" "请新建或清空会话后继续" ) # LangChain 默认按 4000 token 裁剪待摘要消息，单条超长工具结果或用户输入 # （超长无换行文本、非文本多模态块）会被整体丢弃，直接触发"无法压缩"报错。 # 调大该上限可容纳更大单条消息，减少误报；16k 相对常见模型窗口仍然安全， # 摘要模型与主模型同窗口，过大会抬高摘要成本并挤占主模型预算。 _DEFAULT_TRIM_TOKENS_TO_SUMMARIZE = 16000 def __init__( self, model: str | BaseChatModel, *, trigger: ( ContextSize | TriggerClause | list[ContextSize | TriggerClause] | None ) = None, keep: ContextSize = ("messages", 20), token_counter: TokenCounter = count_tokens_approximately, summary_prompt: str = DEFAULT_SUMMARY_PROMPT, trim_tokens_to_summarize: int | None = _DEFAULT_TRIM_TOKENS_TO_SUMMARIZE, **deprecated_kwargs: Any, ) -> None: """按 MoviePilot 的默认压缩策略构建摘要中间件。""" super().__init__( model=model, trigger=trigger, keep=keep, token_counter=token_counter, summary_prompt=summary_prompt, trim_tokens_to_summarize=trim_tokens_to_summarize, **deprecated_kwargs, ) @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) summary_model = getattr(self, "_summary_model", self.model) try: response = summary_model.invoke( self.summary_prompt.format(messages=formatted_messages).rstrip(), config={ "metadata": { "lc_source": "summarization", **_internal_call_metadata(), } }, ) 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) summary_model = getattr(self, "_summary_model", self.model) try: response = await summary_model.ainvoke( self.summary_prompt.format(messages=formatted_messages).rstrip(), config={ "metadata": { "lc_source": "summarization", **_internal_call_metadata(), } }, ) except Exception as err: raise ContextSummarizationError(self._ERROR_MESSAGE) from err return self._require_valid_summary(response.text.strip()) def partition_for_token_limit( self, messages: list[AnyMessage], token_limit: int, *, force: bool = False, minimum_cutoff: int = 1, strict_token_limit: bool = False, ) -> tuple[list[AnyMessage], list[AnyMessage]] | None: """按 token 上限拆分历史，并保持 LangChain 的工具调用事务边界。""" self._ensure_message_ids(messages) if self.token_counter(messages) <= token_limit: return ( self._minimum_safe_partition( messages, minimum_cutoff=minimum_cutoff, token_limit=token_limit if strict_token_limit else None, ) if force else None ) left, right = 0, len(messages) cutoff_candidate = len(messages) while left < right: midpoint = (left + right) // 2 if self._partial_token_counter(messages[midpoint:]) <= token_limit: cutoff_candidate = midpoint right = midpoint else: left = midpoint + 1 if cutoff_candidate >= len(messages): cutoff_candidate = len(messages) cutoff_index = self._find_safe_cutoff_point(messages, cutoff_candidate) if ( cutoff_index <= 0 or cutoff_index >= len(messages) or cutoff_index < minimum_cutoff or not self._contains_unsummarized_message(messages[:cutoff_index]) or ( strict_token_limit and self._partial_token_counter(messages[cutoff_index:]) > token_limit ) ): if cutoff_candidate >= len(messages): if strict_token_limit: return None return self._latest_safe_partition( messages, minimum_cutoff=minimum_cutoff, ) return self._minimum_safe_partition( messages, minimum_cutoff=max(minimum_cutoff, cutoff_candidate), token_limit=token_limit if strict_token_limit else None, ) return self._partition_messages(messages, cutoff_index) def partition_for_retention( self, messages: list[AnyMessage] ) -> tuple[list[AnyMessage], list[AnyMessage]] | None: """按摘要器既有触发和保留策略拆分历史。""" self._ensure_message_ids(messages) total_tokens = self.token_counter(messages) if not self._should_summarize(messages, total_tokens): return None cutoff_index = self._determine_cutoff_index(messages) if cutoff_index <= 0: return None return self._partition_messages(messages, cutoff_index) def _minimum_safe_partition( self, messages: list[AnyMessage], *, minimum_cutoff: int = 1, token_limit: int | None = None, ) -> tuple[list[AnyMessage], list[AnyMessage]] | None: """至少摘要一段旧历史，同时保留最新完整消息事务。""" for candidate in range(max(1, minimum_cutoff), len(messages)): cutoff_index = self._find_safe_cutoff_point(messages, candidate) if ( minimum_cutoff <= cutoff_index < len(messages) and self._contains_unsummarized_message(messages[:cutoff_index]) and ( token_limit is None or self._partial_token_counter(messages[cutoff_index:]) <= token_limit ) ): return self._partition_messages(messages, cutoff_index) return None def _latest_safe_partition( self, messages: list[AnyMessage], *, minimum_cutoff: int, ) -> tuple[list[AnyMessage], list[AnyMessage]] | None: """保留无法满足软预算时的最新完整消息事务。""" for candidate in range(len(messages) - 1, minimum_cutoff - 1, -1): cutoff_index = self._find_safe_cutoff_point(messages, candidate) if ( minimum_cutoff <= cutoff_index < len(messages) and self._contains_unsummarized_message(messages[:cutoff_index]) ): return self._partition_messages(messages, cutoff_index) return None @staticmethod def _contains_unsummarized_message(messages: list[AnyMessage]) -> bool: """确认待摘要段包含可推进上下文的原始消息。""" return any( message.additional_kwargs.get("lc_source") != "summarization" for message in messages ) def build_summary_messages(self, summary: str) -> list[AnyMessage]: """将摘要转换为 LangChain 约定的可识别历史消息。""" return self._build_new_messages(summary) def ensure_message_ids(self, messages: list[AnyMessage]) -> None: """为压缩后消息补齐 LangGraph reducer 所需的稳定 ID。""" self._ensure_message_ids(messages) def create_summary(self, messages_to_summarize: list[AnyMessage]) -> str: """通过 MoviePilot 的失败保护合同生成同步摘要。""" return self._create_summary(messages_to_summarize) async def acreate_summary(self, messages_to_summarize: list[AnyMessage]) -> str: """通过 MoviePilot 的失败保护合同生成异步摘要。""" return await self._acreate_summary(messages_to_summarize) class FinalRequestCompactionMiddleware(AgentMiddleware): """按最终模型请求预算压缩历史，并在模型成功后原子提交新状态。""" _COMPACTION_ANCHOR_KEY = "moviepilot_compaction_anchor_id" _UNCOMPRESSIBLE_REQUEST = ( "最终模型请求压缩后仍超出上下文窗口，原有上下文已保留，" "请减少启用工具或切换更大上下文模型" ) def __init__( self, *, summarizer: ContextPreservingSummarizationMiddleware, trigger_fraction: float = 0.85, keep_fraction: float = 0.10, ) -> None: self.summarizer = summarizer self.trigger_fraction = trigger_fraction self.keep_fraction = keep_fraction def _should_compact(self, budget: dict[str, Any]) -> bool: """以最终请求实际模型窗口判断是否需要压缩。""" estimated_tokens = budget.get("estimated_input_tokens") context_window = budget.get("context_window_tokens") return ( isinstance(estimated_tokens, int) and isinstance(context_window, int) and estimated_tokens >= context_window * self.trigger_fraction ) def _compaction_partition( self, request: ModelRequest ) -> tuple[list[AnyMessage], list[AnyMessage]] | None: """最终输入达到阈值时，拆分需要摘要和需要原样保留的消息。""" messages = list(request.messages) try: budget = UsageMiddleware.estimate_request(request) except Exception as error: logger.debug( "最终模型请求预算评估失败，继续原请求: error_type=%s", type(error).__name__, ) return None context_window = budget.get("context_window_tokens") if self._should_skip_after_current_turn_compaction(messages, budget): return None if not self._should_compact(budget) or not isinstance(context_window, int): return None try: partition = self.summarizer.partition_for_retention(messages) if partition is None: partition = self.summarizer.partition_for_token_limit( messages, max(1, int(context_window * self.keep_fraction)), force=budget["estimated_input_tokens"] > context_window, ) except Exception as error: logger.debug( "最终请求历史拆分失败，继续原请求: error_type=%s", type(error).__name__, ) if budget["estimated_input_tokens"] > context_window: raise ContextSummarizationError( self._UNCOMPRESSIBLE_REQUEST ) from error return None if partition is None: if budget["estimated_input_tokens"] > context_window: raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST) return None messages_to_summarize, preserved_messages = partition if all( message.additional_kwargs.get("lc_source") == "summarization" for message in messages_to_summarize ): if budget["estimated_input_tokens"] > context_window: raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST) return None return messages_to_summarize, preserved_messages @classmethod def _should_skip_after_current_turn_compaction( cls, messages: list[AnyMessage], budget: dict[str, Any] ) -> bool: """同轮只在新工具结果已使请求超窗时再次压缩。""" for message in reversed(messages): anchor_id = message.additional_kwargs.get(cls._COMPACTION_ANCHOR_KEY) if not isinstance(anchor_id, str): continue anchor_index = next( ( index for index, candidate in enumerate(messages) if candidate.id == anchor_id ), None, ) if anchor_index is None: return False messages_after_anchor = messages[anchor_index + 1 :] if any( isinstance(candidate, HumanMessage) for candidate in messages_after_anchor ): return False if any(isinstance(candidate, ToolMessage) for candidate in messages_after_anchor): estimated_tokens = budget.get("estimated_input_tokens") context_window = budget.get("context_window_tokens") return not ( isinstance(estimated_tokens, int) and isinstance(context_window, int) and estimated_tokens > context_window ) return True return False def _build_compacted_messages( self, summary: str, preserved_messages: list[AnyMessage] ) -> list[AnyMessage]: """构造摘要与近期历史，并记录本轮压缩输入边界。""" summary_messages = self.summarizer.build_summary_messages(summary) if summary_messages: self.summarizer.ensure_message_ids(summary_messages) anchor_id = ( preserved_messages[-1].id if preserved_messages else summary_messages[0].id ) first_summary = summary_messages[0] summary_messages[0] = first_summary.model_copy( update={ "additional_kwargs": { **first_summary.additional_kwargs, self._COMPACTION_ANCHOR_KEY: anchor_id, } } ) return [*summary_messages, *preserved_messages] def _validate_or_repartition( self, request: ModelRequest, messages_to_summarize: list[AnyMessage], preserved_messages: list[AnyMessage], summary: str, ) -> tuple[list[AnyMessage], tuple[list[AnyMessage], list[AnyMessage]] | None]: """复核压缩后的最终预算，并计算一次更小的近期历史分区。""" compacted_messages = self._build_compacted_messages(summary, preserved_messages) compacted_budget = UsageMiddleware.estimate_request( request.override(messages=compacted_messages) ) context_window = compacted_budget.get("context_window_tokens") estimated_tokens = compacted_budget.get("estimated_input_tokens") if not isinstance(context_window, int) or not isinstance(estimated_tokens, int): return compacted_messages, None if estimated_tokens <= context_window: return compacted_messages, None summary_messages = self.summarizer.build_summary_messages(summary) fixed_summary_budget = UsageMiddleware.estimate_request( request.override(messages=summary_messages) ) available_recent_tokens = ( context_window - fixed_summary_budget["estimated_input_tokens"] ) if available_recent_tokens <= 0: raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST) repartition = self.summarizer.partition_for_token_limit( list(request.messages), available_recent_tokens, force=True, minimum_cutoff=len(messages_to_summarize) + 1, strict_token_limit=True, ) if ( repartition is None or len(repartition[0]) <= len(messages_to_summarize) ): raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST) return compacted_messages, repartition def _require_within_window( self, request: ModelRequest, compacted_messages: list[AnyMessage] ) -> None: """禁止把已知仍超过主模型窗口的请求发送给 provider。""" budget = UsageMiddleware.estimate_request( request.override(messages=compacted_messages) ) estimated_tokens = budget.get("estimated_input_tokens") context_window = budget.get("context_window_tokens") if ( isinstance(estimated_tokens, int) and isinstance(context_window, int) and estimated_tokens > context_window ): raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST) def _prepare_messages(self, request: ModelRequest) -> list[AnyMessage] | None: """同步生成摘要与需要原样保留的近期消息。""" partition = self._compaction_partition(request) if partition is None: return None messages_to_summarize, preserved_messages = partition summary = self.summarizer.create_summary(messages_to_summarize) compacted_messages, repartition = self._validate_or_repartition( request, messages_to_summarize, preserved_messages, summary, ) if repartition is not None: messages_to_summarize, preserved_messages = repartition summary = self.summarizer.create_summary(messages_to_summarize) compacted_messages = self._build_compacted_messages( summary, preserved_messages ) self._require_within_window(request, compacted_messages) return compacted_messages async def _aprepare_messages( self, request: ModelRequest ) -> list[AnyMessage] | None: """异步生成摘要与需要原样保留的近期消息。""" partition = self._compaction_partition(request) if partition is None: return None messages_to_summarize, preserved_messages = partition summary = await self.summarizer.acreate_summary(messages_to_summarize) compacted_messages, repartition = self._validate_or_repartition( request, messages_to_summarize, preserved_messages, summary, ) if repartition is not None: messages_to_summarize, preserved_messages = repartition summary = await self.summarizer.acreate_summary(messages_to_summarize) compacted_messages = self._build_compacted_messages( summary, preserved_messages ) self._require_within_window(request, compacted_messages) return compacted_messages @staticmethod def _with_state_update( response: ModelResponse, compacted_messages: list[AnyMessage] ) -> ExtendedModelResponse: """主模型成功后一次性替换历史，同时保留本次模型结果。""" return ExtendedModelResponse( model_response=response, command=Command( update={ "messages": [ RemoveMessage(id=REMOVE_ALL_MESSAGES), *compacted_messages, *response.result, ] } ), ) def wrap_model_call( self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse], ) -> ModelResponse | ExtendedModelResponse: """同步压缩最终请求；模型失败时不提交摘要状态。""" compacted_messages = self._prepare_messages(request) if compacted_messages is None: return handler(request) response = handler(request.override(messages=compacted_messages)) return self._with_state_update(response, compacted_messages) async def awrap_model_call( self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]], ) -> ModelResponse | ExtendedModelResponse: """异步压缩最终请求；模型失败时不提交摘要状态。""" compacted_messages = await self._aprepare_messages(request) if compacted_messages is None: return await handler(request) response = await handler(request.override(messages=compacted_messages)) return self._with_state_update(response, compacted_messages)