Add --no-agents to leave the model-backed MCP tools out
haiku-rag mcp --no-agents reaches create_mcp_server(agents=False): ask_question and analyze are not registered and the instructions drop the line describing them. qa.model always has a default, so a server without a usable model cannot be detected from configuration; the flag is how an operator says so. Refs #599
This commit is contained in:
parent
5653e876a8
commit
f45ed90b33
9 changed files with 150 additions and 71 deletions
|
|
@ -4,6 +4,8 @@
|
||||||
|
|
||||||
### Added
|
### 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
|
- MCP tools `get_document_outline` (heading tree with page numbers) and
|
||||||
`get_document_section` (one section's text, subsections included), built
|
`get_document_section` (one section's text, subsections included), built
|
||||||
on `document_items`. `build_toc` in `haiku.rag.context`.
|
on `document_items`. `build_toc` in `haiku.rag.context`.
|
||||||
|
|
|
||||||
|
|
@ -477,6 +477,9 @@ haiku-rag mcp --port 9000
|
||||||
|
|
||||||
# Bind to all interfaces (containers, trusted LAN)
|
# Bind to all interfaces (containers, trusted LAN)
|
||||||
haiku-rag mcp --host 0.0.0.0
|
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
|
See [MCP](mcp.md) for details. For continuous document ingestion
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,9 @@ haiku-rag mcp --host 0.0.0.0 --port 8001
|
||||||
|
|
||||||
# stdio transport (for Claude Desktop)
|
# stdio transport (for Claude Desktop)
|
||||||
haiku-rag mcp --stdio
|
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
|
`--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_outline` | always | `document_id`, `source` |
|
||||||
| `get_document_section` | always | `document_id`, `section_id`, `source` |
|
| `get_document_section` | always | `document_id`, `section_id`, `source` |
|
||||||
| `list_documents` | always | `limit`, `offset`, `filter` |
|
| `list_documents` | always | `limit`, `offset`, `filter` |
|
||||||
| `ask_question` | always | `question`, `images_base64`, `sources` |
|
| `ask_question` | unless `--no-agents` | `question`, `images_base64`, `sources` |
|
||||||
| `analyze` | always | `question`, `filter`, `images_base64`, `sources` |
|
| `analyze` | unless `--no-agents` | `question`, `filter`, `images_base64`, `sources` |
|
||||||
|
|
||||||
`search_documents` runs hybrid search, vector and full-text. Its text content
|
`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
|
is the rendering the in-process agents read: results best first, each with its
|
||||||
|
|
|
||||||
|
|
@ -925,6 +925,7 @@ class HaikuRAGApp:
|
||||||
transport: str | None = None,
|
transport: str | None = None,
|
||||||
host: str = "127.0.0.1",
|
host: str = "127.0.0.1",
|
||||||
port: int = 8001,
|
port: int = 8001,
|
||||||
|
agents: bool = True,
|
||||||
):
|
):
|
||||||
"""Run the MCP server until interrupted.
|
"""Run the MCP server until interrupted.
|
||||||
|
|
||||||
|
|
@ -934,7 +935,7 @@ class HaikuRAGApp:
|
||||||
# The resolved scope: a path overrides a configured URI, and a derived
|
# The resolved scope: a path overrides a configured URI, and a derived
|
||||||
# single-database configuration drops the name results and citations
|
# single-database configuration drops the name results and citations
|
||||||
# carry.
|
# carry.
|
||||||
server = _mcp_server_covering(self.scope, self.config)
|
server = _mcp_server_covering(self.scope, self.config, agents=agents)
|
||||||
try:
|
try:
|
||||||
if transport == "stdio":
|
if transport == "stdio":
|
||||||
await server.run_stdio_async()
|
await server.run_stdio_async()
|
||||||
|
|
|
||||||
|
|
@ -884,13 +884,20 @@ def mcp(
|
||||||
"--port",
|
"--port",
|
||||||
help="Port to bind MCP server to (ignored with --stdio)",
|
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:
|
) -> None:
|
||||||
"""Run the MCP server."""
|
"""Run the MCP server."""
|
||||||
app = create_app(db, covers_set=True)
|
app = create_app(db, covers_set=True)
|
||||||
|
|
||||||
transport = "stdio" if stdio else None
|
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__":
|
if __name__ == "__main__":
|
||||||
|
|
|
||||||
|
|
@ -92,16 +92,22 @@ async def _check_filter(
|
||||||
raise ToolError(f"Invalid filter {filter!r}: {e}") from e
|
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
|
"""What the server is for, naming no tools: the client has every tool's
|
||||||
description from the listing."""
|
description from the listing."""
|
||||||
lines = [
|
lines = [
|
||||||
"haiku-rag is the user's knowledge base: documents they ingested, "
|
"haiku-rag is the user's knowledge base: documents they ingested, "
|
||||||
"searchable by meaning and keyword, readable whole or section by "
|
"searchable by meaning and keyword, readable whole or section by section."
|
||||||
"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.",
|
|
||||||
]
|
]
|
||||||
|
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:
|
if scope.covers_multiple:
|
||||||
lines.append(
|
lines.append(
|
||||||
f"It holds several collections: {', '.join(scope.names)}. Results "
|
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(
|
def create_mcp_server(
|
||||||
db_path: Path | None = None,
|
db_path: Path | None = None,
|
||||||
config: AppConfig | None = None,
|
config: AppConfig | None = None,
|
||||||
|
agents: bool = True,
|
||||||
) -> FastMCP:
|
) -> FastMCP:
|
||||||
"""Create an MCP server over the databases the configuration places.
|
"""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
|
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.
|
||||||
|
agents: Register `ask_question` and `analyze`, which run a model on
|
||||||
|
the server.
|
||||||
"""
|
"""
|
||||||
from haiku.rag.client.scope import DatabaseScope
|
from haiku.rag.client.scope import DatabaseScope
|
||||||
|
|
||||||
config = config if config is not None else get_config()
|
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.
|
"""An MCP server over databases someone already resolved.
|
||||||
|
|
||||||
Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and
|
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.
|
# the traceback goes to the server log. A ToolError reaches the client as is.
|
||||||
mcp = FastMCP(
|
mcp = FastMCP(
|
||||||
"haiku-rag",
|
"haiku-rag",
|
||||||
instructions=_instructions(scope, config),
|
instructions=_instructions(scope, config, agents),
|
||||||
version=metadata.version("haiku.rag-slim"),
|
version=metadata.version("haiku.rag-slim"),
|
||||||
lifespan=lifespan,
|
lifespan=lifespan,
|
||||||
mask_error_details=True,
|
mask_error_details=True,
|
||||||
|
|
@ -453,68 +466,72 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
|
||||||
for doc in documents
|
for doc in documents
|
||||||
]
|
]
|
||||||
|
|
||||||
@mcp.tool(annotations=_read_only("Ask a question"))
|
if agents:
|
||||||
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.
|
|
||||||
|
|
||||||
Use this when the user wants an answer rather than material to read.
|
@mcp.tool(annotations=_read_only("Ask a question"))
|
||||||
It runs a model on the server and is slower than a search. Returns
|
async def ask_question(
|
||||||
the answer, followed by citations to the passages it rests on.
|
question: str,
|
||||||
|
images_base64: list[str] | None = None,
|
||||||
|
sources: Sources = None,
|
||||||
|
) -> str:
|
||||||
|
"""Answer a question from the documents with a retrieval agent.
|
||||||
|
|
||||||
Args:
|
Use this when the user wants an answer rather than material to read.
|
||||||
question: The question, in natural language.
|
It runs a model on the server and is slower than a search. Returns
|
||||||
images_base64: Images to attach to the question, PNG or JPEG
|
the answer, followed by citations to the passages it rests on.
|
||||||
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
|
|
||||||
|
|
||||||
@mcp.tool(annotations=_read_only("Analyze documents"))
|
Args:
|
||||||
async def analyze(
|
question: The question, in natural language.
|
||||||
question: str,
|
images_base64: Images to attach to the question, PNG or JPEG
|
||||||
filter: Filter = None,
|
bytes as base64. Needs a vision-capable model on the server.
|
||||||
images_base64: list[str] | None = None,
|
"""
|
||||||
sources: Sources = None,
|
images = _decode_images(images_base64)
|
||||||
) -> str:
|
rag = await _client()
|
||||||
"""Compute an answer across documents with code.
|
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
|
@mcp.tool(annotations=_read_only("Analyze documents"))
|
||||||
or arithmetic over tables, where reading passages is not enough. A
|
async def analyze(
|
||||||
model writes and runs Python in a sandbox over the selected documents.
|
question: str,
|
||||||
It is the slowest tool. Returns the answer as text.
|
filter: Filter = None,
|
||||||
|
images_base64: list[str] | None = None,
|
||||||
|
sources: Sources = None,
|
||||||
|
) -> str:
|
||||||
|
"""Compute an answer across documents with code.
|
||||||
|
|
||||||
Args:
|
Use this for counting, aggregation, comparison across many documents
|
||||||
question: The question, in natural language.
|
or arithmetic over tables, where reading passages is not enough. A
|
||||||
images_base64: Images to attach to the question, PNG or JPEG
|
model writes and runs Python in a sandbox over the selected documents.
|
||||||
bytes as base64. Needs a vision-capable model on the server.
|
It is the slowest tool. Returns the answer as text.
|
||||||
"""
|
|
||||||
images = _decode_images(images_base64)
|
Args:
|
||||||
rag = await _client()
|
question: The question, in natural language.
|
||||||
try:
|
images_base64: Images to attach to the question, PNG or JPEG
|
||||||
result = await rag.analyze(
|
bytes as base64. Needs a vision-capable model on the server.
|
||||||
question, filter=filter, images=images, sources=sources
|
"""
|
||||||
)
|
images = _decode_images(images_base64)
|
||||||
except UnknownDatabaseError as e:
|
rag = await _client()
|
||||||
raise ToolError(str(e)) from e
|
try:
|
||||||
except Exception as e:
|
result = await rag.analyze(
|
||||||
logger.exception("analyze failed")
|
question, filter=filter, images=images, sources=sources
|
||||||
raise ToolError(f"analyze failed: {type(e).__name__}") from e
|
)
|
||||||
return result.answer
|
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
|
return mcp
|
||||||
|
|
|
||||||
|
|
@ -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):
|
async def test_run_mcp_survives_interruption(app, client, monkeypatch):
|
||||||
server = AsyncMock()
|
server = AsyncMock()
|
||||||
server.run_stdio_async.side_effect = KeyboardInterrupt
|
server.run_stdio_async.side_effect = KeyboardInterrupt
|
||||||
|
|
|
||||||
|
|
@ -1041,6 +1041,14 @@ def test_mcp_stdio_selects_the_transport(app_stub):
|
||||||
app_stub.run_mcp.assert_called_once()
|
app_stub.run_mcp.assert_called_once()
|
||||||
kwargs = app_stub.run_mcp.call_args.kwargs
|
kwargs = app_stub.run_mcp.call_args.kwargs
|
||||||
assert kwargs["transport"] == "stdio"
|
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):
|
def test_mcp_without_stdio_leaves_the_transport_unset(app_stub):
|
||||||
|
|
|
||||||
|
|
@ -632,6 +632,18 @@ class TestMCPDescribesItself:
|
||||||
assert "beta" in covering_both
|
assert "beta" in covering_both
|
||||||
assert "beta" not in covering_one
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_instructions_carry_the_domain_preamble(self, mcp_db):
|
async def test_instructions_carry_the_domain_preamble(self, mcp_db):
|
||||||
from fastmcp import Client
|
from fastmcp import Client
|
||||||
|
|
@ -695,6 +707,18 @@ class TestMCPToolSet:
|
||||||
"analyze",
|
"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:
|
class TestMCPCoversTheConfiguredSet:
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -1235,7 +1259,7 @@ class TestMCPClientLifetime:
|
||||||
async def run_stdio_async(self):
|
async def run_stdio_async(self):
|
||||||
return None
|
return None
|
||||||
|
|
||||||
def fake_covering(scope, config):
|
def fake_covering(scope, config, agents=True):
|
||||||
seen.update(scope=scope, config=config)
|
seen.update(scope=scope, config=config)
|
||||||
return _Server()
|
return _Server()
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue