Accept images on CLI ask/analyze and MCP tools
This commit is contained in:
parent
4c5050d161
commit
c62fd78c7a
8 changed files with 162 additions and 6 deletions
|
|
@ -4,6 +4,8 @@
|
||||||
### Added
|
### Added
|
||||||
|
|
||||||
- `HaikuRAG.ask` and `HaikuRAG.analyze` accept `images: Sequence[bytes]`, attached to the question as model input; requires `vision: true` on the driving model.
|
- `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
|
## [0.69.0] - 2026-07-24
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -161,11 +161,17 @@ Filter to specific documents:
|
||||||
haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
|
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.
|
`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:
|
Flags:
|
||||||
|
|
||||||
- `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results))
|
- `--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
|
## Analyze
|
||||||
|
|
||||||
|
|
@ -184,6 +190,7 @@ haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial%
|
||||||
Flags:
|
Flags:
|
||||||
|
|
||||||
- `--filter` / `-f`: SQL WHERE clause to restrict document access
|
- `--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.
|
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
|
- **`ask_question`** - Ask questions about your documents
|
||||||
- `question` (required): The question to ask
|
- `question` (required): The question to ask
|
||||||
- `cite` (optional): Include source citations (default: false)
|
- `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
|
- **`analyze`** - Answer complex analytical questions via code execution
|
||||||
- `question` (required): The question to answer
|
- `question` (required): The question to answer
|
||||||
- `filter` (optional): SQL WHERE clause to restrict document access
|
- `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
|
- Best for aggregation, computation, and multi-document analysis
|
||||||
|
|
||||||
## Continuous ingestion
|
## Continuous ingestion
|
||||||
|
|
|
||||||
|
|
@ -592,19 +592,25 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
self,
|
self,
|
||||||
question: str,
|
question: str,
|
||||||
filter: str | None = None,
|
filter: str | None = None,
|
||||||
|
images: list[Path] | None = None,
|
||||||
):
|
):
|
||||||
"""Ask a question using the RAG system.
|
"""Ask a question using the RAG system.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
question: The question to ask
|
question: The question to ask
|
||||||
filter: SQL WHERE clause to filter documents
|
filter: SQL WHERE clause to filter documents
|
||||||
|
images: Paths of images to attach to the question
|
||||||
"""
|
"""
|
||||||
async with HaikuRAG(
|
async with HaikuRAG(
|
||||||
db_path=self.db_path,
|
db_path=self.db_path,
|
||||||
config=self.config,
|
config=self.config,
|
||||||
read_only=True,
|
read_only=True,
|
||||||
) as self.client:
|
) 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(f"[bold blue]Question:[/bold blue] {question}")
|
||||||
self.console.print()
|
self.console.print()
|
||||||
|
|
@ -619,12 +625,14 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
self,
|
self,
|
||||||
question: str,
|
question: str,
|
||||||
filter: str | None = None,
|
filter: str | None = None,
|
||||||
|
images: list[Path] | None = None,
|
||||||
):
|
):
|
||||||
"""Answer a question using the analysis capability.
|
"""Answer a question using the analysis capability.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
question: The question to answer
|
question: The question to answer
|
||||||
filter: SQL WHERE clause to filter documents
|
filter: SQL WHERE clause to filter documents
|
||||||
|
images: Paths of images to attach to the question
|
||||||
"""
|
"""
|
||||||
async with HaikuRAG(
|
async with HaikuRAG(
|
||||||
db_path=self.db_path,
|
db_path=self.db_path,
|
||||||
|
|
@ -638,7 +646,11 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
)
|
)
|
||||||
self.console.print()
|
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("[bold green]Answer:[/bold green]")
|
||||||
self.console.print(Markdown(result.answer))
|
self.console.print(Markdown(result.answer))
|
||||||
|
|
|
||||||
|
|
@ -352,12 +352,18 @@ def ask( # pragma: no cover
|
||||||
"-f",
|
"-f",
|
||||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
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)
|
app = create_app(db)
|
||||||
asyncio.run(
|
asyncio.run(
|
||||||
app.ask(
|
app.ask(
|
||||||
question=question,
|
question=question,
|
||||||
filter=filter,
|
filter=filter,
|
||||||
|
images=image,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -378,12 +384,18 @@ def analyze( # pragma: no cover
|
||||||
"-f",
|
"-f",
|
||||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
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)
|
app = create_app(db)
|
||||||
asyncio.run(
|
asyncio.run(
|
||||||
app.analyze(
|
app.analyze(
|
||||||
question=question,
|
question=question,
|
||||||
filter=filter,
|
filter=filter,
|
||||||
|
images=image,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -10,6 +10,14 @@ from haiku.rag.tools.document import DocumentInfo
|
||||||
from haiku.rag.utils import format_citations
|
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(
|
def create_mcp_server(
|
||||||
db_path: Path, config: AppConfig = Config, read_only: bool = False
|
db_path: Path, config: AppConfig = Config, read_only: bool = False
|
||||||
) -> FastMCP:
|
) -> FastMCP:
|
||||||
|
|
@ -186,19 +194,23 @@ def create_mcp_server(
|
||||||
async def ask_question(
|
async def ask_question(
|
||||||
question: str,
|
question: str,
|
||||||
cite: bool = False,
|
cite: bool = False,
|
||||||
|
images_base64: list[str] | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Ask a question using the QA agent.
|
"""Ask a question using the QA agent.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
question: The question to ask.
|
question: The question to ask.
|
||||||
cite: Whether to include citations in the response.
|
cite: Whether to include citations in the response.
|
||||||
|
images_base64: Base64-encoded images attached to the question
|
||||||
|
(requires a vision-capable QA model).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The answer as a string.
|
The answer as a string.
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
|
images = _decode_images(images_base64)
|
||||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
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:
|
if cite and citations:
|
||||||
answer += "\n\n" + format_citations(citations)
|
answer += "\n\n" + format_citations(citations)
|
||||||
return answer
|
return answer
|
||||||
|
|
@ -209,6 +221,7 @@ def create_mcp_server(
|
||||||
async def analyze(
|
async def analyze(
|
||||||
question: str,
|
question: str,
|
||||||
filter: str | None = None,
|
filter: str | None = None,
|
||||||
|
images_base64: list[str] | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""Answer complex questions using the analysis capability.
|
"""Answer complex questions using the analysis capability.
|
||||||
|
|
||||||
|
|
@ -219,13 +232,16 @@ def create_mcp_server(
|
||||||
Args:
|
Args:
|
||||||
question: The question to answer.
|
question: The question to answer.
|
||||||
filter: Optional SQL WHERE clause to filter documents.
|
filter: Optional SQL WHERE clause to filter documents.
|
||||||
|
images_base64: Base64-encoded images attached to the question
|
||||||
|
(requires a vision-capable analysis model).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The answer as a string.
|
The answer as a string.
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
|
images = _decode_images(images_base64)
|
||||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
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
|
return result.answer
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return f"Error running analysis capability: {e!s}" # pragma: no cover
|
return f"Error running analysis capability: {e!s}" # pragma: no cover
|
||||||
|
|
|
||||||
|
|
@ -271,3 +271,61 @@ class TestTagRestore:
|
||||||
result = runner.invoke(cli, ["--help"])
|
result = runner.invoke(cli, ["--help"])
|
||||||
assert "--before" not in result.output
|
assert "--before" not in result.output
|
||||||
assert "--at" 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()]
|
||||||
|
|
|
||||||
|
|
@ -267,3 +267,52 @@ class TestMCPImageQuery:
|
||||||
# in search_documents_by_image rejects it.
|
# in search_documents_by_image rejects it.
|
||||||
results = await search_by_image(image_base64="!!! not base64 !!!")
|
results = await search_by_image(image_base64="!!! not base64 !!!")
|
||||||
assert results == []
|
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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue