Merge pull request #517 from ggozad/feat/image-in-ask-analysis

Attach images to ask/analyze and chat
This commit is contained in:
Yiorgis Gozadinos 2026-07-25 15:09:24 +03:00 committed by GitHub
commit 1e76e66b28
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 996 additions and 32 deletions

View file

@ -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

View file

@ -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?"

View file

@ -39,6 +39,12 @@ You can also render visual grounding from the CLI without launching the TUI:
haiku-rag visualize <chunk_id>
```
## 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.

View file

@ -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.

View file

@ -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

View file

@ -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

View file

@ -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))

View file

@ -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.

View file

@ -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

View file

@ -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:

View file

@ -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")

View file

@ -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")

View file

@ -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,
)
)

View file

@ -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,

View file

@ -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]

View file

@ -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)

View file

@ -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

View file

@ -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],

View file

@ -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

View file

@ -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)

103
tests/test_ask_images.py Normal file
View file

@ -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")])

View file

@ -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()]

View file

@ -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

View file

@ -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