Cover the configured database set from the MCP server

`haiku-rag mcp` passes covers_set=True and _covering no longer refuses a
scope over several databases. search_documents, search_documents_by_image,
ask_question and analyze take `sources`; get_document takes `source`;
DocumentInfo carries `source`. format_citations gains include_source, which
ask_question sets from covers_multiple so citations name their database
only when the server covers several.

Refs #599
This commit is contained in:
Yiorgis Gozadinos 2026-09-04 10:38:48 +03:00
parent ce69a8c989
commit 40d40bcbf2
No known key found for this signature in database
11 changed files with 311 additions and 91 deletions

View file

@ -16,6 +16,12 @@
- `processing.conversion_options.picture_description.model` defaults to - `processing.conversion_options.picture_description.model` defaults to
`enable_thinking: false`, and the field now reaches the VLM: docling's `enable_thinking: false`, and the field now reaches the VLM: docling's
picture-description request carries `reasoning_effort` in `params`. picture-description request carries `reasoning_effort` in `params`.
- `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on
`search_documents`, `search_documents_by_image`, `ask_question` and
`analyze`; `source` on `get_document`; an unknown name is a tool error.
`DocumentInfo.source`; citations name their database when the server
covers several. `format_citations(citations, include_source=False)`.
### Removed ### Removed
- MCP write tools `add_document_from_file`, `add_document_from_url`, - MCP write tools `add_document_from_file`, `add_document_from_url`,

View file

@ -24,7 +24,7 @@ The `haiku-rag` CLI provides complete document management functionality.
haiku-rag add -h haiku-rag add -h
``` ```
With `lancedb.databases` configured, `search`, `ask`, `analyze`, and `chat` use the full set by default. Select one database for other commands with `--db-name` or `--db`. `settings`, `init-config`, and `download-models` do not open a database. See [Multiple Databases](configuration/storage.md#multiple-databases). With `lancedb.databases` configured, `search`, `ask`, `analyze`, `chat`, and `mcp` use the full set by default. Select one database for other commands with `--db-name` or `--db`. `settings`, `init-config`, and `download-models` do not open a database. See [Multiple Databases](configuration/storage.md#multiple-databases).
## Document Management ## Document Management

View file

@ -281,9 +281,9 @@ Conversion, chunking, and title generation do not access a database and remain a
Commands use database sets as follows: Commands use database sets as follows:
- **Set-capable**: `search`, `ask`, `analyze`, and `chat` use the full configured set, or the single database selected by `--db-name`. - **Set-capable**: `search`, `ask`, `analyze`, `chat`, and `mcp` use the full configured set, or the single database selected by `--db-name`.
- **Config-only**: `settings`, `init-config`, and `download-models` do not open a database. - **Config-only**: `settings`, `init-config`, and `download-models` do not open a database.
- **Single-database**: everything else — document writes, `rebuild`, `vacuum`, `migrate`, `init`, `info`, `history`, `tag`, `doctor`, `list`, `inspect`, `visualize`, and `mcp` — works on one database, selected with the global `--db-name` option. - **Single-database**: everything else — document writes, `rebuild`, `vacuum`, `migrate`, `init`, `info`, `history`, `tag`, `doctor`, `list`, `inspect`, and `visualize` — works on one database, selected with the global `--db-name` option.
```bash ```bash
haiku-rag search "query" # every configured database haiku-rag search "query" # every configured database

View file

@ -27,6 +27,16 @@ e.g. inside a Docker container with port mapping, or on a trusted LAN.
The server opens the database read-only. Ingestion goes through the CLI The server opens the database read-only. Ingestion goes through the CLI
(`haiku-rag add`, `add-src`, `delete`) or [`haiku-ingester`](ingester.md). (`haiku-rag add`, `add-src`, `delete`) or [`haiku-ingester`](ingester.md).
## Collections
With several databases in `lancedb.databases`, the server covers all of
them, as `haiku-rag search` does. Results, documents and citations name
theirs in `source`. `sources` on the search and question tools restricts a
call to a subset; `source` on `get_document` names the database holding the
document. A name the server does not cover is an error.
`haiku-rag --db-name NAME mcp` serves one. See
[Multiple Databases](configuration/storage.md#multiple-databases).
## Claude Desktop Integration ## Claude Desktop Integration
Add to your Claude Desktop configuration (`claude_desktop_config.json`): Add to your Claude Desktop configuration (`claude_desktop_config.json`):
@ -63,6 +73,7 @@ After restarting Claude Desktop, you can ask Claude to search your documents or
- **`get_document`** - Retrieve a document by ID - **`get_document`** - Retrieve a document by ID
- `document_id` (required): The document ID - `document_id` (required): The document ID
- `source` (optional): The database holding it
- **`list_documents`** - List documents with pagination and filtering - **`list_documents`** - List documents with pagination and filtering
- `limit` (optional): Maximum number to return - `limit` (optional): Maximum number to return
@ -75,11 +86,13 @@ After restarting Claude Desktop, you can ask Claude to search your documents or
- `query` (required): Search query - `query` (required): Search query
- `limit` (optional): Maximum results (uses config default if not specified) - `limit` (optional): Maximum results (uses config default if not specified)
- `include_images` (optional, default `true`): Attach base64-encoded picture bytes to picture-labeled results - `include_images` (optional, default `true`): Attach base64-encoded picture bytes to picture-labeled results
- `sources` (optional): The databases to search
- **`search_documents_by_image`** - Search using an image as the query (registered only when the configured embedder supports images) - **`search_documents_by_image`** - Search using an image as the query (registered only when the configured embedder supports images)
- `image_base64` (required): Base64-encoded image (PNG/JPEG bytes) - `image_base64` (required): Base64-encoded image (PNG/JPEG bytes)
- `limit` (optional): Maximum results - `limit` (optional): Maximum results
- `include_images` (optional, default `true`) - `include_images` (optional, default `true`)
- `sources` (optional): The databases to search
### Question Answering ### Question Answering
@ -87,11 +100,13 @@ After restarting Claude Desktop, you can ask Claude to search your documents or
- `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)
- `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable QA model) - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable QA model)
- `sources` (optional): The databases to answer from
- **`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
- `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable analysis model) - `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable analysis model)
- `sources` (optional): The databases to analyze
- Best for aggregation, computation, and multi-document analysis - Best for aggregation, computation, and multi-document analysis
## Continuous ingestion ## Continuous ingestion

