diff --git a/CHANGELOG.md b/CHANGELOG.md index 012c59c2..5a90db87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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`, diff --git a/docs/cli.md b/docs/cli.md index d1e083e5..5e0efae8 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -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 diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index baddfbe5..fb0057f9 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -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 diff --git a/docs/mcp.md b/docs/mcp.md index c3862d1a..ca07ce4c 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 836b1efb..c47d3971 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index b9142a6f..44a4f259 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -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}" diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index c4bb7ddb..1458083a 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -27,6 +27,7 @@ class DocumentInfo(BaseModel): title: str uri: str created: str + source: str | None = None class DocumentListResponse(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 321ce3aa..b96bf74b 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -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: ]`` 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) diff --git a/tests/test_cli.py b/tests/test_cli.py index 5c2a7b44..50d389a0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -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"]) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 910f2a2e..06bc32d5 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -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 diff --git a/tests/test_utils.py b/tests/test_utils.py index ca0ae454..5284cb04 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -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) ---