mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-18 20:14:44 +08:00
Implement optimistic UI for file uploads and enhance message handling (#967)
* feat(upload): implement optimistic UI for file uploads and enhance message handling * feat(middleware): enhance file handling by collecting historical uploads from directory * feat(thread-title): update page title handling for new threads and improve loading state * feat(uploads-middleware): enhance file extraction by verifying file existence in uploads directory * feat(thread-stream): update file path reference to use virtual_path for uploads * feat(tests): add core behaviour tests for UploadsMiddleware * feat(tests): remove unused pytest import from test_uploads_middleware_core_logic.py * feat: enhance file upload handling and localization support - Update UploadsMiddleware to validate filenames more robustly. - Modify MessageListItem to parse uploaded files from raw content for backward compatibility. - Add localization for uploading messages in English and Chinese. - Introduce parseUploadedFiles utility to extract uploaded files from message content.
This commit is contained in:
JeffJiang
@@ -1,6 +1,6 @@
|
||||
"""Middleware to inject uploaded files information into agent context."""
|
||||
|
||||
import re
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import NotRequired, override
|
||||
|
||||
@@ -11,6 +11,8 @@ from langgraph.runtime import Runtime
|
||||
|
||||
from src.config.paths import Paths, get_paths
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class UploadsMiddlewareState(AgentState):
|
||||
"""State schema for uploads middleware."""
|
||||
@@ -21,8 +23,9 @@ class UploadsMiddlewareState(AgentState):
|
||||
class UploadsMiddleware(AgentMiddleware[UploadsMiddlewareState]):
|
||||
"""Middleware to inject uploaded files information into the agent context.
|
||||
|
||||
This middleware lists all files in the thread's uploads directory and
|
||||
adds a system message with the file list before the agent processes the request.
|
||||
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
|
||||
@@ -36,111 +39,91 @@ class UploadsMiddleware(AgentMiddleware[UploadsMiddlewareState]):
|
||||
super().__init__()
|
||||
self._paths = Paths(base_dir) if base_dir else get_paths()
|
||||
|
||||
def _get_uploads_dir(self, thread_id: str) -> Path:
|
||||
"""Get the uploads directory for a thread.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID.
|
||||
|
||||
Returns:
|
||||
Path to the uploads directory.
|
||||
"""
|
||||
return self._paths.sandbox_uploads_dir(thread_id)
|
||||
|
||||
def _list_newly_uploaded_files(self, thread_id: str, last_message_files: set[str]) -> list[dict]:
|
||||
"""List only newly uploaded files that weren't in the last message.
|
||||
|
||||
Args:
|
||||
thread_id: The thread ID.
|
||||
last_message_files: Set of filenames that were already shown in previous messages.
|
||||
|
||||
Returns:
|
||||
List of new file information dictionaries.
|
||||
"""
|
||||
uploads_dir = self._get_uploads_dir(thread_id)
|
||||
|
||||
if not uploads_dir.exists():
|
||||
return []
|
||||
|
||||
files = []
|
||||
for file_path in sorted(uploads_dir.iterdir()):
|
||||
if file_path.is_file() and file_path.name not in last_message_files:
|
||||
stat = file_path.stat()
|
||||
files.append(
|
||||
{
|
||||
"filename": file_path.name,
|
||||
"size": stat.st_size,
|
||||
"path": f"/mnt/user-data/uploads/{file_path.name}",
|
||||
"extension": file_path.suffix,
|
||||
}
|
||||
)
|
||||
|
||||
return files
|
||||
|
||||
def _create_files_message(self, files: list[dict]) -> str:
|
||||
def _create_files_message(self, new_files: list[dict], historical_files: list[dict]) -> str:
|
||||
"""Create a formatted message listing uploaded files.
|
||||
|
||||
Args:
|
||||
files: List of file information dictionaries.
|
||||
new_files: Files uploaded in the current message.
|
||||
historical_files: Files uploaded in previous messages.
|
||||
|
||||
Returns:
|
||||
Formatted string listing the files.
|
||||
Formatted string inside <uploaded_files> tags.
|
||||
"""
|
||||
if not files:
|
||||
return "<uploaded_files>\nNo files have been uploaded yet.\n</uploaded_files>"
|
||||
lines = ["<uploaded_files>"]
|
||||
|
||||
lines = ["<uploaded_files>", "The following files have been uploaded and are available for use:", ""]
|
||||
|
||||
for file in files:
|
||||
lines.append("The following files were uploaded in this message:")
|
||||
lines.append("")
|
||||
for file in new_files:
|
||||
size_kb = file["size"] / 1024
|
||||
if size_kb < 1024:
|
||||
size_str = f"{size_kb:.1f} KB"
|
||||
else:
|
||||
size_str = f"{size_kb / 1024:.1f} MB"
|
||||
|
||||
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']}")
|
||||
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:
|
||||
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']}")
|
||||
lines.append("")
|
||||
|
||||
lines.append("You can read these files using the `read_file` tool with the paths shown above.")
|
||||
lines.append("</uploaded_files>")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _extract_files_from_message(self, content: str) -> set[str]:
|
||||
"""Extract filenames from uploaded_files tag in message content.
|
||||
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:
|
||||
content: Message content that may contain <uploaded_files> tag.
|
||||
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:
|
||||
Set of filenames mentioned in the tag.
|
||||
List of file dicts with virtual paths, or None if the field is absent or empty.
|
||||
"""
|
||||
# Match <uploaded_files>...</uploaded_files> tag
|
||||
match = re.search(r"<uploaded_files>([\s\S]*?)</uploaded_files>", content)
|
||||
if not match:
|
||||
return set()
|
||||
kwargs_files = (message.additional_kwargs or {}).get("files")
|
||||
if not isinstance(kwargs_files, list) or not kwargs_files:
|
||||
return None
|
||||
|
||||
files_content = match.group(1)
|
||||
|
||||
# Extract filenames from lines like "- filename.ext (size)"
|
||||
# Need to capture everything before the opening parenthesis, including spaces
|
||||
filenames = set()
|
||||
for line in files_content.split("\n"):
|
||||
# Match pattern: - filename with spaces.ext (size)
|
||||
# Changed from [^\s(]+ to [^(]+ to allow spaces in filename
|
||||
file_match = re.match(r"^-\s+(.+?)\s*\(", line.strip())
|
||||
if file_match:
|
||||
filenames.add(file_match.group(1).strip())
|
||||
|
||||
return filenames
|
||||
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():
|
||||
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.
|
||||
|
||||
Only injects files that weren't already shown in previous messages.
|
||||
Prepends file info to the last human message content.
|
||||
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.
|
||||
@@ -149,72 +132,70 @@ class UploadsMiddleware(AgentMiddleware[UploadsMiddlewareState]):
|
||||
Returns:
|
||||
State updates including uploaded files list.
|
||||
"""
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
thread_id = runtime.context.get("thread_id")
|
||||
if thread_id is None:
|
||||
return None
|
||||
|
||||
messages = list(state.get("messages", []))
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
# Track all filenames that have been shown in previous messages (EXCEPT the last one)
|
||||
shown_files: set[str] = set()
|
||||
for msg in messages[:-1]: # Scan all messages except the last one
|
||||
if isinstance(msg, HumanMessage):
|
||||
content = msg.content if isinstance(msg.content, str) else ""
|
||||
extracted = self._extract_files_from_message(content)
|
||||
shown_files.update(extracted)
|
||||
if extracted:
|
||||
logger.info(f"Found previously shown files: {extracted}")
|
||||
|
||||
logger.info(f"Total shown files from history: {shown_files}")
|
||||
|
||||
# List only newly uploaded files
|
||||
files = self._list_newly_uploaded_files(thread_id, shown_files)
|
||||
logger.info(f"Newly uploaded files to inject: {[f['filename'] for f in files]}")
|
||||
|
||||
if not files:
|
||||
return None
|
||||
|
||||
# Find the last human message and prepend file info to it
|
||||
last_message_index = len(messages) - 1
|
||||
last_message = messages[last_message_index]
|
||||
|
||||
if not isinstance(last_message, HumanMessage):
|
||||
return None
|
||||
|
||||
# Resolve uploads directory for existence checks
|
||||
thread_id = runtime.context.get("thread_id")
|
||||
uploads_dir = self._paths.sandbox_uploads_dir(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 []
|
||||
|
||||
# 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()
|
||||
historical_files.append(
|
||||
{
|
||||
"filename": file_path.name,
|
||||
"size": stat.st_size,
|
||||
"path": f"/mnt/user-data/uploads/{file_path.name}",
|
||||
"extension": file_path.suffix,
|
||||
}
|
||||
)
|
||||
|
||||
if not new_files and not historical_files:
|
||||
return None
|
||||
|
||||
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(files)
|
||||
files_message = self._create_files_message(new_files, historical_files)
|
||||
|
||||
# Extract original content - handle both string and list formats
|
||||
original_content = ""
|
||||
if isinstance(last_message.content, str):
|
||||
original_content = last_message.content
|
||||
elif isinstance(last_message.content, list):
|
||||
# Content is a list of content blocks (e.g., [{"type": "text", "text": "..."}])
|
||||
text_parts = []
|
||||
for block in last_message.content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text_parts.append(block.get("text", ""))
|
||||
original_content = "\n".join(text_parts)
|
||||
|
||||
logger.info(f"Original message content: {original_content[:100] if original_content else '(empty)'}")
|
||||
|
||||
# Create new message with combined 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=f"{files_message}\n\n{original_content}",
|
||||
id=last_message.id,
|
||||
additional_kwargs=last_message.additional_kwargs,
|
||||
)
|
||||
|
||||
# Replace the last message
|
||||
messages[last_message_index] = updated_message
|
||||
|
||||
return {
|
||||
"uploaded_files": files,
|
||||
"uploaded_files": new_files,
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
334
backend/tests/test_uploads_middleware_core_logic.py
Normal file
334
backend/tests/test_uploads_middleware_core_logic.py
Normal file
@@ -0,0 +1,334 @@
|
||||
"""Core behaviour tests for UploadsMiddleware.
|
||||
|
||||
Covers:
|
||||
- _files_from_kwargs: parsing, validation, existence check, virtual-path construction
|
||||
- _create_files_message: output format with new-only and new+historical files
|
||||
- before_agent: full injection pipeline (string & list content, preserved
|
||||
additional_kwargs, historical files from uploads dir, edge-cases)
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from src.agents.middlewares.uploads_middleware import UploadsMiddleware
|
||||
from src.config.paths import Paths
|
||||
|
||||
THREAD_ID = "thread-abc123"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _middleware(tmp_path: Path) -> UploadsMiddleware:
|
||||
return UploadsMiddleware(base_dir=str(tmp_path))
|
||||
|
||||
|
||||
def _runtime(thread_id: str | None = THREAD_ID) -> MagicMock:
|
||||
rt = MagicMock()
|
||||
rt.context = {"thread_id": thread_id}
|
||||
return rt
|
||||
|
||||
|
||||
def _uploads_dir(tmp_path: Path, thread_id: str = THREAD_ID) -> Path:
|
||||
d = Paths(str(tmp_path)).sandbox_uploads_dir(thread_id)
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
return d
|
||||
|
||||
|
||||
def _human(content, files=None, **extra_kwargs):
|
||||
additional_kwargs = dict(extra_kwargs)
|
||||
if files is not None:
|
||||
additional_kwargs["files"] = files
|
||||
return HumanMessage(content=content, additional_kwargs=additional_kwargs)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _files_from_kwargs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFilesFromKwargs:
|
||||
def test_returns_none_when_files_field_absent(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = HumanMessage(content="hello")
|
||||
assert mw._files_from_kwargs(msg) is None
|
||||
|
||||
def test_returns_none_for_empty_files_list(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hello", files=[])
|
||||
assert mw._files_from_kwargs(msg) is None
|
||||
|
||||
def test_returns_none_for_non_list_files(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hello", files="not-a-list")
|
||||
assert mw._files_from_kwargs(msg) is None
|
||||
|
||||
def test_skips_non_dict_entries(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hi", files=["bad", 42, None])
|
||||
assert mw._files_from_kwargs(msg) is None
|
||||
|
||||
def test_skips_entries_with_empty_filename(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hi", files=[{"filename": "", "size": 100, "path": "/mnt/user-data/uploads/x"}])
|
||||
assert mw._files_from_kwargs(msg) is None
|
||||
|
||||
def test_always_uses_virtual_path(self, tmp_path):
|
||||
"""path field must be /mnt/user-data/uploads/<filename> regardless of what the frontend sent."""
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human(
|
||||
"hi",
|
||||
files=[{"filename": "report.pdf", "size": 1024, "path": "/some/arbitrary/path/report.pdf"}],
|
||||
)
|
||||
result = mw._files_from_kwargs(msg)
|
||||
assert result is not None
|
||||
assert result[0]["path"] == "/mnt/user-data/uploads/report.pdf"
|
||||
|
||||
def test_skips_file_that_does_not_exist_on_disk(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
# file is NOT written to disk
|
||||
msg = _human("hi", files=[{"filename": "missing.txt", "size": 50, "path": "/mnt/user-data/uploads/missing.txt"}])
|
||||
assert mw._files_from_kwargs(msg, uploads_dir) is None
|
||||
|
||||
def test_accepts_file_that_exists_on_disk(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "data.csv").write_text("a,b,c")
|
||||
msg = _human("hi", files=[{"filename": "data.csv", "size": 5, "path": "/mnt/user-data/uploads/data.csv"}])
|
||||
result = mw._files_from_kwargs(msg, uploads_dir)
|
||||
assert result is not None
|
||||
assert len(result) == 1
|
||||
assert result[0]["filename"] == "data.csv"
|
||||
assert result[0]["path"] == "/mnt/user-data/uploads/data.csv"
|
||||
|
||||
def test_skips_nonexistent_but_accepts_existing_in_mixed_list(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "present.txt").write_text("here")
|
||||
msg = _human(
|
||||
"hi",
|
||||
files=[
|
||||
{"filename": "present.txt", "size": 4, "path": "/mnt/user-data/uploads/present.txt"},
|
||||
{"filename": "gone.txt", "size": 4, "path": "/mnt/user-data/uploads/gone.txt"},
|
||||
],
|
||||
)
|
||||
result = mw._files_from_kwargs(msg, uploads_dir)
|
||||
assert result is not None
|
||||
assert [f["filename"] for f in result] == ["present.txt"]
|
||||
|
||||
def test_no_existence_check_when_uploads_dir_is_none(self, tmp_path):
|
||||
"""Without an uploads_dir argument the existence check is skipped entirely."""
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hi", files=[{"filename": "phantom.txt", "size": 10, "path": "/mnt/user-data/uploads/phantom.txt"}])
|
||||
result = mw._files_from_kwargs(msg, uploads_dir=None)
|
||||
assert result is not None
|
||||
assert result[0]["filename"] == "phantom.txt"
|
||||
|
||||
def test_size_is_coerced_to_int(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hi", files=[{"filename": "f.txt", "size": "2048", "path": "/mnt/user-data/uploads/f.txt"}])
|
||||
result = mw._files_from_kwargs(msg)
|
||||
assert result is not None
|
||||
assert result[0]["size"] == 2048
|
||||
|
||||
def test_missing_size_defaults_to_zero(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("hi", files=[{"filename": "f.txt", "path": "/mnt/user-data/uploads/f.txt"}])
|
||||
result = mw._files_from_kwargs(msg)
|
||||
assert result is not None
|
||||
assert result[0]["size"] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _create_files_message
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCreateFilesMessage:
|
||||
def _new_file(self, filename="notes.txt", size=1024):
|
||||
return {"filename": filename, "size": size, "path": f"/mnt/user-data/uploads/{filename}"}
|
||||
|
||||
def test_new_files_section_always_present(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = mw._create_files_message([self._new_file()], [])
|
||||
assert "<uploaded_files>" in msg
|
||||
assert "</uploaded_files>" in msg
|
||||
assert "uploaded in this message" in msg
|
||||
assert "notes.txt" in msg
|
||||
assert "/mnt/user-data/uploads/notes.txt" in msg
|
||||
|
||||
def test_historical_section_present_only_when_non_empty(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
|
||||
msg_no_hist = mw._create_files_message([self._new_file()], [])
|
||||
assert "previous messages" not in msg_no_hist
|
||||
|
||||
hist = self._new_file("old.txt")
|
||||
msg_with_hist = mw._create_files_message([self._new_file()], [hist])
|
||||
assert "previous messages" in msg_with_hist
|
||||
assert "old.txt" in msg_with_hist
|
||||
|
||||
def test_size_formatting_kb(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = mw._create_files_message([self._new_file(size=2048)], [])
|
||||
assert "2.0 KB" in msg
|
||||
|
||||
def test_size_formatting_mb(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = mw._create_files_message([self._new_file(size=2 * 1024 * 1024)], [])
|
||||
assert "2.0 MB" in msg
|
||||
|
||||
def test_read_file_instruction_included(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = mw._create_files_message([self._new_file()], [])
|
||||
assert "read_file" in msg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# before_agent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBeforeAgent:
|
||||
def _state(self, *messages):
|
||||
return {"messages": list(messages)}
|
||||
|
||||
def test_returns_none_when_messages_empty(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
assert mw.before_agent({"messages": []}, _runtime()) is None
|
||||
|
||||
def test_returns_none_when_last_message_is_not_human(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
state = self._state(HumanMessage(content="q"), AIMessage(content="a"))
|
||||
assert mw.before_agent(state, _runtime()) is None
|
||||
|
||||
def test_returns_none_when_no_files_in_kwargs(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
state = self._state(_human("plain message"))
|
||||
assert mw.before_agent(state, _runtime()) is None
|
||||
|
||||
def test_returns_none_when_all_files_missing_from_disk(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
_uploads_dir(tmp_path) # directory exists but is empty
|
||||
msg = _human("hi", files=[{"filename": "ghost.txt", "size": 10, "path": "/mnt/user-data/uploads/ghost.txt"}])
|
||||
state = self._state(msg)
|
||||
assert mw.before_agent(state, _runtime()) is None
|
||||
|
||||
def test_injects_uploaded_files_tag_into_string_content(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "report.pdf").write_bytes(b"pdf")
|
||||
|
||||
msg = _human("please analyse", files=[{"filename": "report.pdf", "size": 3, "path": "/mnt/user-data/uploads/report.pdf"}])
|
||||
state = self._state(msg)
|
||||
result = mw.before_agent(state, _runtime())
|
||||
|
||||
assert result is not None
|
||||
updated_msg = result["messages"][-1]
|
||||
assert isinstance(updated_msg.content, str)
|
||||
assert "<uploaded_files>" in updated_msg.content
|
||||
assert "report.pdf" in updated_msg.content
|
||||
assert "please analyse" in updated_msg.content
|
||||
|
||||
def test_injects_uploaded_files_tag_into_list_content(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "data.csv").write_bytes(b"a,b")
|
||||
|
||||
msg = _human(
|
||||
[{"type": "text", "text": "analyse this"}],
|
||||
files=[{"filename": "data.csv", "size": 3, "path": "/mnt/user-data/uploads/data.csv"}],
|
||||
)
|
||||
state = self._state(msg)
|
||||
result = mw.before_agent(state, _runtime())
|
||||
|
||||
assert result is not None
|
||||
updated_msg = result["messages"][-1]
|
||||
assert "<uploaded_files>" in updated_msg.content
|
||||
assert "analyse this" in updated_msg.content
|
||||
|
||||
def test_preserves_additional_kwargs_on_updated_message(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "img.png").write_bytes(b"png")
|
||||
|
||||
files_meta = [{"filename": "img.png", "size": 3, "path": "/mnt/user-data/uploads/img.png", "status": "uploaded"}]
|
||||
msg = _human("check image", files=files_meta, element="task")
|
||||
state = self._state(msg)
|
||||
result = mw.before_agent(state, _runtime())
|
||||
|
||||
assert result is not None
|
||||
updated_kwargs = result["messages"][-1].additional_kwargs
|
||||
assert updated_kwargs.get("files") == files_meta
|
||||
assert updated_kwargs.get("element") == "task"
|
||||
|
||||
def test_uploaded_files_returned_in_state_update(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "notes.txt").write_bytes(b"hello")
|
||||
|
||||
msg = _human("review", files=[{"filename": "notes.txt", "size": 5, "path": "/mnt/user-data/uploads/notes.txt"}])
|
||||
result = mw.before_agent(self._state(msg), _runtime())
|
||||
|
||||
assert result is not None
|
||||
assert result["uploaded_files"] == [
|
||||
{
|
||||
"filename": "notes.txt",
|
||||
"size": 5,
|
||||
"path": "/mnt/user-data/uploads/notes.txt",
|
||||
"extension": ".txt",
|
||||
}
|
||||
]
|
||||
|
||||
def test_historical_files_from_uploads_dir_excluding_new(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "old.txt").write_bytes(b"old")
|
||||
(uploads_dir / "new.txt").write_bytes(b"new")
|
||||
|
||||
msg = _human("go", files=[{"filename": "new.txt", "size": 3, "path": "/mnt/user-data/uploads/new.txt"}])
|
||||
result = mw.before_agent(self._state(msg), _runtime())
|
||||
|
||||
assert result is not None
|
||||
content = result["messages"][-1].content
|
||||
assert "uploaded in this message" in content
|
||||
assert "new.txt" in content
|
||||
assert "previous messages" in content
|
||||
assert "old.txt" in content
|
||||
|
||||
def test_no_historical_section_when_upload_dir_is_empty(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "only.txt").write_bytes(b"x")
|
||||
|
||||
msg = _human("go", files=[{"filename": "only.txt", "size": 1, "path": "/mnt/user-data/uploads/only.txt"}])
|
||||
result = mw.before_agent(self._state(msg), _runtime())
|
||||
|
||||
content = result["messages"][-1].content
|
||||
assert "previous messages" not in content
|
||||
|
||||
def test_no_historical_scan_when_thread_id_is_none(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
msg = _human("go", files=[{"filename": "f.txt", "size": 1, "path": "/mnt/user-data/uploads/f.txt"}])
|
||||
# thread_id=None → _files_from_kwargs skips existence check, no dir scan
|
||||
result = mw.before_agent(self._state(msg), _runtime(thread_id=None))
|
||||
# With no existence check, the file passes through and injection happens
|
||||
assert result is not None
|
||||
content = result["messages"][-1].content
|
||||
assert "previous messages" not in content
|
||||
|
||||
def test_message_id_preserved_on_updated_message(self, tmp_path):
|
||||
mw = _middleware(tmp_path)
|
||||
uploads_dir = _uploads_dir(tmp_path)
|
||||
(uploads_dir / "f.txt").write_bytes(b"x")
|
||||
|
||||
msg = _human("go", files=[{"filename": "f.txt", "size": 1, "path": "/mnt/user-data/uploads/f.txt"}])
|
||||
msg.id = "original-id-42"
|
||||
result = mw.before_agent(self._state(msg), _runtime())
|
||||
|
||||
assert result["messages"][-1].id == "original-id-42"
|
||||
Reference in New Issue
Block a user