View file

@ -886,7 +886,7 @@ def mcp(
), ),
) -> None: ) -> None:
"""Run the MCP server.""" """Run the MCP server."""
app = create_app(db) app = create_app(db, covers_set=True)
transport = "stdio" if stdio else None transport = "stdio" if stdio else None

View file

@ -5,9 +5,11 @@ from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from fastmcp import FastMCP from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, get_config from haiku.rag.config import AppConfig, get_config
from haiku.rag.store.exceptions import UnknownDatabaseError
from haiku.rag.store.models import Document, SearchResult from haiku.rag.store.models import Document, SearchResult
from haiku.rag.tools.document import DocumentInfo from haiku.rag.tools.document import DocumentInfo
from haiku.rag.utils import format_citations from haiku.rag.utils import format_citations
@ -28,11 +30,11 @@ def create_mcp_server(
db_path: Path | None = None, db_path: Path | None = None,
config: AppConfig | None = None, config: AppConfig | None = None,
) -> FastMCP: ) -> FastMCP:
"""Create an MCP server over one database. """Create an MCP server over the databases the configuration places.
Args: Args:
db_path: Path to the database file, where `config` places none; or db_path: Path to the database file, where `config` places none; or
None to serve the database the configuration places. Beside None to serve the databases the configuration places. Beside
`lancedb.databases` a path raises `AmbiguousDatabaseError`. `lancedb.databases` a path raises `AmbiguousDatabaseError`.
config: Configuration to use. config: Configuration to use.
""" """
@ -50,13 +52,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
scope, so the configured name survives, which results and citations carry as scope, so the configured name survives, which results and citations carry as
``source``. ``source``.
""" """
from haiku.rag.store.exceptions import AmbiguousDatabaseError
if scope.covers_multiple:
raise AmbiguousDatabaseError(
"an MCP server serves one database, and this scope covers "
f"{', '.join(scope.names)}; name the one to serve"
)
client: HaikuRAG | None = None client: HaikuRAG | None = None
stack = AsyncExitStack() stack = AsyncExitStack()
client_lock = asyncio.Lock() client_lock = asyncio.Lock()
@ -95,7 +90,10 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
@mcp.tool() @mcp.tool()
async def search_documents( async def search_documents(
query: str, limit: int | None = None, include_images: bool = True query: str,
limit: int | None = None,
include_images: bool = True,
sources: list[str] | None = None,
) -> list[SearchResult]: ) -> list[SearchResult]:
"""Search the RAG system for documents using hybrid search (vector similarity + full-text search). """Search the RAG system for documents using hybrid search (vector similarity + full-text search).
@ -103,10 +101,15 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
in the result set, ``SearchResult.image_data`` carries base64-encoded in the result set, ``SearchResult.image_data`` carries base64-encoded
PNG bytes keyed by self_ref. Set to False to omit the bytes from the PNG bytes keyed by self_ref. Set to False to omit the bytes from the
response (smaller JSON payload for plain-text consumers). response (smaller JSON payload for plain-text consumers).
``sources`` names the databases to search, all of them by default.
""" """
try: try:
rag = await _client() rag = await _client()
return await rag.search(query, limit=limit, include_images=include_images) return await rag.search(
query, limit=limit, include_images=include_images, sources=sources
)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception: except Exception:
return [] return []
@ -123,6 +126,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
image_base64: str, image_base64: str,
limit: int | None = None, limit: int | None = None,
include_images: bool = True, include_images: bool = True,
sources: list[str] | None = None,
) -> list[SearchResult]: ) -> list[SearchResult]:
"""Search the RAG system using an image as the query. """Search the RAG system using an image as the query.
@ -130,6 +134,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
image is embedded via the configured multimodal embedder and the image is embedded via the configured multimodal embedder and the
chunks table is searched vector-only. ``include_images`` controls chunks table is searched vector-only. ``include_images`` controls
whether picture bytes are attached to picture-labeled results. whether picture bytes are attached to picture-labeled results.
``sources`` names the databases to search, all of them by default.
""" """
import base64 import base64
@ -139,16 +144,28 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
return [] return []
try: try:
rag = await _client() rag = await _client()
return await rag.search(raw, limit=limit, include_images=include_images) return await rag.search(
raw, limit=limit, include_images=include_images, sources=sources
)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception: except Exception:
return [] return []
@mcp.tool() @mcp.tool()
async def get_document(document_id: str) -> Document | None: async def get_document(
"""Get a document by its ID.""" document_id: str, source: str | None = None
) -> Document | None:
"""Get a document by its ID.
``source`` names the database holding it; without one every database
is asked.
"""
try: try:
rag = await _client() rag = await _client()
return await rag.get_document_by_id(document_id) return await rag.get_document_by_id(document_id, source)
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception: except Exception:
return None return None
@ -175,6 +192,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
title=doc.title or "Untitled", title=doc.title or "Untitled",
uri=doc.uri or "", uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"), created=doc.created_at.strftime("%Y-%m-%d"),
source=doc.source,
) )
for doc in documents for doc in documents
] ]
@ -186,6 +204,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
question: str, question: str,
cite: bool = False, cite: bool = False,
images_base64: list[str] | None = None, images_base64: list[str] | None = None,
sources: list[str] | None = None,
) -> str: ) -> str:
"""Ask a question using the QA agent. """Ask a question using the QA agent.
@ -194,6 +213,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
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 images_base64: Base64-encoded images attached to the question
(requires a vision-capable QA model). (requires a vision-capable QA model).
sources: The databases to answer from, all of them by default.
Returns: Returns:
The answer as a string. The answer as a string.
@ -201,10 +221,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
try: try:
images = _decode_images(images_base64) images = _decode_images(images_base64)
rag = await _client() rag = await _client()
answer, citations = await rag.ask(question, images=images) answer, citations = await rag.ask(question, images=images, sources=sources)
if cite and citations: if cite and citations:
answer += "\n\n" + format_citations(citations) answer += "\n\n" + format_citations(
citations, include_source=rag.covers_multiple
)
return answer return answer
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception as e: except Exception as e:
return f"Error answering question: {e!s}" return f"Error answering question: {e!s}"
@ -213,6 +237,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
question: str, question: str,
filter: str | None = None, filter: str | None = None,
images_base64: list[str] | None = None, images_base64: list[str] | None = None,
sources: list[str] | None = None,
) -> str: ) -> str:
"""Answer complex questions using the analysis capability. """Answer complex questions using the analysis capability.
@ -225,6 +250,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
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 images_base64: Base64-encoded images attached to the question
(requires a vision-capable analysis model). (requires a vision-capable analysis model).
sources: The databases to analyze, all of them by default.
Returns: Returns:
The answer as a string. The answer as a string.
@ -232,8 +258,12 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
try: try:
images = _decode_images(images_base64) images = _decode_images(images_base64)
rag = await _client() rag = await _client()
result = await rag.analyze(question, filter=filter, images=images) result = await rag.analyze(
question, filter=filter, images=images, sources=sources
)
return result.answer return result.answer
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception as e: except Exception as e:
return f"Error running analysis capability: {e!s}" return f"Error running analysis capability: {e!s}"

