From c62fd78c7ac99f59fcda782010bac3586034b9c8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 25 Jul 2026 10:35:25 +0300 Subject: [PATCH] Accept images on CLI ask/analyze and MCP tools --- CHANGELOG.md | 2 ++ docs/cli.md | 7 ++++ docs/mcp.md | 4 +-- haiku_rag_slim/haiku/rag/app.py | 16 +++++++-- haiku_rag_slim/haiku/rag/cli.py | 12 +++++++ haiku_rag_slim/haiku/rag/mcp.py | 20 ++++++++++-- tests/test_cli.py | 58 +++++++++++++++++++++++++++++++++ tests/test_mcp.py | 49 ++++++++++++++++++++++++++++ 8 files changed, 162 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fe14b011..a1136573 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### 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`. ## [0.69.0] - 2026-07-24 diff --git a/docs/cli.md b/docs/cli.md index 6b1b7970..06c2b726 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -161,11 +161,17 @@ Filter to specific documents: haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'" ``` +Attach images to the question, for example to check an image against indexed documents: +```bash +haiku-rag ask "Does this photo satisfy the spec in the design document?" --image photo.jpg +``` + `ask` runs the [RAG capability](capabilities/rag.md) and always renders citations under the answer. When available, citations use the document title, otherwise they fall back to the URI. Flags: - `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results)) +- `--image`: Path to an image attached to the question (repeatable). Retrieval stays text-based; the model must have `vision: true` configured. ## Analyze @@ -184,6 +190,7 @@ haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial% Flags: - `--filter` / `-f`: SQL WHERE clause to restrict document access +- `--image`: Path to an image attached to the question (repeatable). Requires `vision: true` on the analysis model. See [Analysis capability](capabilities/analysis.md) for details and configuration. diff --git a/docs/mcp.md b/docs/mcp.md index 8d0f74cb..06403bb3 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -107,12 +107,12 @@ After restarting Claude Desktop, you can ask Claude to search your documents, ad - **`ask_question`** - Ask questions about your documents - `question` (required): The question to ask - `cite` (optional): Include source citations (default: false) - - `deep` (optional): Use multi-agent deep QA for complex questions (default: false) + - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable QA model) - **`analyze`** - Answer complex analytical questions via code execution - `question` (required): The question to answer - `filter` (optional): SQL WHERE clause to restrict document access - - `document` (optional): Document title/ID to pre-load (can repeat) + - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable analysis model) - Best for aggregation, computation, and multi-document analysis ## Continuous ingestion diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 23ad6942..47866555 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -592,19 +592,25 @@ class HaikuRAGApp: # pragma: no cover self, question: str, filter: str | None = None, + images: list[Path] | None = None, ): """Ask a question using the RAG system. Args: question: The question to ask filter: SQL WHERE clause to filter documents + images: Paths of images to attach to the question """ async with HaikuRAG( db_path=self.db_path, config=self.config, read_only=True, ) as self.client: - answer, citations = await self.client.ask(question, filter=filter) + answer, citations = await self.client.ask( + question, + filter=filter, + images=[path.read_bytes() for path in images] if images else None, + ) self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() @@ -619,12 +625,14 @@ class HaikuRAGApp: # pragma: no cover self, question: str, filter: str | None = None, + images: list[Path] | None = None, ): """Answer a question using the analysis capability. Args: question: The question to answer filter: SQL WHERE clause to filter documents + images: Paths of images to attach to the question """ async with HaikuRAG( db_path=self.db_path, @@ -638,7 +646,11 @@ class HaikuRAGApp: # pragma: no cover ) self.console.print() - result = await self.client.analyze(question, filter=filter) + result = await self.client.analyze( + question, + filter=filter, + images=[path.read_bytes() for path in images] if images else None, + ) self.console.print("[bold green]Answer:[/bold green]") self.console.print(Markdown(result.answer)) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index b21b9df0..dc7539b0 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -352,12 +352,18 @@ def ask( # pragma: no cover "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + image: list[Path] | None = typer.Option( + None, + "--image", + help="Path to an image to attach to the question (repeatable; requires a vision-capable model)", + ), ): app = create_app(db) asyncio.run( app.ask( question=question, filter=filter, + images=image, ) ) @@ -378,12 +384,18 @@ def analyze( # pragma: no cover "-f", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", ), + image: list[Path] | None = typer.Option( + None, + "--image", + help="Path to an image to attach to the question (repeatable; requires a vision-capable model)", + ), ): app = create_app(db) asyncio.run( app.analyze( question=question, filter=filter, + images=image, ) ) diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 4e91c754..3871fd07 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -10,6 +10,14 @@ from haiku.rag.tools.document import DocumentInfo from haiku.rag.utils import format_citations +def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: + if not images_base64: + return None + import base64 + + return [base64.b64decode(b64, validate=True) for b64 in images_base64] + + def create_mcp_server( db_path: Path, config: AppConfig = Config, read_only: bool = False ) -> FastMCP: @@ -186,19 +194,23 @@ def create_mcp_server( async def ask_question( question: str, cite: bool = False, + images_base64: list[str] | None = None, ) -> str: """Ask a question using the QA agent. Args: question: The question to ask. cite: Whether to include citations in the response. + images_base64: Base64-encoded images attached to the question + (requires a vision-capable QA model). Returns: The answer as a string. """ try: + images = _decode_images(images_base64) async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - answer, citations = await rag.ask(question) + answer, citations = await rag.ask(question, images=images) if cite and citations: answer += "\n\n" + format_citations(citations) return answer @@ -209,6 +221,7 @@ def create_mcp_server( async def analyze( question: str, filter: str | None = None, + images_base64: list[str] | None = None, ) -> str: """Answer complex questions using the analysis capability. @@ -219,13 +232,16 @@ def create_mcp_server( Args: question: The question to answer. filter: Optional SQL WHERE clause to filter documents. + images_base64: Base64-encoded images attached to the question + (requires a vision-capable analysis model). Returns: The answer as a string. """ try: + images = _decode_images(images_base64) async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - result = await rag.analyze(question, filter=filter) + result = await rag.analyze(question, filter=filter, images=images) return result.answer except Exception as e: return f"Error running analysis capability: {e!s}" # pragma: no cover diff --git a/tests/test_cli.py b/tests/test_cli.py index d16e1ff0..994df249 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -271,3 +271,61 @@ class TestTagRestore: result = runner.invoke(cli, ["--help"]) assert "--before" not in result.output assert "--at" not in result.output + + +class TestAskAnalyzeImageOption: + def test_ask_forwards_image_paths(self): + from unittest.mock import AsyncMock + + from haiku.rag.app import HaikuRAGApp + + with patch.object(HaikuRAGApp, "ask", new_callable=AsyncMock) as mock_ask: + result = runner.invoke( + cli, + ["ask", "q", "--image", "/tmp/a.png", "--image", "/tmp/b.jpg"], + ) + assert result.exit_code == 0 + from pathlib import Path + + assert mock_ask.call_args.kwargs["images"] == [ + Path("/tmp/a.png"), + Path("/tmp/b.jpg"), + ] + + def test_analyze_forwards_image_paths(self): + from unittest.mock import AsyncMock + + from haiku.rag.app import HaikuRAGApp + + with patch.object(HaikuRAGApp, "analyze", new_callable=AsyncMock) as mock: + result = runner.invoke(cli, ["analyze", "q", "--image", "/tmp/a.png"]) + assert result.exit_code == 0 + from pathlib import Path + + assert mock.call_args.kwargs["images"] == [Path("/tmp/a.png")] + + @pytest.mark.asyncio + async def test_app_ask_reads_image_bytes(self, temp_db_path, tmp_path): + from io import BytesIO + from unittest.mock import AsyncMock + + from PIL import Image as PILImage + + from haiku.rag.app import HaikuRAGApp + from haiku.rag.client import HaikuRAG + + buffer = BytesIO() + PILImage.new("RGB", (4, 4)).save(buffer, format="PNG") + img_path = tmp_path / "img.png" + img_path.write_bytes(buffer.getvalue()) + + async with HaikuRAG(temp_db_path, create=True): + pass + + with patch.object( + HaikuRAG, "ask", new_callable=AsyncMock, return_value=("answer", []) + ) as mock_ask: + app = HaikuRAGApp(db_path=temp_db_path) + await app.ask("q", images=[img_path]) + + assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()] diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 74bd866a..8f551e51 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -267,3 +267,52 @@ 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