186 lines
7 KiB
Python
186 lines
7 KiB
Python
"""Tests for the multimodal sandbox surface: show_image, picture_refs in
|
|
search results, binary_attachments on SandboxResult, and the absence of
|
|
the old llm() external function.
|
|
"""
|
|
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from haiku.rag.agents.analysis.dependencies import AnalysisContext
|
|
from haiku.rag.agents.analysis.sandbox import Sandbox
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
|
|
|
|
def _png_bytes(color: str = "red", size: tuple[int, int] = (8, 8)) -> bytes:
|
|
"""Generate a real PNG so PIL.Image.verify() accepts it."""
|
|
img = Image.new("RGB", size, color)
|
|
buf = BytesIO()
|
|
img.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
|
|
async def _seed_doc_with_picture(client, *, png: bytes) -> tuple[str, str]:
|
|
"""Create a Document row, replace its items with one picture row carrying
|
|
the given bytes. Returns (doc_id, self_ref)."""
|
|
doc = await client.create_document(content="x", uri="test://pic", title="Pic")
|
|
await client.document_item_repository.delete_by_document_id(doc.id)
|
|
self_ref = "#/pictures/0"
|
|
items = [
|
|
DocumentItem(
|
|
document_id=doc.id,
|
|
position=0,
|
|
self_ref=self_ref,
|
|
label="picture",
|
|
text="",
|
|
page_numbers=[1],
|
|
picture_data=png,
|
|
)
|
|
]
|
|
await client.document_item_repository.create_items(doc.id, items)
|
|
return doc.id, self_ref
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestShowImage:
|
|
"""show_image() appends a BinaryContent attachment when bytes verify."""
|
|
|
|
async def test_appends_binary_attachment(self, temp_db_path):
|
|
png = _png_bytes("red")
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id, ref = await _seed_doc_with_picture(client, png=png)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
result = await sandbox.execute(
|
|
f"await show_image({doc_id!r}, {ref!r})\nprint('done')"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert len(result.binary_attachments) == 1
|
|
att = result.binary_attachments[0]
|
|
assert att.media_type == "image/png"
|
|
assert att.identifier == ref
|
|
assert att.data == png
|
|
|
|
async def test_missing_picture_is_silent_noop(self, temp_db_path):
|
|
png = _png_bytes("red")
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id, _ = await _seed_doc_with_picture(client, png=png)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
result = await sandbox.execute(
|
|
f"await show_image({doc_id!r}, '#/pictures/999')\nprint('ok')"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert result.binary_attachments == []
|
|
|
|
async def test_invalid_bytes_rejected(self, temp_db_path):
|
|
# Garbage bytes — PIL.verify() should refuse, no attachment emitted.
|
|
garbage = b"this is not a PNG"
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id, ref = await _seed_doc_with_picture(client, png=garbage)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
result = await sandbox.execute(
|
|
f"await show_image({doc_id!r}, {ref!r})\nprint('checked')"
|
|
)
|
|
assert result.success, result.stderr
|
|
assert result.binary_attachments == []
|
|
|
|
async def test_attachments_reset_across_executes(self, temp_db_path):
|
|
png = _png_bytes("red")
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
doc_id, ref = await _seed_doc_with_picture(client, png=png)
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
first = await sandbox.execute(
|
|
f"await show_image({doc_id!r}, {ref!r})\nprint('first')"
|
|
)
|
|
second = await sandbox.execute("print('second')")
|
|
assert first.success and len(first.binary_attachments) == 1
|
|
assert second.success and second.binary_attachments == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestSearchPictureRefs:
|
|
"""search() result dicts carry a `picture_refs` list (subset of
|
|
doc_item_refs labeled 'picture'). No `image_data` base64 in the dict."""
|
|
|
|
async def test_picture_refs_extracted_from_labels(self, temp_db_path, monkeypatch):
|
|
# Build a fake SearchResult with mixed labels so we don't need an embedder.
|
|
synthetic = [
|
|
SearchResult(
|
|
chunk_id="c1",
|
|
content="hit",
|
|
document_id="d1",
|
|
document_uri="test://d1",
|
|
document_title="D1",
|
|
score=1.0,
|
|
page_numbers=[1],
|
|
headings=None,
|
|
doc_item_refs=["#/texts/0", "#/pictures/0", "#/pictures/1"],
|
|
labels=["text", "picture", "picture"],
|
|
),
|
|
SearchResult(
|
|
chunk_id="c2",
|
|
content="text only",
|
|
document_id="d1",
|
|
document_uri="test://d1",
|
|
document_title="D1",
|
|
score=0.5,
|
|
page_numbers=[2],
|
|
headings=None,
|
|
doc_item_refs=["#/texts/5"],
|
|
labels=["text"],
|
|
),
|
|
]
|
|
|
|
async def fake_search(self, *args, **kwargs):
|
|
return synthetic
|
|
|
|
async def fake_expand_context(self, results):
|
|
return results
|
|
|
|
# Patch HaikuRAG.search and expand_context so the sandbox closure runs
|
|
# without an embedder. The sandbox opens its own HaikuRAG instance, so
|
|
# we patch on the class.
|
|
monkeypatch.setattr(HaikuRAG, "search", fake_search)
|
|
monkeypatch.setattr(HaikuRAG, "expand_context", fake_expand_context)
|
|
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
pass # ensure the DB exists so the sandbox can open it read-only
|
|
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
external = sandbox._build_external_functions()
|
|
results = await external["search"]("anything")
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["picture_refs"] == ["#/pictures/0", "#/pictures/1"]
|
|
assert results[1]["picture_refs"] == []
|
|
# No raw base64 garbage in the dict.
|
|
assert "image_data" not in results[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
class TestExternalFunctionsShape:
|
|
"""llm() is gone; show_image() is present."""
|
|
|
|
async def test_llm_gone_show_image_present(self, temp_db_path):
|
|
async with HaikuRAG(temp_db_path, create=True):
|
|
pass
|
|
sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
|
|
external = sandbox._build_external_functions()
|
|
assert "llm" not in external
|
|
assert "show_image" in external
|
|
assert "search" in external
|
|
assert "list_documents" in external
|
|
|
|
|
|
# Silence unused-import flake — base64 is reserved for follow-up tests that
|
|
# decode attachment.data and compare. Kept eagerly imported for parity with
|
|
# the QA binary-content tests.
|
|
_ = base64
|