View file

@ -27,6 +27,7 @@ class DocumentInfo(BaseModel):
title: str title: str
uri: str uri: str
created: str created: str
source: str | None = None
class DocumentListResponse(BaseModel): class DocumentListResponse(BaseModel):

View file

@ -393,11 +393,13 @@ def _citation_label(c: "Citation") -> str:
return c.document_title or c.document_uri return c.document_title or c.document_uri
def format_citations(citations: "list[Citation]") -> str: def format_citations(citations: "list[Citation]", include_source: bool = False) -> str:
"""Format citations as plain text with preserved formatting. """Format citations as plain text with preserved formatting.
Used by things like the MCP server where Rich renderables are not available. Used by things like the MCP server where Rich renderables are not available.
Pictures referenced by the chunk are surfaced as ``[Figure: <ref>]`` markers. Pictures referenced by the chunk are surfaced as ``[Figure: <ref>]`` markers.
``include_source`` names each citation's database, for a client covering
several.
""" """
if not citations: if not citations:
return "" return ""
@ -410,6 +412,8 @@ def format_citations(citations: "list[Citation]") -> str:
header = f"[{idx}] {title}" header = f"[{idx}] {title}"
location_parts = [] location_parts = []
if include_source and c.source:
location_parts.append(f"Collection: {c.source}")
pages = _citation_pages(c) pages = _citation_pages(c)
if pages: if pages:
location_parts.append(pages) location_parts.append(pages)

