diff --git a/CHANGELOG.md b/CHANGELOG.md index e84ce856..fe14b011 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # 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. + ## [0.69.0] - 2026-07-24 ### Added diff --git a/docs/python.md b/docs/python.md index 2fb175a6..764f1717 100644 --- a/docs/python.md +++ b/docs/python.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index bff6f32b..ecfd4b01 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index 948f4819..9a739f74 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -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] diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 595e1afb..28c81833 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -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], diff --git a/tests/test_ask_images.py b/tests/test_ask_images.py new file mode 100644 index 00000000..baaab700 --- /dev/null +++ b/tests/test_ask_images.py @@ -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")])