From f87fabe55684c5d97d1114703135667ae3ea1f4d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 24 Jul 2026 16:00:01 +0300 Subject: [PATCH] Make the capability image-attachment gate a vision bool --- .../haiku/rag/capabilities/_base.py | 6 +++--- .../haiku/rag/capabilities/analysis.py | 14 +++++++------ haiku_rag_slim/haiku/rag/capabilities/rag.py | 11 +++++----- haiku_rag_slim/haiku/rag/chat/__init__.py | 4 ++-- tests/chat/test_chat_app.py | 20 ++++++++++--------- 5 files changed, 30 insertions(+), 25 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 04b59775..7131364b 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -22,7 +22,7 @@ from pydantic_ai.toolsets import AgentToolset from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus from haiku.rag.client import HaikuRAG -from haiku.rag.config.models import AppConfig, ModelConfig +from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation, resolve_citations from haiku.rag.tools.search import build_binary_parts_from_results @@ -87,7 +87,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): state_type: type[StateT] state_namespace: str instruction_text: str - model: ModelConfig + vision: bool tool_names: frozenset[str] request_limit: int | None = None state: StateT | None = field(default=None, repr=False) @@ -220,7 +220,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): ) state = cast(Any, self.state) state.searches[query] = results - if self.model.vision and (parts := build_binary_parts_from_results(results)): + if self.vision and (parts := build_binary_parts_from_results(results)): return ToolReturn(return_value=formatted, content=parts) return formatted diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index fa3e0d62..18a3cef9 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -13,7 +13,7 @@ from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, ) -from haiku.rag.config.models import AppConfig, ModelConfig +from haiku.rag.config.models import AppConfig from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation @@ -129,25 +129,27 @@ def create_capability( *, defer_loading: bool = True, request_limit: int | None = 30, - model: ModelConfig | None = None, + vision: bool | None = None, ) -> AnalysisCapability: """Create a native Pydantic AI analysis capability. - ``model`` sets the capability's image-attachment gate and should be the - model the hosting agent actually runs. Defaults to ``config.analysis.model`` - (falling back to ``config.qa.model``). + ``vision`` gates whether picture chunks are attached to search results as + images, and should reflect the model the hosting agent actually runs. + Defaults to ``config.analysis.model.vision`` (falling back to + ``config.qa.model.vision``). """ if config is None: from haiku.rag.config import get_config config = get_config() + analysis_model = config.analysis.model or config.qa.model return AnalysisCapability( db_path=resolve_db_path(db_path, config), config=config, state_type=AnalysisState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), - model=model or config.analysis.model or config.qa.model, + vision=analysis_model.vision if vision is None else vision, tool_names=_TOOL_NAMES, request_limit=request_limit, id=_CAPABILITY_ID, diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 45bd81f3..08931954 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -12,7 +12,7 @@ from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, ) -from haiku.rag.config.models import AppConfig, ModelConfig +from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation @@ -71,12 +71,13 @@ def create_capability( *, defer_loading: bool = True, request_limit: int | None = 20, - model: ModelConfig | None = None, + vision: bool | None = None, ) -> RAGCapability: """Create a native Pydantic AI RAG capability. - ``model`` sets the capability's image-attachment gate and should be the - model the hosting agent actually runs. Defaults to ``config.qa.model``. + ``vision`` gates whether picture chunks are attached to search results as + images, and should reflect the model the hosting agent actually runs. + Defaults to ``config.qa.model.vision``. """ if config is None: from haiku.rag.config import get_config @@ -88,7 +89,7 @@ def create_capability( state_type=RAGState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), - model=model or config.qa.model, + vision=config.qa.model.vision if vision is None else vision, tool_names=_TOOL_NAMES, request_limit=request_limit, id=_CAPABILITY_ID, diff --git a/haiku_rag_slim/haiku/rag/chat/__init__.py b/haiku_rag_slim/haiku/rag/chat/__init__.py index 9ca3bc44..a149d084 100644 --- a/haiku_rag_slim/haiku/rag/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/chat/__init__.py @@ -55,7 +55,7 @@ def run_chat( db_path=db_path, config=config, defer_loading=defer_loading, - model=driving_model, + vision=driving_model.vision, ) ) @@ -67,7 +67,7 @@ def run_chat( db_path=db_path, config=config, defer_loading=defer_loading, - model=driving_model, + vision=driving_model.vision, ) ) diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 5b631f22..7fd6b870 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -47,23 +47,25 @@ def test_run_chat_defers_multiple_capabilities(temp_db_path: Path): @pytest.mark.parametrize( - ("enabled", "expected_model"), + ("enabled", "expected_model", "expected_vision"), [ - (["analysis"], "analysis-model"), - (["rag"], "qa-model"), - (["rag", "analysis"], "qa-model"), + (["analysis"], "analysis-model", False), + (["rag"], "qa-model", True), + (["rag", "analysis"], "qa-model", True), ], ) def test_run_chat_gates_capability_vision_on_driving_model( - temp_db_path: Path, enabled, expected_model + temp_db_path: Path, enabled, expected_model, expected_vision ): """Analysis-only chat runs on analysis.model; otherwise on qa.model. Every - attached capability's vision gate (its ``model``) tracks that one model.""" + attached capability's vision gate tracks that one driving model.""" from haiku.rag.config.models import AppConfig, ModelConfig config = AppConfig() - config.qa.model = ModelConfig(provider="openai", name="qa-model") - config.analysis.model = ModelConfig(provider="openai", name="analysis-model") + config.qa.model = ModelConfig(provider="openai", name="qa-model", vision=True) + config.analysis.model = ModelConfig( + provider="openai", name="analysis-model", vision=False + ) captured: dict[str, str] = {} def fake_get_model(model_config, _config): @@ -81,7 +83,7 @@ def test_run_chat_gates_capability_vision_on_driving_model( assert captured["name"] == expected_model attached = mock_app.call_args.kwargs["capabilities"] - assert {capability.model.name for capability in attached} == {expected_model} + assert {capability.vision for capability in attached} == {expected_vision} def _make_mock_client():