danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cover the configured database set from the MCP server

Browse source 
`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
Show stats Download patch file Download diff file Expand all files Collapse all files

6
CHANGELOG.md
View file

@ -16,6 +16,12 @@
- `processing.conversion_options.picture_description.model` defaults to
`enable_thinking: false`, and the field now reaches the VLM: docling's
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
- MCP write tools `add_document_from_file`, `add_document_from_url`,

2
docs/cli.md
View file

@ -24,7 +24,7 @@ The `haiku-rag` CLI provides complete document management functionality.
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

4
docs/configuration/storage.md
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:
- **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.
- **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
haiku-rag search "query" # every configured database

15
docs/mcp.md
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
(`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
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
- `document_id` (required): The document ID
- `source` (optional): The database holding it
- **`list_documents`** - List documents with pagination and filtering
- `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
- `limit` (optional): Maximum results (uses config default if not specified)
- `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)
- `image_base64` (required): Base64-encoded image (PNG/JPEG bytes)
- `limit` (optional): Maximum results
- `include_images` (optional, default `true`)
- `sources` (optional): The databases to search
### 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
- `cite` (optional): Include source citations (default: false)
- `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
- `question` (required): The question to answer
- `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)
- `sources` (optional): The databases to analyze
- Best for aggregation, computation, and multi-document analysis
## Continuous ingestion

2
haiku_rag_slim/haiku/rag/cli.py
View file

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

