78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from core.artifacts import ArtifactPathError
|
|
from core.attachments import (
|
|
content_for_model,
|
|
materialize_native_images,
|
|
normalize_attachment_refs,
|
|
)
|
|
|
|
|
|
class AttachmentRefTests(unittest.TestCase):
|
|
def test_normalizes_user_relative_path_to_task_relative_metadata(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
wd = root / "实验分析"
|
|
wd.mkdir()
|
|
image = wd / "显微照片.png"
|
|
image.write_bytes(b"png")
|
|
refs = normalize_attachment_refs(
|
|
[{"path": "实验分析/显微照片.png", "kind": "image"}],
|
|
working_dir=wd,
|
|
user_root=root,
|
|
)
|
|
self.assertEqual(refs[0]["path"], "显微照片.png")
|
|
self.assertEqual(refs[0]["kind"], "image")
|
|
self.assertEqual(refs[0]["size_bytes"], 3)
|
|
self.assertEqual(refs[0]["scope"], "working_dir")
|
|
|
|
def test_rejects_attachment_outside_working_dir(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
wd = root / "任务"
|
|
wd.mkdir()
|
|
(root / "outside.txt").write_text("x", encoding="utf-8")
|
|
with self.assertRaises(ArtifactPathError):
|
|
normalize_attachment_refs(
|
|
[{"path": "../outside.txt", "kind": "file"}],
|
|
working_dir=wd,
|
|
user_root=root,
|
|
)
|
|
|
|
def test_builds_provider_content_without_changing_stored_text(self):
|
|
refs = [{"path": "figures/a.png", "kind": "image"}]
|
|
self.assertEqual(
|
|
content_for_model("请分析", refs, "实验分析"),
|
|
"请分析\n\n[用户上传的参考图] 实验分析/figures/a.png",
|
|
)
|
|
self.assertEqual(
|
|
content_for_model("", refs, "实验分析"),
|
|
"[用户上传的参考图] 实验分析/figures/a.png",
|
|
)
|
|
|
|
def test_materializes_native_image_without_persisting_internal_fields(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
root = Path(tmp)
|
|
wd = root / "任务"
|
|
wd.mkdir()
|
|
(wd / "a.png").write_bytes(b"png")
|
|
messages = [{
|
|
"role": "user",
|
|
"content": "请分析\n\n[用户上传的参考图] 任务/a.png",
|
|
"_attachment_refs": [{"path": "a.png", "kind": "image"}],
|
|
}]
|
|
prepared, count = materialize_native_images(
|
|
messages, enabled=True, working_dir=wd, user_root=root,
|
|
)
|
|
self.assertEqual(count, 1)
|
|
self.assertNotIn("_attachment_refs", prepared[0])
|
|
self.assertEqual(prepared[0]["content"][0]["type"], "text")
|
|
self.assertEqual(prepared[0]["content"][1]["type"], "image_url")
|
|
self.assertTrue(
|
|
prepared[0]["content"][1]["image_url"]["url"].startswith(
|
|
"data:image/png;base64,"
|
|
)
|
|
)
|