Accept images on ask/analyze
This commit is contained in:
parent
acbd66afbd
commit
4c5050d161
6 changed files with 173 additions and 7 deletions
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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]
|
||||||
|
|
|
||||||
|
|
@ -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
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")])
|
||||||
Loading…
Reference in a new issue