Accept images on ask/analyze

This commit is contained in:
Yiorgis Gozadinos 2026-07-25 10:18:43 +03:00
parent acbd66afbd
commit 4c5050d161
No known key found for this signature in database
6 changed files with 173 additions and 7 deletions

View file

@ -1,6 +1,10 @@
# Changelog # Changelog
## [Unreleased] ## [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.
## [0.69.0] - 2026-07-24 ## [0.69.0] - 2026-07-24
### Added ### Added

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. `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)). 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` 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. See [Analysis capability](capabilities/analysis.md) for details and configuration.
## Building custom agents ## Building custom agents

View file

@ -484,19 +484,21 @@ class HaikuRAG:
self, self,
question: str, question: str,
filter: str | None = None, filter: str | None = None,
images: Sequence[bytes] | None = None,
) -> "tuple[str, list[Citation]]": ) -> "tuple[str, list[Citation]]":
from haiku.rag.client.agents import ask from haiku.rag.client.agents import ask
return await ask(self, question, filter) return await ask(self, question, filter, images)
async def analyze( async def analyze(
self, self,
question: str, question: str,
filter: str | None = None, filter: str | None = None,
images: Sequence[bytes] | None = None,
) -> "AnalysisResult": ) -> "AnalysisResult":
from haiku.rag.client.agents import analyze from haiku.rag.client.agents import analyze
return await analyze(self, question, filter) return await analyze(self, question, filter, images)
async def visualize_chunk( async def visualize_chunk(
self, self,

View file

@ -1,10 +1,14 @@
from collections.abc import Sequence
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
from pydantic_ai import Agent from pydantic_ai import Agent
if TYPE_CHECKING: if TYPE_CHECKING:
from pydantic_ai.messages import BinaryContent
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import ModelConfig
from haiku.rag.sandbox import AnalysisResult from haiku.rag.sandbox import AnalysisResult
from haiku.rag.store.models.citation import Citation from haiku.rag.store.models.citation import Citation
@ -14,10 +18,28 @@ class _AgentDeps:
state: dict[str, Any] = field(default_factory=dict) 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( async def ask(
client: "HaikuRAG", client: "HaikuRAG",
question: str, question: str,
filter: str | None = None, filter: str | None = None,
images: Sequence[bytes] | None = None,
) -> "tuple[str, list[Citation]]": ) -> "tuple[str, list[Citation]]":
"""Ask a question against the knowledge base via the RAG capability. """Ask a question against the knowledge base via the RAG capability.
@ -25,6 +47,8 @@ async def ask(
client: The HaikuRAG client. client: The HaikuRAG client.
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: Raw image bytes attached to the question (requires a
vision-capable QA model).
Returns: Returns:
Tuple of (answer text, list of resolved citations). Tuple of (answer text, list of resolved citations).
@ -44,6 +68,7 @@ async def ask(
deps = _AgentDeps( deps = _AgentDeps(
state={"rag": RAGState(document_filter=filter).model_dump(mode="json")} 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) model = get_model(client._config.qa.model, client._config)
agent = Agent( agent = Agent(
model, model,
@ -51,7 +76,7 @@ async def ask(
instructions=AGENT_PREAMBLE, instructions=AGENT_PREAMBLE,
capabilities=[capability], capabilities=[capability],
) )
result = await agent.run(question, deps=deps) result = await agent.run(user_prompt, deps=deps)
state = RAGState.model_validate(deps.state["rag"]) state = RAGState.model_validate(deps.state["rag"])
citations = [ citations = [
state.citation_index[cid] state.citation_index[cid]
@ -65,6 +90,7 @@ async def analyze(
client: "HaikuRAG", client: "HaikuRAG",
question: str, question: str,
filter: str | None = None, filter: str | None = None,
images: Sequence[bytes] | None = None,
) -> "AnalysisResult": ) -> "AnalysisResult":
"""Answer a question using the analysis capability. """Answer a question using the analysis capability.
@ -76,6 +102,8 @@ async def analyze(
client: The HaikuRAG client. client: The HaikuRAG client.
question: The question to answer. question: The question to answer.
filter: SQL WHERE clause to filter documents during searches. filter: SQL WHERE clause to filter documents during searches.
images: Raw image bytes attached to the question (requires a
vision-capable analysis model).
Returns: Returns:
AnalysisResult with the answer and resolved citations. AnalysisResult with the answer and resolved citations.
@ -94,15 +122,15 @@ async def analyze(
"analysis": AnalysisState(document_filter=filter).model_dump(mode="json") "analysis": AnalysisState(document_filter=filter).model_dump(mode="json")
} }
) )
model = get_model( model_config = client._config.analysis.model or client._config.qa.model
client._config.analysis.model or client._config.qa.model, client._config user_prompt = _build_user_prompt(question, images, model_config)
) model = get_model(model_config, client._config)
agent = Agent( agent = Agent(
model, model,
deps_type=_AgentDeps, deps_type=_AgentDeps,
capabilities=[capability], capabilities=[capability],
) )
result = await agent.run(question, deps=deps) result = await agent.run(user_prompt, deps=deps)
state = AnalysisState.model_validate(deps.state["analysis"]) state = AnalysisState.model_validate(deps.state["analysis"])
citations = [ citations = [
state.citation_index[cid] state.citation_index[cid]

View file

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, cast
from packaging.version import Version, parse from packaging.version import Version, parse
if TYPE_CHECKING: if TYPE_CHECKING:
from pydantic_ai.messages import BinaryContent
from pydantic_ai.profiles.openai import OpenAIModelProfile from pydantic_ai.profiles.openai import OpenAIModelProfile
from rich.console import RenderableType from rich.console import RenderableType
@ -37,6 +38,21 @@ def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
return dot_product / (norm1 * norm2) 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( def apply_common_settings(
settings: Any | None, settings: Any | None,
settings_class: type[Any], settings_class: type[Any],

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