66
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -5,9 +5,11 @@ from pathlib import Path
from typing import TYPE_CHECKING
from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from haiku.rag.client import HaikuRAG
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.tools.document import DocumentInfo
from haiku.rag.utils import format_citations
@ -28,11 +30,11 @@ def create_mcp_server(
db_path: Path | None = None,
config: AppConfig | None = None,
) -> FastMCP:
"""Create an MCP server over one database.
"""Create an MCP server over the databases the configuration places.
Args:
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`.
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
``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
stack = AsyncExitStack()
client_lock = asyncio.Lock()
@ -95,7 +90,10 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
@mcp.tool()
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]:
"""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
PNG bytes keyed by self_ref. Set to False to omit the bytes from the
response (smaller JSON payload for plain-text consumers).
``sources`` names the databases to search, all of them by default.
"""
try:
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:
return []
@ -123,6 +126,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
image_base64: str,
limit: int | None = None,
include_images: bool = True,
sources: list[str] | None = None,
) -> list[SearchResult]:
"""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
chunks table is searched vector-only. ``include_images`` controls
whether picture bytes are attached to picture-labeled results.
``sources`` names the databases to search, all of them by default.
"""
import base64
@ -139,16 +144,28 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
return []
try:
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:
return []
@mcp.tool()
async def get_document(document_id: str) -> Document | None:
"""Get a document by its ID."""
async def get_document(
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:
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:
return None
@ -175,6 +192,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
source=doc.source,
)
for doc in documents
]
@ -186,6 +204,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
question: str,
cite: bool = False,
images_base64: list[str] | None = None,
sources: list[str] | None = None,
) -> str:
"""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.
images_base64: Base64-encoded images attached to the question
(requires a vision-capable QA model).
sources: The databases to answer from, all of them by default.
Returns:
The answer as a string.
@ -201,10 +221,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
try:
images = _decode_images(images_base64)
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:
answer += "\n\n" + format_citations(citations)
answer += "\n\n" + format_citations(
citations, include_source=rag.covers_multiple
)
return answer
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception as e:
return f"Error answering question: {e!s}"
@ -213,6 +237,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
question: str,
filter: str | None = None,
images_base64: list[str] | None = None,
sources: list[str] | None = None,
) -> str:
"""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.
images_base64: Base64-encoded images attached to the question
(requires a vision-capable analysis model).
sources: The databases to analyze, all of them by default.
Returns:
The answer as a string.
@ -232,8 +258,12 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
try:
images = _decode_images(images_base64)
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
except UnknownDatabaseError as e:
raise ToolError(str(e)) from e
except Exception as e:
return f"Error running analysis capability: {e!s}"

1
haiku_rag_slim/haiku/rag/tools/document.py
View file

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

6
haiku_rag_slim/haiku/rag/utils.py
View file

@ -393,11 +393,13 @@ def _citation_label(c: "Citation") -> str:
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.
Used by things like the MCP server where Rich renderables are not available.
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:
return ""
@ -410,6 +412,8 @@ def format_citations(citations: "list[Citation]") -> str:
header = f"[{idx}] {title}"
location_parts = []
if include_source and c.source:
location_parts.append(f"Collection: {c.source}")
pages = _citation_pages(c)
if pages:
location_parts.append(pages)

15
tests/test_cli.py
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
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():
result = runner.invoke(cli, ["--version"])

269
tests/test_mcp.py
View file

@ -1,3 +1,5 @@
from types import SimpleNamespace
import pytest
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.store.models import Chunk, Document, SearchResult
from haiku.rag.tools.document import DocumentInfo
from tests.multi_db.helpers import _config, _seed
@pytest.fixture(autouse=True)
@ -29,6 +32,22 @@ def mock_embedder(monkeypatch):
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
async def mcp_db(temp_db_path):
"""Create a test database with sample documents."""
@ -46,6 +65,21 @@ async def mcp_db(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):
"""Get a tool function from an MCP server by 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:
"""search_documents_by_image is registered only when the embedder is multimodal."""
@ -195,59 +359,40 @@ class TestMCPImageQuery:
@pytest.mark.asyncio
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
and routes a base64 image through ``client.search``."""
from haiku.rag.embeddings import EmbedderWrapper
and routes the decoded image and the selection through ``client.search``."""
seen = {}
class StubMultimodal(EmbedderWrapper):
supports_images = True
async def fake_search(self, query, **kwargs):
seen.update(query=query, **kwargs)
return []
def __init__(self):
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(),
)
monkeypatch.setattr(HaikuRAG, "search", fake_search)
mcp = create_mcp_server(mcp_db)
names = {t.name for t in await mcp.list_tools()}
assert "search_documents_by_image" in names
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
png_b64 = base64.b64encode(b"\x89PNG\r\n\x1a\n").decode("ascii")
results = await search_by_image(image_base64=png_b64)
# Empty list is fine (the stub vector won't match the toy fixture).
assert isinstance(results, list)
png = b"\x89PNG\r\n\x1a\n"
results = await search_by_image(
image_base64=base64.b64encode(png).decode("ascii"), sources=["alpha"]
)
assert results == []
assert seen["query"] == png
assert seen["sources"] == ["alpha"]
@pytest.mark.asyncio
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
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)
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 !!!")
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:
@pytest.mark.asyncio
@ -264,7 +422,7 @@ class TestMCPImageInput:
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
return ("answer", [])
@ -284,7 +442,7 @@ class TestMCPImageInput:
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
return SimpleNamespace(answer="answer")
@ -309,7 +467,7 @@ class TestMCPImageInput:
async def test_ask_question_without_images_passes_none(self, mcp_db, monkeypatch):
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
return ("answer", [])
@ -356,7 +514,7 @@ class TestMCPToolsDegradeOnError:
@pytest.mark.asyncio
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")
monkeypatch.setattr(HaikuRAG, "analyze", boom)
@ -377,9 +535,10 @@ class TestMCPToolsDegradeOnError:
content="cited text",
document_uri="test://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])
monkeypatch.setattr(HaikuRAG, "ask", fake_ask)
@ -389,6 +548,8 @@ class TestMCPToolsDegradeOnError:
with_cite = await ask(question="q", cite=True)
assert with_cite.startswith("the answer")
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"
@ -499,34 +660,6 @@ class TestMCPClientLifetime:
assert "AI Overview" 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):
"""A path and `lancedb.databases` both place the database."""
from haiku.rag.config.models import AppConfig, LanceDBConfig

16
tests/test_utils.py
View file

@ -753,6 +753,22 @@ def test_format_citations_sequential_indices():
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) ---