diff --git a/CHANGELOG.md b/CHANGELOG.md index 90bf3bf3..7a22246e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### Added +- `haiku-rag mcp --no-agents` leaves `ask_question` and `analyze` + unregistered. `create_mcp_server(agents=)`, `HaikuRAGApp.run_mcp(agents=)`. - MCP tools `get_document_outline` (heading tree with page numbers) and `get_document_section` (one section's text, subsections included), built on `document_items`. `build_toc` in `haiku.rag.context`. diff --git a/docs/cli.md b/docs/cli.md index 5e0efae8..b54cb8fe 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -477,6 +477,9 @@ haiku-rag mcp --port 9000 # Bind to all interfaces (containers, trusted LAN) haiku-rag mcp --host 0.0.0.0 + +# Without the ask_question and analyze tools +haiku-rag mcp --no-agents ``` See [MCP](mcp.md) for details. For continuous document ingestion diff --git a/docs/mcp.md b/docs/mcp.md index 91f636d0..59093b3b 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -18,6 +18,9 @@ haiku-rag mcp --host 0.0.0.0 --port 8001 # stdio transport (for Claude Desktop) haiku-rag mcp --stdio + +# Without ask_question and analyze, which run a model on the server +haiku-rag mcp --stdio --no-agents ``` `--host` defaults to `127.0.0.1` (loopback only). Bind to `0.0.0.0` only @@ -81,8 +84,8 @@ repeating it. | `get_document_outline` | always | `document_id`, `source` | | `get_document_section` | always | `document_id`, `section_id`, `source` | | `list_documents` | always | `limit`, `offset`, `filter` | -| `ask_question` | always | `question`, `images_base64`, `sources` | -| `analyze` | always | `question`, `filter`, `images_base64`, `sources` | +| `ask_question` | unless `--no-agents` | `question`, `images_base64`, `sources` | +| `analyze` | unless `--no-agents` | `question`, `filter`, `images_base64`, `sources` | `search_documents` runs hybrid search, vector and full-text. Its text content is the rendering the in-process agents read: results best first, each with its diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 8bc934ae..d2dbedad 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -925,6 +925,7 @@ class HaikuRAGApp: transport: str | None = None, host: str = "127.0.0.1", port: int = 8001, + agents: bool = True, ): """Run the MCP server until interrupted. @@ -934,7 +935,7 @@ class HaikuRAGApp: # The resolved scope: a path overrides a configured URI, and a derived # single-database configuration drops the name results and citations # carry. - server = _mcp_server_covering(self.scope, self.config) + server = _mcp_server_covering(self.scope, self.config, agents=agents) try: if transport == "stdio": await server.run_stdio_async() diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index c47d3971..06647670 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -884,13 +884,20 @@ def mcp( "--port", help="Port to bind MCP server to (ignored with --stdio)", ), + no_agents: bool = typer.Option( + False, + "--no-agents", + help="Do not register ask_question and analyze, which run a model", + ), ) -> None: """Run the MCP server.""" app = create_app(db, covers_set=True) transport = "stdio" if stdio else None - asyncio.run(app.run_mcp(transport=transport, host=host, port=port)) + asyncio.run( + app.run_mcp(transport=transport, host=host, port=port, agents=not no_agents) + ) if __name__ == "__main__": diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 39130c86..5da35d47 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -92,16 +92,22 @@ async def _check_filter( raise ToolError(f"Invalid filter {filter!r}: {e}") from e -def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: +def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> str: """What the server is for, naming no tools: the client has every tool's description from the listing.""" lines = [ "haiku-rag is the user's knowledge base: documents they ingested, " - "searchable by meaning and keyword, readable whole or section by " - "section, answered with citations, or computed across documents.", - "Use it whenever a question could be answered from those documents, " - "before answering from memory, and say when it had nothing relevant.", + "searchable by meaning and keyword, readable whole or section by section." ] + if agents: + lines.append( + "Questions can be answered from them with citations, or computed " + "across them with code." + ) + lines.append( + "Use it whenever a question could be answered from those documents, " + "before answering from memory, and say when it had nothing relevant." + ) if scope.covers_multiple: lines.append( f"It holds several collections: {', '.join(scope.names)}. Results " @@ -181,6 +187,7 @@ def _find(toc: list["dict[str, Any]"], section_id: str) -> "dict[str, Any] | Non def create_mcp_server( db_path: Path | None = None, config: AppConfig | None = None, + agents: bool = True, ) -> FastMCP: """Create an MCP server over the databases the configuration places. @@ -189,14 +196,20 @@ def create_mcp_server( None to serve the databases the configuration places. Beside `lancedb.databases` a path raises `AmbiguousDatabaseError`. config: Configuration to use. + agents: Register `ask_question` and `analyze`, which run a model on + the server. """ from haiku.rag.client.scope import DatabaseScope config = config if config is not None else get_config() - return _covering(DatabaseScope.resolve(config, database_path=db_path), config) + return _covering( + DatabaseScope.resolve(config, database_path=db_path), config, agents + ) -def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: +def _covering( + scope: "DatabaseScope", config: AppConfig, agents: bool = True +) -> FastMCP: """An MCP server over databases someone already resolved. Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and @@ -242,7 +255,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: # the traceback goes to the server log. A ToolError reaches the client as is. mcp = FastMCP( "haiku-rag", - instructions=_instructions(scope, config), + instructions=_instructions(scope, config, agents), version=metadata.version("haiku.rag-slim"), lifespan=lifespan, mask_error_details=True, @@ -453,68 +466,72 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: for doc in documents ] - @mcp.tool(annotations=_read_only("Ask a question")) - async def ask_question( - question: str, - images_base64: list[str] | None = None, - sources: Sources = None, - ) -> str: - """Answer a question from the documents with a retrieval agent. + if agents: - Use this when the user wants an answer rather than material to read. - It runs a model on the server and is slower than a search. Returns - the answer, followed by citations to the passages it rests on. + @mcp.tool(annotations=_read_only("Ask a question")) + async def ask_question( + question: str, + images_base64: list[str] | None = None, + sources: Sources = None, + ) -> str: + """Answer a question from the documents with a retrieval agent. - Args: - question: The question, in natural language. - images_base64: Images to attach to the question, PNG or JPEG - bytes as base64. Needs a vision-capable model on the server. - """ - images = _decode_images(images_base64) - rag = await _client() - try: - answer, citations = await rag.ask(question, images=images, sources=sources) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e - except Exception as e: - logger.exception("ask_question failed") - raise ToolError(f"ask_question failed: {type(e).__name__}") from e - if citations: - answer += "\n\n" + format_citations( - citations, include_source=rag.covers_multiple - ) - return answer + Use this when the user wants an answer rather than material to read. + It runs a model on the server and is slower than a search. Returns + the answer, followed by citations to the passages it rests on. - @mcp.tool(annotations=_read_only("Analyze documents")) - async def analyze( - question: str, - filter: Filter = None, - images_base64: list[str] | None = None, - sources: Sources = None, - ) -> str: - """Compute an answer across documents with code. + Args: + question: The question, in natural language. + images_base64: Images to attach to the question, PNG or JPEG + bytes as base64. Needs a vision-capable model on the server. + """ + images = _decode_images(images_base64) + rag = await _client() + try: + answer, citations = await rag.ask( + question, images=images, sources=sources + ) + except UnknownDatabaseError as e: + raise ToolError(str(e)) from e + except Exception as e: + logger.exception("ask_question failed") + raise ToolError(f"ask_question failed: {type(e).__name__}") from e + if citations: + answer += "\n\n" + format_citations( + citations, include_source=rag.covers_multiple + ) + return answer - Use this for counting, aggregation, comparison across many documents - or arithmetic over tables, where reading passages is not enough. A - model writes and runs Python in a sandbox over the selected documents. - It is the slowest tool. Returns the answer as text. + @mcp.tool(annotations=_read_only("Analyze documents")) + async def analyze( + question: str, + filter: Filter = None, + images_base64: list[str] | None = None, + sources: Sources = None, + ) -> str: + """Compute an answer across documents with code. - Args: - question: The question, in natural language. - images_base64: Images to attach to the question, PNG or JPEG - bytes as base64. Needs a vision-capable model on the server. - """ - images = _decode_images(images_base64) - rag = await _client() - try: - result = await rag.analyze( - question, filter=filter, images=images, sources=sources - ) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e - except Exception as e: - logger.exception("analyze failed") - raise ToolError(f"analyze failed: {type(e).__name__}") from e - return result.answer + Use this for counting, aggregation, comparison across many documents + or arithmetic over tables, where reading passages is not enough. A + model writes and runs Python in a sandbox over the selected documents. + It is the slowest tool. Returns the answer as text. + + Args: + question: The question, in natural language. + images_base64: Images to attach to the question, PNG or JPEG + bytes as base64. Needs a vision-capable model on the server. + """ + images = _decode_images(images_base64) + rag = await _client() + try: + result = await rag.analyze( + question, filter=filter, images=images, sources=sources + ) + except UnknownDatabaseError as e: + raise ToolError(str(e)) from e + except Exception as e: + logger.exception("analyze failed") + raise ToolError(f"analyze failed: {type(e).__name__}") from e + return result.answer return mcp diff --git a/tests/test_app.py b/tests/test_app.py index cf7104da..589c2eb1 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -705,6 +705,20 @@ async def test_run_mcp_http(app, client, monkeypatch): ) +async def test_run_mcp_hands_the_server_the_agents_switch(app, client, monkeypatch): + seen = {} + + def fake_covering(scope, config, agents=True): + seen["agents"] = agents + return AsyncMock() + + monkeypatch.setattr("haiku.rag.app._mcp_server_covering", fake_covering) + + await app.run_mcp(transport="stdio", agents=False) + + assert seen["agents"] is False + + async def test_run_mcp_survives_interruption(app, client, monkeypatch): server = AsyncMock() server.run_stdio_async.side_effect = KeyboardInterrupt diff --git a/tests/test_cli.py b/tests/test_cli.py index 50d389a0..7d4af634 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1041,6 +1041,14 @@ def test_mcp_stdio_selects_the_transport(app_stub): app_stub.run_mcp.assert_called_once() kwargs = app_stub.run_mcp.call_args.kwargs assert kwargs["transport"] == "stdio" + assert kwargs["agents"] is True + + +def test_mcp_no_agents_leaves_the_agent_tools_out(app_stub): + result = runner.invoke(cli, ["mcp", "--no-agents"] + DB_ARGS) + + assert result.exit_code == 0, result.output + assert app_stub.run_mcp.call_args.kwargs["agents"] is False def test_mcp_without_stdio_leaves_the_transport_unset(app_stub): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 13454930..7bf3183d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -632,6 +632,18 @@ class TestMCPDescribesItself: assert "beta" in covering_both assert "beta" not in covering_one + @pytest.mark.asyncio + async def test_instructions_without_agents_drop_only_their_clause(self, mcp_db): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + full = client.initialize_result.instructions.splitlines() + async with Client(create_mcp_server(mcp_db, agents=False)) as client: + without = client.initialize_result.instructions.splitlines() + + assert set(without) < set(full) + assert len(without) == len(full) - 1 + @pytest.mark.asyncio async def test_instructions_carry_the_domain_preamble(self, mcp_db): from fastmcp import Client @@ -695,6 +707,18 @@ class TestMCPToolSet: "analyze", } + @pytest.mark.asyncio + async def test_without_agents_the_agent_tools_are_not_registered(self, mcp_db): + mcp = create_mcp_server(mcp_db, agents=False) + + assert {t.name for t in await mcp.list_tools()} == { + "search_documents", + "get_document", + "get_document_outline", + "get_document_section", + "list_documents", + } + class TestMCPCoversTheConfiguredSet: @pytest.mark.asyncio @@ -1235,7 +1259,7 @@ class TestMCPClientLifetime: async def run_stdio_async(self): return None - def fake_covering(scope, config): + def fake_covering(scope, config, agents=True): seen.update(scope=scope, config=config) return _Server()