Merge pull request #517 from ggozad/feat/image-in-ask-analysis
Attach images to ask/analyze and chat
This commit is contained in:
commit
1e76e66b28
24 changed files with 996 additions and 32 deletions
11
CHANGELOG.md
11
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
|
||||
|
|
|
|||
|
|
@ -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?"
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
128
haiku_rag_slim/haiku/rag/chat/widgets/image_select.py
Normal file
128
haiku_rag_slim/haiku/rag/chat/widgets/image_select.py
Normal 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")
|
||||
265
haiku_rag_slim/haiku/rag/chat/widgets/prompt.py
Normal file
265
haiku_rag_slim/haiku/rag/chat/widgets/prompt.py
Normal 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")
|
||||
|
|
@ -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,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
162
tests/chat/test_image_input.py
Normal file
162
tests/chat/test_image_input.py
Normal 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
103
tests/test_ask_images.py
Normal 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")])
|
||||
|
|
@ -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()]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue