90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
"""Attachment prompt injection middleware.
|
|
|
|
Read uploaded file metadata from LangGraph state and inject readable paths
|
|
into the system prompt so the model can use `read_file` on demand.
|
|
"""
|
|
|
|
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 langchain_core.messages import SystemMessage
|
|
|
|
from yuxi.utils import logger
|
|
|
|
ATTACHMENT_PROMPT_MARKER = "<!-- attachment_context -->"
|
|
|
|
|
|
class AttachmentState(AgentState):
|
|
"""Extended state schema with uploaded files."""
|
|
|
|
uploads: NotRequired[list[dict]]
|
|
|
|
|
|
def _build_attachment_prompt(uploads: Sequence[dict]) -> str | None:
|
|
"""Render uploads into a concise prompt block."""
|
|
if not uploads:
|
|
return None
|
|
|
|
valid_uploads: list[tuple[str, str]] = []
|
|
for upload in uploads:
|
|
path = upload.get("path")
|
|
if not isinstance(path, str) or not path.strip():
|
|
continue
|
|
file_name = upload.get("file_name", "未知文件")
|
|
valid_uploads.append((str(file_name), path))
|
|
|
|
if not valid_uploads:
|
|
return None
|
|
|
|
upload_infos = [f"- {file_name}: {path}" for file_name, path in valid_uploads]
|
|
lines = [
|
|
"用户上传了以下文件：",
|
|
"",
|
|
*upload_infos,
|
|
"",
|
|
"请优先使用 `read_file` 工具读取这些路径中的文件内容，再回答用户问题。",
|
|
]
|
|
return "\n".join(lines)
|
|
|
|
|
|
class AttachmentMiddleware(AgentMiddleware[AttachmentState]):
|
|
"""Inject upload context from state.uploads into system prompt."""
|
|
|
|
state_schema = AttachmentState
|
|
|
|
async def awrap_model_call(
|
|
self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
|
|
) -> ModelResponse:
|
|
uploads = request.state.get("uploads", [])
|
|
logger.info(f"AttachmentMiddleware: found {len(uploads)} uploads in state")
|
|
|
|
if uploads:
|
|
attachment_prompt = _build_attachment_prompt(uploads)
|
|
if attachment_prompt:
|
|
logger.info("AttachmentMiddleware: injecting attachment prompt")
|
|
existing_blocks = list(request.system_message.content_blocks) if request.system_message else []
|
|
existing_text = "\n".join(
|
|
block.get("text", "")
|
|
for block in existing_blocks
|
|
if isinstance(block, dict) and block.get("type") == "text"
|
|
)
|
|
|
|
if ATTACHMENT_PROMPT_MARKER in existing_text:
|
|
logger.info("AttachmentMiddleware: attachment prompt already injected, skip")
|
|
return await handler(request)
|
|
|
|
merged_blocks = existing_blocks + [
|
|
{"type": "text", "text": f"{ATTACHMENT_PROMPT_MARKER}\n{attachment_prompt}"}
|
|
]
|
|
request = request.override(system_message=SystemMessage(content=merged_blocks))
|
|
|
|
return await handler(request)
|
|
|
|
|
|
save_attachments_to_fs = AttachmentMiddleware()
|
|
inject_attachment_context = save_attachments_to_fs
|