From 75f75ec50522ce76ce261e3a9c60ecf5fd6f9323 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 6 May 2026 12:00:12 +0300 Subject: [PATCH] render attached pictures in inspector context modal under qa.model.vision --- .../rag/inspector/widgets/context_modal.py | 56 +++++++++-- tests/test_inspector.py | 97 ++++++++++++++++++- 2 files changed, 143 insertions(+), 10 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py index ad42edfa..50aeee71 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py @@ -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() diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 1f9c82b4..8c3f6ee8 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -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)) == []