View file

@ -1050,6 +1050,21 @@ def test_mcp_without_stdio_leaves_the_transport_unset(app_stub):
assert app_stub.run_mcp.call_args.kwargs["transport"] is None assert app_stub.run_mcp.call_args.kwargs["transport"] is None
def test_mcp_covers_the_configured_set(monkeypatch):
seen = {}
def create_app(db=None, *, covers_set=False):
seen["covers_set"] = covers_set
return AsyncMock()
monkeypatch.setattr("haiku.rag.cli.create_app", create_app)
result = runner.invoke(cli, ["mcp", "--stdio"])
assert result.exit_code == 0, result.output
assert seen["covers_set"] is True
def test_version_flag_prints_the_version(): def test_version_flag_prints_the_version():
result = runner.invoke(cli, ["--version"]) result = runner.invoke(cli, ["--version"])

View file

@ -1,3 +1,5 @@
from types import SimpleNamespace
import pytest import pytest
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
@ -5,6 +7,7 @@ from haiku.rag.mcp import _covering as _mcp_covering
from haiku.rag.mcp import create_mcp_server from haiku.rag.mcp import create_mcp_server
from haiku.rag.store.models import Chunk, Document, SearchResult from haiku.rag.store.models import Chunk, Document, SearchResult
from haiku.rag.tools.document import DocumentInfo from haiku.rag.tools.document import DocumentInfo
from tests.multi_db.helpers import _config, _seed
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
@ -29,6 +32,22 @@ def mock_embedder(monkeypatch):
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents) monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
@pytest.fixture
def multimodal_embedder(monkeypatch):
"""An embedder reporting image support, so the image-query tool registers."""
from haiku.rag.embeddings import EmbedderWrapper
class StubMultimodal(EmbedderWrapper):
supports_images = True
def __init__(self):
super().__init__(embedder=None, vector_dim=2560)
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder", lambda *a, **kw: StubMultimodal()
)
@pytest.fixture @pytest.fixture
async def mcp_db(temp_db_path): async def mcp_db(temp_db_path):
"""Create a test database with sample documents.""" """Create a test database with sample documents."""
@ -46,6 +65,21 @@ async def mcp_db(temp_db_path):
return temp_db_path return temp_db_path
@pytest.fixture
async def two_dbs(tmp_path):
"""Two configured databases, alpha and beta, one document each."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
return config
def _covering_all(config):
from haiku.rag.client.scope import DatabaseScope
return _mcp_covering(DatabaseScope.resolve(config), config)
async def _get_tool(mcp, name): async def _get_tool(mcp, name):
"""Get a tool function from an MCP server by name.""" """Get a tool function from an MCP server by name."""
tool = await mcp.get_tool(name) tool = await mcp.get_tool(name)
@ -183,6 +217,136 @@ class TestMCPToolSet:
} }
class TestMCPCoversTheConfiguredSet:
@pytest.mark.asyncio
async def test_results_name_the_database_they_came_from(self, two_dbs):
mcp = _covering_all(two_dbs)
search = await _get_tool(mcp, "search_documents")
results = await search(query="cats")
assert {r.source for r in results} == {"alpha", "beta"}
@pytest.mark.asyncio
async def test_sources_narrows_the_search(self, two_dbs):
mcp = _covering_all(two_dbs)
search = await _get_tool(mcp, "search_documents")
results = await search(query="cats", sources=["beta"])
assert results
assert {r.source for r in results} == {"beta"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_name,kwargs",
[
("search_documents", {"query": "cats", "sources": ["nope"]}),
(
"search_documents_by_image",
{"image_base64": "AAAA", "sources": ["nope"]},
),
("get_document", {"document_id": "x", "source": "nope"}),
("ask_question", {"question": "q", "sources": ["nope"]}),
("analyze", {"question": "q", "sources": ["nope"]}),
],
)
async def test_an_unknown_database_is_an_error_not_an_empty_result(
self, two_dbs, multimodal_embedder, tool_name, kwargs
):
from fastmcp.exceptions import ToolError
mcp = _covering_all(two_dbs)
tool = await _get_tool(mcp, tool_name)
with pytest.raises(ToolError, match="nope"):
await tool(**kwargs)
@pytest.mark.asyncio
async def test_the_listing_covers_every_database(self, two_dbs):
mcp = _covering_all(two_dbs)
list_docs = await _get_tool(mcp, "list_documents")
documents = await list_docs()
assert {d.source for d in documents} == {"alpha", "beta"}
@pytest.mark.asyncio
async def test_get_document_reaches_whichever_database_holds_it(self, two_dbs):
mcp = _covering_all(two_dbs)
list_docs = await _get_tool(mcp, "list_documents")
get_doc = await _get_tool(mcp, "get_document")
[beta] = [d for d in await list_docs() if d.source == "beta"]
found = await get_doc(document_id=beta.id)
named = await get_doc(document_id=beta.id, source="beta")
assert found.id == named.id == beta.id
assert found.source == named.source == "beta"
@pytest.mark.asyncio
async def test_the_public_factory_covers_a_configured_set(self, two_dbs):
mcp = create_mcp_server(config=two_dbs)
search = await _get_tool(mcp, "search_documents")
results = await search(query="cats")
assert {r.source for r in results} == {"alpha", "beta"}
@pytest.mark.asyncio
async def test_ask_question_names_each_citations_database(
self, two_dbs, monkeypatch
):
from haiku.rag.store.models.citation import Citation
def cited(source):
return Citation(
chunk_id="c1",
document_id="d1",
content="cited text",
document_uri="test://cats",
document_title="Cats",
source=source,
)
async def fake_ask(self, question, filter=None, images=None, sources=None):
return ("the answer", [cited("alpha"), cited("beta")])
monkeypatch.setattr(HaikuRAG, "ask", fake_ask)
mcp = _covering_all(two_dbs)
ask = await _get_tool(mcp, "ask_question")
answer = await ask(question="q", cite=True)
assert "alpha" in answer
assert "beta" in answer
@pytest.mark.asyncio
@pytest.mark.parametrize(
"tool_name,client_method,returns",
[
("ask_question", "ask", ("answer", [])),
("analyze", "analyze", SimpleNamespace(answer="answer")),
],
)
async def test_agents_search_the_selected_databases(
self, two_dbs, monkeypatch, tool_name, client_method, returns
):
seen = {}
async def fake(self, question, filter=None, images=None, sources=None):
seen["sources"] = sources
return returns
monkeypatch.setattr(HaikuRAG, client_method, fake)
mcp = _covering_all(two_dbs)
tool = await _get_tool(mcp, tool_name)
await tool(question="q", sources=["beta"])
assert seen["sources"] == ["beta"]
class TestMCPImageQuery: class TestMCPImageQuery:
"""search_documents_by_image is registered only when the embedder is multimodal.""" """search_documents_by_image is registered only when the embedder is multimodal."""
@ -195,59 +359,40 @@ class TestMCPImageQuery:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_image_query_tool_registered_for_multimodal_embedder( async def test_image_query_tool_registered_for_multimodal_embedder(
self, mcp_db, monkeypatch self, mcp_db, multimodal_embedder, monkeypatch
): ):
"""When the embedder reports supports_images=True, the tool exists """When the embedder reports supports_images=True, the tool exists
and routes a base64 image through ``client.search``.""" and routes the decoded image and the selection through ``client.search``."""
from haiku.rag.embeddings import EmbedderWrapper seen = {}
class StubMultimodal(EmbedderWrapper): async def fake_search(self, query, **kwargs):
supports_images = True seen.update(query=query, **kwargs)
return []
def __init__(self): monkeypatch.setattr(HaikuRAG, "search", fake_search)
super().__init__(embedder=None, vector_dim=2560)
async def embed_image(self, image):
# Produce a deterministic-ish vector of the right dim.
return [0.0] * 2560
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder",
lambda *a, **kw: StubMultimodal(),
)
mcp = create_mcp_server(mcp_db) mcp = create_mcp_server(mcp_db)
names = {t.name for t in await mcp.list_tools()} names = {t.name for t in await mcp.list_tools()}
assert "search_documents_by_image" in names assert "search_documents_by_image" in names
search_by_image = await _get_tool(mcp, "search_documents_by_image") search_by_image = await _get_tool(mcp, "search_documents_by_image")
# Standalone PNG header (won't decode to a real image but our stub doesn't care).
import base64 import base64
png_b64 = base64.b64encode(b"\x89PNG\r\n\x1a\n").decode("ascii") png = b"\x89PNG\r\n\x1a\n"
results = await search_by_image(image_base64=png_b64) results = await search_by_image(
# Empty list is fine (the stub vector won't match the toy fixture). image_base64=base64.b64encode(png).decode("ascii"), sources=["alpha"]
assert isinstance(results, list) )
assert results == []
assert seen["query"] == png
assert seen["sources"] == ["alpha"]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_image_query_returns_empty_on_invalid_base64( async def test_image_query_returns_empty_on_invalid_base64(
self, mcp_db, monkeypatch self, mcp_db, multimodal_embedder
): ):
"""Garbage base64 from the caller is swallowed, returning an empty """Garbage base64 from the caller is swallowed, returning an empty
list rather than crashing the MCP server.""" list rather than crashing the MCP server."""
from haiku.rag.embeddings import EmbedderWrapper
class StubMultimodal(EmbedderWrapper):
supports_images = True
def __init__(self):
super().__init__(embedder=None, vector_dim=2560)
monkeypatch.setattr(
"haiku.rag.embeddings.get_embedder",
lambda *a, **kw: StubMultimodal(),
)
mcp = create_mcp_server(mcp_db) mcp = create_mcp_server(mcp_db)
search_by_image = await _get_tool(mcp, "search_documents_by_image") search_by_image = await _get_tool(mcp, "search_documents_by_image")
@ -256,6 +401,19 @@ class TestMCPImageQuery:
results = await search_by_image(image_base64="!!! not base64 !!!") results = await search_by_image(image_base64="!!! not base64 !!!")
assert results == [] assert results == []
@pytest.mark.asyncio
async def test_image_query_returns_empty_when_the_search_raises(
self, mcp_db, multimodal_embedder, monkeypatch
):
async def boom(self, *args, **kw):
raise RuntimeError("client exploded")
monkeypatch.setattr(HaikuRAG, "search", boom)
mcp = create_mcp_server(mcp_db)
search_by_image = await _get_tool(mcp, "search_documents_by_image")
assert await search_by_image(image_base64="AAAA") == []
class TestMCPImageInput: class TestMCPImageInput:
@pytest.mark.asyncio @pytest.mark.asyncio
@ -264,7 +422,7 @@ class TestMCPImageInput:
captured = {} captured = {}
async def fake_ask(self, question, filter=None, images=None): async def fake_ask(self, question, filter=None, images=None, sources=None):
captured["images"] = images captured["images"] = images
return ("answer", []) return ("answer", [])
@ -284,7 +442,7 @@ class TestMCPImageInput:
captured = {} captured = {}
async def fake_analyze(self, question, filter=None, images=None): async def fake_analyze(self, question, filter=None, images=None, sources=None):
captured["images"] = images captured["images"] = images
return SimpleNamespace(answer="answer") return SimpleNamespace(answer="answer")
@ -309,7 +467,7 @@ class TestMCPImageInput:
async def test_ask_question_without_images_passes_none(self, mcp_db, monkeypatch): async def test_ask_question_without_images_passes_none(self, mcp_db, monkeypatch):
captured = {} captured = {}
async def fake_ask(self, question, filter=None, images=None): async def fake_ask(self, question, filter=None, images=None, sources=None):
captured["images"] = images captured["images"] = images
return ("answer", []) return ("answer", [])
@ -356,7 +514,7 @@ class TestMCPToolsDegradeOnError:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_analyze_reports_the_error(self, mcp_db, monkeypatch): async def test_analyze_reports_the_error(self, mcp_db, monkeypatch):
async def boom(self, question, filter=None, images=None): async def boom(self, question, filter=None, images=None, sources=None):
raise RuntimeError("sandbox exploded") raise RuntimeError("sandbox exploded")
monkeypatch.setattr(HaikuRAG, "analyze", boom) monkeypatch.setattr(HaikuRAG, "analyze", boom)
@ -377,9 +535,10 @@ class TestMCPToolsDegradeOnError:
content="cited text", content="cited text",
document_uri="test://ai-overview", document_uri="test://ai-overview",
document_title="AI Overview", document_title="AI Overview",
source="alpha",
) )
async def fake_ask(self, question, filter=None, images=None): async def fake_ask(self, question, filter=None, images=None, sources=None):
return ("the answer", [citation]) return ("the answer", [citation])
monkeypatch.setattr(HaikuRAG, "ask", fake_ask) monkeypatch.setattr(HaikuRAG, "ask", fake_ask)
@ -389,6 +548,8 @@ class TestMCPToolsDegradeOnError:
with_cite = await ask(question="q", cite=True) with_cite = await ask(question="q", cite=True)
assert with_cite.startswith("the answer") assert with_cite.startswith("the answer")
assert "AI Overview" in with_cite assert "AI Overview" in with_cite
# One database: its name adds nothing.
assert "alpha" not in with_cite
assert await ask(question="q", cite=False) == "the answer" assert await ask(question="q", cite=False) == "the answer"
@ -499,34 +660,6 @@ class TestMCPClientLifetime:
assert "AI Overview" in titles assert "AI Overview" in titles
assert "Zebras" not in titles assert "Zebras" not in titles
def test_a_scope_covering_a_set_is_refused(self, tmp_path):
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.exceptions import AmbiguousDatabaseError
config = AppConfig(
lancedb=LanceDBConfig(
databases={"alpha": str(tmp_path / "a"), "beta": str(tmp_path / "b")}
)
)
with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
_mcp_covering(DatabaseScope.resolve(config), config)
def test_the_public_factory_refuses_a_configured_set_too(self, tmp_path):
"""It resolves the same scope, so it reaches the same refusal."""
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.exceptions import AmbiguousDatabaseError
config = AppConfig(
lancedb=LanceDBConfig(
databases={"alpha": str(tmp_path / "a"), "beta": str(tmp_path / "b")}
)
)
with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
create_mcp_server(config=config)
def test_the_public_factory_refuses_a_path_beside_a_configured_set(self, tmp_path): def test_the_public_factory_refuses_a_path_beside_a_configured_set(self, tmp_path):
"""A path and `lancedb.databases` both place the database.""" """A path and `lancedb.databases` both place the database."""
from haiku.rag.config.models import AppConfig, LanceDBConfig from haiku.rag.config.models import AppConfig, LanceDBConfig

View file

@ -753,6 +753,22 @@ def test_format_citations_sequential_indices():
assert "[2] Second" in result assert "[2] Second" in result
def test_format_citations_names_the_source_when_asked():
from haiku.rag.store.models.citation import Citation
from haiku.rag.utils import format_citations
citation = Citation(
document_id="doc1",
chunk_id="chunk1",
document_uri="test://doc",
document_title="Test Doc",
content="Content",
source="papers",
)
assert "papers" in format_citations([citation], include_source=True)
assert "papers" not in format_citations([citation])
# --- format_citations tests (pictures) --- # --- format_citations tests (pictures) ---