diff --git a/CHANGELOG.md b/CHANGELOG.md index a1136573..a59f7f7e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - `HaikuRAG.ask` and `HaikuRAG.analyze` accept `images: Sequence[bytes]`, attached to the question as model input; requires `vision: true` on the driving model. - `haiku-rag ask` and `haiku-rag analyze` accept `--image PATH` (repeatable). - MCP `ask_question` and `analyze` tools accept `images_base64`. +- Chat TUI: `Ctrl+I` opens an image picker; attached images insert `[Image #N]` tokens in a multi-line prompt and are sent to the model with the message. ## [0.69.0] - 2026-07-24 diff --git a/README.md b/README.md index ce63dd72..8680d460 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p - **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion - **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query - **Question answering** — RAG capability with citations (page numbers, section headings) -- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text +- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text; attach your own images to questions in `ask`, `analyze`, MCP, and the chat TUI - **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM - **Analysis capability** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis) - **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory @@ -62,6 +62,9 @@ haiku-rag search "attention mechanism" # Ask questions with citations haiku-rag ask "What datasets were used for evaluation?" +# Ask about an image (vision-capable model) +haiku-rag ask "Does this figure match the spec in the design doc?" --image figure.png + # Analyze — complex analytical tasks via code execution haiku-rag analyze "How many documents mention transformers?" diff --git a/docs/chat.md b/docs/chat.md index 2d11128b..091a3eb1 100644 --- a/docs/chat.md +++ b/docs/chat.md @@ -39,6 +39,12 @@ You can also render visual grounding from the CLI without launching the TUI: haiku-rag visualize ``` +## Attaching images + +Press `Ctrl+I` to open the image picker: a directory tree filtered to image files with a live preview. Selecting an image inserts an `[Image #N]` token at the cursor and attaches the image to your next message. Tokens delete as a unit with backspace or delete, and you can place them anywhere in the text to control where each image appears relative to your words. + +Retrieval stays text-based; the images are sent to the model alongside your message, so the driving model needs `vision: true` in its configuration. + ## Command palette `Ctrl+P` opens the palette. diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 6d5f0ea1..53014e3d 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any import textual_image.widget # noqa: F401 - import early for renderer detection from pydantic_ai import Agent from pydantic_ai.messages import ( + BinaryContent, FunctionToolCallEvent, FunctionToolResultEvent, PartDeltaEvent, @@ -19,13 +20,19 @@ from pydantic_ai.messages import ( from pydantic_ai.run import AgentRunResultEvent from textual.app import App, SystemCommand from textual.binding import Binding -from textual.widgets import Footer, Header, Input +from textual.widgets import Footer, Header from textual.worker import Worker from haiku.rag.capabilities._base import RAGCapabilityBase from haiku.rag.capabilities.analysis import AnalysisState from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget +from haiku.rag.chat.widgets.image_select import ImageAdded +from haiku.rag.chat.widgets.prompt import ( + FlexibleInput, + PostableTextArea, + build_user_prompt, +) from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config from haiku.rag.telemetry import configure as configure_telemetry @@ -95,6 +102,7 @@ class ChatApp(App): self._is_processing = False self._current_worker: Worker[None] | None = None self._document_filter: list[str] = [] + self._images: list[bytes] = [] # Stable per-launch id for multi-turn model and telemetry correlation. self._conversation_id = str(uuid.uuid4()) @@ -102,7 +110,7 @@ class ChatApp(App): """Compose the UI layout.""" yield Header() yield ChatHistory(id="chat-history") - yield Input(placeholder="Ask a question...", id="chat-input") + yield FlexibleInput(id="chat-input") yield Footer() def get_system_commands(self, screen: Any) -> Iterable[SystemCommand]: @@ -150,14 +158,14 @@ class ChatApp(App): capability.state_type().model_dump(mode="json") ) - self.query_one(Input).focus() + self.query_one(FlexibleInput).focus() async def on_unmount(self) -> None: """Clean up when unmounting.""" if self.client: await self.client.__aexit__(None, None, None) - async def on_input_submitted(self, event: Input.Submitted) -> None: + async def on_flexible_input_submitted(self, event: FlexibleInput.Submitted) -> None: """Handle user input submission.""" user_message = event.value.strip() if not user_message or self._is_processing: @@ -168,13 +176,24 @@ class ChatApp(App): chat_history = self.query_one(ChatHistory) await chat_history.add_message("user", user_message) + user_prompt = build_user_prompt(user_message, self._images) + self._images = [] + self._is_processing = True - self.query_one(Input).disabled = True + self.query_one(FlexibleInput).disabled = True self._current_worker = self.run_worker( - self._run_agent(user_message), exclusive=True + self._run_agent(user_prompt), exclusive=True ) - async def _run_agent(self, user_message: str) -> None: + def on_image_added(self, event: ImageAdded) -> None: + """Attach a picked image and insert its token into the prompt.""" + self._images.append(event.data) + prompt = self.query_one(FlexibleInput) + prompt.insert_at_cursor(f"[Image #{len(self._images)}]") + prompt.focus() + self.notify(f"Attached {event.path.name}") + + async def _run_agent(self, user_prompt: str | list[str | BinaryContent]) -> None: """Run the agent in a background worker.""" if not self._agent: return @@ -187,7 +206,7 @@ class ChatApp(App): try: async with self._agent.run_stream_events( - user_message, + user_prompt, message_history=self._messages, conversation_id=self._conversation_id, deps=deps, @@ -242,7 +261,7 @@ class ChatApp(App): finally: self._is_processing = False self._current_worker = None - chat_input = self.query_one(Input) + chat_input = self.query_one(FlexibleInput) chat_input.disabled = False chat_input.focus() @@ -304,7 +323,7 @@ class ChatApp(App): """Focus the input field, or cancel if processing.""" if self._is_processing and self._current_worker: self._current_worker.cancel() - self.query_one(Input).focus() + self.query_one(FlexibleInput).focus() def _clear_citation_selection(self) -> None: """Clear citation selection.""" @@ -314,7 +333,7 @@ class ChatApp(App): def on_descendant_focus(self, _event: object) -> None: """Clear citation selection when chat input is focused.""" - if isinstance(self.focused, Input) and self.focused.id == "chat-input": + if isinstance(self.focused, PostableTextArea): self._clear_citation_selection() async def action_show_visual(self) -> None: diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/image_select.py b/haiku_rag_slim/haiku/rag/chat/widgets/image_select.py new file mode 100644 index 00000000..a1d2230f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/widgets/image_select.py @@ -0,0 +1,128 @@ +from collections.abc import Iterable +from io import BytesIO +from pathlib import Path + +import PIL.Image as PILImage +from PIL import UnidentifiedImageError +from textual import on +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Container, Horizontal, Vertical +from textual.message import Message +from textual.screen import ModalScreen +from textual.widgets import DirectoryTree, Input, Label +from textual_image.widget import Image + +IMAGE_EXTENSIONS = PILImage.registered_extensions() + + +def encode_jpeg(path: Path) -> bytes: + """Re-encode an image file as RGB JPEG bytes.""" + image = PILImage.open(path) + if image.mode != "RGB": + image = image.convert("RGB") + buffer = BytesIO() + image.save(buffer, format="JPEG") + return buffer.getvalue() + + +class ImageAdded(Message): + """Emitted when the user picks an image to attach to the prompt.""" + + def __init__(self, path: Path, data: bytes) -> None: + self.path = path + self.data = data + super().__init__() + + +class ImageDirectoryTree(DirectoryTree): + def filter_paths(self, paths: Iterable[Path]) -> Iterable[Path]: + return [ + path for path in paths if path.suffix in IMAGE_EXTENSIONS or path.is_dir() + ] + + +class ImageSelect(ModalScreen[tuple[Path, bytes]]): + """Modal for picking an image file, with a live preview.""" + + BINDINGS = [ + Binding("escape", "cancel", "Cancel", show=False), + ] + + CSS = """ + ImageSelect { + align: center middle; + background: rgba(0, 0, 0, 0.5); + } + + #image-select-container { + width: 80%; + height: 80%; + background: $surface; + border: tall $primary; + padding: 1 2; + } + + #image-directory-tree { + width: 40%; + } + + #image-preview { + width: 60%; + } + + #image-preview #image { + width: auto; + height: auto; + } + + #image-select-container Input { + margin-bottom: 1; + } + """ + + def action_cancel(self) -> None: + self.dismiss() + + async def on_mount(self) -> None: + tree = self.query_one(ImageDirectoryTree) + tree.show_guides = False + tree.focus() + + @on(DirectoryTree.FileSelected) + async def on_image_selected(self, event: DirectoryTree.FileSelected) -> None: + try: + self.dismiss((event.path, encode_jpeg(event.path))) + except UnidentifiedImageError: + self.dismiss() + + @on(DirectoryTree.NodeHighlighted) + async def on_image_highlighted(self, event: DirectoryTree.NodeHighlighted) -> None: + if event.node.data is None: + return + path = event.node.data.path + preview = self.query_one(Image) + if path.suffix in IMAGE_EXTENSIONS: + try: + preview.image = PILImage.open(path.as_posix()) + except UnidentifiedImageError: + preview.image = None + else: + preview.image = None + + @on(Input.Changed) + async def on_root_changed(self, event: Input.Changed) -> None: + path = Path(event.value) + if path.exists() and path.is_dir(): + self.query_one(ImageDirectoryTree).path = path + + def compose(self) -> ComposeResult: + with Container(id="image-select-container"): + with Horizontal(): + with Vertical(id="image-directory-tree"): + yield Label("Select an image:") + yield Label("Root:") + yield Input(Path("./").resolve().as_posix()) + yield ImageDirectoryTree("./") + with Container(id="image-preview"): + yield Image(id="image") diff --git a/haiku_rag_slim/haiku/rag/chat/widgets/prompt.py b/haiku_rag_slim/haiku/rag/chat/widgets/prompt.py new file mode 100644 index 00000000..8aa3e7b6 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/widgets/prompt.py @@ -0,0 +1,261 @@ +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +from pydantic_ai.messages import BinaryContent +from rich.style import Style +from textual import on +from textual.app import ComposeResult +from textual.binding import Binding +from textual.containers import Horizontal +from textual.css.query import NoMatches +from textual.message import Message +from textual.reactive import reactive +from textual.widget import Widget +from textual.widgets import Static, TextArea + +from haiku.rag.chat.widgets.image_select import ImageAdded, ImageSelect +from haiku.rag.utils import image_binary_content + +MAX_PROMPT_LINES = 10 + +IMAGE_TOKEN_RE = re.compile(r"\[Image #(\d+)\]") +_IMAGE_TOKEN_HIGHLIGHT = "image-token" +_IMAGE_TOKEN_STYLE = Style(color="bright_cyan", bold=True) + + +def build_user_prompt( + text: str, images: list[bytes] +) -> str | list[str | BinaryContent]: + """Interleave text and images by ``[Image #N]`` tokens, 1-indexed. + + Without tokens, images are appended after the text. Out-of-range tokens + stay as literal text. + """ + matches = list(IMAGE_TOKEN_RE.finditer(text)) + if not matches: + if not images: + return text + parts: list[str | BinaryContent] = [text] if text else [] + parts.extend(image_binary_content(data) for data in images) + return parts + + parts = [] + last = 0 + for m in matches: + if m.start() > last: + parts.append(text[last : m.start()]) + idx = int(m.group(1)) + if 1 <= idx <= len(images): + parts.append(image_binary_content(images[idx - 1])) + else: + parts.append(m.group(0)) + last = m.end() + if last < len(text): + parts.append(text[last:]) + if not any(isinstance(p, BinaryContent) for p in parts): + return text + return parts + + +class PostableTextArea(TextArea): + """TextArea that auto-grows with content, submits on Enter, newline on Shift+Enter.""" + + BINDINGS = TextArea.BINDINGS + [ + Binding( + key="enter", + action="submit", + description="submit", + show=True, + key_display=None, + priority=True, + ), + Binding( + key="shift+enter", + action="newline", + description="newline", + show=True, + key_display=None, + priority=True, + id="newline", + ), + Binding( + key="ctrl+m", + action="newline", + description="newline", + show=False, + key_display=None, + priority=True, + ), + ] + + @dataclass + class Submitted(Message): + input: "PostableTextArea" + value: str + + @property + def control(self) -> "PostableTextArea": + return self.input + + def on_mount(self) -> None: + self.soft_wrap = True + self._resize_to_content() + if self._theme is not None: # pragma: no branch + self._theme.syntax_styles[_IMAGE_TOKEN_HIGHLIGHT] = _IMAGE_TOKEN_STYLE + self._build_highlight_map() + self.refresh() + + def _resize_to_content(self) -> None: + line_count = max(self.wrapped_document.height, 1) + self.styles.height = min(line_count, MAX_PROMPT_LINES) + + def _build_highlight_map(self) -> None: + super()._build_highlight_map() + for line_idx in range(self.document.line_count): + line = self.document.get_line(line_idx) + for m in IMAGE_TOKEN_RE.finditer(line): + self._highlights[line_idx].append( + (m.start(), m.end(), _IMAGE_TOKEN_HIGHLIGHT) + ) + + def action_submit(self) -> None: + self.post_message(PostableTextArea.Submitted(self, self.text)) + + def action_newline(self) -> None: + self.insert("\n") + + def action_delete_left(self) -> None: + if self.selection.start != self.selection.end: + super().action_delete_left() + return + span = self._image_token_span_at_cursor("left") + if span is not None: + self.delete(*span) + return + super().action_delete_left() + + def action_delete_right(self) -> None: + if self.selection.start != self.selection.end: + super().action_delete_right() + return + span = self._image_token_span_at_cursor("right") + if span is not None: + self.delete(*span) + return + super().action_delete_right() + + def _image_token_span_at_cursor( + self, direction: Literal["left", "right"] + ) -> tuple[tuple[int, int], tuple[int, int]] | None: + row, col = self.cursor_location + line = self.document.get_line(row) + for m in IMAGE_TOKEN_RE.finditer(line): + s, e = m.start(), m.end() + if direction == "left" and s < col <= e: + return (row, s), (row, e) + if direction == "right" and s <= col < e: + return (row, s), (row, e) + return None + + +class FlexibleInput(Widget): + """Prompt input with image attachment via ctrl+i.""" + + text = reactive("") + + BINDINGS = [ + Binding("ctrl+i", "add_image", "add image", id="add.image"), + ] + + DEFAULT_CSS = """ + FlexibleInput { + height: auto; + padding: 0 1; + border-top: solid $primary-darken-1; + } + + FlexibleInput:focus-within { + border-top: solid $primary; + } + + FlexibleInput #promptMarker { + width: 2; + height: 1; + color: $primary; + } + + FlexibleInput #promptArea { + background: transparent; + border: none; + padding: 0; + } + + FlexibleInput #promptArea > .text-area--cursor-line { + background: transparent; + } + """ + + @dataclass + class Submitted(Message): + input: "FlexibleInput" + value: str + + @property + def control(self) -> "FlexibleInput": + return self.input + + def __init__(self, text: str = "", *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.text = text + + def on_mount(self) -> None: + textarea = self.query_one("#promptArea", PostableTextArea) + textarea.show_line_numbers = False + textarea.focus() + + def clear(self) -> None: + self.text = "" + self.query_one("#promptArea", PostableTextArea).text = "" + + def focus(self, scroll_visible: bool = True) -> "FlexibleInput": + self.query_one("#promptArea", PostableTextArea).focus() + return self + + def insert_at_cursor(self, text: str) -> None: + self.query_one("#promptArea", PostableTextArea).insert(text) + + def watch_text(self) -> None: + try: + textarea = self.query_one("#promptArea", PostableTextArea) + if textarea.text != self.text: + textarea.text = self.text + except NoMatches: + pass + + def action_add_image(self) -> None: + async def on_image_selected(image: tuple[Path, bytes] | None) -> None: + if image is None: + return + path, data = image + self.post_message(ImageAdded(path, data)) + + self.app.push_screen(ImageSelect(), on_image_selected) + + @on(PostableTextArea.Submitted, "#promptArea") + def on_textarea_submitted(self, event: PostableTextArea.Submitted) -> None: + self.post_message(self.Submitted(self, event.input.text)) + event.stop() + event.prevent_default() + + @on(TextArea.Changed, "#promptArea") + def on_area_changed(self, event: TextArea.Changed) -> None: + self.text = event.text_area.text + if isinstance(event.text_area, PostableTextArea): # pragma: no branch + event.text_area._resize_to_content() + + def compose(self) -> ComposeResult: + with Horizontal(): + yield Static("❯", id="promptMarker") + yield PostableTextArea(id="promptArea") diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 7fd6b870..94da13ef 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -124,7 +124,7 @@ def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None): @pytest.mark.asyncio async def test_chat_app_has_required_widgets(temp_db_path: Path): - """Test that ChatApp has the required widgets: ChatHistory, Input.""" + """Test that ChatApp has the required widgets: ChatHistory, FlexibleInput.""" from haiku.rag.chat.widgets.chat_history import ChatHistory app, mock_client = _make_app(temp_db_path) @@ -134,9 +134,9 @@ async def test_chat_app_has_required_widgets(temp_db_path: Path): chat_history = app.query_one(ChatHistory) assert chat_history is not None - from textual.widgets import Input + from haiku.rag.chat.widgets.prompt import FlexibleInput - chat_input = app.query_one(Input) + chat_input = app.query_one(FlexibleInput) assert chat_input is not None diff --git a/tests/chat/test_image_input.py b/tests/chat/test_image_input.py new file mode 100644 index 00000000..73a5e6f4 --- /dev/null +++ b/tests/chat/test_image_input.py @@ -0,0 +1,140 @@ +from io import BytesIO +from pathlib import Path + +import pytest +from PIL import Image as PILImage +from pydantic_ai.messages import BinaryContent +from textual.app import App + +from haiku.rag.chat.widgets.image_select import ( + ImageDirectoryTree, + ImageSelect, + encode_jpeg, +) +from haiku.rag.chat.widgets.prompt import ( + FlexibleInput, + PostableTextArea, + build_user_prompt, +) + + +def make_image_bytes(fmt: str = "PNG") -> bytes: + buffer = BytesIO() + PILImage.new("RGB", (4, 4), color="red").save(buffer, format=fmt) + return buffer.getvalue() + + +class TestBuildUserPrompt: + def test_no_images_returns_text(self): + assert build_user_prompt("hello", []) == "hello" + + def test_images_without_tokens_append_at_end(self): + img = make_image_bytes() + prompt = build_user_prompt("hello", [img]) + assert prompt[0] == "hello" + assert isinstance(prompt[1], BinaryContent) + assert prompt[1].data == img + + def test_tokens_interleave_images(self): + first = make_image_bytes("PNG") + second = make_image_bytes("JPEG") + prompt = build_user_prompt( + "compare [Image #1] with [Image #2] please", [first, second] + ) + assert prompt[0] == "compare " + assert isinstance(prompt[1], BinaryContent) + assert prompt[1].data == first + assert prompt[2] == " with " + assert isinstance(prompt[3], BinaryContent) + assert prompt[3].data == second + assert prompt[4] == " please" + + def test_out_of_range_token_stays_literal(self): + assert build_user_prompt("see [Image #2]", [make_image_bytes()]) == ( + "see [Image #2]" + ) + + +class TestImageDirectoryTree: + def test_filter_paths_keeps_images_and_dirs(self, tmp_path): + (tmp_path / "photo.png").write_bytes(make_image_bytes()) + (tmp_path / "notes.txt").write_text("nope") + (tmp_path / "subdir").mkdir() + + tree = ImageDirectoryTree(tmp_path) + kept = {p.name for p in tree.filter_paths(tmp_path.iterdir())} + assert kept == {"photo.png", "subdir"} + + +class TestEncodeJpeg: + def test_reencodes_to_jpeg(self, tmp_path): + path = tmp_path / "img.png" + buffer = BytesIO() + PILImage.new("RGBA", (4, 4)).save(buffer, format="PNG") + path.write_bytes(buffer.getvalue()) + + data = encode_jpeg(path) + assert PILImage.open(BytesIO(data)).format == "JPEG" + + +class PromptApp(App): + def __init__(self) -> None: + super().__init__() + self.submitted: list[str] = [] + + def compose(self): + yield FlexibleInput("", id="chat-input") + + def on_flexible_input_submitted(self, event: FlexibleInput.Submitted) -> None: + self.submitted.append(event.value) + + +class TestFlexibleInput: + @pytest.mark.asyncio + async def test_enter_submits_text(self): + app = PromptApp() + async with app.run_test() as pilot: + area = app.query_one(PostableTextArea) + area.focus() + area.text = "hello" + await pilot.press("enter") + assert app.submitted == ["hello"] + + @pytest.mark.asyncio + async def test_backspace_deletes_whole_image_token(self): + app = PromptApp() + async with app.run_test() as pilot: + area = app.query_one(PostableTextArea) + area.focus() + area.text = "look at [Image #1] now" + area.cursor_location = (0, 18) + await pilot.press("backspace") + assert area.text == "look at now" + + @pytest.mark.asyncio + async def test_ctrl_i_opens_image_select(self): + app = PromptApp() + async with app.run_test() as pilot: + app.query_one(PostableTextArea).focus() + await pilot.press("ctrl+i") + await pilot.pause() + assert isinstance(app.screen, ImageSelect) + + +class TestChatAppImageAttach: + @pytest.mark.asyncio + async def test_image_added_inserts_token_and_stores_bytes(self, temp_db_path): + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.image_select import ImageAdded + from haiku.rag.client import HaikuRAG + + async with HaikuRAG(temp_db_path, create=True): + pass + + app = ChatApp(db_path=temp_db_path, capabilities=[]) + async with app.run_test() as pilot: + data = make_image_bytes() + app.post_message(ImageAdded(Path("img.png"), data)) + await pilot.pause() + assert app._images == [data] + assert "[Image #1]" in app.query_one(PostableTextArea).text