383 lines
16 KiB
Python
383 lines
16 KiB
Python
"""Middleware to inject uploaded files information into agent context."""
|
|
|
|
import logging
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import NotRequired, override
|
|
from urllib.parse import unquote
|
|
|
|
from langchain.agents import AgentState
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.runtime import Runtime
|
|
|
|
from deerflow.config.paths import Paths, get_paths
|
|
from deerflow.runtime.thread_paths import resolve_path_user_id
|
|
from deerflow.utils.file_conversion import extract_outline
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
_OUTLINE_PREVIEW_LINES = 5
|
|
|
|
# Image files are understood via the `view_image` tool (which routes to the
|
|
# configured vision model), never by reading their bytes. Listing them with the
|
|
# generic read_file/grep guidance made the agent try to parse image binaries
|
|
# with bash/python, which fails and wastes turns.
|
|
_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp"}
|
|
_ARTIFACT_URL_MARKER = "/artifacts/"
|
|
|
|
|
|
def _extract_outline_for_file(file_path: Path) -> tuple[list[dict], list[str]]:
|
|
"""Return the document outline and fallback preview for *file_path*.
|
|
|
|
Looks for a sibling ``<stem>.md`` file produced by the upload conversion
|
|
pipeline.
|
|
|
|
Returns:
|
|
(outline, preview) where:
|
|
- outline: list of ``{title, line}`` dicts (plus optional sentinel).
|
|
Empty when no headings are found or no .md exists.
|
|
- preview: first few non-empty lines of the .md, used as a content
|
|
anchor when outline is empty so the agent has some context.
|
|
Empty when outline is non-empty (no fallback needed).
|
|
"""
|
|
md_path = file_path.with_suffix(".md")
|
|
if not md_path.is_file():
|
|
return [], []
|
|
|
|
outline = extract_outline(md_path)
|
|
if outline:
|
|
logger.debug("Extracted %d outline entries from %s", len(outline), file_path.name)
|
|
return outline, []
|
|
|
|
# outline is empty — read the first few non-empty lines as a content preview
|
|
preview: list[str] = []
|
|
try:
|
|
with md_path.open(encoding="utf-8") as f:
|
|
for line in f:
|
|
stripped = line.strip()
|
|
if stripped:
|
|
preview.append(stripped)
|
|
if len(preview) >= _OUTLINE_PREVIEW_LINES:
|
|
break
|
|
except Exception:
|
|
logger.debug("Failed to read preview lines from %s", md_path, exc_info=True)
|
|
return [], preview
|
|
|
|
|
|
class UploadsMiddlewareState(AgentState):
|
|
"""State schema for uploads middleware."""
|
|
|
|
uploaded_files: NotRequired[list[dict] | None]
|
|
|
|
|
|
class UploadsMiddleware(AgentMiddleware[UploadsMiddlewareState]):
|
|
"""Middleware to inject uploaded files information into the agent context.
|
|
|
|
Reads file metadata from the current message's additional_kwargs.files
|
|
(set by the frontend after upload) and prepends an <uploaded_files> block
|
|
to the last human message so the model knows which files are available.
|
|
"""
|
|
|
|
state_schema = UploadsMiddlewareState
|
|
|
|
def __init__(self, base_dir: str | None = None):
|
|
"""Initialize the middleware.
|
|
|
|
Args:
|
|
base_dir: Base directory for thread data. Defaults to Paths resolution.
|
|
"""
|
|
super().__init__()
|
|
self._paths = Paths(base_dir) if base_dir else get_paths()
|
|
|
|
def _format_file_entry(self, file: dict, lines: list[str]) -> None:
|
|
"""Append a single file entry (name, size, path, optional outline) to lines."""
|
|
size_kb = file["size"] / 1024
|
|
size_str = f"{size_kb:.1f} KB" if size_kb < 1024 else f"{size_kb / 1024:.1f} MB"
|
|
lines.append(f"- {file['filename']} ({size_str})")
|
|
lines.append(f" Path: {file['path']}")
|
|
|
|
# Images must be understood through the vision model, not by reading bytes.
|
|
extension = str(file.get("extension") or Path(file["filename"]).suffix).lower()
|
|
if extension in _IMAGE_EXTENSIONS:
|
|
lines.append(f" This is an image. To understand it, call the `view_image` tool with image_path='{file['path']}' (optionally pass a `question`).")
|
|
lines.append(" The configured vision model will analyze it and return a description — trust that result.")
|
|
lines.append(" Do NOT read, open, or parse the image with read_file, grep, bash, python, PIL, or any byte-level inspection; those cannot interpret image content.")
|
|
lines.append("")
|
|
return
|
|
|
|
outline = file.get("outline") or []
|
|
if outline:
|
|
truncated = outline[-1].get("truncated", False)
|
|
visible = [e for e in outline if not e.get("truncated")]
|
|
lines.append(" Document outline (use `read_file` with line ranges to read sections):")
|
|
for entry in visible:
|
|
lines.append(f" L{entry['line']}: {entry['title']}")
|
|
if truncated:
|
|
lines.append(f" ... (showing first {len(visible)} headings; use `read_file` to explore further)")
|
|
else:
|
|
preview = file.get("outline_preview") or []
|
|
if preview:
|
|
lines.append(" No structural headings detected. Document begins with:")
|
|
for text in preview:
|
|
lines.append(f" > {text}")
|
|
lines.append(" Use `grep` to search for keywords (e.g. `grep(pattern='keyword', path='/mnt/user-data/uploads/')`).")
|
|
lines.append("")
|
|
|
|
def _create_files_message(self, new_files: list[dict], historical_files: list[dict]) -> str:
|
|
"""Create a formatted message listing uploaded files.
|
|
|
|
Args:
|
|
new_files: Files uploaded in the current message.
|
|
historical_files: Files uploaded in previous messages.
|
|
Each file dict may contain an optional ``outline`` key — a list of
|
|
``{title, line}`` dicts extracted from the converted Markdown file.
|
|
|
|
Returns:
|
|
Formatted string inside <uploaded_files> tags.
|
|
"""
|
|
lines = ["<uploaded_files>"]
|
|
|
|
lines.append("The following files were uploaded in this message:")
|
|
lines.append("")
|
|
if new_files:
|
|
for file in new_files:
|
|
self._format_file_entry(file, lines)
|
|
else:
|
|
lines.append("(empty)")
|
|
lines.append("")
|
|
|
|
if historical_files:
|
|
lines.append("The following files were uploaded in previous messages and are still available:")
|
|
lines.append("")
|
|
for file in historical_files:
|
|
self._format_file_entry(file, lines)
|
|
|
|
lines.append("To work with these files:")
|
|
lines.append("- Read from the file first — use the outline line numbers and `read_file` to locate relevant sections.")
|
|
lines.append("- Use `grep` to search for keywords when you are not sure which section to look at")
|
|
lines.append(" (e.g. `grep(pattern='revenue', path='/mnt/user-data/uploads/')`).")
|
|
lines.append("- Use `glob` to find files by name pattern")
|
|
lines.append(" (e.g. `glob(pattern='**/*.md', path='/mnt/user-data/uploads/')`).")
|
|
lines.append("- Only fall back to web search if the file content is clearly insufficient to answer the question.")
|
|
lines.append("</uploaded_files>")
|
|
|
|
return "\n".join(lines)
|
|
|
|
def _import_file_from_artifact_url(
|
|
self,
|
|
file_info: dict,
|
|
uploads_dir: Path | None,
|
|
) -> bool:
|
|
"""Copy a previously uploaded file into the current thread when possible.
|
|
|
|
iframe-style callers may pass file metadata that points at an upload
|
|
from another DeerFlow thread. The current run thread still needs a real
|
|
file under its own ``/mnt/user-data/uploads`` directory, otherwise this
|
|
middleware will filter the entry out during existence checks.
|
|
"""
|
|
filename = file_info.get("filename") or ""
|
|
artifact_url = str(file_info.get("artifact_url") or "").strip()
|
|
if not filename or uploads_dir is None or not artifact_url:
|
|
return False
|
|
|
|
if not artifact_url.startswith("/api/threads/") or _ARTIFACT_URL_MARKER not in artifact_url:
|
|
return False
|
|
|
|
prefix, encoded_virtual_path = artifact_url.split(_ARTIFACT_URL_MARKER, 1)
|
|
source_thread_id = prefix.removeprefix("/api/threads/").strip("/")
|
|
if not source_thread_id or not encoded_virtual_path:
|
|
return False
|
|
|
|
virtual_path = "/" + unquote(encoded_virtual_path.lstrip("/"))
|
|
if Path(virtual_path).name != filename:
|
|
return False
|
|
|
|
try:
|
|
source_path = self._paths.resolve_virtual_path(
|
|
source_thread_id,
|
|
virtual_path,
|
|
user_id=resolve_path_user_id(source_thread_id),
|
|
)
|
|
except ValueError:
|
|
logger.debug("Failed to resolve uploaded artifact path: %s", artifact_url, exc_info=True)
|
|
return False
|
|
|
|
if not source_path.is_file():
|
|
return False
|
|
|
|
uploads_dir.mkdir(parents=True, exist_ok=True)
|
|
target_path = uploads_dir / filename
|
|
shutil.copy2(source_path, target_path)
|
|
|
|
source_markdown = source_path.with_suffix(".md")
|
|
if source_markdown.is_file() and source_markdown != source_path:
|
|
shutil.copy2(source_markdown, target_path.with_suffix(".md"))
|
|
|
|
return True
|
|
|
|
def _files_from_kwargs(self, message: HumanMessage, uploads_dir: Path | None = None) -> list[dict] | None:
|
|
"""Extract file info from message additional_kwargs.files.
|
|
|
|
The frontend sends uploaded file metadata in additional_kwargs.files
|
|
after a successful upload. Each entry has: filename, size (bytes),
|
|
path (virtual path), status.
|
|
|
|
Args:
|
|
message: The human message to inspect.
|
|
uploads_dir: Physical uploads directory used to verify file existence.
|
|
When provided, entries whose files no longer exist are skipped.
|
|
|
|
Returns:
|
|
List of file dicts with virtual paths, or None if the field is absent or empty.
|
|
"""
|
|
kwargs_files = (message.additional_kwargs or {}).get("files")
|
|
if not isinstance(kwargs_files, list) or not kwargs_files:
|
|
return None
|
|
|
|
files = []
|
|
for f in kwargs_files:
|
|
if not isinstance(f, dict):
|
|
continue
|
|
filename = f.get("filename") or ""
|
|
if not filename or Path(filename).name != filename:
|
|
continue
|
|
if uploads_dir is not None and not (uploads_dir / filename).is_file():
|
|
if not self._import_file_from_artifact_url(f, uploads_dir):
|
|
continue
|
|
files.append(
|
|
{
|
|
"filename": filename,
|
|
"size": int(f.get("size") or 0),
|
|
"path": f"/mnt/user-data/uploads/{filename}",
|
|
"extension": Path(filename).suffix,
|
|
}
|
|
)
|
|
return files if files else None
|
|
|
|
@override
|
|
def before_agent(self, state: UploadsMiddlewareState, runtime: Runtime) -> dict | None:
|
|
"""Inject uploaded files information before agent execution.
|
|
|
|
New files come from the current message's additional_kwargs.files.
|
|
Historical files are scanned from the thread's uploads directory,
|
|
excluding the new ones.
|
|
|
|
Prepends <uploaded_files> context to the last human message content.
|
|
The original additional_kwargs (including files metadata) is preserved
|
|
on the updated message so the frontend can read it from the stream.
|
|
|
|
Args:
|
|
state: Current agent state.
|
|
runtime: Runtime context containing thread_id.
|
|
|
|
Returns:
|
|
State updates including uploaded files list.
|
|
"""
|
|
messages = list(state.get("messages", []))
|
|
if not messages:
|
|
return None
|
|
|
|
last_message_index = len(messages) - 1
|
|
last_message = messages[last_message_index]
|
|
|
|
if not isinstance(last_message, HumanMessage):
|
|
return None
|
|
|
|
# Only turns with fresh upload metadata need the uploads directory.
|
|
# Avoid resolving the thread owner (which may hit the DB) for ordinary
|
|
# no-file chat turns, including ephemeral open-chat requests.
|
|
kwargs_files = (last_message.additional_kwargs or {}).get("files")
|
|
if not isinstance(kwargs_files, list) or not kwargs_files:
|
|
return None
|
|
|
|
# Resolve uploads directory for existence checks
|
|
thread_id = (runtime.context or {}).get("thread_id")
|
|
if thread_id is None:
|
|
try:
|
|
from langgraph.config import get_config
|
|
|
|
thread_id = get_config().get("configurable", {}).get("thread_id")
|
|
except RuntimeError:
|
|
pass # get_config() raises outside a runnable context (e.g. unit tests)
|
|
uploads_dir = self._paths.sandbox_uploads_dir(thread_id, user_id=resolve_path_user_id(thread_id)) if thread_id else None
|
|
|
|
# Get newly uploaded files from the current message's additional_kwargs.files
|
|
new_files = self._files_from_kwargs(last_message, uploads_dir) or []
|
|
|
|
# Only inject when the *current* message actually carries new uploads.
|
|
# Turns without a fresh upload must not re-inject previously uploaded
|
|
# files: the original upload turn already injected (and persisted) their
|
|
# <uploaded_files> block, and the files remain physically available for
|
|
# read_file/grep/view_image. Re-injecting on every later message made the
|
|
# agent re-reference old images/docs, and — since the frontend parses
|
|
# files out of the <uploaded_files> content when additional_kwargs.files
|
|
# is empty — re-rendered the old attachment on each new question.
|
|
if not new_files:
|
|
return None
|
|
|
|
# Collect historical files from the uploads directory (all except the new ones)
|
|
new_filenames = {f["filename"] for f in new_files}
|
|
historical_files: list[dict] = []
|
|
if uploads_dir and uploads_dir.exists():
|
|
for file_path in sorted(uploads_dir.iterdir()):
|
|
if file_path.is_file() and file_path.name not in new_filenames:
|
|
stat = file_path.stat()
|
|
outline, preview = _extract_outline_for_file(file_path)
|
|
historical_files.append(
|
|
{
|
|
"filename": file_path.name,
|
|
"size": stat.st_size,
|
|
"path": f"/mnt/user-data/uploads/{file_path.name}",
|
|
"extension": file_path.suffix,
|
|
"outline": outline,
|
|
"outline_preview": preview,
|
|
}
|
|
)
|
|
|
|
# Attach outlines to new files as well
|
|
if uploads_dir:
|
|
for file in new_files:
|
|
phys_path = uploads_dir / file["filename"]
|
|
outline, preview = _extract_outline_for_file(phys_path)
|
|
file["outline"] = outline
|
|
file["outline_preview"] = preview
|
|
|
|
logger.debug(f"New files: {[f['filename'] for f in new_files]}, historical: {[f['filename'] for f in historical_files]}")
|
|
|
|
# Create files message and prepend to the last human message content
|
|
files_message = self._create_files_message(new_files, historical_files)
|
|
|
|
# Extract original content - handle both string and list formats
|
|
original_content = last_message.content
|
|
if isinstance(original_content, str):
|
|
# Simple case: string content, just prepend files message
|
|
updated_content = f"{files_message}\n\n{original_content}"
|
|
elif isinstance(original_content, list):
|
|
# Complex case: list content (multimodal), preserve all blocks
|
|
# Prepend files message as the first text block
|
|
files_block = {"type": "text", "text": f"{files_message}\n\n"}
|
|
# Keep all original blocks (including images)
|
|
updated_content = [files_block, *original_content]
|
|
else:
|
|
# Other types, preserve as-is
|
|
updated_content = original_content
|
|
|
|
# Create new message with combined content.
|
|
# Preserve additional_kwargs (including files metadata) so the frontend
|
|
# can read structured file info from the streamed message.
|
|
updated_message = HumanMessage(
|
|
content=updated_content,
|
|
id=last_message.id,
|
|
name=last_message.name,
|
|
additional_kwargs=last_message.additional_kwargs,
|
|
)
|
|
|
|
messages[last_message_index] = updated_message
|
|
|
|
return {
|
|
"uploaded_files": new_files,
|
|
"messages": messages,
|
|
}
|