render attached pictures in inspector context modal under qa.model.vision

This commit is contained in:
Yiorgis Gozadinos 2026-05-06 12:00:12 +03:00
parent 7fac35d2af
commit 75f75ec505
No known key found for this signature in database
2 changed files with 143 additions and 10 deletions

View file

@ -1,10 +1,14 @@
import base64
from io import BytesIO
from typing import TYPE_CHECKING
from PIL import Image as PILImage
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import VerticalScroll
from textual.screen import Screen
from textual.widgets import Markdown, Static
from textual_image.widget import Image as TextualImage
from haiku.rag.store.models import SearchResult
@ -42,6 +46,17 @@ class ContextModal(Screen):
#context-content Markdown {
width: 100%;
}
#context-content Image {
margin: 1 0;
height: auto;
max-height: 30;
}
.picture-caption {
color: $text-muted;
margin-top: 1;
}
"""
def __init__(self, chunk: "Chunk", client: "HaikuRAG"):
@ -56,8 +71,6 @@ class ContextModal(Screen):
yield self._content_widget
async def on_mount(self) -> None:
"""Load and display the expanded context."""
# Create a SearchResult from the chunk
chunk_meta = self.chunk.get_chunk_metadata()
search_result = SearchResult(
content=self.chunk.content,
@ -72,18 +85,45 @@ class ContextModal(Screen):
labels=chunk_meta.labels,
)
# Expand context using the client (this is what agents actually receive)
expanded_results = await self.client.expand_context([search_result])
expanded = expanded_results[0] if expanded_results else search_result
formatted = expanded.format_for_agent()
vision = self.client._config.qa.model.vision
attached = expanded.image_data or {}
content = (
"*This is how the chunk appears to agents after context expansion:*\n\n---\n\n"
f"{formatted}"
if vision and attached:
preface = (
f"*This is what the LLM receives. Vision is enabled — {len(attached)} "
"picture(s) are attached as image content (rendered below).*"
)
elif vision:
preface = (
"*This is what the LLM receives. Vision is enabled, but no pictures "
"are in this expanded context.*"
)
elif attached:
preface = (
f"*This is what the LLM receives. Vision is disabled — {len(attached)} "
"picture(s) are suppressed; any VLM descriptions remain inline below.*"
)
else:
preface = "*This is what the LLM receives.*"
await self._content_widget.update(
f"{preface}\n\n---\n\n{expanded.format_for_agent()}"
)
await self._content_widget.update(content)
if vision and attached:
scroll = self.query_one("#context-content", VerticalScroll)
for self_ref, b64 in attached.items():
try:
pil = PILImage.open(BytesIO(base64.b64decode(b64)))
except Exception:
continue
await scroll.mount(
Static(f"[b]{self_ref}[/b]", classes="picture-caption")
)
await scroll.mount(TextualImage(pil))
async def action_dismiss(self, result=None) -> None:
self.app.pop_screen()

View file

@ -1,14 +1,47 @@
from unittest.mock import AsyncMock, patch
import base64
from io import BytesIO
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from PIL import Image as PILImage
from typer.testing import CliRunner
from haiku.rag.cli import _cli as cli
from haiku.rag.store.models import Document
from haiku.rag.store.models import Chunk, Document, SearchResult
runner = CliRunner()
def _png_b64(color: str = "red", size: tuple[int, int] = (8, 8)) -> str:
img = PILImage.new("RGB", size, color)
buf = BytesIO()
img.save(buf, format="PNG")
return base64.b64encode(buf.getvalue()).decode("ascii")
def _make_client(*, vision: bool, image_data: dict[str, str] | None) -> MagicMock:
"""Stub HaikuRAG that returns one expanded SearchResult with the given attachments."""
from haiku.rag.config import AppConfig
config = AppConfig()
config.qa.model.vision = vision
expanded = SearchResult(
content="expanded text incl. picture descriptions if any",
score=0.5,
chunk_id="chunk-1",
document_id="doc-1",
doc_item_refs=["#/texts/0"] + (list(image_data) if image_data else []),
page_numbers=[1],
labels=["paragraph"] + (["picture"] if image_data else []),
image_data=image_data,
)
client = MagicMock()
client._config = config
client.expand_context = AsyncMock(return_value=[expanded])
return client
def test_inspect_command():
"""Test inspect command launches inspector TUI."""
with patch("haiku.rag.inspector.run_inspector") as mock_inspector:
@ -130,3 +163,63 @@ async def test_document_list_tracks_has_more():
await doc_list.load_more(mock_client)
# After loading partial batch (<50), has_more should be False
assert doc_list.has_more is False
@pytest.mark.asyncio
async def test_context_modal_renders_pictures_when_vision_enabled():
"""ContextModal must mount one TextualImage per attached picture
when qa.model.vision is True that's what the LLM actually sees."""
from textual.app import App
from textual_image.widget import Image as TextualImage
from haiku.rag.inspector.widgets.context_modal import ContextModal
chunk = Chunk(
id="chunk-1", document_id="doc-1", content="raw chunk text", metadata={}
)
client = _make_client(
vision=True,
image_data={"#/pictures/0": _png_b64("red"), "#/pictures/1": _png_b64("blue")},
)
class TestApp(App):
async def on_mount(self) -> None:
await self.push_screen(ContextModal(chunk=chunk, client=client))
app = TestApp()
async with app.run_test() as pilot:
await pilot.pause()
await pilot.pause()
modal = app.screen
images = list(modal.query(TextualImage))
assert len(images) == 2
@pytest.mark.asyncio
async def test_context_modal_suppresses_pictures_when_vision_disabled():
"""ContextModal must NOT mount picture widgets when vision is off,
even if expansion attached image_data text-only models would never
see those bytes, so the inspector shouldn't show them either."""
from textual.app import App
from textual_image.widget import Image as TextualImage
from haiku.rag.inspector.widgets.context_modal import ContextModal
chunk = Chunk(
id="chunk-1", document_id="doc-1", content="raw chunk text", metadata={}
)
client = _make_client(
vision=False,
image_data={"#/pictures/0": _png_b64("red")},
)
class TestApp(App):
async def on_mount(self) -> None:
await self.push_screen(ContextModal(chunk=chunk, client=client))
app = TestApp()
async with app.run_test() as pilot:
await pilot.pause()
await pilot.pause()
modal = app.screen
assert list(modal.query(TextualImage)) == []