diff --git a/CHANGELOG.md b/CHANGELOG.md index e84ce856..0d45a598 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,17 @@ # Changelog ## [Unreleased] +### Added + +- `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. + +### Fixed + +- Chat and inspector TUIs report the actual database-open error instead of an `AttributeError` from teardown. + ## [0.69.0] - 2026-07-24 ### Added 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/docs/cli.md b/docs/cli.md index 6b1b7970..06c2b726 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -161,11 +161,17 @@ Filter to specific documents: haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'" ``` +Attach images to the question, for example to check an image against indexed documents: +```bash +haiku-rag ask "Does this photo satisfy the spec in the design document?" --image photo.jpg +``` + `ask` runs the [RAG capability](capabilities/rag.md) and always renders citations under the answer. When available, citations use the document title, otherwise they fall back to the URI. Flags: - `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results)) +- `--image`: Path to an image attached to the question (repeatable). Retrieval stays text-based; the model must have `vision: true` configured. ## Analyze @@ -184,6 +190,7 @@ haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial% Flags: - `--filter` / `-f`: SQL WHERE clause to restrict document access +- `--image`: Path to an image attached to the question (repeatable). Requires `vision: true` on the analysis model. See [Analysis capability](capabilities/analysis.md) for details and configuration. diff --git a/docs/mcp.md b/docs/mcp.md index 8d0f74cb..06403bb3 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -107,12 +107,12 @@ After restarting Claude Desktop, you can ask Claude to search your documents, ad - **`ask_question`** - Ask questions about your documents - `question` (required): The question to ask - `cite` (optional): Include source citations (default: false) - - `deep` (optional): Use multi-agent deep QA for complex questions (default: false) + - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable QA model) - **`analyze`** - Answer complex analytical questions via code execution - `question` (required): The question to answer - `filter` (optional): SQL WHERE clause to restrict document access - - `document` (optional): Document title/ID to pre-load (can repeat) + - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable analysis model) - Best for aggregation, computation, and multi-document analysis ## Continuous ingestion diff --git a/docs/python.md b/docs/python.md index 2fb175a6..764f1717 100644 --- a/docs/python.md +++ b/docs/python.md @@ -328,6 +328,17 @@ answer, citations = await client.ask( ) ``` +Attach images to the question, for example to check an image against indexed documents: + +```python +answer, citations = await client.ask( + "Does this image satisfy the requirements in the design spec?", + images=[Path("photo.jpg").read_bytes()], +) +``` + +Images are passed to the model alongside the question. Retrieval stays text-based. The QA model must have `vision: true` in its configuration. + `client.ask` runs the [RAG capability](capabilities/rag.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, document references, and the document's metadata (`document_meta`), so UIs can render metadata keys such as a public source URL alongside the citation. The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration/index.md)). @@ -354,6 +365,8 @@ result = await client.analyze( `client.analyze` runs the [analysis capability](capabilities/analysis.md), which writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis. +`client.analyze` also accepts `images=` like `client.ask`, requiring `vision: true` on the analysis model (or the QA model when no analysis model is configured). + See [Analysis capability](capabilities/analysis.md) for details and configuration. ## Building custom agents diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 23ad6942..47866555 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -592,19 +592,25 @@ class HaikuRAGApp: # pragma: no cover self, question: str, filter: str | None = None, + images: list[Path] | None = None, ): """Ask a question using the RAG system. Args: question: The question to ask filter: SQL WHERE clause to filter documents + images: Paths of images to attach to the question """ async with HaikuRAG( db_path=self.db_path, config=self.config, read_only=True, ) as self.client: - answer, citations = await self.client.ask(question, filter=filter) + answer, citations = await self.client.ask( + question, + filter=filter, + images=[path.read_bytes() for path in images] if images else None, + ) self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() @@ -619,12 +625,14 @@ class HaikuRAGApp: # pragma: no cover self, question: str, filter: str | None = None, + images: list[Path] | None = None, ): """Answer a question using the analysis capability. Args: question: The question to answer filter: SQL WHERE clause to filter documents + images: Paths of images to attach to the question """ async with HaikuRAG( db_path=self.db_path, @@ -638,7 +646,11 @@ class HaikuRAGApp: # pragma: no cover ) self.console.print() - result = await self.client.analyze(question, filter=filter) + result = await self.client.analyze( + question, + filter=filter, + images=[path.read_bytes() for path in images] if images else None, + ) self.console.print("[bold green]Answer:[/bold green]") self.console.print(Markdown(result.answer)) diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md index abeade74..92cbf702 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md +++ b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md @@ -92,6 +92,10 @@ Section tree derived from `heading_level`: `{"doc_id", "title", "tree": [...]}` ### Cross-referencing search results with items Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in `items.jsonl`. To find which section a hit lives in: locate the item by `self_ref`, take its line index, and walk `toc.json` to find the deepest node whose `item_range` contains that index. +## Questions with attached images + +The user may attach images to their question. An attached image is part of the question, not knowledge-base content. Search the knowledge base for the criteria, standards, or facts named in the question text, cite them, and apply them to the attached image. Never refuse merely because the image itself is not in the knowledge base. + ## Strategy 1. Search first. diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/rag.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/rag.md index 1ab0c122..93c2ccc1 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/instructions/rag.md +++ b/haiku_rag_slim/haiku/rag/capabilities/instructions/rag.md @@ -21,6 +21,10 @@ Register the chunk IDs that ground your answer. Call this BEFORE writing your fi Use chunk_ids exactly as they appear in the search response — copy the full UUID verbatim. Do not abbreviate, paraphrase, or reconstruct chunk_ids from memory; the tool matches them as opaque strings. +## Questions with attached images + +The user may attach images to their question. An attached image is part of the question, not knowledge-base content. Search the knowledge base for the criteria, standards, or facts named in the question text, cite them, and apply them to the attached image. Never refuse merely because the image itself is not in the knowledge base. + ## How to answer questions 1. Call `rag_search` with relevant keywords from the question diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 6d5f0ea1..0350e5d4 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]: @@ -131,12 +139,15 @@ class ChatApp(App): async def on_mount(self) -> None: """Initialize the app when mounted.""" - self.client = HaikuRAG( + client = HaikuRAG( db_path=self.db_path, config=self.config, read_only=self.read_only, ) - await self.client.__aenter__() + # Assign only after a successful open: on_unmount must not tear down + # a client whose __aenter__ failed. + await client.__aenter__() + self.client = client self._agent = Agent( self._model, @@ -150,14 +161,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 +179,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 +209,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 +264,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 +326,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 +336,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..7d742a91 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/chat/widgets/prompt.py @@ -0,0 +1,265 @@ +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 1 1; + border-top: solid $primary-darken-1; + } + + FlexibleInput:focus-within { + border-top: solid $primary; + } + + FlexibleInput > Horizontal { + height: auto; + } + + 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/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index b21b9df0..dc7539b0 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -352,12 +352,18 @@ def ask( # pragma: no cover "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + image: list[Path] | None = typer.Option( + None, + "--image", + help="Path to an image to attach to the question (repeatable; requires a vision-capable model)", + ), ): app = create_app(db) asyncio.run( app.ask( question=question, filter=filter, + images=image, ) ) @@ -378,12 +384,18 @@ def analyze( # pragma: no cover "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + image: list[Path] | None = typer.Option( + None, + "--image", + help="Path to an image to attach to the question (repeatable; requires a vision-capable model)", + ), ): app = create_app(db) asyncio.run( app.analyze( question=question, filter=filter, + images=image, ) ) diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index bff6f32b..ecfd4b01 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -484,19 +484,21 @@ class HaikuRAG: self, question: str, filter: str | None = None, + images: Sequence[bytes] | None = None, ) -> "tuple[str, list[Citation]]": from haiku.rag.client.agents import ask - return await ask(self, question, filter) + return await ask(self, question, filter, images) async def analyze( self, question: str, filter: str | None = None, + images: Sequence[bytes] | None = None, ) -> "AnalysisResult": from haiku.rag.client.agents import analyze - return await analyze(self, question, filter) + return await analyze(self, question, filter, images) async def visualize_chunk( self, diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index 948f4819..9a739f74 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -1,10 +1,14 @@ +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from pydantic_ai import Agent if TYPE_CHECKING: + from pydantic_ai.messages import BinaryContent + from haiku.rag.client import HaikuRAG + from haiku.rag.config.models import ModelConfig from haiku.rag.sandbox import AnalysisResult from haiku.rag.store.models.citation import Citation @@ -14,10 +18,28 @@ class _AgentDeps: state: dict[str, Any] = field(default_factory=dict) +def _build_user_prompt( + question: str, + images: Sequence[bytes] | None, + model_config: "ModelConfig", +) -> "str | list[str | BinaryContent]": + if not images: + return question + if not model_config.vision: + raise ValueError( + f"Model {model_config.provider}:{model_config.name} is not configured " + "for vision (set `vision: true` on the model config to pass images)." + ) + from haiku.rag.utils import image_binary_content + + return [question, *(image_binary_content(data) for data in images)] + + async def ask( client: "HaikuRAG", question: str, filter: str | None = None, + images: Sequence[bytes] | None = None, ) -> "tuple[str, list[Citation]]": """Ask a question against the knowledge base via the RAG capability. @@ -25,6 +47,8 @@ async def ask( client: The HaikuRAG client. question: The question to ask. filter: SQL WHERE clause to filter documents. + images: Raw image bytes attached to the question (requires a + vision-capable QA model). Returns: Tuple of (answer text, list of resolved citations). @@ -44,6 +68,7 @@ async def ask( deps = _AgentDeps( state={"rag": RAGState(document_filter=filter).model_dump(mode="json")} ) + user_prompt = _build_user_prompt(question, images, client._config.qa.model) model = get_model(client._config.qa.model, client._config) agent = Agent( model, @@ -51,7 +76,7 @@ async def ask( instructions=AGENT_PREAMBLE, capabilities=[capability], ) - result = await agent.run(question, deps=deps) + result = await agent.run(user_prompt, deps=deps) state = RAGState.model_validate(deps.state["rag"]) citations = [ state.citation_index[cid] @@ -65,6 +90,7 @@ async def analyze( client: "HaikuRAG", question: str, filter: str | None = None, + images: Sequence[bytes] | None = None, ) -> "AnalysisResult": """Answer a question using the analysis capability. @@ -76,6 +102,8 @@ async def analyze( client: The HaikuRAG client. question: The question to answer. filter: SQL WHERE clause to filter documents during searches. + images: Raw image bytes attached to the question (requires a + vision-capable analysis model). Returns: AnalysisResult with the answer and resolved citations. @@ -94,15 +122,15 @@ async def analyze( "analysis": AnalysisState(document_filter=filter).model_dump(mode="json") } ) - model = get_model( - client._config.analysis.model or client._config.qa.model, client._config - ) + model_config = client._config.analysis.model or client._config.qa.model + user_prompt = _build_user_prompt(question, images, model_config) + model = get_model(model_config, client._config) agent = Agent( model, deps_type=_AgentDeps, capabilities=[capability], ) - result = await agent.run(question, deps=deps) + result = await agent.run(user_prompt, deps=deps) state = AnalysisState.model_validate(deps.state["analysis"]) citations = [ state.citation_index[cid] diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index 369e4f45..d94b035e 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -83,12 +83,15 @@ class InspectorApp(App): async def on_mount(self) -> None: """Initialize the app when mounted.""" config = get_config() - self.client = HaikuRAG( + client = HaikuRAG( db_path=self.db_path, config=config, read_only=self.read_only, ) - await self.client.__aenter__() + # Assign only after a successful open: on_unmount must not tear down + # a client whose __aenter__ failed. + await client.__aenter__() + self.client = client # Load initial documents doc_list = self.query_one(DocumentList) diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 4e91c754..3871fd07 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -10,6 +10,14 @@ from haiku.rag.tools.document import DocumentInfo from haiku.rag.utils import format_citations +def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: + if not images_base64: + return None + import base64 + + return [base64.b64decode(b64, validate=True) for b64 in images_base64] + + def create_mcp_server( db_path: Path, config: AppConfig = Config, read_only: bool = False ) -> FastMCP: @@ -186,19 +194,23 @@ def create_mcp_server( async def ask_question( question: str, cite: bool = False, + images_base64: list[str] | None = None, ) -> str: """Ask a question using the QA agent. Args: question: The question to ask. cite: Whether to include citations in the response. + images_base64: Base64-encoded images attached to the question + (requires a vision-capable QA model). Returns: The answer as a string. """ try: + images = _decode_images(images_base64) async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - answer, citations = await rag.ask(question) + answer, citations = await rag.ask(question, images=images) if cite and citations: answer += "\n\n" + format_citations(citations) return answer @@ -209,6 +221,7 @@ def create_mcp_server( async def analyze( question: str, filter: str | None = None, + images_base64: list[str] | None = None, ) -> str: """Answer complex questions using the analysis capability. @@ -219,13 +232,16 @@ def create_mcp_server( Args: question: The question to answer. filter: Optional SQL WHERE clause to filter documents. + images_base64: Base64-encoded images attached to the question + (requires a vision-capable analysis model). Returns: The answer as a string. """ try: + images = _decode_images(images_base64) async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - result = await rag.analyze(question, filter=filter) + result = await rag.analyze(question, filter=filter, images=images) return result.answer except Exception as e: return f"Error running analysis capability: {e!s}" # pragma: no cover diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 595e1afb..28c81833 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, cast from packaging.version import Version, parse if TYPE_CHECKING: + from pydantic_ai.messages import BinaryContent from pydantic_ai.profiles.openai import OpenAIModelProfile from rich.console import RenderableType @@ -37,6 +38,21 @@ def cosine_similarity(vec1: list[float], vec2: list[float]) -> float: return dot_product / (norm1 * norm2) +def image_binary_content(data: bytes) -> "BinaryContent": + """Wrap raw image bytes as BinaryContent with the sniffed media type.""" + from io import BytesIO + + from PIL import Image as PILImage + from PIL import UnidentifiedImageError + from pydantic_ai.messages import BinaryContent + + try: + fmt = PILImage.open(BytesIO(data)).format or "PNG" + except UnidentifiedImageError as e: + raise ValueError("data is not a recognizable image") from e + return BinaryContent(data=data, media_type=f"image/{fmt.lower()}") + + def apply_common_settings( settings: Any | None, settings_class: type[Any], diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 7fd6b870..40057501 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 @@ -400,3 +400,15 @@ async def test_document_filter_cleared_when_empty(temp_db_path: Path): rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) assert rag_state.document_filter is None assert app._state["rag"]["document_filter"] is None + + +@pytest.mark.asyncio +async def test_chat_app_open_failure_surfaces_real_error(tmp_path: Path): + """A failed database open must surface its own error, not an + AttributeError from tearing down a client that never opened.""" + from haiku.rag.chat.app import ChatApp + + app = ChatApp(db_path=tmp_path / "missing.lancedb", capabilities=[]) + with pytest.raises(FileNotFoundError): + async with app.run_test(): + pass diff --git a/tests/chat/test_image_input.py b/tests/chat/test_image_input.py new file mode 100644 index 00000000..4a3f6255 --- /dev/null +++ b/tests/chat/test_image_input.py @@ -0,0 +1,162 @@ +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 + + +class TestChatAppLayout: + @pytest.mark.asyncio + async def test_prompt_stays_compact_and_history_visible(self, temp_db_path): + from haiku.rag.chat.app import ChatApp + from haiku.rag.chat.widgets.chat_history import ChatHistory + 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: + await pilot.pause() + prompt = app.query_one(FlexibleInput) + history = app.query_one(ChatHistory) + assert prompt.region.height <= 4 + assert history.region.height > prompt.region.height + assert history.region.y < prompt.region.y + area = app.query_one(PostableTextArea) + assert prompt.region.contains_region(area.region) diff --git a/tests/test_ask_images.py b/tests/test_ask_images.py new file mode 100644 index 00000000..baaab700 --- /dev/null +++ b/tests/test_ask_images.py @@ -0,0 +1,103 @@ +from io import BytesIO +from pathlib import Path + +import pytest +from PIL import Image as PILImage +from pydantic_ai import Agent +from pydantic_ai.messages import BinaryContent + +from haiku.rag.client import HaikuRAG +from haiku.rag.config import AppConfig +from haiku.rag.utils import image_binary_content + + +def make_image_bytes(fmt: str) -> bytes: + buffer = BytesIO() + PILImage.new("RGB", (4, 4), color="red").save(buffer, format=fmt) + return buffer.getvalue() + + +def test_image_binary_content_sniffs_media_type(): + png = image_binary_content(make_image_bytes("PNG")) + assert isinstance(png, BinaryContent) + assert png.media_type == "image/png" + assert image_binary_content(make_image_bytes("JPEG")).media_type == "image/jpeg" + + +def test_image_binary_content_rejects_non_image_bytes(): + with pytest.raises(ValueError, match="not a recognizable image"): + image_binary_content(b"definitely not an image") + + +@pytest.fixture +def captured_run(monkeypatch): + """Capture the user prompt passed to Agent.run without running a model.""" + captured: dict = {} + + async def fake_run(self, user_prompt, **kwargs): + captured["user_prompt"] = user_prompt + + class Result: + output = "answer" + + return Result() + + monkeypatch.setattr(Agent, "run", fake_run) + return captured + + +@pytest.mark.asyncio +async def test_ask_without_images_passes_plain_string(temp_db_path: Path, captured_run): + async with HaikuRAG(temp_db_path, config=AppConfig(), create=True) as client: + await client.ask("What is this?") + assert captured_run["user_prompt"] == "What is this?" + + +@pytest.mark.asyncio +async def test_ask_with_images_passes_binary_content(temp_db_path: Path, captured_run): + config = AppConfig() + config.qa.model.vision = True + png = make_image_bytes("PNG") + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + await client.ask("What is in this image?", images=[png]) + prompt = captured_run["user_prompt"] + assert prompt[0] == "What is in this image?" + assert isinstance(prompt[1], BinaryContent) + assert prompt[1].data == png + assert prompt[1].media_type == "image/png" + + +@pytest.mark.asyncio +async def test_ask_with_images_requires_vision_model(temp_db_path: Path): + config = AppConfig() + config.qa.model.vision = False + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + with pytest.raises(ValueError, match="vision"): + await client.ask("What is this?", images=[make_image_bytes("PNG")]) + + +@pytest.mark.asyncio +async def test_analyze_with_images_passes_binary_content( + temp_db_path: Path, captured_run +): + config = AppConfig() + config.qa.model.vision = True + jpeg = make_image_bytes("JPEG") + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + await client.analyze("Does this image match?", images=[jpeg]) + prompt = captured_run["user_prompt"] + assert prompt[0] == "Does this image match?" + assert isinstance(prompt[1], BinaryContent) + assert prompt[1].media_type == "image/jpeg" + + +@pytest.mark.asyncio +async def test_analyze_with_images_checks_analysis_model_vision(temp_db_path: Path): + from haiku.rag.config.models import ModelConfig + + config = AppConfig() + config.qa.model.vision = True + config.analysis.model = ModelConfig(provider="openai", name="m", vision=False) + async with HaikuRAG(temp_db_path, config=config, create=True) as client: + with pytest.raises(ValueError, match="vision"): + await client.analyze("Does this match?", images=[make_image_bytes("PNG")]) diff --git a/tests/test_cli.py b/tests/test_cli.py index d16e1ff0..994df249 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -271,3 +271,61 @@ class TestTagRestore: result = runner.invoke(cli, ["--help"]) assert "--before" not in result.output assert "--at" not in result.output + + +class TestAskAnalyzeImageOption: + def test_ask_forwards_image_paths(self): + from unittest.mock import AsyncMock + + from haiku.rag.app import HaikuRAGApp + + with patch.object(HaikuRAGApp, "ask", new_callable=AsyncMock) as mock_ask: + result = runner.invoke( + cli, + ["ask", "q", "--image", "/tmp/a.png", "--image", "/tmp/b.jpg"], + ) + assert result.exit_code == 0 + from pathlib import Path + + assert mock_ask.call_args.kwargs["images"] == [ + Path("/tmp/a.png"), + Path("/tmp/b.jpg"), + ] + + def test_analyze_forwards_image_paths(self): + from unittest.mock import AsyncMock + + from haiku.rag.app import HaikuRAGApp + + with patch.object(HaikuRAGApp, "analyze", new_callable=AsyncMock) as mock: + result = runner.invoke(cli, ["analyze", "q", "--image", "/tmp/a.png"]) + assert result.exit_code == 0 + from pathlib import Path + + assert mock.call_args.kwargs["images"] == [Path("/tmp/a.png")] + + @pytest.mark.asyncio + async def test_app_ask_reads_image_bytes(self, temp_db_path, tmp_path): + from io import BytesIO + from unittest.mock import AsyncMock + + from PIL import Image as PILImage + + from haiku.rag.app import HaikuRAGApp + from haiku.rag.client import HaikuRAG + + buffer = BytesIO() + PILImage.new("RGB", (4, 4)).save(buffer, format="PNG") + img_path = tmp_path / "img.png" + img_path.write_bytes(buffer.getvalue()) + + async with HaikuRAG(temp_db_path, create=True): + pass + + with patch.object( + HaikuRAG, "ask", new_callable=AsyncMock, return_value=("answer", []) + ) as mock_ask: + app = HaikuRAGApp(db_path=temp_db_path) + await app.ask("q", images=[img_path]) + + assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()] diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 8c3f6ee8..9ee6d476 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -223,3 +223,15 @@ async def test_context_modal_suppresses_pictures_when_vision_disabled(): await pilot.pause() modal = app.screen assert list(modal.query(TextualImage)) == [] + + +@pytest.mark.asyncio +async def test_inspector_open_failure_surfaces_real_error(tmp_path): + """A failed database open must surface its own error, not an + AttributeError from tearing down a client that never opened.""" + from haiku.rag.inspector.app import InspectorApp + + app = InspectorApp(db_path=tmp_path / "missing.lancedb", read_only=True) + with pytest.raises(FileNotFoundError): + async with app.run_test(): + pass diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 74bd866a..eda2e2e2 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -267,3 +267,68 @@ class TestMCPImageQuery: # in search_documents_by_image rejects it. results = await search_by_image(image_base64="!!! not base64 !!!") assert results == [] + + +class TestMCPImageInput: + @pytest.mark.asyncio + async def test_ask_question_decodes_images(self, mcp_db, monkeypatch): + from base64 import b64encode + + captured = {} + + async def fake_ask(self, question, filter=None, images=None): + captured["images"] = images + return ("answer", []) + + monkeypatch.setattr(HaikuRAG, "ask", fake_ask) + mcp = create_mcp_server(mcp_db, read_only=True) + ask = await _get_tool(mcp, "ask_question") + + png = b"fake image bytes" + result = await ask(question="q", images_base64=[b64encode(png).decode()]) + assert result == "answer" + assert captured["images"] == [png] + + @pytest.mark.asyncio + async def test_analyze_decodes_images(self, mcp_db, monkeypatch): + from base64 import b64encode + from types import SimpleNamespace + + captured = {} + + async def fake_analyze(self, question, filter=None, images=None): + captured["images"] = images + return SimpleNamespace(answer="answer") + + monkeypatch.setattr(HaikuRAG, "analyze", fake_analyze) + mcp = create_mcp_server(mcp_db, read_only=True) + analyze = await _get_tool(mcp, "analyze") + + jpeg = b"fake jpeg bytes" + result = await analyze(question="q", images_base64=[b64encode(jpeg).decode()]) + assert result == "answer" + assert captured["images"] == [jpeg] + + @pytest.mark.asyncio + async def test_ask_question_rejects_invalid_base64(self, mcp_db): + mcp = create_mcp_server(mcp_db, read_only=True) + ask = await _get_tool(mcp, "ask_question") + + result = await ask(question="q", images_base64=["!!! not base64 !!!"]) + assert "Error" in result + + @pytest.mark.asyncio + async def test_ask_question_without_images_passes_none(self, mcp_db, monkeypatch): + captured = {} + + async def fake_ask(self, question, filter=None, images=None): + captured["images"] = images + return ("answer", []) + + monkeypatch.setattr(HaikuRAG, "ask", fake_ask) + mcp = create_mcp_server(mcp_db, read_only=True) + ask = await _get_tool(mcp, "ask_question") + + result = await ask(question="q") + assert result == "answer" + assert captured["images"] is None