ForcePilot/src/agents/common/middlewares/attachment_middleware.py

"""附件注入中间件 - 使用 LangChain 标准中间件实现

从 State 中读取附件信息，注入提示词让模型使用 read_file 工具读取附件内容。
"""

from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import NotRequired

from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse

from src.utils import logger


class AttachmentState(AgentState):
    """扩展 AgentState 以支持附件"""

    attachments: NotRequired[list[dict]]


def _build_attachment_prompt(attachments: Sequence[dict]) -> str | None:
    """Render attachments into a system prompt block with file paths.

    提示模型使用 read_file 工具读取附件内容。
    """
    if not attachments:
        return None

    valid_attachments = [a for a in attachments if a.get("status") == "parsed"]

    if not valid_attachments:
        return None

    attachment_infos: list[str] = []
    for attachment in valid_attachments:
        file_name = attachment.get("file_name", "未知文件")
        file_path = attachment.get("file_path", "")
        truncated = "（已截断）" if attachment.get("truncated") else ""

        if file_path:
            attachment_infos.append(f"- {file_name}{truncated}: {file_path}")
        else:
            attachment_infos.append(f"- {file_name}{truncated}")

    lines = [
        "用户上传了以下附件：",
        "",
        *attachment_infos,
        "",
        "请使用 read_file 工具读取附件内容后，再回答用户的问题。",
    ]

    return "\n".join(lines)


class AttachmentMiddleware(AgentMiddleware[AttachmentState]):
    """
    LangChain 标准中间件：从 State 中读取附件并注入提示词。

    LangGraph 会自动从 checkpointer 恢复 state，包括 attachments。
    从 request.state 中读取附件，将其转换为 SystemMessage 并注入到消息列表开头。
    """

    state_schema = AttachmentState

    async def awrap_model_call(
        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
    ) -> ModelResponse:
        # 从 state 获取附件（LangGraph 自动从 checkpointer 恢复）
        attachments = request.state.get("attachments", [])
        logger.info(f"AttachmentMiddleware: found {len(attachments)} attachments in state")

        if attachments:
            # 构建附件提示
            attachment_prompt = _build_attachment_prompt(attachments)

            if attachment_prompt:
                logger.info("AttachmentMiddleware: injecting attachment prompt")

                messages = list(request.messages)
                insert_idx = 0
                for idx, msg in enumerate(messages):
                    if isinstance(msg, dict):
                        role = msg.get("role") or msg.get("type")
                        is_system = role == "system"
                    else:
                        msg_type = getattr(msg, "type", None) or getattr(msg, "role", None)
                        is_system = msg_type == "system"

                    if not is_system:
                        break
                    insert_idx = idx + 1

                messages.insert(insert_idx, {"role": "system", "content": attachment_prompt})
                request = request.override(messages=messages)

        return await handler(request)


# 创建中间件实例，供其他模块使用
save_attachments_to_fs = AttachmentMiddleware()

# 保留旧名称以保持向后兼容（已废弃）
inject_attachment_context = save_attachments_to_fs