From 7eeddd1a7c0259f3a5373747dd31dea32e036d84 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 10:00:12 +0300 Subject: [PATCH 01/18] Start MCP server revisited (#599) From ce69a8c989b9de78482df760eeb039761dc8f9a4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 10:27:32 +0300 Subject: [PATCH 02/18] Remove the MCP write tools; the server opens read-only add_document_from_file, add_document_from_url, add_document_from_text and delete_document are gone. create_mcp_server and _covering lose read_only; the client always opens with read_only=True, so the --read-only flag before `mcp` is redundant and the docs, Dockerfiles and compose example drop it. Ingestion is `haiku-rag add`/`add-src`/`delete` and haiku-ingester. No known consumer used the write tools; one stdio server per client window made multi-writer the accidental default, and streamable HTTP has no auth. The test_app stub drops a comment and __init__ that described run_mcp constructing the client positionally; every path goes through _covering. Refs #599 --- CHANGELOG.md | 6 + README.md | 2 +- docker/Dockerfile | 2 +- docker/Dockerfile.slim | 2 +- docs/cli.md | 3 - docs/configuration/storage.md | 2 +- docs/mcp.md | 29 +--- examples/docker/docker-compose.yml | 1 - haiku_rag_slim/haiku/rag/app.py | 2 +- haiku_rag_slim/haiku/rag/mcp.py | 78 +--------- tests/test_app.py | 4 - tests/test_mcp.py | 232 ++++++----------------------- 12 files changed, 61 insertions(+), 302 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8834d6de..012c59c2 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`. +### Removed + +- MCP write tools `add_document_from_file`, `add_document_from_url`, + `add_document_from_text` and `delete_document`. The server opens the + database read-only; ingest with `haiku-rag add`, `add-src`, `delete` or + `haiku-ingester`. `create_mcp_server` loses `read_only`. ## [0.82.1] - 2026-09-03 diff --git a/README.md b/README.md index 3829c502..9016e322 100644 --- a/README.md +++ b/README.md @@ -129,7 +129,7 @@ Add to your Claude Desktop configuration: } ``` -Provides tools for document management, search, QA, and analysis directly in your AI assistant. +Provides search, document, QA, and analysis tools directly in your AI assistant. ## Examples diff --git a/docker/Dockerfile b/docker/Dockerfile index 51e4f336..fa5713d1 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -40,4 +40,4 @@ EXPOSE 8001 8765 # Default command: read-only MCP server. The companion ingester service is # launched via docker-compose against the same image. -CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "--read-only", "mcp", "--port", "8001"] +CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "mcp", "--port", "8001"] diff --git a/docker/Dockerfile.slim b/docker/Dockerfile.slim index 529cb6c0..72844915 100644 --- a/docker/Dockerfile.slim +++ b/docker/Dockerfile.slim @@ -39,4 +39,4 @@ EXPOSE 8001 8765 # Default command: read-only MCP server. The companion ingester service is # launched via docker-compose against the same image. -CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "--read-only", "mcp", "--port", "8001"] +CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "mcp", "--port", "8001"] diff --git a/docs/cli.md b/docs/cli.md index f9baf5d6..d1e083e5 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -477,9 +477,6 @@ haiku-rag mcp --port 9000 # Bind to all interfaces (containers, trusted LAN) haiku-rag mcp --host 0.0.0.0 - -# Read-only mode (no write tools) -haiku-rag --read-only mcp ``` See [MCP](mcp.md) for details. For continuous document ingestion diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index fd35edb3..baddfbe5 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -196,7 +196,7 @@ writing process per database URI, any number of read-only consumers. The recommended layout for production is "different buckets, same account, separate IAM roles per process": - **Ingestion process** — IAM role with `s3:Get/List` on the documents bucket and `s3:Get/Put/Delete` on the LanceDB bucket. Runs `haiku-ingester serve` (with `ingester.sources[type=s3]` pointing at the documents bucket). Exactly one such process per LanceDB URI. -- **Consumer processes** (1..N) — IAM role with `s3:Get/List` on the LanceDB bucket only. Run `haiku-rag --read-only mcp`, the chat TUI, etc. They never see the documents bucket. +- **Consumer processes** (1..N) — IAM role with `s3:Get/List` on the LanceDB bucket only. Run `haiku-rag mcp`, the chat TUI, etc. They never see the documents bucket. Each process picks up its own credentials from the AWS default chain (env vars, IAM instance role, AWS profile), so no credentials are hard-coded in the configuration files. diff --git a/docs/mcp.md b/docs/mcp.md index 06403bb3..c3862d1a 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -18,16 +18,14 @@ haiku-rag mcp --host 0.0.0.0 --port 8001 # stdio transport (for Claude Desktop) haiku-rag mcp --stdio - -# Read-only mode (excludes write tools) -haiku-rag --read-only mcp --stdio ``` `--host` defaults to `127.0.0.1` (loopback only). Bind to `0.0.0.0` only when you want the MCP server reachable from outside the local machine — e.g. inside a Docker container with port mapping, or on a trusted LAN. -**Read-only mode:** When `--read-only` is specified, write tools (`add_document_from_file`, `add_document_from_url`, `add_document_from_text`, `delete_document`) are not registered. Only search and query tools remain available. +The server opens the database read-only. Ingestion goes through the CLI +(`haiku-rag add`, `add-src`, `delete`) or [`haiku-ingester`](ingester.md). ## Claude Desktop Integration @@ -57,27 +55,11 @@ With a custom database path: } ``` -After restarting Claude Desktop, you can ask Claude to search your documents, add new content, or answer questions using your knowledge base. +After restarting Claude Desktop, you can ask Claude to search your documents or answer questions using your knowledge base. ## Available Tools -### Document Management - -- **`add_document_from_file`** - Add documents from local file paths - - `file_path` (required): Path to the file - - `metadata` (optional): Key-value metadata - - `title` (optional): Human-readable title - -- **`add_document_from_url`** - Add documents from URLs - - `url` (required): URL to fetch - - `metadata` (optional): Key-value metadata - - `title` (optional): Human-readable title - -- **`add_document_from_text`** - Add documents from raw text content - - `content` (required): Text content - - `uri` (optional): URI identifier - - `metadata` (optional): Key-value metadata - - `title` (optional): Human-readable title +### Documents - **`get_document`** - Retrieve a document by ID - `document_id` (required): The document ID @@ -87,9 +69,6 @@ After restarting Claude Desktop, you can ask Claude to search your documents, ad - `offset` (optional): Number to skip - `filter` (optional): SQL WHERE clause for filtering -- **`delete_document`** - Delete a document by ID - - `document_id` (required): The document ID - ### Search - **`search_documents`** - Search using hybrid search (vector + full-text) diff --git a/examples/docker/docker-compose.yml b/examples/docker/docker-compose.yml index a889bf7d..45f068ac 100644 --- a/examples/docker/docker-compose.yml +++ b/examples/docker/docker-compose.yml @@ -103,7 +103,6 @@ services: "haiku-rag", "--config", "/app/haiku.rag.yaml", - "--read-only", "mcp", "--host", "0.0.0.0", diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 0c71ec47..8bc934ae 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -934,7 +934,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, self.read_only) + server = _mcp_server_covering(self.scope, self.config) try: if transport == "stdio": await server.run_stdio_async() diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index f5f808a3..b9142a6f 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -2,7 +2,7 @@ import asyncio from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from fastmcp import FastMCP @@ -27,7 +27,6 @@ def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: def create_mcp_server( db_path: Path | None = None, config: AppConfig | None = None, - read_only: bool = False, ) -> FastMCP: """Create an MCP server over one database. @@ -36,17 +35,14 @@ def create_mcp_server( None to serve the database the configuration places. Beside `lancedb.databases` a path raises `AmbiguousDatabaseError`. config: Configuration to use. - read_only: If True, write tools (add_document_*, delete_document) are not registered. """ 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, read_only - ) + return _covering(DatabaseScope.resolve(config, database_path=db_path), config) -def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> FastMCP: +def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: """An MCP server over databases someone already resolved. Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and @@ -76,7 +72,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> Fas async with client_lock: if client is None: client = await stack.enter_async_context( - HaikuRAG._covering(scope, config, read_only=read_only) + HaikuRAG._covering(scope, config, read_only=True) ) return client @@ -97,72 +93,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> Fas mcp = FastMCP("haiku-rag", lifespan=lifespan) - # Write tools - only registered when not in read-only mode - if not read_only: - - @mcp.tool() - async def add_document_from_file( - file_path: str, - metadata: dict[str, Any] | None = None, - title: str | None = None, - ) -> str | None: - """Add a document to the RAG system from a file path.""" - try: - rag = await _client() - result = await rag.create_document_from_source( - Path(file_path), title=title, metadata=metadata or {} - ) - # Handle both single document and list of documents (directories) - if isinstance(result, list): - return result[0].id if result else None - return result.id - except Exception: - return None - - @mcp.tool() - async def add_document_from_url( - url: str, metadata: dict[str, Any] | None = None, title: str | None = None - ) -> str | None: - """Add a document to the RAG system from a URL.""" - try: - rag = await _client() - result = await rag.create_document_from_source( - url, title=title, metadata=metadata or {} - ) - # Handle both single document and list of documents - if isinstance(result, list): - return result[0].id if result else None - return result.id - except Exception: - return None - - @mcp.tool() - async def add_document_from_text( - content: str, - uri: str | None = None, - metadata: dict[str, Any] | None = None, - title: str | None = None, - ) -> str | None: - """Add a document to the RAG system from text content.""" - try: - rag = await _client() - document = await rag.create_document( - content, uri, title=title, metadata=metadata or {} - ) - return document.id - except Exception: - return None - - @mcp.tool() - async def delete_document(document_id: str) -> bool: - """Delete a document by its ID.""" - try: - rag = await _client() - return await rag.delete_document(document_id) - except Exception: - return False - - # Read tools - always registered @mcp.tool() async def search_documents( query: str, limit: int | None = None, include_images: bool = True diff --git a/tests/test_app.py b/tests/test_app.py index ebe43f9b..cf7104da 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -31,10 +31,6 @@ def client(): @pytest.fixture def app(tmp_path, client, monkeypatch): class StubHaikuRAG: - # run_mcp passes db_path positionally; every other caller uses kwargs. - def __init__(self, *args, **kwargs): - pass - @classmethod def _covering(cls, *args, **kwargs): return cls() diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 7b465b45..910f2a2e 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -55,7 +55,7 @@ async def _get_tool(mcp, name): class TestMCPReadTools: @pytest.mark.asyncio async def test_search_documents(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") results = await search(query="artificial intelligence") @@ -64,7 +64,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_search_documents_with_limit(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") results = await search(query="artificial intelligence", limit=1) @@ -93,7 +93,7 @@ class TestMCPReadTools: ) await rag.store.chunks_table.optimize() - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) async with Client(mcp) as client: result = await client.call_tool( "search_documents", {"query": "artificial intelligence"} @@ -107,7 +107,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_get_document(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) get_doc = await _get_tool(mcp, "get_document") # First get the ID via list @@ -122,7 +122,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_get_document_excludes_docling_fields(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) get_doc = await _get_tool(mcp, "get_document") list_docs = await _get_tool(mcp, "list_documents") @@ -136,7 +136,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_get_document_not_found(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) get_doc = await _get_tool(mcp, "get_document") result = await get_doc(document_id="nonexistent-id") @@ -144,7 +144,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_list_documents(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) list_docs = await _get_tool(mcp, "list_documents") results = await list_docs() @@ -153,7 +153,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_list_documents_with_limit(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) list_docs = await _get_tool(mcp, "list_documents") results = await list_docs(limit=1) @@ -161,7 +161,7 @@ class TestMCPReadTools: @pytest.mark.asyncio async def test_list_documents_with_filter(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) list_docs = await _get_tool(mcp, "list_documents") results = await list_docs(filter="title = 'AI Overview'") @@ -169,66 +169,18 @@ class TestMCPReadTools: assert results[0].title == "AI Overview" -class TestMCPWriteTools: +class TestMCPToolSet: @pytest.mark.asyncio - async def test_write_tools_registered_when_not_read_only(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True): - pass - mcp = create_mcp_server(temp_db_path, read_only=False) - tools = await mcp.list_tools() - tool_names = [t.name for t in tools] - assert "add_document_from_text" in tool_names - assert "add_document_from_file" in tool_names - assert "add_document_from_url" in tool_names - assert "delete_document" in tool_names + async def test_the_server_registers_read_tools_only(self, mcp_db): + mcp = create_mcp_server(mcp_db) - @pytest.mark.asyncio - async def test_write_tools_not_registered_when_read_only(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True): - pass - mcp = create_mcp_server(temp_db_path, read_only=True) - tools = await mcp.list_tools() - tool_names = [t.name for t in tools] - assert "add_document_from_text" not in tool_names - assert "delete_document" not in tool_names - - @pytest.mark.asyncio - async def test_add_document_from_text(self, temp_db_path): - async with HaikuRAG(temp_db_path, create=True): - pass - mcp = create_mcp_server(temp_db_path, read_only=False) - add_text = await _get_tool(mcp, "add_document_from_text") - - doc_id = await add_text(content="Test content for MCP", title="MCP Test Doc") - assert doc_id is not None - - get_doc = await _get_tool(mcp, "get_document") - doc = await get_doc(document_id=doc_id) - assert doc.title == "MCP Test Doc" - assert doc.content == "Test content for MCP" - - @pytest.mark.asyncio - async def test_delete_document(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=False) - list_docs = await _get_tool(mcp, "list_documents") - delete_doc = await _get_tool(mcp, "delete_document") - - docs = await list_docs() - assert len(docs) == 2 - - result = await delete_doc(document_id=docs[0].id) - assert result is True - - docs_after = await list_docs() - assert len(docs_after) == 1 - - @pytest.mark.asyncio - async def test_delete_document_not_found(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=False) - delete_doc = await _get_tool(mcp, "delete_document") - - result = await delete_doc(document_id="nonexistent-id") - assert result is False + assert {t.name for t in await mcp.list_tools()} == { + "search_documents", + "get_document", + "list_documents", + "ask_question", + "analyze", + } class TestMCPImageQuery: @@ -237,7 +189,7 @@ class TestMCPImageQuery: @pytest.mark.asyncio async def test_image_query_tool_absent_for_text_only_embedder(self, mcp_db): """Default text-only embedder must not expose the image-query tool.""" - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) names = {t.name for t in await mcp.list_tools()} assert "search_documents_by_image" not in names @@ -264,7 +216,7 @@ class TestMCPImageQuery: lambda *a, **kw: StubMultimodal(), ) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) names = {t.name for t in await mcp.list_tools()} assert "search_documents_by_image" in names @@ -296,7 +248,7 @@ class TestMCPImageQuery: lambda *a, **kw: StubMultimodal(), ) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) search_by_image = await _get_tool(mcp, "search_documents_by_image") # Not valid base64 (contains non-base64 chars) — the strict decoder @@ -317,7 +269,7 @@ class TestMCPImageInput: return ("answer", []) monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) ask = await _get_tool(mcp, "ask_question") png = b"fake image bytes" @@ -337,7 +289,7 @@ class TestMCPImageInput: return SimpleNamespace(answer="answer") monkeypatch.setattr(HaikuRAG, "analyze", fake_analyze) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) analyze = await _get_tool(mcp, "analyze") jpeg = b"fake jpeg bytes" @@ -347,7 +299,7 @@ class TestMCPImageInput: @pytest.mark.asyncio async def test_ask_question_rejects_invalid_base64(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) ask = await _get_tool(mcp, "ask_question") result = await ask(question="q", images_base64=["!!! not base64 !!!"]) @@ -362,7 +314,7 @@ class TestMCPImageInput: return ("answer", []) monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) ask = await _get_tool(mcp, "ask_question") result = await ask(question="q") @@ -370,70 +322,6 @@ class TestMCPImageInput: assert captured["images"] is None -class TestMCPFileAndUrlIngestion: - @pytest.mark.asyncio - async def test_add_document_from_file(self, temp_db_path, tmp_path): - async with HaikuRAG(temp_db_path, create=True): - pass - source = tmp_path / "note.txt" - source.write_text("Ingested from a file path.") - - mcp = create_mcp_server(temp_db_path, read_only=False) - add_file = await _get_tool(mcp, "add_document_from_file") - - doc_id = await add_file(file_path=str(source), title="File Doc") - assert doc_id is not None - - get_doc = await _get_tool(mcp, "get_document") - doc = await get_doc(document_id=doc_id) - assert doc.title == "File Doc" - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "tool_name,kwargs", - [ - ("add_document_from_file", {"file_path": "/tmp/x.txt"}), - ("add_document_from_url", {"url": "https://example.com/x.txt"}), - ], - ) - @pytest.mark.parametrize( - "results,expected", - [ - ( - [Document(id="first", content="a"), Document(id="second", content="b")], - "first", - ), - ([], None), - ], - ids=["directory_reports_first_id", "empty_directory_reports_none"], - ) - async def test_add_tools_handle_multi_document_sources( - self, mcp_db, monkeypatch, tool_name, kwargs, results, expected - ): - """A source resolving to several documents reports the first id.""" - - async def fake_from_source(self, source, title=None, metadata=None, **kw): - return results - - monkeypatch.setattr(HaikuRAG, "create_document_from_source", fake_from_source) - mcp = create_mcp_server(mcp_db, read_only=False) - add = await _get_tool(mcp, tool_name) - - assert await add(**kwargs) == expected - - @pytest.mark.asyncio - async def test_add_document_from_url(self, mcp_db, monkeypatch): - async def fake_from_source(self, source, title=None, metadata=None, **kwargs): - assert source == "https://example.com/doc.txt" - return Document(id="url-doc", content="fetched") - - monkeypatch.setattr(HaikuRAG, "create_document_from_source", fake_from_source) - mcp = create_mcp_server(mcp_db, read_only=False) - add_url = await _get_tool(mcp, "add_document_from_url") - - assert await add_url(url="https://example.com/doc.txt") == "url-doc" - - class TestMCPToolsDegradeOnError: """Every tool swallows client failures and returns its empty value rather than propagating an exception to the MCP transport.""" @@ -442,20 +330,6 @@ class TestMCPToolsDegradeOnError: @pytest.mark.parametrize( "client_method,tool_name,kwargs,expected", [ - ( - "create_document_from_source", - "add_document_from_file", - {"file_path": "/tmp/x.txt"}, - None, - ), - ( - "create_document_from_source", - "add_document_from_url", - {"url": "https://example.com/x"}, - None, - ), - ("create_document", "add_document_from_text", {"content": "x"}, None), - ("delete_document", "delete_document", {"document_id": "x"}, False), ("search", "search_documents", {"query": "x"}, []), ("get_document_by_id", "get_document", {"document_id": "x"}, None), ("list_documents", "list_documents", {}, []), @@ -468,14 +342,14 @@ class TestMCPToolsDegradeOnError: raise RuntimeError("client exploded") monkeypatch.setattr(HaikuRAG, client_method, boom) - mcp = create_mcp_server(mcp_db, read_only=False) + mcp = create_mcp_server(mcp_db) tool = await _get_tool(mcp, tool_name) assert await tool(**kwargs) == expected @pytest.mark.asyncio async def test_list_documents_returns_empty_for_invalid_filter(self, mcp_db): - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) list_docs = await _get_tool(mcp, "list_documents") assert await list_docs(filter="no_such_column = 1") == [] @@ -486,7 +360,7 @@ class TestMCPToolsDegradeOnError: raise RuntimeError("sandbox exploded") monkeypatch.setattr(HaikuRAG, "analyze", boom) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) analyze = await _get_tool(mcp, "analyze") assert "sandbox exploded" in await analyze(question="q") @@ -509,7 +383,7 @@ class TestMCPToolsDegradeOnError: return ("the answer", [citation]) monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) ask = await _get_tool(mcp, "ask_question") with_cite = await ask(question="q", cite=True) @@ -534,7 +408,7 @@ class TestMCPClientLifetime: monkeypatch.setattr(Store, "_initialize", counted) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") list_docs = await _get_tool(mcp, "list_documents") await search(query="artificial intelligence") @@ -559,7 +433,7 @@ class TestMCPClientLifetime: monkeypatch.setattr(Store, "_initialize", counted) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) list_docs = await _get_tool(mcp, "list_documents") results = await asyncio.gather(*(list_docs() for _ in range(5))) @@ -567,18 +441,6 @@ class TestMCPClientLifetime: assert opens == 1 assert all(len(r) == 2 for r in results) - @pytest.mark.asyncio - async def test_a_write_is_visible_to_the_next_read(self, mcp_db): - """One connection sees its own writes, whatever the consistency interval.""" - mcp = create_mcp_server(mcp_db, read_only=False) - list_docs = await _get_tool(mcp, "list_documents") - delete_doc = await _get_tool(mcp, "delete_document") - - docs = await list_docs() - assert await delete_doc(document_id=docs[0].id) is True - - assert len(await list_docs()) == len(docs) - 1 - @pytest.mark.asyncio async def test_lifespan_opens_and_closes_once(self, mcp_db, monkeypatch): from haiku.rag.store.engine import Store @@ -593,7 +455,7 @@ class TestMCPClientLifetime: monkeypatch.setattr(Store, "_initialize", counted) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) # _lifespan_manager is what every transport enters; the public # lifespan() combines provider lifespans only. async with mcp._lifespan_manager(): @@ -624,7 +486,7 @@ class TestMCPClientLifetime: ) scope = DatabaseScope.resolve(config, database_name="alpha") - mcp = _mcp_covering(scope, config, read_only=True) + mcp = _mcp_covering(scope, config) async with mcp._lifespan_manager(): search = await _get_tool(mcp, "search_documents") results = await search(query="artificial intelligence") @@ -649,7 +511,7 @@ class TestMCPClientLifetime: ) with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"): - _mcp_covering(DatabaseScope.resolve(config), config, read_only=True) + _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.""" @@ -696,7 +558,7 @@ class TestMCPClientLifetime: async def run_stdio_async(self): return None - def fake_covering(scope, config, read_only): + def fake_covering(scope, config): seen.update(scope=scope, config=config) return _Server() @@ -715,7 +577,7 @@ class TestMCPClientLifetime: @pytest.mark.asyncio async def test_startup_fails_when_the_database_cannot_open(self, tmp_path): - mcp = create_mcp_server(tmp_path / "does-not-exist.lancedb", read_only=True) + mcp = create_mcp_server(tmp_path / "does-not-exist.lancedb") with pytest.raises(FileNotFoundError): async with mcp._lifespan_manager(): @@ -737,7 +599,7 @@ class TestMCPClientLifetime: monkeypatch.setattr(Store, "_initialize", counted) - mcp = create_mcp_server(mcp_db, read_only=True) + mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") async with mcp._lifespan_manager(): @@ -750,23 +612,13 @@ class TestMCPClientLifetime: assert len(results) > 0 @pytest.mark.asyncio - async def test_same_dim_drift_starts_read_only_but_not_writable(self, mcp_db): - """Validation is unchanged: same-dimension identity drift warns in - read-only mode and raises in writable mode. The MCP server no longer - opts out of it for deletion.""" + async def test_same_dim_drift_starts(self, mcp_db): + """Same-dimension identity drift warns on a read-only open and raises + on a writable one; the server starts, so it opened read-only.""" from haiku.rag.config import get_config - from haiku.rag.store.exceptions import ConfigMismatchError drifted = get_config().model_copy(deep=True) drifted.embeddings.model.name = "a-different-model" - async with create_mcp_server( - mcp_db, config=drifted, read_only=True - )._lifespan_manager(): + async with create_mcp_server(mcp_db, config=drifted)._lifespan_manager(): pass - - with pytest.raises(ConfigMismatchError): - async with create_mcp_server( - mcp_db, config=drifted, read_only=False - )._lifespan_manager(): - pass From 40d40bcbf2ca08c3841f3ac968400c65f0d599ee Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 10:38:48 +0300 Subject: [PATCH 03/18] 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 --- CHANGELOG.md | 6 + docs/cli.md | 2 +- docs/configuration/storage.md | 4 +- docs/mcp.md | 15 ++ haiku_rag_slim/haiku/rag/cli.py | 2 +- haiku_rag_slim/haiku/rag/mcp.py | 66 +++-- haiku_rag_slim/haiku/rag/tools/document.py | 1 + haiku_rag_slim/haiku/rag/utils.py | 6 +- tests/test_cli.py | 15 ++ tests/test_mcp.py | 269 +++++++++++++++------ tests/test_utils.py | 16 ++ 11 files changed, 311 insertions(+), 91 deletions(-) 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) --- From 88cd2b1d6cf6b08150a1012db1b4d80e6792c5d5 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 11:31:45 +0300 Subject: [PATCH 04/18] Describe the MCP server and its tools to clients FastMCP gets instructions, the haiku.rag-slim version and a lifespan. Every tool carries read-only ToolAnnotations with a title, a description that says when to use it, and a description on every parameter. The filter description lists the document columns from DocumentMetaRecord and how to match metadata; filter is accepted by both search tools. One strict base64 decoder serves every image parameter. DocumentInfo carries metadata. Refs #599 --- CHANGELOG.md | 6 + docs/mcp.md | 71 +++---- haiku_rag_slim/haiku/rag/mcp.py | 209 ++++++++++++++------- haiku_rag_slim/haiku/rag/tools/document.py | 1 + tests/test_mcp.py | 137 +++++++++++++- 5 files changed, 327 insertions(+), 97 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a90db87..bc98984f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## [Unreleased] +### Added + +- MCP server `instructions`, `version`, and read-only `ToolAnnotations` on + every tool; every parameter carries a description. `filter` on + `search_documents` and `search_documents_by_image`. `DocumentInfo.metadata`. + ### Changed - Default models are `ollama:qwen3.8`: `ModelConfig`, `qa.model`, diff --git a/docs/mcp.md b/docs/mcp.md index ca07ce4c..fd500d07 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -67,47 +67,54 @@ With a custom database path: After restarting Claude Desktop, you can ask Claude to search your documents or answer questions using your knowledge base. -## Available Tools +## Tools -### Documents +Every tool is read-only and says so in its annotations. Each parameter carries +a description in the tool schema, so the listing below names them without +repeating it. -- **`get_document`** - Retrieve a document by ID - - `document_id` (required): The document ID - - `source` (optional): The database holding it +| Tool | Registered | Parameters | +|---|---|---| +| `search_documents` | always | `query`, `limit`, `include_images`, `filter`, `sources` | +| `search_documents_by_image` | multimodal embedder only | `image_base64`, `limit`, `include_images`, `filter`, `sources` | +| `get_document` | always | `document_id`, `source` | +| `list_documents` | always | `limit`, `offset`, `filter` | +| `ask_question` | always | `question`, `cite`, `images_base64`, `sources` | +| `analyze` | always | `question`, `filter`, `images_base64`, `sources` | -- **`list_documents`** - List documents with pagination and filtering - - `limit` (optional): Maximum number to return - - `offset` (optional): Number to skip - - `filter` (optional): SQL WHERE clause for filtering +`search_documents` runs hybrid search, vector and full-text, and returns +results best first. Scores are not comparable across queries or search types. +Rank is the signal. `include_images` attaches picture bytes as base64 PNG under +`image_data`. `search_documents_by_image` embeds the query image and searches +by vector similarity alone. -### Search +`get_document` returns a document whole, in reading order. `list_documents` +returns titles, URIs and metadata, which is how a client learns what a filter +can match. -- **`search_documents`** - Search using hybrid search (vector + full-text) - - `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 +`ask_question` runs the RAG agent on the server and returns an answer, with +citations when `cite` is set. `analyze` writes and runs Python in a sandbox +over the documents, for counting, aggregation and computation across +documents. Both cost a model call. -- **`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 +### Filters -### Question Answering +`filter` is a SQL WHERE clause over the document columns `id`, `uri`, `title`, +`metadata`, `created_at`, `updated_at`. `metadata` is a JSON string, so match +its keys with LIKE: -- **`ask_question`** - Ask questions about your documents - - `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 +```sql +metadata LIKE '%"author": "Smith"%' +uri LIKE '%.pdf' +title = 'Q3 report' +``` -- **`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 +### Instructions + +The server publishes `instructions` describing the knowledge base: what it +holds, when to reach for it, the collection names when it covers several, and +`prompts.domain_preamble` when set. Claude Code shows them to the model. Claude +Desktop does not, so every tool description stands on its own. ## Continuous ingestion diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 44a4f259..e94b7624 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,29 +1,79 @@ import asyncio from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager +from importlib import metadata from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Annotated from fastmcp import FastMCP from fastmcp.exceptions import ToolError +from mcp.types import ToolAnnotations +from pydantic import Field 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.store.schema import DocumentMetaRecord from haiku.rag.tools.document import DocumentInfo from haiku.rag.utils import format_citations if TYPE_CHECKING: from haiku.rag.client.scope import DatabaseScope +_FILTER_COLUMNS = ", ".join(DocumentMetaRecord.model_fields) + +Filter = Annotated[ + str | None, + Field( + description=( + f"SQL WHERE clause over the document columns {_FILTER_COLUMNS}, " + "restricting which documents are used. `metadata` is a JSON string, " + 'so match its keys with LIKE: metadata LIKE \'%"author": "Smith"%\'. ' + "Also uri LIKE '%.pdf', title = 'Q3 report'." + ) + ), +] +Sources = Annotated[ + list[str] | None, + Field(description="Collections to use, by name. All of them by default."), +] + + +def _read_only(title: str) -> ToolAnnotations: + return ToolAnnotations(title=title, readOnlyHint=True, openWorldHint=False) + + +def _decode_image(image_base64: str) -> bytes: + import base64 + + return base64.b64decode(image_base64, validate=True) + def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: if not images_base64: return None - import base64 + return [_decode_image(b64) for b64 in images_base64] - return [base64.b64decode(b64, validate=True) for b64 in images_base64] + +def _instructions(scope: "DatabaseScope", config: AppConfig) -> 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, 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 scope.covers_multiple: + lines.append( + f"It holds several collections: {', '.join(scope.names)}. Results " + "and citations name theirs in `source`; pass `sources` to use a subset." + ) + if config.prompts.domain_preamble: + lines.append(config.prompts.domain_preamble) + return "\n".join(lines) def create_mcp_server( @@ -86,27 +136,45 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: finally: client = None - mcp = FastMCP("haiku-rag", lifespan=lifespan) + mcp = FastMCP( + "haiku-rag", + instructions=_instructions(scope, config), + version=metadata.version("haiku.rag-slim"), + lifespan=lifespan, + ) - @mcp.tool() + @mcp.tool(annotations=_read_only("Search documents")) async def search_documents( query: str, limit: int | None = None, include_images: bool = True, - sources: list[str] | None = None, + filter: Filter = None, + sources: Sources = None, ) -> list[SearchResult]: - """Search the RAG system for documents using hybrid search (vector similarity + full-text search). + """Search the knowledge base by meaning and keyword. - When include_images is True (default) and a picture-labeled chunk is - 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. + Use this first for any question the documents might answer; it needs + no model and is the cheapest call. Results come best first, each with + the document's id, title and collection, the section headings and the + matching passage. Scores are not comparable across queries, so read + the order, not the numbers. If nothing relevant comes back, rephrase + once or narrow with `filter` before concluding the material is absent. + + Args: + query: What to look for, in natural language or keywords. + limit: How many results to return; the server's configured default + when omitted. + include_images: Attach the bytes of pictures in the results as + base64 PNG under `image_data`. False for a smaller response. """ try: rag = await _client() return await rag.search( - query, limit=limit, include_images=include_images, sources=sources + query, + limit=limit, + filter=filter, + include_images=include_images, + sources=sources, ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e @@ -121,45 +189,57 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: if get_embedder(config).supports_images: - @mcp.tool() + @mcp.tool(annotations=_read_only("Search documents by image")) async def search_documents_by_image( image_base64: str, limit: int | None = None, include_images: bool = True, - sources: list[str] | None = None, + filter: Filter = None, + sources: Sources = None, ) -> list[SearchResult]: - """Search the RAG system using an image as the query. + """Search the knowledge base with an image as the query. - ``image_base64`` is a base64-encoded image (PNG/JPEG bytes). The - 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. + Use this when the question is about a picture rather than words. + The image is embedded and matched against document text and + figures by vector similarity alone. Results have the shape of + `search_documents` results. + + Args: + image_base64: The query image, PNG or JPEG bytes as base64. + limit: How many results to return; the server's configured + default when omitted. + include_images: Attach the bytes of pictures in the results as + base64 PNG under `image_data`. False for a smaller response. """ - import base64 - - try: - raw = base64.b64decode(image_base64) - except Exception: - return [] try: + raw = _decode_image(image_base64) rag = await _client() return await rag.search( - raw, limit=limit, include_images=include_images, sources=sources + raw, + limit=limit, + filter=filter, + include_images=include_images, + sources=sources, ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e except Exception: return [] - @mcp.tool() + @mcp.tool(annotations=_read_only("Get document")) async def get_document( document_id: str, source: str | None = None ) -> Document | None: - """Get a document by its ID. + """Read one document whole, in reading order. - ``source`` names the database holding it; without one every database - is asked. + Use this after a search when a passage is not enough. Returns the + document's content, title, uri and metadata. Ids come from search + results and `list_documents`. + + Args: + document_id: The document's id. + source: The collection holding it. Without one every collection + is asked. """ try: rag = await _client() @@ -169,18 +249,21 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: except Exception: return None - @mcp.tool() + @mcp.tool(annotations=_read_only("List documents")) async def list_documents( limit: int | None = None, offset: int | None = None, - filter: str | None = None, + filter: Filter = None, ) -> list[DocumentInfo]: - """List all documents with optional pagination and filtering. + """List what the knowledge base holds. + + Use this to see which documents exist, their titles, URIs and + metadata, and so what a `filter` can match. Not a search: it returns + no passages. Args: - limit: Maximum number of documents to return. - offset: Number of documents to skip. - filter: Optional SQL WHERE clause to filter documents. + limit: How many documents to return. + offset: How many documents to skip, for paging. """ try: rag = await _client() @@ -193,30 +276,32 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: uri=doc.uri or "", created=doc.created_at.strftime("%Y-%m-%d"), source=doc.source, + metadata=doc.metadata, ) for doc in documents ] except Exception: return [] - @mcp.tool() + @mcp.tool(annotations=_read_only("Ask a question")) async def ask_question( question: str, cite: bool = False, images_base64: list[str] | None = None, - sources: list[str] | None = None, + sources: Sources = None, ) -> str: - """Ask a question using the QA agent. + """Answer a question from the documents with a retrieval agent. + + 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 when + `cite` is set. Args: - question: The question to ask. - 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. + question: The question, in natural language. + cite: Append citations to the answer. + images_base64: Images to attach to the question, PNG or JPEG + bytes as base64. Needs a vision-capable model on the server. """ try: images = _decode_images(images_base64) @@ -232,28 +317,24 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: except Exception as e: return f"Error answering question: {e!s}" - @mcp.tool() + @mcp.tool(annotations=_read_only("Analyze documents")) async def analyze( question: str, - filter: str | None = None, + filter: Filter = None, images_base64: list[str] | None = None, - sources: list[str] | None = None, + sources: Sources = None, ) -> str: - """Answer complex questions using the analysis capability. + """Compute an answer across documents with code. - Use this for questions requiring computation, aggregation, or - structural traversal across documents. The capability can write and - execute Python code in a sandboxed interpreter. + 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 to answer. - 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. + 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. """ try: images = _decode_images(images_base64) diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index 1458083a..960d5194 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -28,6 +28,7 @@ class DocumentInfo(BaseModel): uri: str created: str source: str | None = None + metadata: dict = {} class DocumentListResponse(BaseModel): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 06bc32d5..82799789 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -56,6 +56,7 @@ async def mcp_db(temp_db_path): "Artificial intelligence is transforming industries worldwide.", title="AI Overview", uri="test://ai-overview", + metadata={"author": "Ada"}, ) await rag.create_document( "Machine learning is a subset of artificial intelligence.", @@ -104,6 +105,21 @@ class TestMCPReadTools: results = await search(query="artificial intelligence", limit=1) assert len(results) == 1 + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_search_documents_with_filter(self, mcp_db): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + result = await client.call_tool( + "search_documents", + {"query": "artificial intelligence", "filter": "title = 'ML Basics'"}, + ) + + results = result.structured_content["result"] + assert results + assert {r["document_title"] for r in results} == {"ML Basics"} + @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") async def test_search_documents_preserves_chunk_meta_through_serialization( @@ -202,6 +218,103 @@ class TestMCPReadTools: assert len(results) == 1 assert results[0].title == "AI Overview" + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_list_documents_carries_metadata(self, mcp_db): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + result = await client.call_tool("list_documents", {}) + + [overview] = [ + d + for d in result.structured_content["result"] + if d["title"] == "AI Overview" + ] + assert overview["metadata"] == {"author": "Ada"} + + +@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") +class TestMCPDescribesItself: + """What a client learns from initialize and list_tools, over the wire.""" + + @pytest.mark.asyncio + async def test_instructions_and_version_are_set(self, mcp_db): + from importlib import metadata + + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + init = client.initialize_result + + assert init.instructions + assert init.serverInfo.version == metadata.version("haiku.rag-slim") + + @pytest.mark.asyncio + async def test_instructions_name_the_collections_when_covering_several( + self, two_dbs + ): + from fastmcp import Client + + from haiku.rag.client.scope import DatabaseScope + + async with Client(_covering_all(two_dbs)) as client: + covering_both = client.initialize_result.instructions + one = DatabaseScope.resolve(two_dbs, database_name="alpha") + async with Client(_mcp_covering(one, two_dbs)) as client: + covering_one = client.initialize_result.instructions + + assert "alpha" in covering_both + assert "beta" in covering_both + assert "beta" not in covering_one + + @pytest.mark.asyncio + async def test_instructions_carry_the_domain_preamble(self, mcp_db): + from fastmcp import Client + + from haiku.rag.config import get_config + + config = get_config().model_copy(deep=True) + config.prompts.domain_preamble = "Everything here is about zebras." + + async with Client(create_mcp_server(mcp_db, config=config)) as client: + with_preamble = client.initialize_result.instructions + async with Client(create_mcp_server(mcp_db)) as client: + without = client.initialize_result.instructions + + assert "Everything here is about zebras." in with_preamble + assert "zebras" not in without + + @pytest.mark.asyncio + async def test_every_tool_is_annotated_read_only(self, mcp_db, multimodal_embedder): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + tools = await client.list_tools() + + assert len(tools) == 6 + for tool in tools: + assert tool.annotations is not None, tool.name + assert tool.annotations.readOnlyHint is True, tool.name + assert tool.annotations.openWorldHint is False, tool.name + assert tool.annotations.title, tool.name + + @pytest.mark.asyncio + async def test_every_parameter_is_described(self, mcp_db, multimodal_embedder): + from fastmcp import Client + + async with Client(create_mcp_server(mcp_db)) as client: + tools = await client.list_tools() + + undescribed = [ + f"{tool.name}.{name}" + for tool in tools + for name, schema in tool.inputSchema.get("properties", {}).items() + if not schema.get("description") + ] + assert len(tools) == 6 + assert undescribed == [] + class TestMCPToolSet: @pytest.mark.asyncio @@ -380,13 +493,35 @@ class TestMCPImageQuery: png = b"\x89PNG\r\n\x1a\n" results = await search_by_image( - image_base64=base64.b64encode(png).decode("ascii"), sources=["alpha"] + image_base64=base64.b64encode(png).decode("ascii"), + filter="uri LIKE 'x%'", + sources=["alpha"], ) assert results == [] assert seen["query"] == png + assert seen["filter"] == "uri LIKE 'x%'" assert seen["sources"] == ["alpha"] + @pytest.mark.asyncio + async def test_image_query_rejects_characters_outside_the_alphabet( + self, mcp_db, multimodal_embedder, monkeypatch + ): + """A lenient decoder would drop the stray characters and search.""" + searched = False + + async def fake_search(self, query, **kwargs): + nonlocal searched + searched = True + return [] + + monkeypatch.setattr(HaikuRAG, "search", fake_search) + 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!!!!") == [] + assert not searched + @pytest.mark.asyncio async def test_image_query_returns_empty_on_invalid_base64( self, mcp_db, multimodal_embedder From 2a6d72171d1d8fae156e42695f8ce383e78e5bd5 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 12:03:06 +0300 Subject: [PATCH 05/18] Make MCP failures errors on the wire FastMCP masks unexpected exceptions and logs the traceback server-side, so paths and provider URLs never cross the transport. Expected failures raise ToolError with a message: unknown document, unknown collection, invalid filter, invalid base64, and agent failures naming only the exception type. No tool returns an empty value or an error string on failure any more. A filter is validated on its own before the read that would use it, with a filtered count on one of the selected databases: that is the query engine rejecting the filter and nothing else, so its message (columns and the statement) can be forwarded, while a ValueError raised later in the read stays masked and no database outside the selection is opened. Refs #599 --- CHANGELOG.md | 5 + docs/mcp.md | 9 + haiku_rag_slim/haiku/rag/mcp.py | 117 ++++++++----- tests/test_mcp.py | 294 ++++++++++++++++++++------------ 4 files changed, 271 insertions(+), 154 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bc98984f..7940d86c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,11 @@ - `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`. +- MCP tools raise on failure; an empty result no longer doubles as an error. + Unknown document, unknown collection, invalid filter and invalid base64 + carry a message; `ask_question` and `analyze` failures name the exception + type. Anything else is masked (`mask_error_details=True`) and logged + server-side. - `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. diff --git a/docs/mcp.md b/docs/mcp.md index fd500d07..d63f5880 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -109,6 +109,15 @@ uri LIKE '%.pdf' title = 'Q3 report' ``` +### Errors + +A failure is an MCP error, never an empty result. Expected failures carry a +message: a document id that matches nothing, a collection the server does not +cover, a filter the query engine rejects (with its message), invalid base64, +and an `ask_question` or `analyze` failure naming only the exception type. +Anything else reaches the client as `Error calling tool 'name'` and its +traceback goes to the server log. + ### Instructions The server publishes `instructions` describing the knowledge base: what it diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index e94b7624..5c285952 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,4 +1,5 @@ import asyncio +import logging from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from importlib import metadata @@ -21,6 +22,8 @@ from haiku.rag.utils import format_citations if TYPE_CHECKING: from haiku.rag.client.scope import DatabaseScope +logger = logging.getLogger(__name__) + _FILTER_COLUMNS = ", ".join(DocumentMetaRecord.model_fields) Filter = Annotated[ @@ -47,7 +50,12 @@ def _read_only(title: str) -> ToolAnnotations: def _decode_image(image_base64: str) -> bytes: import base64 - return base64.b64decode(image_base64, validate=True) + try: + return base64.b64decode(image_base64, validate=True) + except ValueError as e: + # binascii.Error for characters outside the alphabet or bad padding, + # ValueError itself for non-ASCII input. + raise ToolError("Invalid base64 image") from e def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: @@ -56,6 +64,28 @@ def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: return [_decode_image(b64) for b64 in images_base64] +async def _check_filter( + rag: HaikuRAG, filter: str | None, sources: list[str] | None = None +) -> None: + """Evaluate a filter on its own before the read that would use it. + + A filtered count on one selected database runs the same predicate on the + same table and nothing else, so a ValueError here is the query engine + rejecting the filter; its message names columns and the statement, never + a location. A ValueError raised later in the read stays masked. Only the + selection is touched: every database shares the schema, so one suffices. + """ + if filter is None: + return + selected = await rag.clients_covering(sources) + if not selected: + return + try: + await selected[0].count_documents(filter=filter) + except ValueError as e: + raise ToolError(f"Invalid filter {filter!r}: {e}") from e + + def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: """What the server is for, naming no tools: the client has every tool's description from the listing.""" @@ -136,11 +166,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: finally: client = None + # Masking keeps paths and provider URLs out of an unexpected error's text; + # the traceback goes to the server log. A ToolError reaches the client as is. mcp = FastMCP( "haiku-rag", instructions=_instructions(scope, config), version=metadata.version("haiku.rag-slim"), lifespan=lifespan, + mask_error_details=True, ) @mcp.tool(annotations=_read_only("Search documents")) @@ -167,8 +200,9 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: include_images: Attach the bytes of pictures in the results as base64 PNG under `image_data`. False for a smaller response. """ + rag = await _client() try: - rag = await _client() + await _check_filter(rag, filter, sources) return await rag.search( query, limit=limit, @@ -178,8 +212,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - except Exception: - return [] # Image-as-query tool, only registered when the configured embedder # supports image embeddings. Probed at server-build time when no Store is @@ -211,9 +243,10 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: include_images: Attach the bytes of pictures in the results as base64 PNG under `image_data`. False for a smaller response. """ + raw = _decode_image(image_base64) + rag = await _client() try: - raw = _decode_image(image_base64) - rag = await _client() + await _check_filter(rag, filter, sources) return await rag.search( raw, limit=limit, @@ -223,13 +256,9 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - except Exception: - return [] @mcp.tool(annotations=_read_only("Get document")) - async def get_document( - document_id: str, source: str | None = None - ) -> Document | None: + async def get_document(document_id: str, source: str | None = None) -> Document: """Read one document whole, in reading order. Use this after a search when a passage is not enough. Returns the @@ -241,13 +270,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: source: The collection holding it. Without one every collection is asked. """ + rag = await _client() try: - rag = await _client() - return await rag.get_document_by_id(document_id, source) + document = await rag.get_document_by_id(document_id, source) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - except Exception: - return None + if document is None: + raise ToolError(f"No document with id {document_id!r}") + return document @mcp.tool(annotations=_read_only("List documents")) async def list_documents( @@ -265,23 +295,20 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: limit: How many documents to return. offset: How many documents to skip, for paging. """ - try: - rag = await _client() - documents = await rag.list_documents(limit, offset, filter) - - return [ - DocumentInfo( - id=doc.id, - title=doc.title or "Untitled", - uri=doc.uri or "", - created=doc.created_at.strftime("%Y-%m-%d"), - source=doc.source, - metadata=doc.metadata, - ) - for doc in documents - ] - except Exception: - return [] + rag = await _client() + await _check_filter(rag, filter) + documents = await rag.list_documents(limit, offset, filter) + return [ + DocumentInfo( + id=doc.id, + title=doc.title or "Untitled", + uri=doc.uri or "", + created=doc.created_at.strftime("%Y-%m-%d"), + source=doc.source, + metadata=doc.metadata, + ) + for doc in documents + ] @mcp.tool(annotations=_read_only("Ask a question")) async def ask_question( @@ -303,19 +330,20 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: 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: - images = _decode_images(images_base64) - rag = await _client() answer, citations = await rag.ask(question, images=images, sources=sources) - if cite and 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}" + logger.exception("ask_question failed") + raise ToolError(f"ask_question failed: {type(e).__name__}") from e + if cite and citations: + answer += "\n\n" + format_citations( + citations, include_source=rag.covers_multiple + ) + return answer @mcp.tool(annotations=_read_only("Analyze documents")) async def analyze( @@ -336,16 +364,17 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: 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: - images = _decode_images(images_base64) - rag = await _client() 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}" + logger.exception("analyze failed") + raise ToolError(f"analyze failed: {type(e).__name__}") from e + return result.answer return mcp diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 82799789..b21b24cd 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,6 +1,8 @@ +import logging from types import SimpleNamespace import pytest +from fastmcp.exceptions import ToolError from haiku.rag.client import HaikuRAG from haiku.rag.mcp import _covering as _mcp_covering @@ -87,6 +89,14 @@ async def _get_tool(mcp, name): return tool.fn +async def _call(mcp, name, **kwargs): + """Call a tool over the wire, returning the result whether or not it errored.""" + from fastmcp import Client + + async with Client(mcp) as client: + return await client.call_tool(name, kwargs, raise_on_error=False) + + class TestMCPReadTools: @pytest.mark.asyncio async def test_search_documents(self, mcp_db): @@ -184,14 +194,6 @@ class TestMCPReadTools: assert "docling_document" not in serialized assert "docling_version" not in serialized - @pytest.mark.asyncio - async def test_get_document_not_found(self, mcp_db): - mcp = create_mcp_server(mcp_db) - get_doc = await _get_tool(mcp, "get_document") - - result = await get_doc(document_id="nonexistent-id") - assert result is None - @pytest.mark.asyncio async def test_list_documents(self, mcp_db): mcp = create_mcp_server(mcp_db) @@ -233,6 +235,36 @@ class TestMCPReadTools: ] assert overview["metadata"] == {"author": "Ada"} + @pytest.mark.asyncio + async def test_ask_question_appends_citations_when_requested( + self, mcp_db, monkeypatch + ): + from haiku.rag.store.models.citation import Citation + + citation = Citation( + chunk_id="c1", + document_id="d1", + content="cited text", + document_uri="test://ai-overview", + document_title="AI Overview", + source="alpha", + ) + + async def fake_ask(self, question, filter=None, images=None, sources=None): + return ("the answer", [citation]) + + monkeypatch.setattr(HaikuRAG, "ask", fake_ask) + mcp = create_mcp_server(mcp_db) + ask = await _get_tool(mcp, "ask_question") + + 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" + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") class TestMCPDescribesItself: @@ -367,14 +399,28 @@ class TestMCPCoversTheConfiguredSet: 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_a_filtered_search_touches_only_the_selected_databases(self, two_dbs): + """alpha is gone; a filtered search selecting beta must not notice.""" + import shutil + + shutil.rmtree(two_dbs.lancedb.databases["alpha"]) + mcp = _covering_all(two_dbs) + search = await _get_tool(mcp, "search_documents") + + results = await search( + query="cats", filter="uri LIKE '%beta%'", sources=["beta"] + ) + assert results + assert {r.source for r in results} == {"beta"} + assert await search(query="cats", filter="uri LIKE '%beta%'", sources=[]) == [] + @pytest.mark.asyncio async def test_the_listing_covers_every_database(self, two_dbs): mcp = _covering_all(two_dbs) @@ -495,13 +541,13 @@ class TestMCPImageQuery: results = await search_by_image( image_base64=base64.b64encode(png).decode("ascii"), filter="uri LIKE 'x%'", - sources=["alpha"], + sources=[], ) assert results == [] assert seen["query"] == png assert seen["filter"] == "uri LIKE 'x%'" - assert seen["sources"] == ["alpha"] + assert seen["sources"] == [] @pytest.mark.asyncio async def test_image_query_rejects_characters_outside_the_alphabet( @@ -519,36 +565,10 @@ class TestMCPImageQuery: 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!!!!") == [] + with pytest.raises(ToolError): + await search_by_image(image_base64="AAAA!!!!") assert not searched - @pytest.mark.asyncio - async def test_image_query_returns_empty_on_invalid_base64( - self, mcp_db, multimodal_embedder - ): - """Garbage base64 from the caller is swallowed, returning an empty - list rather than crashing the MCP server.""" - mcp = create_mcp_server(mcp_db) - search_by_image = await _get_tool(mcp, "search_documents_by_image") - - # Not valid base64 (contains non-base64 chars) — the strict decoder - # in search_documents_by_image rejects it. - 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 @@ -590,14 +610,6 @@ class TestMCPImageInput: assert result == "answer" assert captured["images"] == [jpeg] - @pytest.mark.asyncio - async def test_ask_question_rejects_invalid_base64(self, mcp_db): - mcp = create_mcp_server(mcp_db) - ask = await _get_tool(mcp, "ask_question") - - result = await ask(question="q", images_base64=["!!! not base64 !!!"]) - assert "Error" in result - @pytest.mark.asyncio async def test_ask_question_without_images_passes_none(self, mcp_db, monkeypatch): captured = {} @@ -615,78 +627,140 @@ class TestMCPImageInput: assert captured["images"] is None -class TestMCPToolsDegradeOnError: - """Every tool swallows client failures and returns its empty value rather - than propagating an exception to the MCP transport.""" +@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") +class TestMCPErrorContract: + """A failure is an error on the wire, never an empty result. Expected + failures say what went wrong; anything else is masked and logged on the + server.""" + + @pytest.mark.asyncio + async def test_an_unknown_document_is_an_error(self, mcp_db): + result = await _call( + create_mcp_server(mcp_db), "get_document", document_id="nonexistent-id" + ) + + assert result.is_error + assert "nonexistent-id" in result.content[0].text @pytest.mark.asyncio @pytest.mark.parametrize( - "client_method,tool_name,kwargs,expected", - [ - ("search", "search_documents", {"query": "x"}, []), - ("get_document_by_id", "get_document", {"document_id": "x"}, None), - ("list_documents", "list_documents", {}, []), - ], + "tool_name,kwargs", + [("search_documents", {"query": "x"}), ("list_documents", {})], ) - async def test_tool_returns_empty_value_when_client_raises( - self, mcp_db, monkeypatch, client_method, tool_name, kwargs, expected + async def test_an_invalid_filter_is_an_error_naming_the_filter( + self, mcp_db, tool_name, kwargs ): - async def boom(self, *args, **kw): - raise RuntimeError("client exploded") - - monkeypatch.setattr(HaikuRAG, client_method, boom) - mcp = create_mcp_server(mcp_db) - tool = await _get_tool(mcp, tool_name) - - assert await tool(**kwargs) == expected - - @pytest.mark.asyncio - async def test_list_documents_returns_empty_for_invalid_filter(self, mcp_db): - mcp = create_mcp_server(mcp_db) - list_docs = await _get_tool(mcp, "list_documents") - - assert await list_docs(filter="no_such_column = 1") == [] - - @pytest.mark.asyncio - async def test_analyze_reports_the_error(self, mcp_db, monkeypatch): - async def boom(self, question, filter=None, images=None, sources=None): - raise RuntimeError("sandbox exploded") - - monkeypatch.setattr(HaikuRAG, "analyze", boom) - mcp = create_mcp_server(mcp_db) - analyze = await _get_tool(mcp, "analyze") - - assert "sandbox exploded" in await analyze(question="q") - - @pytest.mark.asyncio - async def test_ask_question_appends_citations_when_requested( - self, mcp_db, monkeypatch - ): - from haiku.rag.store.models.citation import Citation - - citation = Citation( - chunk_id="c1", - document_id="d1", - content="cited text", - document_uri="test://ai-overview", - document_title="AI Overview", - source="alpha", + result = await _call( + create_mcp_server(mcp_db), tool_name, filter="no_such_column = 1", **kwargs ) - async def fake_ask(self, question, filter=None, images=None, sources=None): - return ("the answer", [citation]) + assert result.is_error + assert "no_such_column = 1" in result.content[0].text - monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db) - ask = await _get_tool(mcp, "ask_question") + @pytest.mark.asyncio + @pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"]) + async def test_a_value_error_from_the_read_is_not_an_invalid_filter( + self, mcp_db, monkeypatch, filter + ): + """Only the filter check translates ValueError; one raised by the read + itself, with or without a valid filter, stays masked.""" - 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 + async def boom(self, *args, **kw): + raise ValueError("boom at /secret/path") - assert await ask(question="q", cite=False) == "the answer" + monkeypatch.setattr(HaikuRAG, "search", boom) + result = await _call( + create_mcp_server(mcp_db), "search_documents", query="x", filter=filter + ) + + assert result.is_error + assert "filter" not in result.content[0].text + assert "/secret/path" not in result.content[0].text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "payload", ["!!! not base64 !!!", "é"], ids=["outside_alphabet", "non_ascii"] + ) + @pytest.mark.parametrize( + "tool_name,image_param,many", + [ + ("search_documents_by_image", "image_base64", False), + ("ask_question", "images_base64", True), + ("analyze", "images_base64", True), + ], + ) + async def test_invalid_base64_is_an_error( + self, mcp_db, multimodal_embedder, tool_name, image_param, many, payload + ): + kwargs: dict[str, object] = {"question": "q"} if many else {} + kwargs[image_param] = [payload] if many else payload + + result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs) + + assert result.is_error + assert "base64" in result.content[0].text + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_method,tool_name", + [("ask", "ask_question"), ("analyze", "analyze")], + ) + async def test_an_agent_failure_names_only_its_type( + self, mcp_db, monkeypatch, caplog, client_method, tool_name + ): + async def boom(self, question, filter=None, images=None, sources=None): + raise RuntimeError("boom at /secret/path") + + monkeypatch.setattr(HaikuRAG, client_method, boom) + with caplog.at_level(logging.ERROR, logger="haiku.rag.mcp"): + result = await _call(create_mcp_server(mcp_db), tool_name, question="q") + + assert result.is_error + assert "RuntimeError" in result.content[0].text + assert "/secret/path" not in result.content[0].text + assert any( + r.exc_info and "boom at /secret/path" in str(r.exc_info[1]) + for r in caplog.records + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_method,tool_name,kwargs", + [ + ("search", "search_documents", {"query": "x"}), + ("search", "search_documents_by_image", {"image_base64": "AAAA"}), + ("get_document_by_id", "get_document", {"document_id": "x"}), + ("list_documents", "list_documents", {}), + ], + ) + async def test_an_unexpected_failure_is_masked_and_logged( + self, + mcp_db, + multimodal_embedder, + monkeypatch, + caplog, + client_method, + tool_name, + kwargs, + ): + async def boom(self, *args, **kw): + raise RuntimeError("boom at /secret/path") + + monkeypatch.setattr(HaikuRAG, client_method, boom) + # fastmcp's logger does not propagate, so listen to it directly. + fastmcp_logger = logging.getLogger("fastmcp") + fastmcp_logger.addHandler(caplog.handler) + try: + result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs) + finally: + fastmcp_logger.removeHandler(caplog.handler) + + assert result.is_error + assert "/secret/path" not in result.content[0].text + assert any( + r.exc_info and "boom at /secret/path" in str(r.exc_info[1]) + for r in caplog.records + ) class TestMCPClientLifetime: From 15afb97a6e2a73f44213eb481342e196269be813 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 12:16:00 +0300 Subject: [PATCH 06/18] Navigate documents by outline and section from the MCP server build_toc moves from the sandbox into haiku.rag.context; the sandbox keeps its toc.json unchanged. get_document_outline returns the heading tree with page numbers and get_document_section one section's text, subsections included, both resolved in the database holding the document. Chunk ids never leave the server. ask_question drops `cite` and always appends its citations. Refs #599 --- CHANGELOG.md | 4 + docs/mcp.md | 22 ++- haiku_rag_slim/haiku/rag/context.py | 77 ++++++++ haiku_rag_slim/haiku/rag/mcp.py | 103 ++++++++++- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 78 +------- haiku_rag_slim/haiku/rag/tools/document.py | 19 ++ tests/test_mcp.py | 191 ++++++++++++++++++-- 7 files changed, 391 insertions(+), 103 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7940d86c..9f6fb5c6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ ### Added +- 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`. - MCP server `instructions`, `version`, and read-only `ToolAnnotations` on every tool; every parameter carries a description. `filter` on `search_documents` and `search_documents_by_image`. `DocumentInfo.metadata`. @@ -35,6 +38,7 @@ ### Removed +- `cite` on the MCP `ask_question` tool; citations are always appended. - MCP write tools `add_document_from_file`, `add_document_from_url`, `add_document_from_text` and `delete_document`. The server opens the database read-only; ingest with `haiku-rag add`, `add-src`, `delete` or diff --git a/docs/mcp.md b/docs/mcp.md index d63f5880..c4815434 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -78,8 +78,10 @@ repeating it. | `search_documents` | always | `query`, `limit`, `include_images`, `filter`, `sources` | | `search_documents_by_image` | multimodal embedder only | `image_base64`, `limit`, `include_images`, `filter`, `sources` | | `get_document` | always | `document_id`, `source` | +| `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`, `cite`, `images_base64`, `sources` | +| `ask_question` | always | `question`, `images_base64`, `sources` | | `analyze` | always | `question`, `filter`, `images_base64`, `sources` | `search_documents` runs hybrid search, vector and full-text, and returns @@ -88,12 +90,15 @@ Rank is the signal. `include_images` attaches picture bytes as base64 PNG under `image_data`. `search_documents_by_image` embeds the query image and searches by vector similarity alone. -`get_document` returns a document whole, in reading order. `list_documents` -returns titles, URIs and metadata, which is how a client learns what a filter -can match. +`get_document` returns a document whole, in reading order. For a long one, +`get_document_outline` returns the heading tree with page numbers and +`get_document_section` the text of one section, subsections included; a +node's `id` in the outline is the `section_id`. A document without headings +has an empty outline. `list_documents` returns titles, URIs and metadata, +which is how a client learns what a filter can match. -`ask_question` runs the RAG agent on the server and returns an answer, with -citations when `cite` is set. `analyze` writes and runs Python in a sandbox +`ask_question` runs the RAG agent on the server and returns an answer +followed by its citations. `analyze` writes and runs Python in a sandbox over the documents, for counting, aggregation and computation across documents. Both cost a model call. @@ -112,8 +117,9 @@ title = 'Q3 report' ### Errors A failure is an MCP error, never an empty result. Expected failures carry a -message: a document id that matches nothing, a collection the server does not -cover, a filter the query engine rejects (with its message), invalid base64, +message: a document or section id that matches nothing, a collection the +server does not cover, a filter the query engine rejects (with its message), +invalid base64, and an `ask_question` or `analyze` failure naming only the exception type. Anything else reaches the client as `Error calling tool 'name'` and its traceback goes to the server log. diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index b4214da5..17fb44c5 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -32,6 +32,8 @@ In both cases: - Results without doc_item_refs pass through unexpanded """ +from typing import Any + from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.document_item import DocumentItem @@ -488,3 +490,78 @@ def expand_with_items( final_results.append(built) return final_results + passthrough + + +def build_toc( + items: list["DocumentItem"], + chunk_index: dict[str, list[str]], +) -> list[dict[str, Any]]: + """Build a nested section tree from items in position order. + + Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting + follows the explicit levels: a header pops the stack until the top is at + a strictly shallower level, then becomes a child of that top (or a root). + + ``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the + position of the next header whose level is the same or shallower (i.e. + the next sibling or ancestor that ends this section), or the total item + count if no such header exists. + + ``chunk_ids`` aggregates the chunks covered by all items in the section's + ``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to + ground a section-scoped answer without a corpus-wide ``search()`` call. + + Items without a section_header label (or with ``heading_level == 0``) are + skipped. When all section_headers carry the same level the output is a + flat sibling list (see docling-project/docling#2121 for an upstream case + where every PDF section_header is emitted at level=1). + """ + # Defensive: every consumer is supposed to pass items in position order, + # but the end_exclusive lookahead below silently miscomputes section + # boundaries if it's not — better to sort once than trust the caller. + items = sorted(items, key=lambda i: i.position) + headers: list[DocumentItem] = [ + i for i in items if i.label == "section_header" and i.heading_level > 0 + ] + if not headers: + return [] + + total = max((i.position for i in items), default=-1) + 1 + items_by_position: dict[int, DocumentItem] = {i.position: i for i in items} + + ends: list[int] = [] + for idx, h in enumerate(headers): + end = total + for j in range(idx + 1, len(headers)): + if headers[j].heading_level <= h.heading_level: + end = headers[j].position + break + ends.append(end) + + roots: list[dict[str, Any]] = [] + stack: list[tuple[int, dict[str, Any]]] = [] + for h, end in zip(headers, ends, strict=True): + seen: set[str] = set() + chunk_ids: list[str] = [] + for pos in range(h.position, end): + item = items_by_position.get(pos) + if item is None: + continue + for cid in chunk_index.get(item.self_ref, []): + if cid not in seen: + seen.add(cid) + chunk_ids.append(cid) + node: dict[str, Any] = { + "self_ref": h.self_ref, + "level": h.heading_level, + "title": h.text, + "page_numbers": list(h.page_numbers), + "item_range": [h.position, end], + "chunk_ids": chunk_ids, + "children": [], + } + while stack and stack[-1][0] >= h.heading_level: + stack.pop() + (stack[-1][1]["children"] if stack else roots).append(node) + stack.append((h.heading_level, node)) + return roots diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 5c285952..ec97476f 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -13,13 +13,17 @@ from pydantic import Field from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, get_config +from haiku.rag.context import build_toc from haiku.rag.store.exceptions import UnknownDatabaseError from haiku.rag.store.models import Document, SearchResult +from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.schema import DocumentMetaRecord -from haiku.rag.tools.document import DocumentInfo +from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode from haiku.rag.utils import format_citations if TYPE_CHECKING: + from typing import Any + from haiku.rag.client.scope import DatabaseScope logger = logging.getLogger(__name__) @@ -91,8 +95,8 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: description from the listing.""" lines = [ "haiku-rag is the user's knowledge base: documents they ingested, " - "searchable by meaning and keyword, readable whole, answered with " - "citations, or computed across documents.", + "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.", ] @@ -106,6 +110,26 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: return "\n".join(lines) +def _node(toc: "dict[str, Any]") -> OutlineNode: + return OutlineNode( + id=toc["self_ref"], + title=toc["title"], + level=toc["level"], + page_numbers=toc["page_numbers"], + children=[_node(child) for child in toc["children"]], + ) + + +def _find(toc: list["dict[str, Any]"], section_id: str) -> "dict[str, Any] | None": + for node in toc: + if node["self_ref"] == section_id: + return node + found = _find(node["children"], section_id) + if found is not None: + return found + return None + + def create_mcp_server( db_path: Path | None = None, config: AppConfig | None = None, @@ -279,6 +303,72 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: raise ToolError(f"No document with id {document_id!r}") return document + async def _items_of(document_id: str, source: str | None) -> list[DocumentItem]: + """A document's items in reading order, from the database holding it.""" + rag = await _client() + try: + document = await rag.get_document_by_id(document_id, source) + if document is None: + raise ToolError(f"No document with id {document_id!r}") + owner = await rag.reader_for(source or document.source) + except UnknownDatabaseError as e: + raise ToolError(str(e)) from e + assert owner is not None, "a stored document names its database" + return await owner.document_item_repository.get_all_items(document_id) + + @mcp.tool(annotations=_read_only("Document outline")) + async def get_document_outline( + document_id: str, source: str | None = None + ) -> list[OutlineNode]: + """The heading tree of a document, with page numbers. + + Use this on a long document to see its structure before reading, then + pass a node's `id` to `get_document_section`. Returns the headings + nested by level; an empty list means the document has no headings, + so read it with `get_document`. + + Args: + document_id: The document's id. + source: The collection holding it. Without one every collection + is asked. + """ + return [ + _node(toc) for toc in build_toc(await _items_of(document_id, source), {}) + ] + + @mcp.tool(annotations=_read_only("Document section")) + async def get_document_section( + document_id: str, section_id: str, source: str | None = None + ) -> DocumentSection: + """The text of one section of a document, subsections included. + + Use this to read a part of a long document instead of the whole. + `section_id` is a node `id` from `get_document_outline`. Returns the + section's heading, page numbers and text in reading order, up to the + next heading of the same or a higher level. + + Args: + document_id: The document's id. + section_id: The `id` of a node in the document's outline. + source: The collection holding it. Without one every collection + is asked. + """ + items = await _items_of(document_id, source) + node = _find(build_toc(items, {}), section_id) + if node is None: + raise ToolError(f"No section {section_id!r} in document {document_id!r}") + start, end = node["item_range"] + return DocumentSection( + id=node["self_ref"], + title=node["title"], + page_numbers=node["page_numbers"], + content="\n\n".join( + item.text + for item in items + if start <= item.position < end and item.text + ), + ) + @mcp.tool(annotations=_read_only("List documents")) async def list_documents( limit: int | None = None, @@ -313,7 +403,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: @mcp.tool(annotations=_read_only("Ask a question")) async def ask_question( question: str, - cite: bool = False, images_base64: list[str] | None = None, sources: Sources = None, ) -> str: @@ -321,12 +410,10 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: 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 when - `cite` is set. + the answer, followed by citations to the passages it rests on. Args: question: The question, in natural language. - cite: Append citations to the answer. images_base64: Images to attach to the question, PNG or JPEG bytes as base64. Needs a vision-capable model on the server. """ @@ -339,7 +426,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: except Exception as e: logger.exception("ask_question failed") raise ToolError(f"ask_question failed: {type(e).__name__}") from e - if cite and citations: + if citations: answer += "\n\n" + format_citations( citations, include_source=rag.covers_multiple ) diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 4102d961..cddbdd46 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -17,6 +17,7 @@ from pydantic_monty import ( ) from haiku.rag.config.models import AppConfig +from haiku.rag.context import build_toc from haiku.rag.sandbox.dependencies import AnalysisContext from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem @@ -38,81 +39,6 @@ class SandboxResult: success: bool -def _build_toc( - items: list["DocumentItem"], - chunk_index: dict[str, list[str]], -) -> list[dict[str, Any]]: - """Build a nested section tree from items in position order. - - Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting - follows the explicit levels: a header pops the stack until the top is at - a strictly shallower level, then becomes a child of that top (or a root). - - ``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the - position of the next header whose level is the same or shallower (i.e. - the next sibling or ancestor that ends this section), or the total item - count if no such header exists. - - ``chunk_ids`` aggregates the chunks covered by all items in the section's - ``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to - ground a section-scoped answer without a corpus-wide ``search()`` call. - - Items without a section_header label (or with ``heading_level == 0``) are - skipped. When all section_headers carry the same level the output is a - flat sibling list (see docling-project/docling#2121 for an upstream case - where every PDF section_header is emitted at level=1). - """ - # Defensive: every consumer is supposed to pass items in position order, - # but the end_exclusive lookahead below silently miscomputes section - # boundaries if it's not — better to sort once than trust the caller. - items = sorted(items, key=lambda i: i.position) - headers: list[DocumentItem] = [ - i for i in items if i.label == "section_header" and i.heading_level > 0 - ] - if not headers: - return [] - - total = max((i.position for i in items), default=-1) + 1 - items_by_position: dict[int, DocumentItem] = {i.position: i for i in items} - - ends: list[int] = [] - for idx, h in enumerate(headers): - end = total - for j in range(idx + 1, len(headers)): - if headers[j].heading_level <= h.heading_level: - end = headers[j].position - break - ends.append(end) - - roots: list[dict[str, Any]] = [] - stack: list[tuple[int, dict[str, Any]]] = [] - for h, end in zip(headers, ends, strict=True): - seen: set[str] = set() - chunk_ids: list[str] = [] - for pos in range(h.position, end): - item = items_by_position.get(pos) - if item is None: - continue - for cid in chunk_index.get(item.self_ref, []): - if cid not in seen: - seen.add(cid) - chunk_ids.append(cid) - node: dict[str, Any] = { - "self_ref": h.self_ref, - "level": h.heading_level, - "title": h.text, - "page_numbers": list(h.page_numbers), - "item_range": [h.position, end], - "chunk_ids": chunk_ids, - "children": [], - } - while stack and stack[-1][0] >= h.heading_level: - stack.pop() - (stack[-1][1]["children"] if stack else roots).append(node) - stack.append((h.heading_level, node)) - return roots - - class Sandbox: """Execute code in a sandboxed Python interpreter. @@ -520,7 +446,7 @@ class Sandbox: { "doc_id": did, "title": doc_titles.get(did), - "tree": _build_toc(items, chunk_index), + "tree": build_toc(items, chunk_index), }, ensure_ascii=False, ) diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index 960d5194..89648127 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -31,6 +31,25 @@ class DocumentInfo(BaseModel): metadata: dict = {} +class OutlineNode(BaseModel): + """A heading in a document's outline. `id` is the heading item's self_ref.""" + + id: str + title: str + level: int + page_numbers: list[int] = [] + children: list["OutlineNode"] = [] + + +class DocumentSection(BaseModel): + """One section's text in reading order, subsections included.""" + + id: str + title: str + page_numbers: list[int] = [] + content: str + + class DocumentListResponse(BaseModel): """Response from list_documents tool.""" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index b21b24cd..b908a190 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -236,9 +236,7 @@ class TestMCPReadTools: assert overview["metadata"] == {"author": "Ada"} @pytest.mark.asyncio - async def test_ask_question_appends_citations_when_requested( - self, mcp_db, monkeypatch - ): + async def test_ask_question_appends_the_citations(self, mcp_db, monkeypatch): from haiku.rag.store.models.citation import Citation citation = Citation( @@ -257,13 +255,182 @@ class TestMCPReadTools: mcp = create_mcp_server(mcp_db) ask = await _get_tool(mcp, "ask_question") - with_cite = await ask(question="q", cite=True) - assert with_cite.startswith("the answer") - assert "AI Overview" in with_cite + answer = await ask(question="q") + assert answer.startswith("the answer") + assert "AI Overview" in answer # One database: its name adds nothing. - assert "alpha" not in with_cite + assert "alpha" not in answer - assert await ask(question="q", cite=False) == "the answer" + +@pytest.fixture +async def outlined_db(temp_db_path): + """A database with one document whose items carry a heading hierarchy. + + Rows are written through the repositories, so no embedder is involved. + Returns the path and the document id.""" + from haiku.rag.store.models.document import Document as DocumentModel + from haiku.rag.store.models.document_item import DocumentItem + + def header(pos, level, text): + return DocumentItem( + document_id="", + position=pos, + self_ref=f"#/texts/{pos}", + label="section_header", + text=text, + page_numbers=[pos // 4 + 1], + heading_level=level, + ) + + def para(pos): + return DocumentItem( + document_id="", + position=pos, + self_ref=f"#/texts/{pos}", + label="paragraph", + text=f"para{pos}", + page_numbers=[pos // 4 + 1], + ) + + async with HaikuRAG(temp_db_path, create=True) as rag: + doc = await rag.document_repository.create( + DocumentModel(content="x", uri="test://outlined", title="Outlined") + ) + items = [ + header(0, 1, "Intro"), + para(1), + header(2, 2, "Background"), + para(3), + header(4, 3, "Prior Work"), + para(5), + header(6, 2, "Approach"), + para(7), + header(8, 1, "Methods"), + para(9), + ] + for item in items: + item.document_id = doc.id + await rag.document_item_repository.create_items(doc.id, items) + return temp_db_path, doc.id + + +class TestMCPDocumentNavigation: + @pytest.mark.asyncio + async def test_the_outline_nests_headings_by_level(self, outlined_db): + db, doc_id = outlined_db + outline = await _get_tool(create_mcp_server(db), "get_document_outline") + + roots = await outline(document_id=doc_id) + + assert [n.title for n in roots] == ["Intro", "Methods"] + intro = roots[0] + assert (intro.id, intro.level, intro.page_numbers) == ("#/texts/0", 1, [1]) + assert [c.title for c in intro.children] == ["Background", "Approach"] + assert [c.title for c in intro.children[0].children] == ["Prior Work"] + assert intro.children[0].children[0].level == 3 + assert roots[1].children == [] + + @pytest.mark.asyncio + async def test_a_document_without_headings_has_an_empty_outline(self, mcp_db): + mcp = create_mcp_server(mcp_db) + [doc] = await (await _get_tool(mcp, "list_documents"))(limit=1) + outline = await _get_tool(mcp, "get_document_outline") + + assert await outline(document_id=doc.id) == [] + + @pytest.mark.asyncio + async def test_a_section_covers_its_subsections_and_stops_at_its_sibling( + self, outlined_db + ): + db, doc_id = outlined_db + section = await _get_tool(create_mcp_server(db), "get_document_section") + + background = await section(document_id=doc_id, section_id="#/texts/2") + + assert background.title == "Background" + assert background.content.split("\n\n") == [ + "Background", + "para3", + "Prior Work", + "para5", + ] + assert background.page_numbers == [1] + + intro = await section(document_id=doc_id, section_id="#/texts/0") + assert intro.content.startswith("Intro") + assert "para7" in intro.content + assert "Methods" not in intro.content + + @pytest.mark.asyncio + async def test_an_unknown_section_or_document_is_an_error(self, outlined_db): + db, doc_id = outlined_db + mcp = create_mcp_server(db) + section = await _get_tool(mcp, "get_document_section") + outline = await _get_tool(mcp, "get_document_outline") + + with pytest.raises(ToolError, match="#/texts/99"): + await section(document_id=doc_id, section_id="#/texts/99") + with pytest.raises(ToolError, match="nonexistent-id"): + await outline(document_id="nonexistent-id") + with pytest.raises(ToolError, match="nonexistent-id"): + await section(document_id="nonexistent-id", section_id="#/texts/0") + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_outline_and_section_serialize_over_the_wire(self, outlined_db): + db, doc_id = outlined_db + mcp = create_mcp_server(db) + + outline = await _call(mcp, "get_document_outline", document_id=doc_id) + section = await _call( + mcp, "get_document_section", document_id=doc_id, section_id="#/texts/8" + ) + + assert not outline.is_error and not section.is_error + [intro, methods] = outline.structured_content["result"] + assert set(intro) == {"id", "title", "level", "page_numbers", "children"} + assert intro["children"][0]["children"][0]["title"] == "Prior Work" + assert set(section.structured_content) == { + "id", + "title", + "page_numbers", + "content", + } + assert section.structured_content["content"] == "Methods\n\npara9" + + @pytest.mark.asyncio + async def test_source_routes_to_the_database_holding_the_document(self, two_dbs): + from haiku.rag.store.models.document_item import DocumentItem + + async with HaikuRAG(config=two_dbs, sources=["beta"]) as beta: + [doc] = await beta.list_documents() + await beta.document_item_repository.create_items( + doc.id, + [ + DocumentItem( + document_id=doc.id, + position=0, + self_ref="#/texts/0", + label="section_header", + text="Only in beta", + heading_level=1, + ) + ], + ) + mcp = _covering_all(two_dbs) + outline = await _get_tool(mcp, "get_document_outline") + section = await _get_tool(mcp, "get_document_section") + + named = await outline(document_id=doc.id, source="beta") + found = await outline(document_id=doc.id) + assert [n.title for n in named] == [n.title for n in found] == ["Only in beta"] + assert ( + await section(document_id=doc.id, section_id="#/texts/0", source="beta") + ).title == "Only in beta" + with pytest.raises(ToolError, match="nope"): + await outline(document_id=doc.id, source="nope") + with pytest.raises(ToolError, match=doc.id): + await outline(document_id=doc.id, source="alpha") @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @@ -324,7 +491,7 @@ class TestMCPDescribesItself: async with Client(create_mcp_server(mcp_db)) as client: tools = await client.list_tools() - assert len(tools) == 6 + assert len(tools) == 8 for tool in tools: assert tool.annotations is not None, tool.name assert tool.annotations.readOnlyHint is True, tool.name @@ -344,7 +511,7 @@ class TestMCPDescribesItself: for name, schema in tool.inputSchema.get("properties", {}).items() if not schema.get("description") ] - assert len(tools) == 6 + assert len(tools) == 8 assert undescribed == [] @@ -356,6 +523,8 @@ class TestMCPToolSet: assert {t.name for t in await mcp.list_tools()} == { "search_documents", "get_document", + "get_document_outline", + "get_document_section", "list_documents", "ask_question", "analyze", @@ -475,7 +644,7 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) ask = await _get_tool(mcp, "ask_question") - answer = await ask(question="q", cite=True) + answer = await ask(question="q") assert "alpha" in answer assert "beta" in answer From 5653e876a812f05a5a2dcfb84c4070663cadedd4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 12:27:11 +0300 Subject: [PATCH 07/18] Return search results as agent text, images and structured content search_documents and search_documents_by_image return a ToolResult: the format_for_agent rendering with rank, Document ID and Collection so the text alone drives the document tools; one ImageContent per distinct picture, labelled with its result; and the SearchResult list without image_data as structured content. format_for_agent gains an opt-in include_document_id, so the capabilities' rendering is unchanged. collect_pictures is the one place pictures are deduplicated and validated for both wire formats. Refs #599 --- CHANGELOG.md | 6 + docs/mcp.md | 14 +- haiku_rag_slim/haiku/rag/mcp.py | 79 +++++-- .../haiku/rag/store/models/chunk.py | 7 +- haiku_rag_slim/haiku/rag/tools/search.py | 59 +++--- tests/test_chunk.py | 11 + tests/test_mcp.py | 200 ++++++++++++++++-- 7 files changed, 316 insertions(+), 60 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f6fb5c6..90bf3bf3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,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`. +- MCP `search_documents` and `search_documents_by_image` return the agent + rendering as text (rank, `Document ID`, `Collection` over several + databases, title, headings, passage), pictures as `ImageContent` blocks, + and the `SearchResult` list without `image_data` as structured content. + `SearchResult.format_for_agent(include_document_id=)`; + `collect_pictures` in `haiku.rag.tools.search`. - MCP tools raise on failure; an empty result no longer doubles as an error. Unknown document, unknown collection, invalid filter and invalid base64 carry a message; `ask_question` and `analyze` failures name the exception diff --git a/docs/mcp.md b/docs/mcp.md index c4815434..91f636d0 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -84,11 +84,15 @@ repeating it. | `ask_question` | always | `question`, `images_base64`, `sources` | | `analyze` | always | `question`, `filter`, `images_base64`, `sources` | -`search_documents` runs hybrid search, vector and full-text, and returns -results best first. Scores are not comparable across queries or search types. -Rank is the signal. `include_images` attaches picture bytes as base64 PNG under -`image_data`. `search_documents_by_image` embeds the query image and searches -by vector similarity alone. +`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 +rank, `Document ID`, `Collection` when the server covers several, the document +title, section headings and the passage. Pictures in the results follow as +image blocks, one per distinct picture, each preceded by a line naming its +result; `include_images: false` leaves them out. The structured content is the +`SearchResult` list without picture bytes. Scores are not comparable across +queries or search types, so rank is the signal. `search_documents_by_image` +embeds the query image and searches by vector similarity alone. `get_document` returns a document whole, in reading order. For a long one, `get_document_outline` returns the heading tree with page numbers and diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index ec97476f..39130c86 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -8,7 +8,8 @@ from typing import TYPE_CHECKING, Annotated from fastmcp import FastMCP from fastmcp.exceptions import ToolError -from mcp.types import ToolAnnotations +from fastmcp.tools import ToolResult +from mcp.types import ContentBlock, ImageContent, TextContent, ToolAnnotations from pydantic import Field from haiku.rag.client import HaikuRAG @@ -19,6 +20,7 @@ from haiku.rag.store.models import Document, SearchResult from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.schema import DocumentMetaRecord from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode +from haiku.rag.tools.search import collect_pictures from haiku.rag.utils import format_citations if TYPE_CHECKING: @@ -110,6 +112,52 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: return "\n".join(lines) +def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult: + """Results as the in-process agents read them, then each distinct picture + as an image block labelled with its result, and the results as structured + content without the picture bytes.""" + import base64 + + total = len(results) + text = "\n\n".join( + result.format_for_agent( + rank=rank, + total=total, + include_collection=covers_multiple, + include_document_id=True, + ) + for rank, result in enumerate(results, 1) + ) + content: list[ContentBlock] = [ + TextContent(type="text", text=text or "No results found.") + ] + pictures, _ = collect_pictures(results) + for source, chunk_id, self_ref, picture in pictures: + collection = f" in {source}" if covers_multiple and source else "" + content.append( + TextContent( + type="text", + text=f"Picture {self_ref} of search result [{chunk_id}]{collection}", + ) + ) + content.append( + ImageContent( + type="image", + data=base64.b64encode(picture.data).decode("ascii"), + mimeType="image/png", + ) + ) + return ToolResult( + content=content, + structured_content={ + "result": [ + result.model_dump(mode="json", exclude={"image_data"}) + for result in results + ] + }, + ) + + def _node(toc: "dict[str, Any]") -> OutlineNode: return OutlineNode( id=toc["self_ref"], @@ -207,27 +255,30 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: include_images: bool = True, filter: Filter = None, sources: Sources = None, - ) -> list[SearchResult]: + ) -> ToolResult: """Search the knowledge base by meaning and keyword. Use this first for any question the documents might answer; it needs no model and is the cheapest call. Results come best first, each with - the document's id, title and collection, the section headings and the - matching passage. Scores are not comparable across queries, so read - the order, not the numbers. If nothing relevant comes back, rephrase - once or narrow with `filter` before concluding the material is absent. + its rank, `Document ID`, `Collection` when the server covers several, + the document title, section headings and the matching passage; pass + the id and collection to the document tools. Pictures in the results + follow as images, each labelled with its result. Ranks, not scores, + are the signal: scores are not comparable across queries. If nothing + relevant comes back, rephrase once or narrow with `filter` before + concluding the material is absent. Args: query: What to look for, in natural language or keywords. limit: How many results to return; the server's configured default when omitted. - include_images: Attach the bytes of pictures in the results as - base64 PNG under `image_data`. False for a smaller response. + include_images: Return the pictures in the results as images. + False for a smaller response. """ rag = await _client() try: await _check_filter(rag, filter, sources) - return await rag.search( + results = await rag.search( query, limit=limit, filter=filter, @@ -236,6 +287,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e + return _search_result(results, rag.covers_multiple) # Image-as-query tool, only registered when the configured embedder # supports image embeddings. Probed at server-build time when no Store is @@ -252,7 +304,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: include_images: bool = True, filter: Filter = None, sources: Sources = None, - ) -> list[SearchResult]: + ) -> ToolResult: """Search the knowledge base with an image as the query. Use this when the question is about a picture rather than words. @@ -264,14 +316,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: image_base64: The query image, PNG or JPEG bytes as base64. limit: How many results to return; the server's configured default when omitted. - include_images: Attach the bytes of pictures in the results as - base64 PNG under `image_data`. False for a smaller response. + include_images: Return the pictures in the results as images. + False for a smaller response. """ raw = _decode_image(image_base64) rag = await _client() try: await _check_filter(rag, filter, sources) - return await rag.search( + results = await rag.search( raw, limit=limit, filter=filter, @@ -280,6 +332,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e + return _search_result(results, rag.covers_multiple) @mcp.tool(annotations=_read_only("Get document")) async def get_document(document_id: str, source: str | None = None) -> Document: diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 6c2b4215..b21231e1 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -202,6 +202,7 @@ class SearchResult(BaseModel): total: int | None = None, *, include_collection: bool = False, + include_document_id: bool = False, ) -> str: """Format this search result for inclusion in agent context. @@ -215,7 +216,8 @@ class SearchResult(BaseModel): `include_collection` is the caller's decision, not this result's: a search spanning one collection has nothing to distinguish, whether or - not that collection is named. + not that collection is named. `include_document_id` is for a reader + that will fetch the document by id from the text alone. """ if rank is not None and total is not None: parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"] @@ -224,6 +226,9 @@ class SearchResult(BaseModel): else: parts = [f"[{self.chunk_id}] (score: {self.score:.2f})"] + if include_document_id and self.document_id: + parts.append(f"Document ID: {self.document_id}") + if include_collection and self.source: parts.append(f"Collection: {self.source}") diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index c779a19a..54bd1b99 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -52,31 +52,16 @@ def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None: return BinaryContent(data=data, media_type="image/png", identifier=self_ref) -def build_image_content_from_results( - results: list[SearchResult], - include_collection: bool = False, - exclude: AbstractSet[PictureKey] = frozenset(), -) -> tuple[list[str | BinaryContent], set[PictureKey]]: - """Decode and validate picture bytes attached to search results, labelled. +def collect_pictures( + results: list[SearchResult], exclude: AbstractSet[PictureKey] = frozenset() +) -> tuple[list[tuple[str | None, str | None, str, BinaryContent]], set[PictureKey]]: + """Every distinct, decodable picture attached to ``results``, in order. - Returns the labelled content and the ``PictureKey`` of every picture it - emitted. Dedup keyed on ``PictureKey`` so the same picture in - different chunks is sent once, and a copy in another collection is its - own; ``exclude`` seeds that dedup with pictures already sent. Pictures that fail - ``PIL.Image.verify()`` are skipped — the model adapter renders one - vision placeholder per ``BinaryContent``, so emitting one for an - image the server can't decode leaves the processor with an - off-by-one count. - - Every picture is preceded by a line naming the result it belongs to. - ``ToolReturn.content`` reaches the model as a user-role message, so - retrieved pictures are otherwise indistinguishable from ones the user - attached, and models narrate them as part of the question: unlabelled, - gemma4-26b answered about a figure from an unrelated document, and with a - single note ahead of the batch it still called them "images in the prompt". - The label also names the chunk to cite for a figure, which - ``BinaryContent.identifier`` cannot do — it does not survive serialization - to the vision API. + Returns ``(source, chunk_id, self_ref, picture)`` per picture and the + ``PictureKey`` of each. Dedup keyed on ``PictureKey`` so the same picture in + different chunks is emitted once, and a copy in another collection is its + own; ``exclude`` seeds that dedup with pictures already sent. Pictures that + fail ``PIL.Image.verify()`` are skipped. """ collected: list[tuple[str | None, str | None, str, BinaryContent]] = [] seen: set[PictureKey] = set(exclude) @@ -94,7 +79,33 @@ def build_image_content_from_results( collected.append((result.source, result.chunk_id, self_ref, picture)) seen.add(key) emitted.add(key) + return collected, emitted + +def build_image_content_from_results( + results: list[SearchResult], + include_collection: bool = False, + exclude: AbstractSet[PictureKey] = frozenset(), +) -> tuple[list[str | BinaryContent], set[PictureKey]]: + """Decode and validate picture bytes attached to search results, labelled. + + Returns the labelled content and the ``PictureKey`` of every picture it + emitted, as ``collect_pictures`` decides them. An undecodable picture is + skipped because the model adapter renders one vision placeholder per + ``BinaryContent``, so emitting one for an image the server can't decode + leaves the processor with an off-by-one count. + + Every picture is preceded by a line naming the result it belongs to. + ``ToolReturn.content`` reaches the model as a user-role message, so + retrieved pictures are otherwise indistinguishable from ones the user + attached, and models narrate them as part of the question: unlabelled, + gemma4-26b answered about a figure from an unrelated document, and with a + single note ahead of the batch it still called them "images in the prompt". + The label also names the chunk to cite for a figure, which + ``BinaryContent.identifier`` cannot do — it does not survive serialization + to the vision API. + """ + collected, emitted = collect_pictures(results, exclude) content: list[str | BinaryContent] = [] total = len(collected) for position, (source, chunk_id, self_ref, picture) in enumerate(collected, 1): diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 98b24674..1bcf6022 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -414,6 +414,17 @@ def test_search_result_format_for_agent_source_line(fields, expected_source): assert expected_source in result.format_for_agent() +def test_search_result_format_for_agent_document_id_is_opt_in(): + """The capabilities' rendering is unchanged; only a caller that asks gets + the id it will fetch the document by.""" + result = SearchResult(content="x", score=0.5, chunk_id="c1", document_id="doc-1") + + assert "Document ID" not in result.format_for_agent(rank=1, total=1) + assert "Document ID: doc-1" in result.format_for_agent( + rank=1, total=1, include_document_id=True + ) + + @pytest.mark.parametrize( "labels,expected", [ diff --git a/tests/test_mcp.py b/tests/test_mcp.py index b908a190..13454930 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -97,22 +97,38 @@ async def _call(mcp, name, **kwargs): return await client.call_tool(name, kwargs, raise_on_error=False) +def _results(search_result) -> list[dict]: + """The search results a tool returned, as the client sees them.""" + return search_result.structured_content["result"] + + +def _png_b64() -> str: + import base64 + from io import BytesIO + + from PIL import Image as PILImage + + buf = BytesIO() + PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode() + + class TestMCPReadTools: @pytest.mark.asyncio async def test_search_documents(self, mcp_db): mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") - results = await search(query="artificial intelligence") + results = _results(await search(query="artificial intelligence")) assert len(results) > 0 - assert all(isinstance(r, SearchResult) for r in results) + assert all(r["chunk_id"] and r["content"] for r in results) @pytest.mark.asyncio async def test_search_documents_with_limit(self, mcp_db): mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") - results = await search(query="artificial intelligence", limit=1) + results = _results(await search(query="artificial intelligence", limit=1)) assert len(results) == 1 @pytest.mark.asyncio @@ -433,6 +449,155 @@ class TestMCPDocumentNavigation: await outline(document_id=doc.id, source="alpha") +@pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") +class TestMCPSearchResultShape: + """Text as the in-process agents read it, one image per distinct picture, + and the results as structured content without picture bytes.""" + + @staticmethod + def _serve(monkeypatch, results): + async def fake_search(self, *args, **kwargs): + return results + + monkeypatch.setattr(HaikuRAG, "search", fake_search) + + @pytest.mark.asyncio + async def test_text_ranks_then_one_image_per_distinct_picture( + self, mcp_db, monkeypatch + ): + from mcp.types import ImageContent, TextContent + + shared = {"#/pictures/0": _png_b64()} + self._serve( + monkeypatch, + [ + SearchResult( + content="a", + score=0.9, + chunk_id="c1", + document_id="d1", + image_data=shared, + ), + SearchResult( + content="b", + score=0.8, + chunk_id="c2", + document_id="d1", + image_data=shared, + ), + SearchResult( + content="c", + score=0.7, + chunk_id="c3", + document_id="d2", + image_data={"#/pictures/3": _png_b64()}, + ), + ], + ) + + result = await _call(create_mcp_server(mcp_db), "search_documents", query="q") + + text, *rest = result.content + assert isinstance(text, TextContent) + assert "[rank 1 of 3]" in text.text and "[rank 3 of 3]" in text.text + assert "score" not in text.text + assert "Document ID: d1" in text.text + images = [block for block in rest if isinstance(block, ImageContent)] + labels = [block.text for block in rest if isinstance(block, TextContent)] + assert len(images) == 2 + assert all(image.mimeType == "image/png" for image in images) + assert [ + label for label in labels if "[c1]" in label and "#/pictures/0" in label + ] + assert [ + label for label in labels if "[c3]" in label and "#/pictures/3" in label + ] + structured = _results(result) + assert [r["chunk_id"] for r in structured] == ["c1", "c2", "c3"] + assert all("image_data" not in r for r in structured) + + @pytest.mark.asyncio + async def test_an_undecodable_picture_yields_no_image(self, mcp_db, monkeypatch): + import base64 + + self._serve( + monkeypatch, + [ + SearchResult( + content="a", + score=0.9, + chunk_id="c1", + document_id="d1", + image_data={ + "#/pictures/0": base64.b64encode(b"not a png").decode() + }, + ) + ], + ) + + result = await _call(create_mcp_server(mcp_db), "search_documents", query="q") + + assert len(result.content) == 1 + assert "[rank 1 of 1]" in result.content[0].text + + @pytest.mark.asyncio + async def test_no_results_says_so(self, mcp_db, monkeypatch): + self._serve(monkeypatch, []) + + result = await _call(create_mcp_server(mcp_db), "search_documents", query="q") + + assert [block.text for block in result.content] == ["No results found."] + assert _results(result) == [] + + @pytest.mark.asyncio + async def test_search_text_alone_drives_the_document_tools(self, two_dbs): + """Over two databases, every result's `Document ID` and `Collection` + parsed from the text are working arguments for the outline and + section tools.""" + import re + + from haiku.rag.store.models.document_item import DocumentItem + + for name in ("alpha", "beta"): + async with HaikuRAG(config=two_dbs, sources=[name]) as rag: + [doc] = await rag.list_documents() + await rag.document_item_repository.create_items( + doc.id, + [ + DocumentItem( + document_id=doc.id, + position=0, + self_ref="#/texts/0", + label="section_header", + text=f"Heading in {name}", + heading_level=1, + ) + ], + ) + mcp = _covering_all(two_dbs) + + search = await _call(mcp, "search_documents", query="cats") + pairs = re.findall( + r"Document ID: (\S+)\nCollection: (\S+)", search.content[0].text + ) + + assert len(pairs) == len(_results(search)) == 2 + assert {source for _, source in pairs} == {"alpha", "beta"} + for document_id, source in pairs: + outline = await _call( + mcp, "get_document_outline", document_id=document_id, source=source + ) + [node] = _results(outline) + section = await _call( + mcp, + "get_document_section", + document_id=document_id, + section_id=node["id"], + source=source, + ) + assert section.structured_content["title"] == f"Heading in {source}" + + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") class TestMCPDescribesItself: """What a client learns from initialize and list_tools, over the wire.""" @@ -537,19 +702,19 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = await search(query="cats") + results = _results(await search(query="cats")) - assert {r.source for r in results} == {"alpha", "beta"} + 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"]) + results = _results(await search(query="cats", sources=["beta"])) assert results - assert {r.source for r in results} == {"beta"} + assert {r["source"] for r in results} == {"beta"} @pytest.mark.asyncio @pytest.mark.parametrize( @@ -583,12 +748,13 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = await search( - query="cats", filter="uri LIKE '%beta%'", sources=["beta"] + results = _results( + await search(query="cats", filter="uri LIKE '%beta%'", sources=["beta"]) ) assert results - assert {r.source for r in results} == {"beta"} - assert await search(query="cats", filter="uri LIKE '%beta%'", sources=[]) == [] + assert {r["source"] for r in results} == {"beta"} + none = await search(query="cats", filter="uri LIKE '%beta%'", sources=[]) + assert _results(none) == [] @pytest.mark.asyncio async def test_the_listing_covers_every_database(self, two_dbs): @@ -617,9 +783,9 @@ class TestMCPCoversTheConfiguredSet: mcp = create_mcp_server(config=two_dbs) search = await _get_tool(mcp, "search_documents") - results = await search(query="cats") + results = _results(await search(query="cats")) - assert {r.source for r in results} == {"alpha", "beta"} + assert {r["source"] for r in results} == {"alpha", "beta"} @pytest.mark.asyncio async def test_ask_question_names_each_citations_database( @@ -713,7 +879,7 @@ class TestMCPImageQuery: sources=[], ) - assert results == [] + assert _results(results) == [] assert seen["query"] == png assert seen["filter"] == "uri LIKE 'x%'" assert seen["sources"] == [] @@ -1028,12 +1194,12 @@ class TestMCPClientLifetime: mcp = _mcp_covering(scope, config) async with mcp._lifespan_manager(): search = await _get_tool(mcp, "search_documents") - results = await search(query="artificial intelligence") + results = _results(await search(query="artificial intelligence")) listing = await _get_tool(mcp, "list_documents") documents = await listing() assert results - assert {r.source for r in results} == {"alpha"} + assert {r["source"] for r in results} == {"alpha"} titles = {d.title for d in documents} assert "AI Overview" in titles assert "Zebras" not in titles @@ -1118,7 +1284,7 @@ class TestMCPClientLifetime: assert opens == 1 async with mcp._lifespan_manager(): - results = await search(query="artificial intelligence") + results = _results(await search(query="artificial intelligence")) assert opens == 2 assert len(results) > 0 From f45ed90b338e54f1ebc252f2464c1bc54a9d7872 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 12:33:32 +0300 Subject: [PATCH 08/18] 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 --- CHANGELOG.md | 2 + docs/cli.md | 3 + docs/mcp.md | 7 +- haiku_rag_slim/haiku/rag/app.py | 3 +- haiku_rag_slim/haiku/rag/cli.py | 9 +- haiku_rag_slim/haiku/rag/mcp.py | 149 ++++++++++++++++++-------------- tests/test_app.py | 14 +++ tests/test_cli.py | 8 ++ tests/test_mcp.py | 26 +++++- 9 files changed, 150 insertions(+), 71 deletions(-) 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() From eb9995e9b73228da471ebcaaa81d40bb7bb5728d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 12:54:30 +0300 Subject: [PATCH 09/18] Ship a Claude Code plugin with the haiku-rag skill claude-plugin/ holds the plugin manifest, the server configuration (haiku-rag mcp --stdio, the configuration decides the database) and a skill that says when to reach for the knowledge base and how to move from a search result to a document, a section, an answer or a computation. A repo-root marketplace manifest makes `claude plugin marketplace add ggozad/haiku.rag` work. The skill pre-approves every tool the server registers, and a test keeps the two in step. The manifest carries the package version, which bump_version.py now rewrites: a versioned plugin updates only on a bump, so the installed skill stays in step with the haiku-rag release the user has. Refs #599 --- .claude-plugin/marketplace.json | 14 +++++ CHANGELOG.md | 3 ++ README.md | 9 +++- claude-plugin/.claude-plugin/plugin.json | 12 +++++ claude-plugin/.mcp.json | 8 +++ claude-plugin/skills/haiku-rag/SKILL.md | 65 ++++++++++++++++++++++++ docs/mcp.md | 19 +++++++ scripts/bump_version.py | 23 ++++++++- tests/test_bump_version.py | 38 ++++++++++++++ tests/test_mcp.py | 43 ++++++++++++++++ 10 files changed, 231 insertions(+), 3 deletions(-) create mode 100644 .claude-plugin/marketplace.json create mode 100644 claude-plugin/.claude-plugin/plugin.json create mode 100644 claude-plugin/.mcp.json create mode 100644 claude-plugin/skills/haiku-rag/SKILL.md create mode 100644 tests/test_bump_version.py diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json new file mode 100644 index 00000000..24d16eee --- /dev/null +++ b/.claude-plugin/marketplace.json @@ -0,0 +1,14 @@ +{ + "name": "haiku-rag", + "description": "The haiku.rag knowledge base as Claude Code tools and a skill.", + "owner": { + "name": "Yiorgis Gozadinos" + }, + "plugins": [ + { + "name": "haiku-rag", + "source": "./claude-plugin", + "description": "Search, read and question your haiku.rag knowledge base from Claude Code." + } + ] +} diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a22246e..41411c13 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,9 @@ ### Added +- Claude Code plugin under `claude-plugin/`: the server configuration and the + `haiku-rag` skill. `claude plugin marketplace add ggozad/haiku.rag`, then + `claude plugin install haiku-rag`. - `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 diff --git a/README.md b/README.md index 9016e322..fb3c8d5b 100644 --- a/README.md +++ b/README.md @@ -110,12 +110,19 @@ For direct agent composition, see the [capabilities documentation](https://ggoza ## MCP Server -Use with AI assistants like Claude Desktop: +Use with AI assistants like Claude Code and Claude Desktop: ```bash haiku-rag mcp --stdio ``` +In Claude Code, install the plugin, which registers the server and a skill: + +```bash +claude plugin marketplace add ggozad/haiku.rag +claude plugin install haiku-rag +``` + Add to your Claude Desktop configuration: ```json diff --git a/claude-plugin/.claude-plugin/plugin.json b/claude-plugin/.claude-plugin/plugin.json new file mode 100644 index 00000000..93282056 --- /dev/null +++ b/claude-plugin/.claude-plugin/plugin.json @@ -0,0 +1,12 @@ +{ + "name": "haiku-rag", + "version": "0.82.1", + "description": "Search, read and question your haiku.rag knowledge base from Claude Code.", + "author": { + "name": "Yiorgis Gozadinos", + "email": "ggozadinos@gmail.com" + }, + "homepage": "https://ggozad.github.io/haiku.rag/mcp/", + "repository": "https://github.com/ggozad/haiku.rag", + "license": "MIT" +} diff --git a/claude-plugin/.mcp.json b/claude-plugin/.mcp.json new file mode 100644 index 00000000..f6fb25f9 --- /dev/null +++ b/claude-plugin/.mcp.json @@ -0,0 +1,8 @@ +{ + "mcpServers": { + "haiku-rag": { + "command": "haiku-rag", + "args": ["mcp", "--stdio"] + } + } +} diff --git a/claude-plugin/skills/haiku-rag/SKILL.md b/claude-plugin/skills/haiku-rag/SKILL.md new file mode 100644 index 00000000..f93631a3 --- /dev/null +++ b/claude-plugin/skills/haiku-rag/SKILL.md @@ -0,0 +1,65 @@ +--- +name: haiku-rag +description: Search, read and question the user's haiku.rag knowledge base + through the haiku-rag MCP tools. Use whenever a request could be answered + from the user's ingested documents, when asked to find, look up, check or + cite something in their documents or knowledge base, or when the question is + about the user's own material rather than general knowledge. +allowed-tools: + - mcp__plugin_haiku-rag_haiku-rag__search_documents + - mcp__plugin_haiku-rag_haiku-rag__search_documents_by_image + - mcp__plugin_haiku-rag_haiku-rag__get_document + - mcp__plugin_haiku-rag_haiku-rag__get_document_outline + - mcp__plugin_haiku-rag_haiku-rag__get_document_section + - mcp__plugin_haiku-rag_haiku-rag__list_documents + - mcp__plugin_haiku-rag_haiku-rag__ask_question + - mcp__plugin_haiku-rag_haiku-rag__analyze +--- + +# Working with the knowledge base + +Check the knowledge base before answering from memory whenever the question +could be about the user's documents. Say so when it has nothing relevant. + +## Find + +`search_documents` is the first call. Results come best first with the document +title, section headings and the matching passage. `filter` restricts which +documents are searched, `limit` how many results come back. If it misses, +rephrase once or narrow with a filter before concluding the material is not +there. + +## Read + +Every search result shows its `Document ID` (and `Collection` when there are +several); pass them to the read tools. `get_document` returns a document's +whole text in reading order. For a long one, `get_document_outline` gives the +heading tree with page numbers and `get_document_section` the text of one +section, subsections included. + +## Answer or compute + +`ask_question` runs the RAG agent on the server and returns an answer with +citations; use it when the user wants an answer rather than material. +`analyze` runs code in a sandbox over the documents; use it for counting, +aggregation, comparison across many documents or computation over tables. Both +cost a model call and are slower than a search. + +## Explore + +`list_documents` shows what is stored: titles, URIs and metadata. It is how you +learn what a filter can match. + +## Filters + +A SQL WHERE clause over the document columns `id`, `uri`, `title`, +`created_at`, `updated_at`, `metadata`. `metadata` is a JSON string, so match +it with LIKE: `metadata LIKE '%"author": "Smith"%'`. Also `uri LIKE '%.pdf'`, +`title = 'Q3 report'`. + +## Results and citations + +Rank is the signal; scores are not comparable across queries and are never +confidence. Cite the document title or URI, the section heading and page +numbers when present. When results carry `source`, the server covers several +collections: name it, and pass `sources` to search a subset. diff --git a/docs/mcp.md b/docs/mcp.md index 59093b3b..cbf08999 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -40,6 +40,25 @@ 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 Code + +The repository ships a plugin that registers the server and a skill telling +Claude when and how to use it: + +```bash +claude plugin marketplace add ggozad/haiku.rag +claude plugin install haiku-rag +``` + +The plugin runs `haiku-rag mcp --stdio`, so `haiku-rag` must be on the PATH +and the configuration decides the database. The skill pre-approves every tool +and is also invocable as `/haiku-rag`. To register the server without the +plugin: + +```bash +claude mcp add haiku-rag -- haiku-rag mcp --stdio +``` + ## Claude Desktop Integration Add to your Claude Desktop configuration (`claude_desktop_config.json`): diff --git a/scripts/bump_version.py b/scripts/bump_version.py index 5c5a1e2e..9593cd89 100755 --- a/scripts/bump_version.py +++ b/scripts/bump_version.py @@ -2,7 +2,8 @@ """ Version bumping script for haiku.rag workspace. -Updates version in all pyproject.toml files and CHANGELOG.md. +Updates version in all pyproject.toml files, the Claude Code plugin manifest +and CHANGELOG.md. """ import re @@ -54,6 +55,19 @@ def update_example_dependencies(file_path: Path, new_version: str) -> None: print(f"✓ Updated example dependencies in {file_path.relative_to(Path.cwd())}") +def update_plugin_version(file_path: Path, new_version: str) -> None: + """Update the version in the Claude Code plugin manifest.""" + content = file_path.read_text() + updated = re.sub( + r'^(\s*"version": )"[^"]+"', + rf'\1"{new_version}"', + content, + flags=re.MULTILINE, + ) + file_path.write_text(updated) + print(f"✓ Updated {file_path.relative_to(Path.cwd())}") + + def update_changelog(changelog_path: Path, new_version: str) -> None: """Update CHANGELOG.md with new version.""" content = changelog_path.read_text() @@ -122,10 +136,13 @@ def main(): root / "app" / "backend" / "pyproject.toml", ] + plugin_file = root / "claude-plugin" / ".claude-plugin" / "plugin.json" changelog_file = root / "CHANGELOG.md" # Check all files exist - for file in pyproject_files + example_pyproject_files + [changelog_file]: + for file in ( + pyproject_files + example_pyproject_files + [plugin_file, changelog_file] + ): if not file.exists(): print(f"Error: {file} not found") sys.exit(1) @@ -155,6 +172,8 @@ def main(): for file in example_pyproject_files: update_example_dependencies(file, new_version) + update_plugin_version(plugin_file, new_version) + # Update CHANGELOG.md update_changelog(changelog_file, new_version) diff --git a/tests/test_bump_version.py b/tests/test_bump_version.py new file mode 100644 index 00000000..4f8fcf4f --- /dev/null +++ b/tests/test_bump_version.py @@ -0,0 +1,38 @@ +import importlib.util +import json +from pathlib import Path + +_spec = importlib.util.spec_from_file_location( + "bump_version", Path(__file__).resolve().parents[1] / "scripts" / "bump_version.py" +) +assert _spec is not None and _spec.loader is not None +bump_version = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(bump_version) + + +def test_update_plugin_version_rewrites_only_the_version_field(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + manifest = tmp_path / "plugin.json" + manifest.write_text( + '{\n "name": "haiku-rag",\n "version": "0.1.0",\n "license": "MIT"\n}\n' + ) + + bump_version.update_plugin_version(manifest, "0.2.0") + + assert json.loads(manifest.read_text()) == { + "name": "haiku-rag", + "version": "0.2.0", + "license": "MIT", + } + assert manifest.read_text().endswith("}\n") + + +def test_the_shipped_plugin_manifest_carries_the_package_version(): + root = Path(__file__).resolve().parents[1] + plugin = json.loads( + (root / "claude-plugin" / ".claude-plugin" / "plugin.json").read_text() + ) + + assert plugin["version"] == bump_version.get_current_version( + root / "haiku_rag_slim" / "pyproject.toml" + ) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 7bf3183d..4f1160b8 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import logging +from pathlib import Path from types import SimpleNamespace import pytest @@ -1122,6 +1123,48 @@ class TestMCPErrorContract: ) +class TestClaudeCodePlugin: + """The plugin under claude-plugin/ points at the server this module builds.""" + + root = Path(__file__).resolve().parents[1] + + def test_the_manifests_name_the_plugin_and_its_server(self): + import json + + plugin = json.loads( + (self.root / "claude-plugin/.claude-plugin/plugin.json").read_text() + ) + marketplace = json.loads( + (self.root / ".claude-plugin/marketplace.json").read_text() + ) + servers = json.loads((self.root / "claude-plugin/.mcp.json").read_text()) + + assert plugin["name"] == "haiku-rag" + assert plugin["description"] + [entry] = marketplace["plugins"] + assert entry["name"] == plugin["name"] + assert entry["source"] == "./claude-plugin" + assert servers["mcpServers"]["haiku-rag"]["args"] == ["mcp", "--stdio"] + + @pytest.mark.asyncio + async def test_the_skill_pre_approves_every_tool_the_server_registers( + self, mcp_db, multimodal_embedder + ): + import yaml + + text = (self.root / "claude-plugin/skills/haiku-rag/SKILL.md").read_text() + _, frontmatter, _ = text.split("---", 2) + skill = yaml.safe_load(frontmatter) + prefix = "mcp__plugin_haiku-rag_haiku-rag__" + + assert skill["name"] == "haiku-rag" + assert skill["description"] + assert all(tool.startswith(prefix) for tool in skill["allowed-tools"]) + approved = {tool.removeprefix(prefix) for tool in skill["allowed-tools"]} + registered = {t.name for t in await create_mcp_server(mcp_db).list_tools()} + assert approved == registered + + class TestMCPClientLifetime: @pytest.mark.asyncio async def test_tool_calls_share_one_database_open(self, mcp_db, monkeypatch): From 494294046f82901c05d0173300786514759f62bd Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 13:11:06 +0300 Subject: [PATCH 10/18] Move to fastmcp 4 and MCP SDK 2 fastmcp>=4.0.2,<5.0.0. Protocol types are snake_case (read_only_hint, mime_type, input_schema, server_info); the camelCase names warn and go in fastmcp 5, so the next major is an explicit upgrade. The client defaults to the sessionless protocol, where initialize_result is None, so the tests read client.instructions and client.server_info, which both protocol modes populate. The server answers either mode. --- CHANGELOG.md | 2 ++ haiku_rag_slim/haiku/rag/mcp.py | 4 +-- haiku_rag_slim/pyproject.toml | 2 +- tests/test_mcp.py | 28 ++++++++------- uv.lock | 62 +++++++++++++++++++-------------- 5 files changed, 55 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 41411c13..848bf844 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,8 @@ ### Changed +- `fastmcp>=4.0.2,<5.0.0`, on MCP Python SDK 2. The MCP server answers both the + session-based and the sessionless (2026-07-28) protocol. - Default models are `ollama:qwen3.8`: `ModelConfig`, `qa.model`, `processing.title_model` (was `ollama:gpt-oss`) and `processing.conversion_options.picture_description.model` (was diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 5da35d47..b3e8d680 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -50,7 +50,7 @@ Sources = Annotated[ def _read_only(title: str) -> ToolAnnotations: - return ToolAnnotations(title=title, readOnlyHint=True, openWorldHint=False) + return ToolAnnotations(title=title, read_only_hint=True, open_world_hint=False) def _decode_image(image_base64: str) -> bytes: @@ -150,7 +150,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe ImageContent( type="image", data=base64.b64encode(picture.data).decode("ascii"), - mimeType="image/png", + mime_type="image/png", ) ) return ToolResult( diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index e438dbfd..b3e5dbc6 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "docling-core>=2.82.0,<3.0.0", "httpx>=0.28.1", "jinja2>=3.1.0", - "fastmcp>=3.3.0", + "fastmcp>=4.0.2,<5.0.0", "lancedb==0.37.1", "pathspec>=1.0.4", "pydantic>=2.12.5", diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 4f1160b8..002723cc 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -506,7 +506,7 @@ class TestMCPSearchResultShape: images = [block for block in rest if isinstance(block, ImageContent)] labels = [block.text for block in rest if isinstance(block, TextContent)] assert len(images) == 2 - assert all(image.mimeType == "image/png" for image in images) + assert all(image.mime_type == "image/png" for image in images) assert [ label for label in labels if "[c1]" in label and "#/pictures/0" in label ] @@ -610,10 +610,12 @@ class TestMCPDescribesItself: from fastmcp import Client async with Client(create_mcp_server(mcp_db)) as client: - init = client.initialize_result + instructions = client.instructions + server_info = client.server_info - assert init.instructions - assert init.serverInfo.version == metadata.version("haiku.rag-slim") + assert instructions + assert server_info is not None + assert server_info.version == metadata.version("haiku.rag-slim") @pytest.mark.asyncio async def test_instructions_name_the_collections_when_covering_several( @@ -624,10 +626,10 @@ class TestMCPDescribesItself: from haiku.rag.client.scope import DatabaseScope async with Client(_covering_all(two_dbs)) as client: - covering_both = client.initialize_result.instructions + covering_both = client.instructions one = DatabaseScope.resolve(two_dbs, database_name="alpha") async with Client(_mcp_covering(one, two_dbs)) as client: - covering_one = client.initialize_result.instructions + covering_one = client.instructions assert "alpha" in covering_both assert "beta" in covering_both @@ -638,9 +640,9 @@ class TestMCPDescribesItself: from fastmcp import Client async with Client(create_mcp_server(mcp_db)) as client: - full = client.initialize_result.instructions.splitlines() + full = client.instructions.splitlines() async with Client(create_mcp_server(mcp_db, agents=False)) as client: - without = client.initialize_result.instructions.splitlines() + without = client.instructions.splitlines() assert set(without) < set(full) assert len(without) == len(full) - 1 @@ -655,9 +657,9 @@ class TestMCPDescribesItself: config.prompts.domain_preamble = "Everything here is about zebras." async with Client(create_mcp_server(mcp_db, config=config)) as client: - with_preamble = client.initialize_result.instructions + with_preamble = client.instructions async with Client(create_mcp_server(mcp_db)) as client: - without = client.initialize_result.instructions + without = client.instructions assert "Everything here is about zebras." in with_preamble assert "zebras" not in without @@ -672,8 +674,8 @@ class TestMCPDescribesItself: assert len(tools) == 8 for tool in tools: assert tool.annotations is not None, tool.name - assert tool.annotations.readOnlyHint is True, tool.name - assert tool.annotations.openWorldHint is False, tool.name + assert tool.annotations.read_only_hint is True, tool.name + assert tool.annotations.open_world_hint is False, tool.name assert tool.annotations.title, tool.name @pytest.mark.asyncio @@ -686,7 +688,7 @@ class TestMCPDescribesItself: undescribed = [ f"{tool.name}.{name}" for tool in tools - for name, schema in tool.inputSchema.get("properties", {}).items() + for name, schema in tool.input_schema.get("properties", {}).items() if not schema.get("description") ] assert len(tools) == 8 diff --git a/uv.lock b/uv.lock index 2754fbb2..4d9636e2 100644 --- a/uv.lock +++ b/uv.lock @@ -1209,21 +1209,22 @@ wheels = [ [[package]] name = "fastmcp" -version = "3.3.1" +version = "4.0.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "fastmcp-slim", extra = ["client", "server"] }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3b/a9/5c5a01b6abd5346bf60b97cfd29e4a86661940c27dd562bfcda07fd03519/fastmcp-3.3.1.tar.gz", hash = "sha256:979362ea557de42a5f40342563c7e4b236bcc8e7cd192715f50030695d1a71cd", size = 28681699, upload-time = "2026-05-15T15:50:39.673Z" } +sdist = { url = "https://files.pythonhosted.org/packages/37/1c/981a1854f91a08872f4b8b9a627d5d751cafc1340d29b21b747f0b520b0a/fastmcp-4.0.2.tar.gz", hash = "sha256:60d5c5ead3b6a117bfada5c0f95fe5c1aba53d1577079ecbdf42eeff0cd9b931", size = 42306015, upload-time = "2026-09-02T23:28:08.386Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/11/6b1bdada6ccfe647d615ae63f9106f8136aec17971e9361546af01c7d38e/fastmcp-3.3.1-py3-none-any.whl", hash = "sha256:862440c5c4d281363a5995eee59d77f0f7cac1f18869038729cecf03b02fc522", size = 7903, upload-time = "2026-05-15T15:50:36.424Z" }, + { url = "https://files.pythonhosted.org/packages/58/3f/b97cfb92e0d6db8232c67c258117cd0dd9def86c8b472270bd7196d5cd9d/fastmcp-4.0.2-py3-none-any.whl", hash = "sha256:9075e64a94634ad660971ed14374c87be06f2a16a921028ca87987e6aa2f3bfa", size = 8078, upload-time = "2026-09-02T23:28:03.777Z" }, ] [[package]] name = "fastmcp-slim" -version = "3.3.1" +version = "4.0.2" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "mcp-types" }, { name = "platformdirs" }, { name = "pydantic", extra = ["email"] }, { name = "pydantic-settings" }, @@ -1231,26 +1232,28 @@ dependencies = [ { name = "rich" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d1/a0/627103e517e1d0d6f1eec633d5662d13e776f01b45ad188e4f5f7478b438/fastmcp_slim-3.3.1.tar.gz", hash = "sha256:0957835fc59452e143ab2f4b7836d2d2df9b2d9958408edc79ba8b56232b2a88", size = 567007, upload-time = "2026-05-15T15:50:10.426Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/7d/c2597734e3a0859d62c9d8f6f35067d1e296537512e280db3af19204be64/fastmcp_slim-4.0.2.tar.gz", hash = "sha256:86b99bdcb872b52d964c79bc6d43ce79f40ed5538b589d102792b4a7cf3947f4", size = 684052, upload-time = "2026-09-02T23:27:39.868Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/ee/97047f4cc2d7b1d46670d08d8ad01a96e7a748cc01c0b4b351ad8eddbc7a/fastmcp_slim-3.3.1-py3-none-any.whl", hash = "sha256:6cf1c2d77e3adb0d409d6825ed6b0b2a999062973e00b8eea03bd48bf9b4c043", size = 738644, upload-time = "2026-05-15T15:50:08.336Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c0/c022eba3a25ebb56111de1b5f76fbfca81925def58880464263582acfcd8/fastmcp_slim-4.0.2-py3-none-any.whl", hash = "sha256:6bd5b5885628f73263fa2247ea1d26e4a514499a6e079ee3e340cd03a7fe5ed8", size = 858100, upload-time = "2026-09-02T23:27:38.459Z" }, ] [package.optional-dependencies] client = [ { name = "authlib" }, { name = "exceptiongroup" }, - { name = "httpx" }, + { name = "httpx2" }, { name = "mcp" }, { name = "opentelemetry-api" }, { name = "py-key-value-aio", extra = ["filetree", "keyring", "memory"] }, + { name = "starlette" }, ] server = [ { name = "authlib" }, { name = "cyclopts" }, { name = "exceptiongroup" }, { name = "griffelib" }, - { name = "httpx" }, + { name = "httpx2" }, + { name = "joserfc" }, { name = "jsonref" }, { name = "jsonschema-path" }, { name = "mcp" }, @@ -1261,6 +1264,7 @@ server = [ { name = "pyperclip" }, { name = "python-multipart" }, { name = "pyyaml" }, + { name = "starlette" }, { name = "uncalled-for" }, { name = "uvicorn" }, { name = "watchfiles" }, @@ -1755,7 +1759,7 @@ requires-dist = [ { name = "docling", marker = "extra == 'docling'", specifier = ">=2.102.2,<3.0.0" }, { name = "docling-core", specifier = ">=2.82.0,<3.0.0" }, { name = "fastapi", marker = "extra == 'ingester'", specifier = ">=0.125" }, - { name = "fastmcp", specifier = ">=3.3.0" }, + { name = "fastmcp", specifier = ">=4.0.2,<5.0.0" }, { name = "haiku-rag-slim", extras = ["s3"], marker = "extra == 'ingester'", editable = "haiku_rag_slim" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "jinja2", specifier = ">=3.1.0" }, @@ -1887,15 +1891,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, ] -[[package]] -name = "httpx-sse" -version = "0.4.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, -] - [[package]] name = "httpx2" version = "2.8.0" @@ -2590,15 +2585,15 @@ wheels = [ [[package]] name = "mcp" -version = "1.28.1" +version = "2.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, + { name = "httpx2" }, { name = "jsonschema" }, + { name = "mcp-types" }, + { name = "opentelemetry-api" }, { name = "pydantic" }, - { name = "pydantic-settings" }, { name = "pyjwt", extra = ["crypto"] }, { name = "python-multipart" }, { name = "pywin32", marker = "sys_platform == 'win32'" }, @@ -2608,9 +2603,22 @@ dependencies = [ { name = "typing-inspection" }, { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6e/77/9450b8f251a13affb6281997d0523c4615f8a8b35d0b21ff30db3a5aac9d/mcp-1.28.1.tar.gz", hash = "sha256:d51e36a5f5644faea4f85ea649bfffa6bc6c26770d42798ad6a3de3d2ba69683", size = 638501, upload-time = "2026-06-26T12:57:29.093Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/6e/21fb8e5d579dbe21d96ea4d5034200d46d8bdf2261053b5bd041f3c2f612/mcp-2.1.1.tar.gz", hash = "sha256:50b7ba1ebbe117008ea7bdd288234043e69c20b403d6851d19661e6d431a75ef", size = 3984589, upload-time = "2026-08-25T16:14:02.376Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/5e/d118fce19f87a2e7d8101c35c8ae0ec289098a4df0ff244cec23e415aca0/mcp-1.28.1-py3-none-any.whl", hash = "sha256:2726bca5e7193f61c5dde8b12500a6de2d9acf6d1a1c0be9e8c2e706437991df", size = 222620, upload-time = "2026-06-26T12:57:27.218Z" }, + { url = "https://files.pythonhosted.org/packages/50/af/8644cc5fa26a59afd2df2e98eeb19e72926887fa4b7441aba4ff661140db/mcp-2.1.1-py3-none-any.whl", hash = "sha256:1c6c31c5d6471c58db76af3af8af67f46d11d01f0a59077d0a308cbdb3d3e915", size = 357912, upload-time = "2026-08-25T16:13:59.024Z" }, +] + +[[package]] +name = "mcp-types" +version = "2.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6a/dd/1c4417dc0b722c23a1669032d5f044e41170fe5d4773b488a50fcce98c32/mcp_types-2.1.1.tar.gz", hash = "sha256:77dcbe48fba73cca71a673f2646a5f037a017b7a0a07ac89cec1113028890eda", size = 66674, upload-time = "2026-08-25T16:14:03.861Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/d0/242e63c510f4a17381f55b1549a3f94f5687a0595984febd2b6f87a687a0/mcp_types-2.1.1-py3-none-any.whl", hash = "sha256:26f9f7f03f2a5730717a5b98e2ab7eb640ac352d05a00cdc725c311864778295", size = 69656, upload-time = "2026-08-25T16:14:00.667Z" }, ] [[package]] @@ -5543,11 +5551,11 @@ wheels = [ [[package]] name = "uncalled-for" -version = "0.2.0" +version = "0.4.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/02/7c/b5b7d8136f872e3f13b0584e576886de0489d7213a12de6bebf29ff6ebfc/uncalled_for-0.2.0.tar.gz", hash = "sha256:b4f8fdbcec328c5a113807d653e041c5094473dd4afa7c34599ace69ccb7e69f", size = 49488, upload-time = "2026-02-27T17:40:58.137Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6b/5a/92ce0b3ea5481915f55da994c2c2c5f7a3c09949afde196ee89f8ab961aa/uncalled_for-0.4.0.tar.gz", hash = "sha256:335b95bd2422332ec210d518f314a16e4c640921c39fc8bf2ad095bd3538f4af", size = 56979, upload-time = "2026-08-10T14:51:46.247Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/7f/4320d9ce3be404e6310b915c3629fe27bf1e2f438a1a7a3cb0396e32e9a9/uncalled_for-0.2.0-py3-none-any.whl", hash = "sha256:2c0bd338faff5f930918f79e7eb9ff48290df2cb05fcc0b40a7f334e55d4d85f", size = 11351, upload-time = "2026-02-27T17:40:56.804Z" }, + { url = "https://files.pythonhosted.org/packages/a2/40/97cec87c077eb3291fc7905e6633e08b7ca593c57d30238444bcb6bb3d53/uncalled_for-0.4.0-py3-none-any.whl", hash = "sha256:16c4bb3337532e4bd5569adc192285976f3ad5305402256d34c67a12b5c968bd", size = 15502, upload-time = "2026-08-10T14:51:45.068Z" }, ] [[package]] From 2582f2c05acdda5690b4b9c323a46417bc78e769 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 14:12:12 +0300 Subject: [PATCH 11/18] Expand MCP search results and render the matched chunk's metadata Both search tools pass their results through HaikuRAG.expand_context, as every other consumer of search results already did, so a client reads the hit in its section rather than the chunk that matched. The rendering gains an opt-in include_chunk_meta that shows the metadata stored with the matched chunk beyond haiku.rag's structural keys, labelled as the matched chunk's because an expanded passage spans several chunks and only the anchor's metadata survives expansion. The capabilities' rendering is unchanged. Refs #599 --- CHANGELOG.md | 13 ++++---- claude-plugin/skills/haiku-rag/SKILL.md | 3 +- docs/mcp.md | 4 ++- haiku_rag_slim/haiku/rag/mcp.py | 20 +++++++----- .../haiku/rag/store/models/chunk.py | 20 ++++++++++-- tests/test_chunk.py | 31 +++++++++++++++++++ tests/test_mcp.py | 20 +++++++++++- 7 files changed, 92 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 848bf844..9716736d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,12 +32,13 @@ - `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`. -- MCP `search_documents` and `search_documents_by_image` return the agent - rendering as text (rank, `Document ID`, `Collection` over several - databases, title, headings, passage), pictures as `ImageContent` blocks, - and the `SearchResult` list without `image_data` as structured content. - `SearchResult.format_for_agent(include_document_id=)`; - `collect_pictures` in `haiku.rag.tools.search`. +- MCP `search_documents` and `search_documents_by_image` expand results to + their section (`HaikuRAG.expand_context`) and return the agent rendering + as text (rank, `Document ID`, `Collection` over several databases, title, + headings, the matched chunk's metadata, passage), pictures as + `ImageContent` blocks, and the `SearchResult` list without `image_data` as + structured content. `SearchResult.format_for_agent(include_document_id=, + include_chunk_meta=)`; `collect_pictures` in `haiku.rag.tools.search`. - MCP tools raise on failure; an empty result no longer doubles as an error. Unknown document, unknown collection, invalid filter and invalid base64 carry a message; `ask_question` and `analyze` failures name the exception diff --git a/claude-plugin/skills/haiku-rag/SKILL.md b/claude-plugin/skills/haiku-rag/SKILL.md index f93631a3..a0a0d57b 100644 --- a/claude-plugin/skills/haiku-rag/SKILL.md +++ b/claude-plugin/skills/haiku-rag/SKILL.md @@ -24,7 +24,8 @@ could be about the user's documents. Say so when it has nothing relevant. ## Find `search_documents` is the first call. Results come best first with the document -title, section headings and the matching passage. `filter` restricts which +title, section headings, the matched chunk's metadata when it has any, and the +passage in its section. `filter` restricts which documents are searched, `limit` how many results come back. If it misses, rephrase once or narrow with a filter before concluding the material is not there. diff --git a/docs/mcp.md b/docs/mcp.md index cbf08999..ab3b8aa7 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -109,7 +109,9 @@ repeating it. `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 rank, `Document ID`, `Collection` when the server covers several, the document -title, section headings and the passage. Pictures in the results follow as +title, section headings, the matched chunk's metadata when it has any, and the +passage expanded to its section the way the agents get it +(`search.max_context_chars` caps it). Pictures in the results follow as image blocks, one per distinct picture, each preceded by a line naming its result; `include_images: false` leaves them out. The structured content is the `SearchResult` list without picture bytes. Scores are not comparable across diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index b3e8d680..c80ea1ec 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -119,9 +119,9 @@ def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> st def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult: - """Results as the in-process agents read them, then each distinct picture - as an image block labelled with its result, and the results as structured - content without the picture bytes.""" + """Results as the in-process agents read them, plus the matched chunk's + metadata, then each distinct picture as an image block labelled with its + result, and the results as structured content without the picture bytes.""" import base64 total = len(results) @@ -131,6 +131,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe total=total, include_collection=covers_multiple, include_document_id=True, + include_chunk_meta=True, ) for rank, result in enumerate(results, 1) ) @@ -274,9 +275,10 @@ def _covering( Use this first for any question the documents might answer; it needs no model and is the cheapest call. Results come best first, each with its rank, `Document ID`, `Collection` when the server covers several, - the document title, section headings and the matching passage; pass - the id and collection to the document tools. Pictures in the results - follow as images, each labelled with its result. Ranks, not scores, + the document title, section headings, the matched chunk's metadata + when it has any, and the matching passage expanded to its section; + pass the id and collection to the document tools. Pictures in the + results follow as images, each labelled with its result. Ranks, not scores, are the signal: scores are not comparable across queries. If nothing relevant comes back, rephrase once or narrow with `filter` before concluding the material is absent. @@ -300,7 +302,7 @@ def _covering( ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - return _search_result(results, rag.covers_multiple) + return _search_result(await rag.expand_context(results), rag.covers_multiple) # Image-as-query tool, only registered when the configured embedder # supports image embeddings. Probed at server-build time when no Store is @@ -345,7 +347,9 @@ def _covering( ) except UnknownDatabaseError as e: raise ToolError(str(e)) from e - return _search_result(results, rag.covers_multiple) + return _search_result( + await rag.expand_context(results), rag.covers_multiple + ) @mcp.tool(annotations=_read_only("Get document")) async def get_document(document_id: str, source: str | None = None) -> Document: diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index b21231e1..fc48966f 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -1,3 +1,4 @@ +import json from typing import TYPE_CHECKING, Literal from pydantic import BaseModel, PrivateAttr @@ -143,8 +144,9 @@ class SearchResult(BaseModel): consumers (UIs). Never part of ``format_for_agent`` output. ``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not - include the metadata of any other chunks merged with it. Never part of - ``format_for_agent`` output. + include the metadata of any other chunks merged with it. Left out of + ``format_for_agent`` output unless ``include_chunk_meta`` asks for its custom + keys. ``source`` names the database a result came from: the name from ``lancedb.databases`` or a path's stem, never a path or URI, so a location @@ -203,6 +205,7 @@ class SearchResult(BaseModel): *, include_collection: bool = False, include_document_id: bool = False, + include_chunk_meta: bool = False, ) -> str: """Format this search result for inclusion in agent context. @@ -218,6 +221,9 @@ class SearchResult(BaseModel): search spanning one collection has nothing to distinguish, whether or not that collection is named. `include_document_id` is for a reader that will fetch the document by id from the text alone. + `include_chunk_meta` renders the metadata stored with the matched + chunk beyond haiku.rag's own structural keys; on an expanded result it + locates the hit, not the whole passage. """ if rank is not None and total is not None: parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"] @@ -247,6 +253,16 @@ class SearchResult(BaseModel): if primary_label: parts.append(f"Type: {primary_label}") + if include_chunk_meta: + custom = { + key: value + for key, value in self.chunk_meta.items() + if key not in ChunkMetadata.model_fields + } + if custom: + rendered = json.dumps(custom, ensure_ascii=False, sort_keys=True) + parts.append(f"Matched chunk metadata: {rendered}") + # Surface picture captions when present. Order matches the binary # attachments emitted by build_image_content_from_results, so the model # can correlate caption ↔ attached image by position (BinaryContent diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 1bcf6022..a24ff97a 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -264,6 +264,37 @@ def test_search_result_format_for_agent_omits_chunk_meta(): assert "para_no" not in formatted +def test_search_result_format_for_agent_chunk_meta_is_opt_in(): + """A caller that asks sees the chunk's own metadata, never the structural + keys haiku.rag stores beside it.""" + result = SearchResult( + content="Some content.", + score=0.9, + chunk_id="chunk-1", + chunk_meta={ + "para_no": "12", + "doc_item_refs": ["#/texts/0"], + "page_numbers": [1], + "headings": ["Intro"], + "labels": ["paragraph"], + }, + ) + + opted = result.format_for_agent(rank=1, total=1, include_chunk_meta=True) + + assert "para_no" in opted + assert "12" in opted + assert "doc_item_refs" not in opted + assert "#/texts/0" not in opted + + structural_only = result.model_copy( + update={"chunk_meta": {"doc_item_refs": ["#/texts/0"], "page_numbers": [1]}} + ) + assert structural_only.format_for_agent( + rank=1, total=1, include_chunk_meta=True + ) == structural_only.format_for_agent(rank=1, total=1) + + def test_search_result_format_for_agent_omits_document_meta(): """Document metadata is UI plumbing, never shown to the model.""" result = SearchResult( diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 002723cc..ddc46fdb 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -10,7 +10,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 +from tests.multi_db.helpers import _config, _seed, _seed_expandable @pytest.fixture(autouse=True) @@ -181,6 +181,24 @@ class TestMCPReadTools: assert any( r["chunk_meta"] == {"fake-metadata-for-testing": "42"} for r in results ) + text = result.content[0].text + assert "fake-metadata-for-testing" in text + assert "42" in text + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") + async def test_search_results_come_expanded(self, tmp_path): + """The passage is the hit in its section, as the in-process agents read + it, not the chunk that matched.""" + config = _config(tmp_path, ["alpha"]) + sentences = ["Gardens need water.", "Roses need pruning.", "Tulips need sun."] + await _seed_expandable(config, "alpha", sentences) + + result = await _call(_covering_all(config), "search_documents", query="gardens") + + [hit] = _results(result) + assert all(sentence in hit["content"] for sentence in sentences) + assert all(sentence in result.content[0].text for sentence in sentences) @pytest.mark.asyncio async def test_get_document(self, mcp_db): From 0026142e7d033da067113cd9042397eb71a9afeb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Sep 2026 15:50:24 +0300 Subject: [PATCH 12/18] Give MCP search results one channel and tidy two messages Search results carry text and image blocks only. Claude Code and the Agent SDK do not forward text blocks when structuredContent is present and Desktop forwards both, so sending both either hid the rendering or doubled it. The invalid-filter error keeps the engine's diagnosis and lists our columns instead of lance's internals. format_citations no longer repeats the URI of an untitled document. Refs #599 --- CHANGELOG.md | 12 ++-- docs/mcp.md | 9 ++- haiku_rag_slim/haiku/rag/mcp.py | 21 +++--- haiku_rag_slim/haiku/rag/utils.py | 7 +- tests/test_mcp.py | 109 +++++++++++++++++------------- tests/test_utils.py | 17 +++++ 6 files changed, 107 insertions(+), 68 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9716736d..64343190 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,10 +35,10 @@ - MCP `search_documents` and `search_documents_by_image` expand results to their section (`HaikuRAG.expand_context`) and return the agent rendering as text (rank, `Document ID`, `Collection` over several databases, title, - headings, the matched chunk's metadata, passage), pictures as - `ImageContent` blocks, and the `SearchResult` list without `image_data` as - structured content. `SearchResult.format_for_agent(include_document_id=, - include_chunk_meta=)`; `collect_pictures` in `haiku.rag.tools.search`. + headings, the matched chunk's metadata, passage) and pictures as + `ImageContent` blocks, with no structured content. + `SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`; + `collect_pictures` in `haiku.rag.tools.search`. - MCP tools raise on failure; an empty result no longer doubles as an error. Unknown document, unknown collection, invalid filter and invalid base64 carry a message; `ask_question` and `analyze` failures name the exception @@ -50,6 +50,10 @@ `DocumentInfo.source`; citations name their database when the server covers several. `format_citations(citations, include_source=False)`. +### Fixed + +- MCP citations no longer repeat the URI of an untitled document. + ### Removed - `cite` on the MCP `ask_question` tool; citations are always appended. diff --git a/docs/mcp.md b/docs/mcp.md index ab3b8aa7..df7b4448 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -113,8 +113,9 @@ title, section headings, the matched chunk's metadata when it has any, and the passage expanded to its section the way the agents get it (`search.max_context_chars` caps it). Pictures in the results follow as image blocks, one per distinct picture, each preceded by a line naming its -result; `include_images: false` leaves them out. The structured content is the -`SearchResult` list without picture bytes. Scores are not comparable across +result; `include_images: false` leaves them out. Search results carry no +structured content, so every client shows the model the same text and +images. Scores are not comparable across queries or search types, so rank is the signal. `search_documents_by_image` embeds the query image and searches by vector similarity alone. @@ -128,7 +129,9 @@ which is how a client learns what a filter can match. `ask_question` runs the RAG agent on the server and returns an answer followed by its citations. `analyze` writes and runs Python in a sandbox over the documents, for counting, aggregation and computation across -documents. Both cost a model call. +documents. Both cost a model call. Claude Code moves a call still running +after about two minutes to a background task, which a slow local model can +trigger; `--no-agents` leaves both tools out. ### Filters diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index c80ea1ec..112a156c 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,5 +1,6 @@ import asyncio import logging +import re from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from importlib import metadata @@ -89,7 +90,12 @@ async def _check_filter( try: await selected[0].count_documents(filter=filter) except ValueError as e: - raise ToolError(f"Invalid filter {filter!r}: {e}") from e + # The engine lists its own columns too, lance internals among them. + reason = re.sub(r"\s*Valid fields are .*", "", str(e), flags=re.DOTALL) + raise ToolError( + f"Invalid filter {filter!r}: {reason.rstrip('. ')}. " + f"Columns: {_FILTER_COLUMNS}." + ) from e def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> str: @@ -121,7 +127,8 @@ def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> st def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult: """Results as the in-process agents read them, plus the matched chunk's metadata, then each distinct picture as an image block labelled with its - result, and the results as structured content without the picture bytes.""" + result. No structured content: a client given both shows the model the + JSON and drops the text, or shows both.""" import base64 total = len(results) @@ -154,15 +161,7 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe mime_type="image/png", ) ) - return ToolResult( - content=content, - structured_content={ - "result": [ - result.model_dump(mode="json", exclude={"image_data"}) - for result in results - ] - }, - ) + return ToolResult(content=content) def _node(toc: "dict[str, Any]") -> OutlineNode: diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index b96bf74b..2d71b47c 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -421,11 +421,12 @@ def format_citations(citations: "list[Citation]", include_source: bool = False) if section: location_parts.append(f"Section: {section}") - source = c.document_uri + # The URI is the header when there is no title; do not repeat it. + line = f"{header} {c.document_uri}" if c.document_title else header if location_parts: - source += f" - {', '.join(location_parts)}" + line += f" - {', '.join(location_parts)}" - lines.append(f"{header} {source}") + lines.append(line) for ref in c.picture_refs: lines.append(f"[Figure: {ref}]") lines.append(c.content) diff --git a/tests/test_mcp.py b/tests/test_mcp.py index ddc46fdb..0fada9a3 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,5 @@ import logging +import re from pathlib import Path from types import SimpleNamespace @@ -98,9 +99,25 @@ async def _call(mcp, name, **kwargs): return await client.call_tool(name, kwargs, raise_on_error=False) -def _results(search_result) -> list[dict]: - """The search results a tool returned, as the client sees them.""" - return search_result.structured_content["result"] +def _results(tool_result) -> list[dict]: + """A tool's structured result list, as the client sees it.""" + return tool_result.structured_content["result"] + + +_HEADER = re.compile(r"^\[[^\]]+\] \[rank \d+ of \d+\]$", re.MULTILINE) + + +def _rendered(search_result) -> list[str]: + """The result blocks of a search, split from the text the model reads.""" + text = search_result.content[0].text + starts = [match.start() for match in _HEADER.finditer(text)] + return [text[a:b].strip() for a, b in zip(starts, starts[1:] + [len(text)])] + + +def _line(block: str, name: str) -> str | None: + """The value of a `Name: value` line in a rendered result, if present.""" + match = re.search(rf"^{re.escape(name)}: (.+)$", block, re.MULTILINE) + return match.group(1) if match else None def _png_b64() -> str: @@ -120,17 +137,17 @@ class TestMCPReadTools: mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="artificial intelligence")) - assert len(results) > 0 - assert all(r["chunk_id"] and r["content"] for r in results) + blocks = _rendered(await search(query="artificial intelligence")) + assert blocks + assert all("Content:" in block for block in blocks) @pytest.mark.asyncio async def test_search_documents_with_limit(self, mcp_db): mcp = create_mcp_server(mcp_db) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="artificial intelligence", limit=1)) - assert len(results) == 1 + blocks = _rendered(await search(query="artificial intelligence", limit=1)) + assert len(blocks) == 1 @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @@ -143,18 +160,15 @@ class TestMCPReadTools: {"query": "artificial intelligence", "filter": "title = 'ML Basics'"}, ) - results = result.structured_content["result"] - assert results - assert {r["document_title"] for r in results} == {"ML Basics"} + blocks = _rendered(result) + assert blocks + assert all('"ML Basics"' in block for block in blocks) @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") - async def test_search_documents_preserves_chunk_meta_through_serialization( - self, mcp_db - ): - """Chunk_meta must survive FastMCP's actual wire serialization. - - Calling the tool function directly bypasses that serialization step entirely.""" + async def test_search_documents_carries_the_matched_chunks_metadata(self, mcp_db): + """The chunk's own metadata reaches the text the model reads, over the + wire, without haiku.rag's structural keys.""" from fastmcp import Client async with HaikuRAG(mcp_db, create=True) as rag: @@ -176,14 +190,11 @@ class TestMCPReadTools: "search_documents", {"query": "artificial intelligence"} ) - results = result.structured_content["result"] - assert results - assert any( - r["chunk_meta"] == {"fake-metadata-for-testing": "42"} for r in results - ) text = result.content[0].text assert "fake-metadata-for-testing" in text assert "42" in text + assert "doc_item_refs" not in text + assert result.structured_content is None @pytest.mark.asyncio @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") @@ -196,9 +207,8 @@ class TestMCPReadTools: result = await _call(_covering_all(config), "search_documents", query="gardens") - [hit] = _results(result) - assert all(sentence in hit["content"] for sentence in sentences) - assert all(sentence in result.content[0].text for sentence in sentences) + [hit] = _rendered(result) + assert all(sentence in hit for sentence in sentences) @pytest.mark.asyncio async def test_get_document(self, mcp_db): @@ -531,9 +541,12 @@ class TestMCPSearchResultShape: assert [ label for label in labels if "[c3]" in label and "#/pictures/3" in label ] - structured = _results(result) - assert [r["chunk_id"] for r in structured] == ["c1", "c2", "c3"] - assert all("image_data" not in r for r in structured) + assert [block.split("]")[0] for block in _rendered(result)] == [ + "[c1", + "[c2", + "[c3", + ] + assert result.structured_content is None @pytest.mark.asyncio async def test_an_undecodable_picture_yields_no_image(self, mcp_db, monkeypatch): @@ -566,7 +579,7 @@ class TestMCPSearchResultShape: result = await _call(create_mcp_server(mcp_db), "search_documents", query="q") assert [block.text for block in result.content] == ["No results found."] - assert _results(result) == [] + assert result.structured_content is None @pytest.mark.asyncio async def test_search_text_alone_drives_the_document_tools(self, two_dbs): @@ -600,7 +613,7 @@ class TestMCPSearchResultShape: r"Document ID: (\S+)\nCollection: (\S+)", search.content[0].text ) - assert len(pairs) == len(_results(search)) == 2 + assert len(pairs) == len(_rendered(search)) == 2 assert {source for _, source in pairs} == {"alpha", "beta"} for document_id, source in pairs: outline = await _call( @@ -747,19 +760,19 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="cats")) + blocks = _rendered(await search(query="cats")) - assert {r["source"] for r in results} == {"alpha", "beta"} + assert {_line(block, "Collection") for block in blocks} == {"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 = _results(await search(query="cats", sources=["beta"])) + blocks = _rendered(await search(query="cats", sources=["beta"])) - assert results - assert {r["source"] for r in results} == {"beta"} + assert blocks + assert {_line(block, "Collection") for block in blocks} == {"beta"} @pytest.mark.asyncio @pytest.mark.parametrize( @@ -793,13 +806,13 @@ class TestMCPCoversTheConfiguredSet: mcp = _covering_all(two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results( + blocks = _rendered( await search(query="cats", filter="uri LIKE '%beta%'", sources=["beta"]) ) - assert results - assert {r["source"] for r in results} == {"beta"} + assert blocks + assert {_line(block, "Collection") for block in blocks} == {"beta"} none = await search(query="cats", filter="uri LIKE '%beta%'", sources=[]) - assert _results(none) == [] + assert _rendered(none) == [] @pytest.mark.asyncio async def test_the_listing_covers_every_database(self, two_dbs): @@ -828,9 +841,9 @@ class TestMCPCoversTheConfiguredSet: mcp = create_mcp_server(config=two_dbs) search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="cats")) + blocks = _rendered(await search(query="cats")) - assert {r["source"] for r in results} == {"alpha", "beta"} + assert {_line(block, "Collection") for block in blocks} == {"alpha", "beta"} @pytest.mark.asyncio async def test_ask_question_names_each_citations_database( @@ -924,7 +937,7 @@ class TestMCPImageQuery: sources=[], ) - assert _results(results) == [] + assert _rendered(results) == [] assert seen["query"] == png assert seen["filter"] == "uri LIKE 'x%'" assert seen["sources"] == [] @@ -1036,6 +1049,8 @@ class TestMCPErrorContract: assert result.is_error assert "no_such_column = 1" in result.content[0].text + assert "created_at" in result.content[0].text + assert "_rowid" not in result.content[0].text @pytest.mark.asyncio @pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"]) @@ -1281,12 +1296,12 @@ class TestMCPClientLifetime: mcp = _mcp_covering(scope, config) async with mcp._lifespan_manager(): search = await _get_tool(mcp, "search_documents") - results = _results(await search(query="artificial intelligence")) + blocks = _rendered(await search(query="artificial intelligence")) listing = await _get_tool(mcp, "list_documents") documents = await listing() - assert results - assert {r["source"] for r in results} == {"alpha"} + assert blocks + assert {_line(block, "Collection") for block in blocks} == {None} titles = {d.title for d in documents} assert "AI Overview" in titles assert "Zebras" not in titles @@ -1371,9 +1386,9 @@ class TestMCPClientLifetime: assert opens == 1 async with mcp._lifespan_manager(): - results = _results(await search(query="artificial intelligence")) + blocks = _rendered(await search(query="artificial intelligence")) assert opens == 2 - assert len(results) > 0 + assert blocks @pytest.mark.asyncio async def test_same_dim_drift_starts(self, mcp_db): diff --git a/tests/test_utils.py b/tests/test_utils.py index 5284cb04..9d4cba45 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -769,6 +769,23 @@ def test_format_citations_names_the_source_when_asked(): assert "papers" not in format_citations([citation]) +def test_format_citations_names_an_untitled_document_once(): + 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", + content="Content", + page_numbers=[3], + ) + result = format_citations([citation]) + + assert result.count("test://doc") == 1 + assert "[1] test://doc - p. 3" in result + + # --- format_citations tests (pictures) --- From b374d5eb83e7ffa54925afcca62f0b8a8f8c5d50 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 7 Sep 2026 10:02:42 +0300 Subject: [PATCH 13/18] Replace ask_question and analyze with execute_code In Claude Code the client is the model, so the server no longer runs one. execute_code runs a Python program per call in the analysis sandbox over the selected documents and returns what it printed; the sandbox is created and closed per call so Monty's cumulative budget and a frozen mount never outlive a program. --no-agents goes with the two tools, and format_citations in haiku.rag.utils goes with its only caller. The sandbox exposes chunk metadata to code: chunk_meta on search results, metadata on list_documents rows and in metadata.json, and chunks.jsonl per document. A host-side failure inside a program, a document read or an in-code search raising, reaches the program by exception type only and is logged with its traceback. recovery_hint moves to haiku.rag.sandbox. Closes #604. --- CHANGELOG.md | 32 +- README.md | 2 +- claude-plugin/skills/haiku-rag/SKILL.md | 21 +- docs/capabilities/analysis.md | 2 +- docs/cli.md | 3 - docs/configuration/qa.md | 2 +- docs/mcp.md | 33 +- haiku_rag_slim/haiku/rag/app.py | 3 +- .../haiku/rag/capabilities/analysis.py | 19 +- .../rag/capabilities/instructions/analysis.md | 12 +- haiku_rag_slim/haiku/rag/cli.py | 9 +- haiku_rag_slim/haiku/rag/mcp.py | 141 +++----- haiku_rag_slim/haiku/rag/sandbox/__init__.py | 3 +- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 95 ++++- haiku_rag_slim/haiku/rag/utils.py | 42 --- tests/sandbox/test_sandbox.py | 231 +++++++++++- tests/sandbox/test_sandbox_toc.py | 33 ++ tests/test_app.py | 14 - tests/test_cli.py | 8 - tests/test_mcp.py | 336 ++++++++---------- tests/test_utils.py | 160 +-------- 21 files changed, 631 insertions(+), 570 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 64343190..fc815f0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,8 +7,13 @@ - Claude Code plugin under `claude-plugin/`: the server configuration and the `haiku-rag` skill. `claude plugin marketplace add ggozad/haiku.rag`, then `claude plugin install haiku-rag`. -- `haiku-rag mcp --no-agents` leaves `ask_question` and `analyze` - unregistered. `create_mcp_server(agents=)`, `HaikuRAGApp.run_mcp(agents=)`. +- MCP tool `execute_code(code, filter, sources)`: runs a program in the + analysis sandbox over the selected documents and returns what it printed; + one sandbox per call. +- In the analysis sandbox, `search()` results carry `chunk_meta`, + `list_documents()` rows and `metadata.json` carry the document `metadata`, + and `/documents/{id}/chunks.jsonl` lists chunk ids with their metadata. + `recovery_hint` in `haiku.rag.sandbox`. - 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`. @@ -40,23 +45,20 @@ `SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`; `collect_pictures` in `haiku.rag.tools.search`. - MCP tools raise on failure; an empty result no longer doubles as an error. - Unknown document, unknown collection, invalid filter and invalid base64 - carry a message; `ask_question` and `analyze` failures name the exception - type. Anything else is masked (`mask_error_details=True`) and logged - server-side. + Unknown document, unknown collection, invalid filter, invalid base64 and a + failing program carry a message. Anything else is masked + (`mask_error_details=True`) and logged server-side. +- A host-side failure inside the analysis sandbox (a document read or an + in-code `search()` raising) reaches the program as + `RuntimeError(" failed: ")`; the traceback is logged. - `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)`. - -### Fixed - -- MCP citations no longer repeat the URI of an untitled document. + `search_documents`, `search_documents_by_image` and `execute_code`; `source` + on `get_document`; an unknown name is a tool error. `DocumentInfo.source`. ### Removed -- `cite` on the MCP `ask_question` tool; citations are always appended. +- MCP tools `ask_question` and `analyze`. +- `format_citations` in `haiku.rag.utils`; `format_citations_rich` stays. - MCP write tools `add_document_from_file`, `add_document_from_url`, `add_document_from_text` and `delete_document`. The server opens the database read-only; ingest with `haiku-rag add`, `add-src`, `delete` or diff --git a/README.md b/README.md index fb3c8d5b..24d832e3 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/ - **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion - **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query - **Question answering** — RAG capability with citations (page numbers, section headings) -- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text; attach your own images to questions in `ask`, `analyze`, MCP, and the chat TUI +- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text; attach your own images to questions in `ask`, `analyze` and the chat TUI - **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM - **Analysis capability** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis) - **Evidence compaction** — Optional capability that replaces earlier questions' search results on the request with the evidence they cited, so long conversations stop resending everything they retrieved diff --git a/claude-plugin/skills/haiku-rag/SKILL.md b/claude-plugin/skills/haiku-rag/SKILL.md index a0a0d57b..981baf43 100644 --- a/claude-plugin/skills/haiku-rag/SKILL.md +++ b/claude-plugin/skills/haiku-rag/SKILL.md @@ -1,6 +1,6 @@ --- name: haiku-rag -description: Search, read and question the user's haiku.rag knowledge base +description: Search, read and compute over the user's haiku.rag knowledge base through the haiku-rag MCP tools. Use whenever a request could be answered from the user's ingested documents, when asked to find, look up, check or cite something in their documents or knowledge base, or when the question is @@ -12,8 +12,7 @@ allowed-tools: - mcp__plugin_haiku-rag_haiku-rag__get_document_outline - mcp__plugin_haiku-rag_haiku-rag__get_document_section - mcp__plugin_haiku-rag_haiku-rag__list_documents - - mcp__plugin_haiku-rag_haiku-rag__ask_question - - mcp__plugin_haiku-rag_haiku-rag__analyze + - mcp__plugin_haiku-rag_haiku-rag__execute_code --- # Working with the knowledge base @@ -38,13 +37,17 @@ whole text in reading order. For a long one, `get_document_outline` gives the heading tree with page numbers and `get_document_section` the text of one section, subsections included. -## Answer or compute +## Compute -`ask_question` runs the RAG agent on the server and returns an answer with -citations; use it when the user wants an answer rather than material. -`analyze` runs code in a sandbox over the documents; use it for counting, -aggregation, comparison across many documents or computation over tables. Both -cost a model call and are slower than a search. +`execute_code` runs a Python program on the server over the same documents. +Under `/documents/{id}/` each has `metadata.json`, `content.txt`, `items.jsonl`, +`chunks.jsonl` and `toc.json`, and the program can `await search(query)` and +`await list_documents()`. Write code when the answer is a count, an aggregate, a +comparison across many documents, a lookup by document or chunk metadata, or a +pattern over whole documents: whatever search cannot rank. Each call is one +program and variables do not carry over, so gather, compute and `print` a +compact result in the same program. `filter` and `sources` select the documents +it sees. Answer and cite from what it printed. ## Explore diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md index b3a58138..604eb3d3 100644 --- a/docs/capabilities/analysis.md +++ b/docs/capabilities/analysis.md @@ -16,7 +16,7 @@ When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool | `analysis_execute_code(code)` | Run Python against the virtual document filesystem. | | `analysis_cite(chunk_ids)` | Register retrieved or filesystem-derived chunk IDs. | -The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`. +The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, `chunks.jsonl` (chunk ids with their metadata) and `toc.json`. In code, `await search()` results carry `chunk_meta` and `await list_documents()` rows carry `metadata`. ## Compose an agent diff --git a/docs/cli.md b/docs/cli.md index b54cb8fe..5e0efae8 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -477,9 +477,6 @@ 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/configuration/qa.md b/docs/configuration/qa.md index 7b059bd8..32c4c34b 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -20,7 +20,7 @@ Context expansion is automatic and section-aware. For structured documents (with ## Question Answering Configuration -Configure the RAG capability (used by `client.ask`, `haiku-rag ask`, and the MCP `ask_question` tool): +Configure the RAG capability (used by `client.ask` and `haiku-rag ask`): ```yaml qa: diff --git a/docs/mcp.md b/docs/mcp.md index df7b4448..9fef2c7d 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -19,8 +19,6 @@ 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 @@ -59,6 +57,10 @@ plugin: claude mcp add haiku-rag -- haiku-rag mcp --stdio ``` +The skill works with that registration too: copy `claude-plugin/skills/haiku-rag` +into `~/.claude/skills/` and change the tool prefix in its `allowed-tools` from +`mcp__plugin_haiku-rag_haiku-rag__` to `mcp__haiku-rag__`. + ## Claude Desktop Integration Add to your Claude Desktop configuration (`claude_desktop_config.json`): @@ -103,8 +105,7 @@ 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` | unless `--no-agents` | `question`, `images_base64`, `sources` | -| `analyze` | unless `--no-agents` | `question`, `filter`, `images_base64`, `sources` | +| `execute_code` | always | `code`, `filter`, `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 @@ -126,12 +127,18 @@ node's `id` in the outline is the `section_id`. A document without headings has an empty outline. `list_documents` returns titles, URIs and metadata, which is how a client learns what a filter can match. -`ask_question` runs the RAG agent on the server and returns an answer -followed by its citations. `analyze` writes and runs Python in a sandbox -over the documents, for counting, aggregation and computation across -documents. Both cost a model call. Claude Code moves a call still running -after about two minutes to a background task, which a slow local model can -trigger; `--no-agents` leaves both tools out. +`execute_code` runs a Python program in the sandbox of the +[analysis capability](capabilities/analysis.md), over the documents `filter` +and `sources` select, and returns what it printed. The program reads +`/documents/{document_id}/` (`metadata.json`, `content.txt`, `items.jsonl`, +`chunks.jsonl`, `toc.json`) and can `await search()` and +`await list_documents()`; the tool description spells out the fields and the +interpreter's limits. Each call is one program: nothing carries over between +calls, and the sandbox is created and closed per call. A failing program is a +tool error carrying the interpreter's message and any output printed before +it. `analysis.code_timeout` bounds a call and `analysis.max_output_chars` its +output; no model runs on the server. Claude Code moves a call still running +after about two minutes to a background task. ### Filters @@ -150,8 +157,10 @@ title = 'Q3 report' A failure is an MCP error, never an empty result. Expected failures carry a message: a document or section id that matches nothing, a collection the server does not cover, a filter the query engine rejects (with its message), -invalid base64, -and an `ask_question` or `analyze` failure naming only the exception type. +invalid base64, and a program that fails in `execute_code`. A failure on the +server inside a program, a database read or an in-code search raising, reaches +the program and the client as its exception type only; the traceback goes to +the server log. Anything else reaches the client as `Error calling tool 'name'` and its traceback goes to the server log. diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index d2dbedad..8bc934ae 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -925,7 +925,6 @@ class HaikuRAGApp: transport: str | None = None, host: str = "127.0.0.1", port: int = 8001, - agents: bool = True, ): """Run the MCP server until interrupted. @@ -935,7 +934,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, agents=agents) + server = _mcp_server_covering(self.scope, self.config) try: if transport == "stdio": await server.run_stdio_async() diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index fb88a081..2dafbbe8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -19,7 +19,7 @@ from haiku.rag.capabilities._base import ( ) from haiku.rag.capabilities._tools import merge_results from haiku.rag.config.models import AppConfig -from haiku.rag.sandbox import AnalysisContext, Sandbox +from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint STATE_NAMESPACE = "analysis" _CAPABILITY_ID = "haiku-rag-analysis" @@ -49,21 +49,6 @@ def multiple_collections_instructions() -> str: return _multiple_collections_path.read_text().rstrip() -def _recovery_hint(stderr: str) -> str: - """Name the workaround for sandbox limits models trip over repeatedly. - - The instructions already say file objects are not iterable, and models write - ``for line in open(...)`` regardless. Carrying the fix in the error gives - them something to act on for the retry. - """ - if "TextIOWrapper" in stderr and "not iterable" in stderr: - return ( - "\n\nHint: file objects cannot be iterated here. Read lines with " - '.readlines() or .read().split("\\n").' - ) - return "" - - @dataclass class AnalysisCapability(RAGCapabilityBase[AnalysisState]): """Deferred capability for sandboxed computation over a RAG corpus.""" @@ -139,7 +124,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ) if not result.success: raise ToolFailed( - f"{result.stderr}{_recovery_hint(result.stderr)}" + f"{result.stderr}{recovery_hint(result.stderr)}" f"\n\nOutput: {result.stdout}" ) return result.stdout or "No output." diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md index cd0a6bea..37dc8807 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md +++ b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md @@ -13,8 +13,8 @@ You can mix the two. The rule: always call `analysis_cite` before answering — Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results. Inside the code, these functions are available (use `await`): -- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`) -- `await list_documents()` → list of dicts with keys: id, title, uri, created_at +- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`), chunk_meta (the matched chunk's stored metadata, custom keys included) +- `await list_documents()` → list of dicts with keys: id, title, uri, created_at, metadata Available modules: `json`, `re`, `math`, `pathlib` Not supported: class inheritance and metaclasses, generators/yield, match statements, decorators, `collections`, iterating a file object (`for line in f`) @@ -39,9 +39,10 @@ All documents are mounted as a virtual filesystem at `/documents/`: ``` /documents/{document_id}/ - metadata.json # {"id", "title", "uri", "created_at"} + metadata.json # {"id", "title", "uri", "created_at", "metadata"} content.txt # Full document text items.jsonl # Structured items (one JSON object per line) + chunks.jsonl # Chunks in order with their metadata (one JSON object per line) toc.json # Section tree derived from heading_level ``` @@ -70,7 +71,7 @@ for line in Path(f'/documents/{doc_id}/items.jsonl').read_text().strip().split(" ``` ### metadata.json -Document metadata: `id`, `title`, `uri`, `created_at`. +Document metadata: `id`, `title`, `uri`, `created_at`, and `metadata`, the keys stored with the document. ### content.txt Full text content. Use for regex or keyword search across a whole document. @@ -86,6 +87,9 @@ Each row carries: - `chunk_ids`: chunks that contain this item — pass to `analysis_cite()` to ground an answer that read this item directly - `heading_level`: H-level for `section_header` rows; `0` on non-header rows +### chunks.jsonl +The document's chunks in order, one JSON object per line: `chunk_id` and `metadata`, the chunk's stored metadata (`doc_item_refs`, `headings`, `labels`, `page_numbers`, and any custom keys such as paragraph or footnote numbers). To read by chunk metadata, keep the matching rows and take the `items.jsonl` rows whose `chunk_ids` name them. + ### toc.json Section tree derived from `heading_level`: `{"doc_id", "title", "tree": [...]}` where each node has `{self_ref, level, title, page_numbers, item_range: [start, end_exclusive], chunk_ids, children}`. `item_range` is a line slice into `items.jsonl` — `items[start:end]`. `chunk_ids` aggregates the citable chunks across all items in the section — pass directly to `analysis_cite()` to ground a section-scoped answer without a corpus-wide `search()` call. `tree: []` for docs with no headers. diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 06647670..c47d3971 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -884,20 +884,13 @@ 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, agents=not no_agents) - ) + asyncio.run(app.run_mcp(transport=transport, host=host, port=port)) if __name__ == "__main__": diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 112a156c..ebfc533b 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -16,13 +16,13 @@ from pydantic import Field from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, get_config from haiku.rag.context import build_toc +from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint from haiku.rag.store.exceptions import UnknownDatabaseError from haiku.rag.store.models import Document, SearchResult from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.schema import DocumentMetaRecord from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode from haiku.rag.tools.search import collect_pictures -from haiku.rag.utils import format_citations if TYPE_CHECKING: from typing import Any @@ -65,12 +65,6 @@ def _decode_image(image_base64: str) -> bytes: raise ToolError("Invalid base64 image") from e -def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: - if not images_base64: - return None - return [_decode_image(b64) for b64 in images_base64] - - async def _check_filter( rag: HaikuRAG, filter: str | None, sources: list[str] | None = None ) -> None: @@ -98,18 +92,14 @@ async def _check_filter( ) from e -def _instructions(scope: "DatabaseScope", config: AppConfig, agents: bool) -> str: +def _instructions(scope: "DatabaseScope", config: AppConfig) -> 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." + "searchable by meaning and keyword, readable whole or section by section, " + "or computed across with code." ] - 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." @@ -185,9 +175,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, + db_path: Path | None = None, config: AppConfig | None = None ) -> FastMCP: """Create an MCP server over the databases the configuration places. @@ -196,20 +184,14 @@ 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, agents - ) + return _covering(DatabaseScope.resolve(config, database_path=db_path), config) -def _covering( - scope: "DatabaseScope", config: AppConfig, agents: bool = True -) -> FastMCP: +def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: """An MCP server over databases someone already resolved. Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and @@ -255,7 +237,7 @@ def _covering( # the traceback goes to the server log. A ToolError reaches the client as is. mcp = FastMCP( "haiku-rag", - instructions=_instructions(scope, config, agents), + instructions=_instructions(scope, config), version=metadata.version("haiku.rag-slim"), lifespan=lifespan, mask_error_details=True, @@ -469,72 +451,51 @@ def _covering( for doc in documents ] - if agents: + @mcp.tool(annotations=_read_only("Run code over the documents")) + async def execute_code( + code: str, filter: Filter = None, sources: Sources = None + ) -> str: + """Run a Python program over the documents and return what it printed. - @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. + Use this when the answer is a count, an aggregate, a comparison across + many documents, a lookup by document or chunk metadata, or a pattern + over whole documents: whatever a search cannot rank. The program runs + in a sandboxed interpreter on the server. Each call is one program, + nothing carries over between calls, and `print` is the only output. - 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. + Inside the program, `/documents/{document_id}/` holds `metadata.json` + (id, title, uri, created_at, metadata), `content.txt` (the whole text), + `items.jsonl` (one item per line: self_ref, label, text, page_numbers, + heading_level, chunk_ids), `chunks.jsonl` (one chunk per line: chunk_id, + metadata) and `toc.json` (the section tree, each node with an item_range + slice into items.jsonl). Read files with `Path.read_text()` or `open()`; + a file object cannot be iterated, use `.readlines()`. + `await search(query, limit=10)` returns dicts with chunk_id, content, + document_id, document_title, document_uri, source, score, page_numbers, + headings, doc_item_refs, labels and chunk_meta. `await list_documents()` + returns dicts with id, title, uri, created_at, source and metadata. + Modules: json, re, math, pathlib. Not available: generators, class + inheritance, match statements, decorators, collections. - 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 - - @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. - - 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 + Args: + code: The program. Use `await` on search and list_documents. + """ + rag = await _client() + sandbox = Sandbox._covering( + scope, config, AnalysisContext(filter=filter, sources=sources), rag=rag + ) + try: + await _check_filter(rag, filter, sources) + result = await sandbox.execute(code) + except UnknownDatabaseError as e: + raise ToolError(str(e)) from e + finally: + await sandbox.close() + if not result.success: + raise ToolError( + f"{result.stderr}{recovery_hint(result.stderr)}" + f"\n\nOutput: {result.stdout}" + ) + return result.stdout or "No output." return mcp diff --git a/haiku_rag_slim/haiku/rag/sandbox/__init__.py b/haiku_rag_slim/haiku/rag/sandbox/__init__.py index 1ca4b1de..0e7576f3 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/__init__.py +++ b/haiku_rag_slim/haiku/rag/sandbox/__init__.py @@ -1,10 +1,11 @@ from haiku.rag.sandbox.dependencies import AnalysisContext from haiku.rag.sandbox.models import AnalysisResult -from haiku.rag.sandbox.sandbox import Sandbox, SandboxResult +from haiku.rag.sandbox.sandbox import Sandbox, SandboxResult, recovery_hint __all__ = [ "AnalysisContext", "AnalysisResult", "Sandbox", "SandboxResult", + "recovery_hint", ] diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index cddbdd46..47149f1e 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -1,5 +1,6 @@ import asyncio import json +import logging import os from collections.abc import AsyncIterator, Callable, Coroutine from contextlib import asynccontextmanager, suppress @@ -19,7 +20,7 @@ from pydantic_monty import ( from haiku.rag.config.models import AppConfig from haiku.rag.context import build_toc from haiku.rag.sandbox.dependencies import AnalysisContext -from haiku.rag.store.models.chunk import SearchResult +from haiku.rag.store.models.chunk import Chunk, SearchResult from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem from haiku.rag.utils import gather_all @@ -30,6 +31,19 @@ if TYPE_CHECKING: from haiku.rag.client.scope import DatabaseScope +logger = logging.getLogger(__name__) + + +def _host_failure(where: str, e: Exception) -> RuntimeError: + """The error a program gets for a failure on the host side of a call. + + The message and traceback go to the log. The program, and through the MCP + server its client, learn the exception type only. + """ + logger.exception("%s failed inside the sandbox", where) + return RuntimeError(f"{where} failed: {type(e).__name__}") + + @dataclass class SandboxResult: """Result of executing code in the sandbox.""" @@ -39,6 +53,21 @@ class SandboxResult: success: bool +def recovery_hint(stderr: str) -> str: + """Name the workaround for sandbox limits models trip over repeatedly. + + The instructions already say file objects are not iterable, and models write + ``for line in open(...)`` regardless. Carrying the fix in the error gives + them something to act on for the retry. + """ + if "TextIOWrapper" in stderr and "not iterable" in stderr: + return ( + "\n\nHint: file objects cannot be iterated here. Read lines with " + '.readlines() or .read().split("\\n").' + ) + return "" + + class Sandbox: """Execute code in a sandboxed Python interpreter. @@ -46,7 +75,8 @@ class Sandbox: The interpreter runs in a subprocess worker checked out of an ``AsyncMonty`` pool. External functions (search, list_documents) are called by Monty code using ``await`` and resolved asynchronously on the host. Documents are - exposed via a virtual filesystem at ``/documents/{id}/``. + exposed via a virtual filesystem at ``/documents/{id}/``: ``metadata.json``, + ``content.txt``, ``items.jsonl``, ``chunks.jsonl`` and ``toc.json``. The session persists across ``execute()`` calls within the same Sandbox instance — variables carry over. Call ``close()`` to return the worker to @@ -76,6 +106,7 @@ class Sandbox: _doc_items: dict[str, list["DocumentItem"]] _doc_chunk_index: dict[str, dict[str, list[str]]] _items_jsonl_cache: dict[str, str] + _chunks_jsonl_cache: dict[str, str] _toc_json_cache: dict[str, str] _opened: "HaikuRAG | None" _pool: AsyncMonty | None @@ -142,6 +173,7 @@ class Sandbox: self._doc_items = {} self._doc_chunk_index = {} self._items_jsonl_cache = {} + self._chunks_jsonl_cache = {} self._toc_json_cache = {} self._pool = None self._session = None @@ -249,7 +281,7 @@ class Sandbox: loop overruns it by however long the outstanding reads take. Raising from inside the callback answers the worker's suspension, which keeps the session usable — cancelling ``feed_run`` from outside does not, and wedges - the protocol. + the protocol. A failed read reaches the program by type only. """ assert self._loop is not None, ( "VFS reads happen during execute(); the loop must be captured first." @@ -260,7 +292,10 @@ class Sandbox: "time limit exceeded: no further document reads after " f"{self._config.analysis.code_timeout}s" ) - return asyncio.run_coroutine_threadsafe(coro, self._loop).result() + try: + return asyncio.run_coroutine_threadsafe(coro, self._loop).result() + except Exception as e: + raise _host_failure("document read", e) from None async def _discard_session(self) -> None: """Drop a session whose worker is gone. @@ -330,6 +365,7 @@ class Sandbox: "doc_item_refs": r.doc_item_refs, "labels": r.labels, "picture_refs": picture_refs, + "chunk_meta": r.chunk_meta, } ) return out @@ -343,15 +379,28 @@ class Sandbox: "uri": d.uri, "created_at": str(d.created_at), "source": d.source, + "metadata": d.metadata, } for d in docs ] return { - "search": search, - "list_documents": list_documents, + "search": self._guarded("search()", search), + "list_documents": self._guarded("list_documents()", list_documents), } + @staticmethod + def _guarded( + where: str, fn: Callable[..., Coroutine[Any, Any, Any]] + ) -> Callable[..., Coroutine[Any, Any, Any]]: + async def call(*args: Any, **kwargs: Any) -> Any: + try: + return await fn(*args, **kwargs) + except Exception as e: + raise _host_failure(where, e) from None + + return call + async def _build_vfs(self) -> OSAccess: """Build the virtual filesystem with document data. @@ -359,6 +408,7 @@ class Sandbox: - metadata.json: CallbackFile (eager, small) - content.txt: CallbackFile (lazy, can be large) - items.jsonl: CallbackFile (lazy, bulk-cached) + - chunks.jsonl: CallbackFile (lazy, bulk-cached) - toc.json: CallbackFile (lazy, bulk-cached) """ files: list[CallbackFile] = [] @@ -433,6 +483,31 @@ class Sandbox: return read_items + def _make_chunks_reader( + did: str, + ) -> Callable[["PurePosixPath"], str]: + def read_chunks(_path: "PurePosixPath") -> str: + cached = sandbox._chunks_jsonl_cache.get(did) + if cached is not None: + return cached + + async def _fetch() -> list[Chunk]: + async with sandbox._connection(sandbox._owners.get(did)) as rag: + return await rag.chunk_repository.get_by_document_id(did) + + chunks = sandbox._run_on_loop(_fetch()) + jsonl = "\n".join( + json.dumps( + {"chunk_id": chunk.id, "metadata": chunk.metadata}, + ensure_ascii=False, + ) + for chunk in chunks + ) + sandbox._chunks_jsonl_cache[did] = jsonl + return jsonl + + return read_chunks + def _make_toc_reader( did: str, ) -> Callable[["PurePosixPath"], str]: @@ -467,6 +542,7 @@ class Sandbox: "title": doc.title, "uri": doc.uri, "created_at": str(doc.created_at), + "metadata": doc.metadata, }, ensure_ascii=False, ) @@ -508,6 +584,13 @@ class Sandbox: write=_deny_write, ) ) + files.append( + CallbackFile( + f"{doc_dir}/chunks.jsonl", + read=_make_chunks_reader(doc_id), + write=_deny_write, + ) + ) # HAIKU_RAG_DISABLE_TOC is an evaluation-time toggle for measuring # whether toc.json's outline view earns its place in the VFS. # Production callers should leave it unset. diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 2d71b47c..6327b808 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -393,48 +393,6 @@ def _citation_label(c: "Citation") -> str: return c.document_title or c.document_uri -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 "" - - lines = ["## Citations\n"] - - for i, c in enumerate(citations): - idx = c.index if c.index is not None else (i + 1) - title = c.document_title or c.document_uri - 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) - section = _citation_section(c) - if section: - location_parts.append(f"Section: {section}") - - # The URI is the header when there is no title; do not repeat it. - line = f"{header} {c.document_uri}" if c.document_title else header - if location_parts: - line += f" - {', '.join(location_parts)}" - - lines.append(line) - for ref in c.picture_refs: - lines.append(f"[Figure: {ref}]") - lines.append(c.content) - lines.append("") - - return "\n".join(lines) - - def truncated(text: str, limit: int) -> str: """The first `limit` characters of `text`, with `…` appended when anything was dropped. A cut result is `limit` characters plus the mark.""" diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index ed4a8a2d..cb6a1946 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -1,4 +1,5 @@ import asyncio +import logging import threading from pathlib import Path @@ -112,6 +113,41 @@ class TestSandboxListDocuments: assert "Test Document" in result.stdout assert temp_db_path.stem in result.stdout + @pytest.mark.asyncio + async def test_list_documents_carries_metadata(self, temp_db_path): + """Rows carry the document's metadata, so a corpus-wide pass over it is + one call rather than a file read per document.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + config = AppConfig() + docling = DoclingDocument(name="d") + docling.add_text(label=DocItemLabel.TEXT, text="Test content") + async with HaikuRAG(temp_db_path, create=True) as client: + await client.import_document( + docling, + [ + Chunk( + content="Test content", + embedding=[0.1] * config.embeddings.model.vector_dim, + order=0, + ) + ], + uri="test://doc1", + title="Test Document", + metadata={"author": "Ada"}, + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + "docs = await list_documents()\nprint(docs[0]['metadata']['author'])" + ) + finally: + await sb.close() + assert result.success, result.stderr + assert "Ada" in result.stdout + class TestSandboxSearch: """Test search function in sandbox.""" @@ -188,6 +224,51 @@ class TestSandboxSearch: assert "str" in result.stdout assert "True" in result.stdout + @pytest.mark.asyncio + async def test_search_returns_the_matched_chunks_metadata( + self, temp_db_path, monkeypatch + ): + """Results carry the stored metadata of the chunk that matched, custom + keys included.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + from haiku.rag.embeddings import EmbedderWrapper + + config = AppConfig() + dim = config.embeddings.model.vector_dim + + async def embed_query(self, text): + return [0.1] * dim + + monkeypatch.setattr(EmbedderWrapper, "embed_query", embed_query) + docling = DoclingDocument(name="d") + docling.add_text(label=DocItemLabel.TEXT, text="Paragraph fourteen.") + async with HaikuRAG(temp_db_path, create=True) as client: + await client.import_document( + docling, + [ + Chunk( + content="Paragraph fourteen.", + embedding=[0.1] * dim, + order=0, + metadata={"para_no": "14"}, + ) + ], + uri="test://paras", + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + "results = await search('fourteen', limit=1)\n" + "print(results[0]['chunk_meta']['para_no'])" + ) + finally: + await sb.close() + assert result.success, result.stderr + assert "14" in result.stdout + class TestSandboxExternalFunctionEdgeCases: """Test edge cases in external function dispatch.""" @@ -240,6 +321,71 @@ class TestSandboxExternalFunctionEdgeCases: assert not result.success assert "external error" in result.stderr + @pytest.mark.asyncio + async def test_a_failing_search_reaches_the_program_by_type_only( + self, sandbox, monkeypatch, caplog + ): + """A host-side failure inside search() names its exception type to + the program; the message and traceback go to the log.""" + + async def boom(self, *args, **kwargs): + raise ValueError("failed at /secret/path") + + monkeypatch.setattr(HaikuRAG, "search", boom) + + with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): + result = await sandbox.execute("await search('hello')") + + assert not result.success + assert "search() failed: ValueError" in result.stderr + assert "/secret/path" not in result.stderr + assert any( + r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) + for r in caplog.records + ) + + @pytest.mark.asyncio + async def test_a_failing_document_read_reaches_the_program_by_type_only( + self, temp_db_path, monkeypatch, caplog + ): + """A program can catch a failed file read, and what it catches names + the exception type only.""" + from haiku.rag.store.models.document import Document + + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.document_repository.create( + Document(content="x", uri="test://read", title="Read") + ) + repository = type(client.document_repository) + + async def boom(self, *args, **kwargs): + raise ValueError("failed at /secret/path") + + monkeypatch.setattr(repository, "get_content", boom) + sb = Sandbox( + db_path=temp_db_path, config=AppConfig(), context=AnalysisContext() + ) + try: + with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): + result = await sb.execute( + "from pathlib import Path\n" + "try:\n" + f" Path('/documents/{doc.id}/content.txt').read_text()\n" + "except Exception as e:\n" + " print('caught:', e)" + ) + finally: + await sb.close() + + assert result.success, result.stderr + assert "caught:" in result.stdout + assert "ValueError" in result.stdout + assert "/secret/path" not in result.stdout + assert any( + r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) + for r in caplog.records + ) + class TestSandboxOutputTruncation: """Test output truncation behavior.""" @@ -312,13 +458,14 @@ class TestSandboxVFS: @pytest.mark.asyncio @pytest.mark.vcr() async def test_metadata_json(self, temp_db_path): - """metadata.json contains document title and uri.""" + """metadata.json contains document title, uri and stored metadata.""" config = AppConfig() async with HaikuRAG(temp_db_path, create=True) as client: doc = await client.create_document( content="Test content", uri="test://doc1", title="Test Document", + metadata={"author": "Ada"}, ) context = AnalysisContext() @@ -328,11 +475,13 @@ class TestSandboxVFS: "import json\n" f"meta = json.loads(Path('/documents/{doc.id}/metadata.json').read_text())\n" "print(meta['title'])\n" - "print(meta['uri'])" + "print(meta['uri'])\n" + "print(meta['metadata']['author'])" ) - assert result.success + assert result.success, result.stderr assert "Test Document" in result.stdout assert "test://doc1" in result.stdout + assert "Ada" in result.stdout @pytest.mark.asyncio @pytest.mark.vcr() @@ -386,6 +535,59 @@ class TestSandboxVFS: assert result.success assert result.stdout.count("True") == 6 + @pytest.mark.asyncio + async def test_chunks_jsonl(self, temp_db_path): + """chunks.jsonl lists a document's chunks in order with their stored + metadata; a chunk found by its metadata leads to its items through + their chunk_ids.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + config = AppConfig() + dim = config.embeddings.model.vector_dim + docling = DoclingDocument(name="d") + docling.add_text(label=DocItemLabel.TEXT, text="Paragraph thirteen.") + docling.add_text(label=DocItemLabel.TEXT, text="Paragraph fourteen.") + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.import_document( + docling, + [ + Chunk( + content="Paragraph thirteen.", + embedding=[0.1] * dim, + order=0, + metadata={"para_no": "13", "doc_item_refs": ["#/texts/0"]}, + ), + Chunk( + content="Paragraph fourteen.", + embedding=[0.1] * dim, + order=1, + metadata={"para_no": "14", "doc_item_refs": ["#/texts/1"]}, + ), + ], + uri="test://paras", + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + "from pathlib import Path\n" + "import json\n" + f"root = Path('/documents/{doc.id}')\n" + "def rows(name):\n" + " return [json.loads(l) for l in (root / name).read_text().strip().split('\\n')]\n" + "chunks = rows('chunks.jsonl')\n" + "print(len(chunks))\n" + "hit = [c for c in chunks if c['metadata'].get('para_no') == '14']\n" + "print(len(hit))\n" + "items = rows('items.jsonl')\n" + "print([i['text'] for i in items if hit[0]['chunk_id'] in i['chunk_ids']])" + ) + finally: + await sb.close() + assert result.success, result.stderr + assert result.stdout.splitlines() == ["2", "1", "['Paragraph fourteen.']"] + @pytest.mark.asyncio @pytest.mark.vcr() async def test_open_read(self, temp_db_path): @@ -433,7 +635,8 @@ class TestSandboxVFS: @pytest.mark.asyncio @pytest.mark.parametrize( - "filename", ["content.txt", "items.jsonl", "toc.json", "metadata.json"] + "filename", + ["content.txt", "items.jsonl", "chunks.jsonl", "toc.json", "metadata.json"], ) async def test_write_denied_for_every_document_file(self, temp_db_path, filename): """Every file in the document VFS is read-only, metadata.json included.""" @@ -797,6 +1000,26 @@ class TestSandboxReadDeadline: cannot check its duration budget while one is in flight. The sandbox enforces the budget itself, before each read.""" + @pytest.mark.asyncio + async def test_a_failed_read_reaches_the_program_by_type_only( + self, sandbox, caplog + ): + """The bridged read hands the program the exception type, not the + message, and logs the traceback.""" + sandbox._loop = asyncio.get_running_loop() + + async def failing_read(): + raise ValueError("failed at /secret/path") + + with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): + with pytest.raises(RuntimeError, match="document read failed: ValueError"): + await asyncio.to_thread(sandbox._run_on_loop, failing_read()) + + assert any( + r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) + for r in caplog.records + ) + @pytest.mark.asyncio async def test_read_after_deadline_raises_without_scheduling(self, sandbox): """A read attempted past the deadline fails instead of querying.""" diff --git a/tests/sandbox/test_sandbox_toc.py b/tests/sandbox/test_sandbox_toc.py index 2d359212..d08b9340 100644 --- a/tests/sandbox/test_sandbox_toc.py +++ b/tests/sandbox/test_sandbox_toc.py @@ -16,6 +16,7 @@ import pytest from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.sandbox import AnalysisContext, Sandbox +from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document from haiku.rag.store.models.document_item import DocumentItem @@ -434,6 +435,38 @@ class TestVfsReadPaths: "nope", ) + async def test_chunks_jsonl_lists_chunks_in_order_with_their_metadata( + self, temp_db_path + ): + """One row per chunk, in chunk order, carrying the stored metadata as + is; the second read of a document is served from the sandbox's cache.""" + config = AppConfig() + dim = config.embeddings.model.vector_dim + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://paras", title="Paras") + for order, para_no in enumerate(["13", "14"]): + await client.chunk_repository.create( + Chunk( + document_id=doc_id, + content=f"Paragraph {para_no}.", + embedding=[0.1] * dim, + order=order, + metadata={"para_no": para_no, "doc_item_refs": []}, + ) + ) + + sandbox = Sandbox(temp_db_path, config, AnalysisContext()) + first = await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl") + rows = [json.loads(line) for line in first.split("\n")] + + assert [row["metadata"]["para_no"] for row in rows] == ["13", "14"] + assert all(set(row) == {"chunk_id", "metadata"} for row in rows) + assert rows[0]["metadata"] == {"para_no": "13", "doc_item_refs": []} + assert sandbox._chunks_jsonl_cache[doc_id] == first + assert ( + await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl") == first + ) + async def test_toc_skips_gaps_in_item_positions(self, temp_db_path): """Positions need not be contiguous — a heading's span may cover positions that carry no item.""" diff --git a/tests/test_app.py b/tests/test_app.py index 589c2eb1..cf7104da 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -705,20 +705,6 @@ 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 7d4af634..50d389a0 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1041,14 +1041,6 @@ 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 0fada9a3..3954625a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,7 +1,6 @@ import logging import re from pathlib import Path -from types import SimpleNamespace import pytest from fastmcp.exceptions import ToolError @@ -280,32 +279,6 @@ class TestMCPReadTools: ] assert overview["metadata"] == {"author": "Ada"} - @pytest.mark.asyncio - async def test_ask_question_appends_the_citations(self, mcp_db, monkeypatch): - from haiku.rag.store.models.citation import Citation - - citation = Citation( - chunk_id="c1", - document_id="d1", - content="cited text", - document_uri="test://ai-overview", - document_title="AI Overview", - source="alpha", - ) - - async def fake_ask(self, question, filter=None, images=None, sources=None): - return ("the answer", [citation]) - - monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db) - ask = await _get_tool(mcp, "ask_question") - - answer = await ask(question="q") - assert answer.startswith("the answer") - assert "AI Overview" in answer - # One database: its name adds nothing. - assert "alpha" not in answer - @pytest.fixture async def outlined_db(temp_db_path): @@ -666,18 +639,6 @@ 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.instructions.splitlines() - async with Client(create_mcp_server(mcp_db, agents=False)) as client: - without = client.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 @@ -702,7 +663,7 @@ class TestMCPDescribesItself: async with Client(create_mcp_server(mcp_db)) as client: tools = await client.list_tools() - assert len(tools) == 8 + assert len(tools) == 7 for tool in tools: assert tool.annotations is not None, tool.name assert tool.annotations.read_only_hint is True, tool.name @@ -722,7 +683,7 @@ class TestMCPDescribesItself: for name, schema in tool.input_schema.get("properties", {}).items() if not schema.get("description") ] - assert len(tools) == 8 + assert len(tools) == 7 assert undescribed == [] @@ -737,21 +698,138 @@ class TestMCPToolSet: "get_document_outline", "get_document_section", "list_documents", - "ask_question", - "analyze", + "execute_code", } + +_COUNT_DOCUMENTS = ( + "from pathlib import Path\n" + "n = 0\n" + "for d in Path('/documents').iterdir():\n" + " n += 1\n" + "print(n)" +) + + +class TestMCPExecuteCode: + """`execute_code` runs one program per call in the analysis sandbox over + the documents the filter and sources select, and returns what it printed.""" + @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) + async def test_a_program_reads_the_documents_and_returns_what_it_printed( + self, mcp_db + ): + result = await _call( + create_mcp_server(mcp_db), "execute_code", code=_COUNT_DOCUMENTS + ) - assert {t.name for t in await mcp.list_tools()} == { - "search_documents", - "get_document", - "get_document_outline", - "get_document_section", - "list_documents", - } + assert not result.is_error + assert result.content[0].text.strip() == "2" + + @pytest.mark.asyncio + async def test_a_silent_program_says_so(self, mcp_db): + result = await _call(create_mcp_server(mcp_db), "execute_code", code="x = 1") + + assert not result.is_error + assert result.content[0].text == "No output." + + @pytest.mark.asyncio + async def test_filter_narrows_the_documents_a_program_sees(self, mcp_db): + result = await _call( + create_mcp_server(mcp_db), + "execute_code", + code=_COUNT_DOCUMENTS, + filter="title = 'AI Overview'", + ) + + assert result.content[0].text.strip() == "1" + + @pytest.mark.asyncio + async def test_sources_narrows_the_documents_a_program_sees(self, two_dbs): + mcp = _covering_all(two_dbs) + + both = await _call(mcp, "execute_code", code=_COUNT_DOCUMENTS) + beta = await _call(mcp, "execute_code", code=_COUNT_DOCUMENTS, sources=["beta"]) + + assert both.content[0].text.strip() == "2" + assert beta.content[0].text.strip() == "1" + + @pytest.mark.asyncio + async def test_a_failing_program_is_an_error_carrying_the_cause_and_its_output( + self, mcp_db + ): + code = ( + "from pathlib import Path\n" + "print('before')\n" + "for d in Path('/documents').iterdir():\n" + " for line in open(d / 'items.jsonl'):\n" + " pass" + ) + + result = await _call(create_mcp_server(mcp_db), "execute_code", code=code) + + assert result.is_error + text = result.content[0].text + assert "not iterable" in text + assert ".readlines()" in text + assert "Output: before" in text + + @pytest.mark.asyncio + async def test_calls_share_no_state(self, mcp_db): + mcp = create_mcp_server(mcp_db) + + first = await _call(mcp, "execute_code", code="x = 1\nprint(x)") + second = await _call(mcp, "execute_code", code="print(x)") + + assert first.content[0].text.strip() == "1" + assert second.is_error + assert "NameError" in second.content[0].text + + @pytest.mark.asyncio + async def test_every_call_closes_its_sandbox(self, mcp_db, monkeypatch): + from haiku.rag.sandbox import Sandbox + + closed = [] + close = Sandbox.close + + async def closing(self): + closed.append(self) + await close(self) + + monkeypatch.setattr(Sandbox, "close", closing) + mcp = create_mcp_server(mcp_db) + + await _call(mcp, "execute_code", code="print(1)") + await _call(mcp, "execute_code", code="raise ValueError('x')") + + assert len(closed) == 2 + assert closed[0] is not closed[1] + + @pytest.mark.asyncio + async def test_a_program_reaches_chunk_metadata(self, mcp_db): + async with HaikuRAG(mcp_db, create=True) as rag: + doc = await rag.get_document_by_uri("test://ai-overview") + embedding = (await rag.embedder.embed_documents(["x"]))[0] + await rag.chunk_repository.create( + Chunk( + document_id=doc.id, + content="Paragraph fourteen.", + metadata={"para_no": "14"}, + embedding=embedding, + ) + ) + code = ( + "from pathlib import Path\n" + "import json\n" + f"text = Path('/documents/{doc.id}/chunks.jsonl').read_text()\n" + "rows = [json.loads(line) for line in text.strip().split('\\n')]\n" + "print(len([r for r in rows if r['metadata'].get('para_no') == '14']))" + ) + + result = await _call(create_mcp_server(mcp_db), "execute_code", code=code) + + assert not result.is_error, result.content[0].text + assert result.content[0].text.strip() == "1" class TestMCPCoversTheConfiguredSet: @@ -784,8 +862,7 @@ class TestMCPCoversTheConfiguredSet: {"image_base64": "AAAA", "sources": ["nope"]}, ), ("get_document", {"document_id": "x", "source": "nope"}), - ("ask_question", {"question": "q", "sources": ["nope"]}), - ("analyze", {"question": "q", "sources": ["nope"]}), + ("execute_code", {"code": "print(1)", "sources": ["nope"]}), ], ) async def test_an_unknown_database_is_an_error_not_an_empty_result( @@ -845,59 +922,6 @@ class TestMCPCoversTheConfiguredSet: assert {_line(block, "Collection") for block in blocks} == {"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") - - 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.""" @@ -963,63 +987,6 @@ class TestMCPImageQuery: assert not searched -class TestMCPImageInput: - @pytest.mark.asyncio - async def test_ask_question_decodes_images(self, mcp_db, monkeypatch): - from base64 import b64encode - - captured = {} - - async def fake_ask(self, question, filter=None, images=None, sources=None): - captured["images"] = images - return ("answer", []) - - monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db) - ask = await _get_tool(mcp, "ask_question") - - png = b"fake image bytes" - result = await ask(question="q", images_base64=[b64encode(png).decode()]) - assert result == "answer" - assert captured["images"] == [png] - - @pytest.mark.asyncio - async def test_analyze_decodes_images(self, mcp_db, monkeypatch): - from base64 import b64encode - from types import SimpleNamespace - - captured = {} - - async def fake_analyze(self, question, filter=None, images=None, sources=None): - captured["images"] = images - return SimpleNamespace(answer="answer") - - monkeypatch.setattr(HaikuRAG, "analyze", fake_analyze) - mcp = create_mcp_server(mcp_db) - analyze = await _get_tool(mcp, "analyze") - - jpeg = b"fake jpeg bytes" - result = await analyze(question="q", images_base64=[b64encode(jpeg).decode()]) - assert result == "answer" - assert captured["images"] == [jpeg] - - @pytest.mark.asyncio - async def test_ask_question_without_images_passes_none(self, mcp_db, monkeypatch): - captured = {} - - async def fake_ask(self, question, filter=None, images=None, sources=None): - captured["images"] = images - return ("answer", []) - - monkeypatch.setattr(HaikuRAG, "ask", fake_ask) - mcp = create_mcp_server(mcp_db) - ask = await _get_tool(mcp, "ask_question") - - result = await ask(question="q") - assert result == "answer" - assert captured["images"] is None - - @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") class TestMCPErrorContract: """A failure is an error on the wire, never an empty result. Expected @@ -1038,7 +1005,11 @@ class TestMCPErrorContract: @pytest.mark.asyncio @pytest.mark.parametrize( "tool_name,kwargs", - [("search_documents", {"query": "x"}), ("list_documents", {})], + [ + ("search_documents", {"query": "x"}), + ("list_documents", {}), + ("execute_code", {"code": "print(1)"}), + ], ) async def test_an_invalid_filter_is_an_error_naming_the_filter( self, mcp_db, tool_name, kwargs @@ -1076,39 +1047,28 @@ class TestMCPErrorContract: @pytest.mark.parametrize( "payload", ["!!! not base64 !!!", "é"], ids=["outside_alphabet", "non_ascii"] ) - @pytest.mark.parametrize( - "tool_name,image_param,many", - [ - ("search_documents_by_image", "image_base64", False), - ("ask_question", "images_base64", True), - ("analyze", "images_base64", True), - ], - ) async def test_invalid_base64_is_an_error( - self, mcp_db, multimodal_embedder, tool_name, image_param, many, payload + self, mcp_db, multimodal_embedder, payload ): - kwargs: dict[str, object] = {"question": "q"} if many else {} - kwargs[image_param] = [payload] if many else payload - - result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs) + result = await _call( + create_mcp_server(mcp_db), "search_documents_by_image", image_base64=payload + ) assert result.is_error assert "base64" in result.content[0].text @pytest.mark.asyncio - @pytest.mark.parametrize( - "client_method,tool_name", - [("ask", "ask_question"), ("analyze", "analyze")], - ) - async def test_an_agent_failure_names_only_its_type( - self, mcp_db, monkeypatch, caplog, client_method, tool_name + async def test_a_host_failure_inside_a_program_names_only_its_type( + self, mcp_db, monkeypatch, caplog ): - async def boom(self, question, filter=None, images=None, sources=None): + async def boom(self, *args, **kwargs): raise RuntimeError("boom at /secret/path") - monkeypatch.setattr(HaikuRAG, client_method, boom) - with caplog.at_level(logging.ERROR, logger="haiku.rag.mcp"): - result = await _call(create_mcp_server(mcp_db), tool_name, question="q") + monkeypatch.setattr(HaikuRAG, "search", boom) + with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): + result = await _call( + create_mcp_server(mcp_db), "execute_code", code="await search('x')" + ) assert result.is_error assert "RuntimeError" in result.content[0].text diff --git a/tests/test_utils.py b/tests/test_utils.py index 9d4cba45..f63bf07e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -662,150 +662,6 @@ def test_format_bytes(): assert format_bytes(1125899906842624) == "1.0 PB" -# --- format_citations tests --- - - -def test_format_citations_empty(): - from haiku.rag.utils import format_citations - - assert format_citations([]) == "" - - -def test_format_citations_with_citation(): - 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="Some content", - page_numbers=[1], - headings=["Intro"], - ) - result = format_citations([citation]) - assert "[1] Test Doc" in result - assert "doc1" not in result - assert "chunk1" not in result - assert "test://doc" in result - assert "p. 1" in result - assert "Section: Intro" in result - assert "Some content" in result - - -def test_format_citations_multiple_pages(): - 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", - content="Content", - page_numbers=[1, 2, 3], - ) - result = format_citations([citation]) - assert "[1] test://doc" in result - assert "pp. 1-3" in result - # No title: the URI stands in, and the document id never leaks. - assert "doc1" not in result - - -def test_format_citations_with_index(): - from haiku.rag.store.models.citation import Citation - from haiku.rag.utils import format_citations - - citation = Citation( - index=5, - document_id="doc1", - chunk_id="chunk1", - document_uri="test://doc", - document_title="Test Doc", - content="Content", - ) - result = format_citations([citation]) - assert "[5] Test Doc" in result - - -def test_format_citations_sequential_indices(): - from haiku.rag.store.models.citation import Citation - from haiku.rag.utils import format_citations - - citations = [ - Citation( - document_id="doc1", - chunk_id="chunk1", - document_uri="test://doc1", - document_title="First", - content="Content 1", - ), - Citation( - document_id="doc2", - chunk_id="chunk2", - document_uri="test://doc2", - document_title="Second", - content="Content 2", - ), - ] - result = format_citations(citations) - assert "[1] First" 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]) - - -def test_format_citations_names_an_untitled_document_once(): - 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", - content="Content", - page_numbers=[3], - ) - result = format_citations([citation]) - - assert result.count("test://doc") == 1 - assert "[1] test://doc - p. 3" in result - - -# --- format_citations tests (pictures) --- - - -def test_format_citations_picture_refs_render_as_markers(): - 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="text body", - picture_refs=["#/pictures/0", "#/pictures/3"], - ) - result = format_citations([citation]) - assert "[Figure: #/pictures/0]" in result - assert "[Figure: #/pictures/3]" in result - - # --- format_citations_rich tests --- @@ -846,6 +702,22 @@ async def test_format_citations_rich_header_and_footer(): assert "chunk: chunk-uuid-1" in output +async def test_format_citations_rich_names_a_single_page(): + from haiku.rag.store.models.citation import Citation + from haiku.rag.utils import format_citations_rich + + citation = Citation( + document_id="doc1", + chunk_id="chunk1", + document_uri="test://doc", + content="Body", + page_numbers=[3], + ) + output = _render_rich(await format_citations_rich([citation])) + assert "p. 3" in output + assert "pp." not in output + + async def test_format_citations_rich_names_the_database_when_federating(): """Across databases, a citation has to say which one it came from.""" from unittest.mock import AsyncMock From 110adf23eeb898b299fc7f975f49adbcda6a74f8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 7 Sep 2026 10:27:02 +0300 Subject: [PATCH 14/18] Make toc.json item_range the line slice it documents build_toc stored item positions, while the sandbox instructions describe item_range as a slice into items.jsonl; a gap in positions pulled the next heading into a section. Ranges are now indices into the position-ordered items, and get_document_section slices by index too. docs/mcp.md names the tools that take sources. --- CHANGELOG.md | 5 ++++ docs/mcp.md | 4 +-- haiku_rag_slim/haiku/rag/context.py | 41 ++++++++++++++--------------- haiku_rag_slim/haiku/rag/mcp.py | 7 ++--- tests/sandbox/test_sandbox_toc.py | 28 ++++++++++++++++++++ tests/test_mcp.py | 40 ++++++++++++++++++++++++++++ 6 files changed, 97 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc815f0c..81e41e43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,11 @@ `search_documents`, `search_documents_by_image` and `execute_code`; `source` on `get_document`; an unknown name is a tool error. `DocumentInfo.source`. +### Fixed + +- `toc.json` `item_range` in the analysis sandbox is a line slice into + `items.jsonl`, as documented; it held item positions. + ### Removed - MCP tools `ask_question` and `analyze`. diff --git a/docs/mcp.md b/docs/mcp.md index 9fef2c7d..5079ee26 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -32,8 +32,8 @@ The server opens the database read-only. Ingestion goes through the CLI 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 +theirs in `source`. `sources` on `search_documents`, `search_documents_by_image` +and `execute_code` 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). diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 17fb44c5..82fe6ed7 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -502,10 +502,12 @@ def build_toc( follows the explicit levels: a header pops the stack until the top is at a strictly shallower level, then becomes a child of that top (or a root). - ``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the - position of the next header whose level is the same or shallower (i.e. - the next sibling or ancestor that ends this section), or the total item - count if no such header exists. + ``item_range = [start, end_exclusive]`` indexes the position-ordered item + list, which is the line numbering of the sandbox's ``items.jsonl``: ``start`` + is the header's index and ``end_exclusive`` the index of the next header + whose level is the same or shallower (the next sibling or ancestor that + ends this section), or the item count if no such header exists. Indices, + not positions: positions may have gaps. ``chunk_ids`` aggregates the chunks covered by all items in the section's ``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to @@ -520,33 +522,30 @@ def build_toc( # but the end_exclusive lookahead below silently miscomputes section # boundaries if it's not — better to sort once than trust the caller. items = sorted(items, key=lambda i: i.position) - headers: list[DocumentItem] = [ - i for i in items if i.label == "section_header" and i.heading_level > 0 + header_indices = [ + idx + for idx, i in enumerate(items) + if i.label == "section_header" and i.heading_level > 0 ] - if not headers: + if not header_indices: return [] - total = max((i.position for i in items), default=-1) + 1 - items_by_position: dict[int, DocumentItem] = {i.position: i for i in items} - ends: list[int] = [] - for idx, h in enumerate(headers): - end = total - for j in range(idx + 1, len(headers)): - if headers[j].heading_level <= h.heading_level: - end = headers[j].position + for n, idx in enumerate(header_indices): + end = len(items) + for later in header_indices[n + 1 :]: + if items[later].heading_level <= items[idx].heading_level: + end = later break ends.append(end) roots: list[dict[str, Any]] = [] stack: list[tuple[int, dict[str, Any]]] = [] - for h, end in zip(headers, ends, strict=True): + for idx, end in zip(header_indices, ends, strict=True): + h = items[idx] seen: set[str] = set() chunk_ids: list[str] = [] - for pos in range(h.position, end): - item = items_by_position.get(pos) - if item is None: - continue + for item in items[idx:end]: for cid in chunk_index.get(item.self_ref, []): if cid not in seen: seen.add(cid) @@ -556,7 +555,7 @@ def build_toc( "level": h.heading_level, "title": h.text, "page_numbers": list(h.page_numbers), - "item_range": [h.position, end], + "item_range": [idx, end], "chunk_ids": chunk_ids, "children": [], } diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index ebfc533b..b9eb2bb1 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -409,15 +409,12 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: if node is None: raise ToolError(f"No section {section_id!r} in document {document_id!r}") start, end = node["item_range"] + ordered = sorted(items, key=lambda item: item.position) return DocumentSection( id=node["self_ref"], title=node["title"], page_numbers=node["page_numbers"], - content="\n\n".join( - item.text - for item in items - if start <= item.position < end and item.text - ), + content="\n\n".join(item.text for item in ordered[start:end] if item.text), ) @mcp.tool(annotations=_read_only("List documents")) diff --git a/tests/sandbox/test_sandbox_toc.py b/tests/sandbox/test_sandbox_toc.py index d08b9340..9ac5c464 100644 --- a/tests/sandbox/test_sandbox_toc.py +++ b/tests/sandbox/test_sandbox_toc.py @@ -467,6 +467,34 @@ class TestVfsReadPaths: await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl") == first ) + async def test_item_range_is_a_line_slice_into_items_jsonl(self, temp_db_path): + """`item_range` indexes lines of items.jsonl, as documented, not item + positions: a gap in positions must not pull the next heading into a + section.""" + async with HaikuRAG(temp_db_path, create=True) as client: + doc_id = await _empty_doc(client, uri="test://slice", title="Slice") + items = [ + _header(doc_id, 0, 1, "Intro"), + _para(doc_id, 1), + _header(doc_id, 3, 1, "Methods"), + _para(doc_id, 4), + ] + await client.document_item_repository.create_items(doc_id, items) + + sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext()) + toc = await _read_toc(sandbox, doc_id) + raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl") + lines = raw.split("\n") + + intro, methods = toc["tree"] + assert intro["item_range"] == [0, 2] + assert methods["item_range"] == [2, 4] + start, end = intro["item_range"] + assert [json.loads(line)["self_ref"] for line in lines[start:end]] == [ + "#/texts/0", + "#/texts/1", + ] + async def test_toc_skips_gaps_in_item_positions(self, temp_db_path): """Positions need not be contiguous — a heading's span may cover positions that carry no item.""" diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 3954625a..9a9faee6 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -379,6 +379,46 @@ class TestMCPDocumentNavigation: assert "para7" in intro.content assert "Methods" not in intro.content + @pytest.mark.asyncio + async def test_a_section_stops_at_the_next_heading_across_a_position_gap( + self, temp_db_path + ): + from haiku.rag.store.models.document import Document as DocumentModel + from haiku.rag.store.models.document_item import DocumentItem + + def item(pos, label, text, level=0): + return DocumentItem( + document_id="", + position=pos, + self_ref=f"#/texts/{pos}", + label=label, + text=text, + heading_level=level, + ) + + async with HaikuRAG(temp_db_path, create=True) as rag: + doc = await rag.document_repository.create( + DocumentModel(content="x", uri="test://gapped", title="Gapped") + ) + items = [ + item(0, "section_header", "Intro", 1), + item(1, "paragraph", "para1"), + item(3, "section_header", "Methods", 1), + item(4, "paragraph", "para4"), + ] + for i in items: + i.document_id = doc.id + await rag.document_item_repository.create_items(doc.id, items) + + section = await _call( + create_mcp_server(temp_db_path), + "get_document_section", + document_id=doc.id, + section_id="#/texts/0", + ) + + assert section.structured_content["content"] == "Intro\n\npara1" + @pytest.mark.asyncio async def test_an_unknown_section_or_document_is_an_error(self, outlined_db): db, doc_id = outlined_db From d7c131bc2801b8532222f6281f20365ec2292974 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 7 Sep 2026 11:21:25 +0300 Subject: [PATCH 15/18] Package the plugin for Codex from the same bundle The Claude Code plugin moves to plugins/haiku-rag/ and gains a Codex manifest at .codex-plugin/plugin.json; both share one .mcp.json and one haiku-rag Agent Skill. .agents/plugins/marketplace.json serves Codex, the Claude marketplace points at the new path, and scripts/bump_version.py bumps both manifests. The skill declares its compatibility. --- .agents/plugins/marketplace.json | 20 +++++++++++ .claude-plugin/marketplace.json | 4 +-- CHANGELOG.md | 5 ++- README.md | 11 ++++-- docs/mcp.md | 30 ++++++++++++++-- .../haiku-rag}/.claude-plugin/plugin.json | 2 +- plugins/haiku-rag/.codex-plugin/plugin.json | 26 ++++++++++++++ .../haiku-rag}/.mcp.json | 0 .../haiku-rag}/skills/haiku-rag/SKILL.md | 1 + scripts/bump_version.py | 16 +++++---- tests/test_bump_version.py | 15 ++++---- tests/test_mcp.py | 36 ++++++++++++------- 12 files changed, 131 insertions(+), 35 deletions(-) create mode 100644 .agents/plugins/marketplace.json rename {claude-plugin => plugins/haiku-rag}/.claude-plugin/plugin.json (74%) create mode 100644 plugins/haiku-rag/.codex-plugin/plugin.json rename {claude-plugin => plugins/haiku-rag}/.mcp.json (100%) rename {claude-plugin => plugins/haiku-rag}/skills/haiku-rag/SKILL.md (97%) diff --git a/.agents/plugins/marketplace.json b/.agents/plugins/marketplace.json new file mode 100644 index 00000000..ed9452d9 --- /dev/null +++ b/.agents/plugins/marketplace.json @@ -0,0 +1,20 @@ +{ + "name": "haiku-rag", + "interface": { + "displayName": "haiku.rag" + }, + "plugins": [ + { + "name": "haiku-rag", + "source": { + "source": "local", + "path": "./plugins/haiku-rag" + }, + "policy": { + "installation": "AVAILABLE", + "authentication": "ON_INSTALL" + }, + "category": "Productivity" + } + ] +} diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 24d16eee..93de1bc0 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -7,8 +7,8 @@ "plugins": [ { "name": "haiku-rag", - "source": "./claude-plugin", - "description": "Search, read and question your haiku.rag knowledge base from Claude Code." + "source": "./plugins/haiku-rag", + "description": "Search, read and analyze your haiku.rag knowledge base from Claude Code." } ] } diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e41e43..3ace83d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,9 +4,8 @@ ### Added -- Claude Code plugin under `claude-plugin/`: the server configuration and the - `haiku-rag` skill. `claude plugin marketplace add ggozad/haiku.rag`, then - `claude plugin install haiku-rag`. +- Claude Code and Codex plugin under `plugins/haiku-rag/`: two client manifests + sharing the server configuration and the `haiku-rag` Agent Skill. - MCP tool `execute_code(code, filter, sources)`: runs a program in the analysis sandbox over the selected documents and returns what it printed; one sandbox per call. diff --git a/README.md b/README.md index 24d832e3..5970c003 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,7 @@ For direct agent composition, see the [capabilities documentation](https://ggoza ## MCP Server -Use with AI assistants like Claude Code and Claude Desktop: +Use with AI assistants like Claude Code, Codex, and Claude Desktop: ```bash haiku-rag mcp --stdio @@ -123,6 +123,13 @@ claude plugin marketplace add ggozad/haiku.rag claude plugin install haiku-rag ``` +In Codex, install the same plugin from its marketplace: + +```bash +codex plugin marketplace add ggozad/haiku.rag +codex plugin add haiku-rag@haiku-rag +``` + Add to your Claude Desktop configuration: ```json @@ -136,7 +143,7 @@ Add to your Claude Desktop configuration: } ``` -Provides search, document, QA, and analysis tools directly in your AI assistant. +Provides search, document reading, and analysis tools directly in your AI assistant. ## Examples diff --git a/docs/mcp.md b/docs/mcp.md index 5079ee26..4d268d23 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -57,10 +57,34 @@ plugin: claude mcp add haiku-rag -- haiku-rag mcp --stdio ``` -The skill works with that registration too: copy `claude-plugin/skills/haiku-rag` +The skill works with that registration too: copy `plugins/haiku-rag/skills/haiku-rag` into `~/.claude/skills/` and change the tool prefix in its `allowed-tools` from `mcp__plugin_haiku-rag_haiku-rag__` to `mcp__haiku-rag__`. +## Codex + +The repository's Codex plugin registers the server and installs the same Agent +Skill: + +```bash +codex plugin marketplace add ggozad/haiku.rag +codex plugin add haiku-rag@haiku-rag +``` + +The plugin runs `haiku-rag mcp --stdio`, so `haiku-rag` must be on the PATH. +Invoke the skill as `$haiku-rag`. Codex can also select it automatically from +its description. To register the server without the plugin: + +```bash +codex mcp add haiku-rag -- haiku-rag mcp --stdio +``` + +The skill works with that registration too: copy +`plugins/haiku-rag/skills/haiku-rag` into `~/.agents/skills/`. +The `allowed-tools` field supplies Claude Code's tool pre-approval and may be +ignored by other Agent Skills clients. Codex configures MCP tool approvals +separately in `config.toml`. + ## Claude Desktop Integration Add to your Claude Desktop configuration (`claude_desktop_config.json`): @@ -168,8 +192,8 @@ traceback goes to the server log. The server publishes `instructions` describing the knowledge base: what it holds, when to reach for it, the collection names when it covers several, and -`prompts.domain_preamble` when set. Claude Code shows them to the model. Claude -Desktop does not, so every tool description stands on its own. +`prompts.domain_preamble` when set. Claude Code and Codex show them to the +model. Claude Desktop does not, so every tool description stands on its own. ## Continuous ingestion diff --git a/claude-plugin/.claude-plugin/plugin.json b/plugins/haiku-rag/.claude-plugin/plugin.json similarity index 74% rename from claude-plugin/.claude-plugin/plugin.json rename to plugins/haiku-rag/.claude-plugin/plugin.json index 93282056..cf30ea82 100644 --- a/claude-plugin/.claude-plugin/plugin.json +++ b/plugins/haiku-rag/.claude-plugin/plugin.json @@ -1,7 +1,7 @@ { "name": "haiku-rag", "version": "0.82.1", - "description": "Search, read and question your haiku.rag knowledge base from Claude Code.", + "description": "Search, read and analyze your haiku.rag knowledge base from Claude Code.", "author": { "name": "Yiorgis Gozadinos", "email": "ggozadinos@gmail.com" diff --git a/plugins/haiku-rag/.codex-plugin/plugin.json b/plugins/haiku-rag/.codex-plugin/plugin.json new file mode 100644 index 00000000..f27c32ac --- /dev/null +++ b/plugins/haiku-rag/.codex-plugin/plugin.json @@ -0,0 +1,26 @@ +{ + "name": "haiku-rag", + "version": "0.82.1", + "description": "Search, read and analyze your haiku.rag knowledge base from Codex.", + "author": { + "name": "Yiorgis Gozadinos", + "email": "ggozadinos@gmail.com", + "url": "https://github.com/ggozad" + }, + "homepage": "https://ggozad.github.io/haiku.rag/mcp/", + "repository": "https://github.com/ggozad/haiku.rag", + "license": "MIT", + "keywords": ["rag", "knowledge-base", "search", "documents", "mcp"], + "skills": "./skills/", + "mcpServers": "./.mcp.json", + "interface": { + "displayName": "haiku.rag", + "shortDescription": "Search and analyze your haiku.rag knowledge base", + "longDescription": "Search, read, and compute over documents in your local haiku.rag knowledge base through MCP tools.", + "developerName": "Yiorgis Gozadinos", + "category": "Productivity", + "capabilities": ["Interactive", "Read"], + "websiteURL": "https://ggozad.github.io/haiku.rag/", + "defaultPrompt": "Search my haiku.rag knowledge base and cite the relevant documents." + } +} diff --git a/claude-plugin/.mcp.json b/plugins/haiku-rag/.mcp.json similarity index 100% rename from claude-plugin/.mcp.json rename to plugins/haiku-rag/.mcp.json diff --git a/claude-plugin/skills/haiku-rag/SKILL.md b/plugins/haiku-rag/skills/haiku-rag/SKILL.md similarity index 97% rename from claude-plugin/skills/haiku-rag/SKILL.md rename to plugins/haiku-rag/skills/haiku-rag/SKILL.md index 981baf43..c513fd2b 100644 --- a/claude-plugin/skills/haiku-rag/SKILL.md +++ b/plugins/haiku-rag/skills/haiku-rag/SKILL.md @@ -5,6 +5,7 @@ description: Search, read and compute over the user's haiku.rag knowledge base from the user's ingested documents, when asked to find, look up, check or cite something in their documents or knowledge base, or when the question is about the user's own material rather than general knowledge. +compatibility: Requires the haiku-rag MCP server to be registered in the client. allowed-tools: - mcp__plugin_haiku-rag_haiku-rag__search_documents - mcp__plugin_haiku-rag_haiku-rag__search_documents_by_image diff --git a/scripts/bump_version.py b/scripts/bump_version.py index 9593cd89..f696c2f0 100755 --- a/scripts/bump_version.py +++ b/scripts/bump_version.py @@ -2,8 +2,8 @@ """ Version bumping script for haiku.rag workspace. -Updates version in all pyproject.toml files, the Claude Code plugin manifest -and CHANGELOG.md. +Updates version in all pyproject.toml files, both plugin manifests, and +CHANGELOG.md. """ import re @@ -56,7 +56,7 @@ def update_example_dependencies(file_path: Path, new_version: str) -> None: def update_plugin_version(file_path: Path, new_version: str) -> None: - """Update the version in the Claude Code plugin manifest.""" + """Update the version in a plugin manifest.""" content = file_path.read_text() updated = re.sub( r'^(\s*"version": )"[^"]+"', @@ -136,12 +136,15 @@ def main(): root / "app" / "backend" / "pyproject.toml", ] - plugin_file = root / "claude-plugin" / ".claude-plugin" / "plugin.json" + plugin_files = [ + root / "plugins" / "haiku-rag" / ".claude-plugin" / "plugin.json", + root / "plugins" / "haiku-rag" / ".codex-plugin" / "plugin.json", + ] changelog_file = root / "CHANGELOG.md" # Check all files exist for file in ( - pyproject_files + example_pyproject_files + [plugin_file, changelog_file] + pyproject_files + example_pyproject_files + plugin_files + [changelog_file] ): if not file.exists(): print(f"Error: {file} not found") @@ -172,7 +175,8 @@ def main(): for file in example_pyproject_files: update_example_dependencies(file, new_version) - update_plugin_version(plugin_file, new_version) + for file in plugin_files: + update_plugin_version(file, new_version) # Update CHANGELOG.md update_changelog(changelog_file, new_version) diff --git a/tests/test_bump_version.py b/tests/test_bump_version.py index 4f8fcf4f..4b6fdc27 100644 --- a/tests/test_bump_version.py +++ b/tests/test_bump_version.py @@ -27,12 +27,15 @@ def test_update_plugin_version_rewrites_only_the_version_field(tmp_path, monkeyp assert manifest.read_text().endswith("}\n") -def test_the_shipped_plugin_manifest_carries_the_package_version(): +def test_the_shipped_plugin_manifests_carry_the_package_version(): root = Path(__file__).resolve().parents[1] - plugin = json.loads( - (root / "claude-plugin" / ".claude-plugin" / "plugin.json").read_text() - ) - - assert plugin["version"] == bump_version.get_current_version( + package_version = bump_version.get_current_version( root / "haiku_rag_slim" / "pyproject.toml" ) + for client in ("claude", "codex"): + manifest = json.loads( + ( + root / "plugins" / "haiku-rag" / f".{client}-plugin" / "plugin.json" + ).read_text() + ) + assert manifest["version"] == package_version diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 9a9faee6..18d7106d 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1158,27 +1158,39 @@ class TestMCPErrorContract: ) -class TestClaudeCodePlugin: - """The plugin under claude-plugin/ points at the server this module builds.""" +class TestAgentPlugins: + """The shared Claude Code and Codex plugin points at this MCP server.""" root = Path(__file__).resolve().parents[1] def test_the_manifests_name_the_plugin_and_its_server(self): import json - plugin = json.loads( - (self.root / "claude-plugin/.claude-plugin/plugin.json").read_text() + claude_plugin = json.loads( + (self.root / "plugins/haiku-rag/.claude-plugin/plugin.json").read_text() ) - marketplace = json.loads( + codex_plugin = json.loads( + (self.root / "plugins/haiku-rag/.codex-plugin/plugin.json").read_text() + ) + claude_marketplace = json.loads( (self.root / ".claude-plugin/marketplace.json").read_text() ) - servers = json.loads((self.root / "claude-plugin/.mcp.json").read_text()) + codex_marketplace = json.loads( + (self.root / ".agents/plugins/marketplace.json").read_text() + ) + servers = json.loads((self.root / "plugins/haiku-rag/.mcp.json").read_text()) - assert plugin["name"] == "haiku-rag" - assert plugin["description"] - [entry] = marketplace["plugins"] - assert entry["name"] == plugin["name"] - assert entry["source"] == "./claude-plugin" + assert claude_plugin["name"] == codex_plugin["name"] == "haiku-rag" + assert claude_plugin["description"] + assert codex_plugin["description"] + assert codex_plugin["skills"] == "./skills/" + assert codex_plugin["mcpServers"] == "./.mcp.json" + [claude_entry] = claude_marketplace["plugins"] + assert claude_entry["name"] == claude_plugin["name"] + assert claude_entry["source"] == "./plugins/haiku-rag" + [codex_entry] = codex_marketplace["plugins"] + assert codex_entry["name"] == codex_plugin["name"] + assert codex_entry["source"]["path"] == "./plugins/haiku-rag" assert servers["mcpServers"]["haiku-rag"]["args"] == ["mcp", "--stdio"] @pytest.mark.asyncio @@ -1187,7 +1199,7 @@ class TestClaudeCodePlugin: ): import yaml - text = (self.root / "claude-plugin/skills/haiku-rag/SKILL.md").read_text() + text = (self.root / "plugins/haiku-rag/skills/haiku-rag/SKILL.md").read_text() _, frontmatter, _ = text.split("---", 2) skill = yaml.safe_load(frontmatter) prefix = "mcp__plugin_haiku-rag_haiku-rag__" From bae359e3bf02102d3e62e9e93bde4d5ae7f92303 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 7 Sep 2026 12:02:33 +0300 Subject: [PATCH 16/18] Align the MCP guidance with the analysis instructions and move to Monty 0.0.23 The execute_code description gains the toc.json node shape, the patterns the evals put into the analysis instructions (one list_documents call to map a title to an id, files carry no source, toc before search for a known document, doc_item_refs are items self_refs, chunk ids join files and are not citations) and the sandbox's read-only, no-network, time-limit and output facts, which the analysis instructions now state too. The skill gains pictures as images, image search when offered, and citing chunk metadata locators. docs/mcp.md lists the interpreter's limits under Code. pydantic-monty>=0.0.23 brings collections, itertools, functools, dataclasses, function decorators and str.format into the sandbox; every layer names the same modules, and a test imports them. Monty now caps host callbacks per checkout at 1000 by default; the sandbox raises it out of reach, since the time budgets govern. --- CHANGELOG.md | 3 + docs/capabilities/analysis.md | 2 +- docs/mcp.md | 19 ++- .../rag/capabilities/instructions/analysis.md | 8 +- haiku_rag_slim/haiku/rag/mcp.py | 34 +++-- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 11 +- haiku_rag_slim/pyproject.toml | 2 +- plugins/haiku-rag/skills/haiku-rag/SKILL.md | 21 ++- tests/sandbox/test_sandbox.py | 58 ++++++- uv.lock | 142 ++++++++---------- 10 files changed, 195 insertions(+), 105 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ace83d2..3070696d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,9 @@ ### Changed +- `pydantic-monty>=0.0.23`. The analysis sandbox gains `collections`, + `itertools`, `functools`, `dataclasses`, function decorators and + `str.format`. - `fastmcp>=4.0.2,<5.0.0`, on MCP Python SDK 2. The MCP server answers both the session-based and the sessionless (2026-07-28) protocol. - Default models are `ollama:qwen3.8`: `ModelConfig`, `qa.model`, diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md index 604eb3d3..18cd12bd 100644 --- a/docs/capabilities/analysis.md +++ b/docs/capabilities/analysis.md @@ -16,7 +16,7 @@ When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool | `analysis_execute_code(code)` | Run Python against the virtual document filesystem. | | `analysis_cite(chunk_ids)` | Register retrieved or filesystem-derived chunk IDs. | -The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, `chunks.jsonl` (chunk ids with their metadata) and `toc.json`. In code, `await search()` results carry `chunk_meta` and `await list_documents()` rows carry `metadata`. +The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, `chunks.jsonl` (chunk ids with their metadata) and `toc.json`. In code, `await search()` results carry `chunk_meta` and `await list_documents()` rows carry `metadata`. The interpreter's limits and the per-call budgets are listed under [MCP, Code](../mcp.md#code). ## Compose an agent diff --git a/docs/mcp.md b/docs/mcp.md index 4d268d23..d1cb4bac 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -151,18 +151,29 @@ node's `id` in the outline is the `section_id`. A document without headings has an empty outline. `list_documents` returns titles, URIs and metadata, which is how a client learns what a filter can match. +### Code + `execute_code` runs a Python program in the sandbox of the [analysis capability](capabilities/analysis.md), over the documents `filter` and `sources` select, and returns what it printed. The program reads `/documents/{document_id}/` (`metadata.json`, `content.txt`, `items.jsonl`, `chunks.jsonl`, `toc.json`) and can `await search()` and `await list_documents()`; the tool description spells out the fields and the -interpreter's limits. Each call is one program: nothing carries over between +patterns that matter. Each call is one program: nothing carries over between calls, and the sandbox is created and closed per call. A failing program is a tool error carrying the interpreter's message and any output printed before -it. `analysis.code_timeout` bounds a call and `analysis.max_output_chars` its -output; no model runs on the server. Claude Code moves a call still running -after about two minutes to a background task. +it. No model runs on the server. Claude Code moves a call still running after +about two minutes to a background task. + +The interpreter is [Monty](https://github.com/pydantic/monty), a Python subset. +Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`, +`collections`, `itertools`, `functools` and `dataclasses`. Absent, and often +reached for: `decimal` and `statistics`. No generator functions, class +inheritance or `match` statements, and a file object cannot be iterated. Files are read-only, and +there is no network and no filesystem +beyond `/documents`. `analysis.code_timeout` bounds a call, counted separately +for compute and for document reads, and `analysis.max_output_chars` bounds its +output. ### Filters diff --git a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md index 37dc8807..808edf59 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md +++ b/haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md @@ -16,8 +16,8 @@ Inside the code, these functions are available (use `await`): - `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`), chunk_meta (the matched chunk's stored metadata, custom keys included) - `await list_documents()` → list of dicts with keys: id, title, uri, created_at, metadata -Available modules: `json`, `re`, `math`, `pathlib` -Not supported: class inheritance and metaclasses, generators/yield, match statements, decorators, `collections`, iterating a file object (`for line in f`) +Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`, `collections`, `itertools`, `functools` and `dataclasses`. `decimal` and `statistics` do not exist. +Not supported: class inheritance and metaclasses, generators/yield, match statements, iterating a file object (`for line in f`) ### analysis_search Search the knowledge base directly (outside code execution). Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots. @@ -49,7 +49,7 @@ All documents are mounted as a virtual filesystem at `/documents/`: `{document_id}` is an internal identifier, not the user-facing `uri` (filename, URL, etc.). When you only know a document by its URI or title, use `await list_documents()` to enumerate ids and match against `uri` / `title` — that's a single call to the host. Iterating `/documents/` and reading every `metadata.json` works too but is much slower on portal-scale corpora. ### Reading files -Read with `Path.read_text()` or `open()` (including `with` blocks); file objects support `.read()`, `.readline()`, and `.readlines()`. A file object cannot be iterated, so read line-wise with `.readlines()` or `.read().split("\n")` instead of `for line in f`. Files are read-only; writing raises `PermissionError`. +Read with `Path.read_text()` or `open()` (including `with` blocks); file objects support `.read()`, `.readline()`, and `.readlines()`. A file object cannot be iterated, so read line-wise with `.readlines()` or `.read().split("\n")` instead of `for line in f`. Files are read-only; writing raises `PermissionError`. There is no network. A call has a time limit, named in the error when it is hit, and output past a size is cut with an `... (output truncated)` marker. ```python from pathlib import Path @@ -116,6 +116,6 @@ You MUST call `analysis_cite` before producing your final answer, every time, wi - Use `print()` to output results — the output is your only feedback - When you write code, execute it — don't describe what code would do. But not every question needs code; simple lookups are best answered by `analysis_search → analysis_cite`. - Use `await` for all async functions inside `analysis_execute_code` (`search`, `list_documents`) -- Read files with `Path.read_text()` or `open()`/`with`. For lines use `.readlines()` or `.read().split("\n")`, never `for line in f`. The `collections` module is unavailable. +- Read files with `Path.read_text()` or `open()`/`with`. For lines use `.readlines()` or `.read().split("\n")`, never `for line in f`. - Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `analysis_cite` tool separately to register citations. `cite{...}` markdown-style inline references do nothing; only an actual `analysis_cite` tool call registers a citation. - **Before you write your final answer, invoke the `analysis_cite` tool with the supporting chunk_ids, or with an empty list if there are none.** This is the last tool call before answering, every time. diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index b9eb2bb1..58a94e9b 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -464,15 +464,31 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: (id, title, uri, created_at, metadata), `content.txt` (the whole text), `items.jsonl` (one item per line: self_ref, label, text, page_numbers, heading_level, chunk_ids), `chunks.jsonl` (one chunk per line: chunk_id, - metadata) and `toc.json` (the section tree, each node with an item_range - slice into items.jsonl). Read files with `Path.read_text()` or `open()`; - a file object cannot be iterated, use `.readlines()`. - `await search(query, limit=10)` returns dicts with chunk_id, content, - document_id, document_title, document_uri, source, score, page_numbers, - headings, doc_item_refs, labels and chunk_meta. `await list_documents()` - returns dicts with id, title, uri, created_at, source and metadata. - Modules: json, re, math, pathlib. Not available: generators, class - inheritance, match statements, decorators, collections. + metadata) and `toc.json` (`doc_id`, `title`, `tree`; each node has + self_ref, level, title, page_numbers, item_range as a slice into + items.jsonl, chunk_ids and children; an empty tree means no headings). + Read files with `Path.read_text()` or `open()`; a file object cannot be + iterated, use `.readlines()`. `await search(query, limit=10)` returns + dicts with chunk_id, content, document_id, document_title, document_uri, + source, score, page_numbers, headings, doc_item_refs, labels, + picture_refs (the doc_item_refs that are pictures) and chunk_meta. + `await list_documents()` returns dicts with id, title, uri, created_at, + source and metadata. Both see the documents `filter` and `sources` + select. Useful modules include json, re, math, pathlib, datetime, + collections, itertools, functools and dataclasses; decimal and + statistics do not exist. No generator functions, match statements or + class inheritance. + Files are read-only, there is no network, a call has a time limit named + in the error when it is hit, and output past a size is truncated. + + Map a title or URI to a document id with one `list_documents()` call + rather than reading every `metadata.json`. The files carry no `source`, + so over several collections group by the `source` of `list_documents()` + rows. For a known document's structure read its `toc.json` before + searching: `search()` ranks across every document. A hit's + `doc_item_refs` are `self_ref` values in `items.jsonl`, which places it + in its section. `chunk_ids` on items and `chunk_id` in `chunks.jsonl` + join the two files; they are not citations. Args: code: The program. Use `await` on search and list_documents. diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 47149f1e..823f0040 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -33,6 +33,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +_MAX_HOST_CALLS = 10_000_000 + def _host_failure(where: str, e: Exception) -> RuntimeError: """The error a program gets for a failure on the host side of a call. @@ -613,9 +615,16 @@ class Sandbox: covers the whole run. ``code_timeout`` is enforced per call elsewhere: the read deadline in ``_run_on_loop`` bounds a call that reads, and the pool's ``request_timeout`` bounds one that computes. + + ``max_suspensions`` counts host callbacks per session, document reads + included, defaults to 1000 and cannot be disabled. The time budgets are + the governors here, so it is set where no program reaches it. """ analysis = self._config.analysis - return {"max_duration_secs": analysis.code_timeout * analysis.max_executions} + return { + "max_duration_secs": analysis.code_timeout * analysis.max_executions, + "max_suspensions": _MAX_HOST_CALLS, + } async def _ensure_initialized(self) -> tuple[AsyncMontySession, OSAccess]: """Check out a worker session and build the VFS on first use.""" diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index b3e5dbc6..a304e941 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -45,7 +45,7 @@ dependencies = [ "pathspec>=1.0.4", "pydantic>=2.12.5", "pydantic-ai-slim[openai,logfire,ag-ui]>=2.18.0,<3.0.0", - "pydantic-monty>=0.0.19", + "pydantic-monty>=0.0.23", "pypdfium2>=5.0", "python-dotenv>=1.2.2", "pyyaml>=6.0.3", diff --git a/plugins/haiku-rag/skills/haiku-rag/SKILL.md b/plugins/haiku-rag/skills/haiku-rag/SKILL.md index c513fd2b..532186dc 100644 --- a/plugins/haiku-rag/skills/haiku-rag/SKILL.md +++ b/plugins/haiku-rag/skills/haiku-rag/SKILL.md @@ -25,10 +25,12 @@ could be about the user's documents. Say so when it has nothing relevant. `search_documents` is the first call. Results come best first with the document title, section headings, the matched chunk's metadata when it has any, and the -passage in its section. `filter` restricts which -documents are searched, `limit` how many results come back. If it misses, -rephrase once or narrow with a filter before concluding the material is not -there. +passage in its section. Pictures in the results arrive as images: answer +figure questions from them. `filter` restricts which documents are searched, +`limit` how many results come back. If it misses, rephrase once or narrow with +a filter before concluding the material is not there. When the question is +about an image rather than words and the server offers +`search_documents_by_image`, it takes the image as the query. ## Read @@ -48,7 +50,11 @@ comparison across many documents, a lookup by document or chunk metadata, or a pattern over whole documents: whatever search cannot rank. Each call is one program and variables do not carry over, so gather, compute and `print` a compact result in the same program. `filter` and `sources` select the documents -it sees. Answer and cite from what it printed. +it sees. For a known document's structure read its `toc.json` first; `search()` +ranks across every document. Map a title or URI to an id with one +`list_documents()` call rather than reading every `metadata.json`; the files +carry no `source`, so over several collections group by its rows. Answer and +cite from what it printed. ## Explore @@ -66,5 +72,6 @@ it with LIKE: `metadata LIKE '%"author": "Smith"%'`. Also `uri LIKE '%.pdf'`, Rank is the signal; scores are not comparable across queries and are never confidence. Cite the document title or URI, the section heading and page -numbers when present. When results carry `source`, the server covers several -collections: name it, and pass `sources` to search a subset. +numbers when present, and the matched chunk's metadata when it carries locators +such as paragraph or footnote numbers. When results carry `source`, the server +covers several collections: name it, and pass `sources` to search a subset. diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index cb6a1946..98465329 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -19,6 +19,18 @@ def vcr_cassette_dir(): class TestSandboxBasics: """Test basic sandbox functionality.""" + @pytest.mark.asyncio + async def test_the_documented_modules_import(self, sandbox): + """The modules the instructions and the MCP description promise.""" + result = await sandbox.execute( + "import json, re, math, pathlib, datetime\n" + "import collections, itertools, functools, dataclasses\n" + "print(collections.Counter('aab').most_common(1)," + " list(itertools.islice(itertools.count(), 2)))" + ) + assert result.success, result.stderr + assert "[('a', 2)] [0, 1]" in result.stdout + @pytest.mark.asyncio async def test_execute_simple_code(self, sandbox): """Test executing simple code in the sandbox.""" @@ -1048,7 +1060,51 @@ class TestSandboxReadDeadline: sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) - assert sb._session_limits() == {"max_duration_secs": 15.0} + limits = sb._session_limits() + + assert limits["max_duration_secs"] == 15.0 + cap = limits["max_suspensions"] + assert cap is not None + assert cap >= 1_000_000 + + @pytest.mark.asyncio + async def test_a_program_may_read_more_than_a_thousand_times(self, temp_db_path): + """Monty caps host callbacks per checkout at 1000 unless told otherwise; + a corpus-wide pass over documents reads far more than that.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel + + config = AppConfig() + docling = DoclingDocument(name="d") + docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.") + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.import_document( + docling, + [ + Chunk( + content="Foxes and dogs.", + embedding=[0.1] * config.embeddings.model.vector_dim, + order=0, + ) + ], + uri="test://many-reads", + ) + + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + "from pathlib import Path\n" + f"p = Path('/documents/{doc.id}/content.txt')\n" + "n = 0\n" + "for i in range(1100):\n" + " n += len(p.read_text())\n" + "print(n)" + ) + finally: + await sb.close() + + assert result.success, result.stderr + assert result.stdout.strip() == str(1100 * len("Foxes and dogs.")) @pytest.mark.asyncio async def test_refused_read_fails_the_execution(self, temp_db_path, monkeypatch): diff --git a/uv.lock b/uv.lock index 4d9636e2..eba87258 100644 --- a/uv.lock +++ b/uv.lock @@ -1776,7 +1776,7 @@ requires-dist = [ { name = "pydantic-ai-slim", extras = ["mistral"], marker = "extra == 'mistral'" }, { name = "pydantic-ai-slim", extras = ["openai", "logfire", "ag-ui"], specifier = ">=2.18.0,<3.0.0" }, { name = "pydantic-ai-slim", extras = ["voyageai"], marker = "extra == 'voyageai'" }, - { name = "pydantic-monty", specifier = ">=0.0.19" }, + { name = "pydantic-monty", specifier = ">=0.0.23" }, { name = "pypdfium2", specifier = ">=5.0" }, { name = "python-dotenv", specifier = ">=1.2.2" }, { name = "pyyaml", specifier = ">=6.0.3" }, @@ -3991,94 +3991,82 @@ wheels = [ [[package]] name = "pydantic-monty" -version = "0.0.19" +version = "0.0.23" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "pydantic-monty-client" }, { name = "pydantic-monty-runtime" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/59/01/cd7927f51500a13c1661e2db6488f6ff668194c4378d0fb218928892aaba/pydantic_monty-0.0.23.tar.gz", hash = "sha256:ee674b81ed12f81cfbe0db210fe5e803c754d0682634331931f4d70d5742058d", size = 6713, upload-time = "2026-09-05T19:26:03.24Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/b3/a259a600df30a54ee9e576c90bb7cb6c2b2d4750774c91748922cc202c5e/pydantic_monty-0.0.23-py3-none-any.whl", hash = "sha256:cddcf7d4d7dd163b56e4411f450ea0244a7205988c77cd85441ec89fedaab91a", size = 6302, upload-time = "2026-09-05T19:24:01.829Z" }, +] + +[[package]] +name = "pydantic-monty-client" +version = "0.0.23" +source = { registry = "https://pypi.org/simple" } +dependencies = [ { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/51/f8/c04df414086488834d102ab85cf8172c114108f842fb142ac92a490213fd/pydantic_monty-0.0.19.tar.gz", hash = "sha256:f3f9e256058b4085349dd4ad347795d5203a8310e9eaad9a2bff93890ac5b86d", size = 1484032, upload-time = "2026-07-24T10:00:13.194Z" } +sdist = { url = "https://files.pythonhosted.org/packages/29/e9/dcbf9a90ccd195be0db5257c2234d002d551d47f9d8987543eddd3095812/pydantic_monty_client-0.0.23.tar.gz", hash = "sha256:8b31f75afebb60c0416c869c8d6667a75e84b0f05b7a0b44e51b77ef98f294d7", size = 1983536, upload-time = "2026-09-05T19:26:04.167Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/36/bf/152bbb3315dfa46d4e4aae71779230e50c67a34d859a3470fd75c01b795c/pydantic_monty-0.0.19-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:b073e64edfd62cca918d792d6fe559512472f981f949e53a7aec673201f5f554", size = 2492733, upload-time = "2026-07-24T09:56:49.612Z" }, - { url = "https://files.pythonhosted.org/packages/46/16/9d37f1bf94c47c593a06ecb918bb64a2c8a38af24d6fa2b0d0e031938edf/pydantic_monty-0.0.19-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5825ae6c40270166f0f31b6e62d59b0fe47201d12b1be5842efc22d3b7dadcee", size = 2234610, upload-time = "2026-07-24T09:56:51.257Z" }, - { url = "https://files.pythonhosted.org/packages/82/4f/31732f4b9c2b6574eb564e420593b9be3f80d92440211970fab23e882bc5/pydantic_monty-0.0.19-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5b6821c0035ff2c02cf0f37823fb905890f7238ce923bbc3ca937740f9554836", size = 2303116, upload-time = "2026-07-24T09:56:52.581Z" }, - { url = "https://files.pythonhosted.org/packages/21/2c/c21e6179361100dd8b9ad410df775d176a29e0b85f2ffc3144fd29840ee9/pydantic_monty-0.0.19-cp312-cp312-manylinux_2_28_armv7l.whl", hash = "sha256:5fb4514d99a10e304237a82baeb6072e934ce8462dcd5fb5c3fea0ee0ef16aeb", size = 2007947, upload-time = "2026-07-24T09:56:54.006Z" }, - { url = "https://files.pythonhosted.org/packages/7b/f0/ad64f4894334499f689bdce7e5b5dda6b680d98989424092dcbf21666564/pydantic_monty-0.0.19-cp312-cp312-manylinux_2_28_i686.whl", hash = "sha256:2b0f154ac2e6450befa337a99a8519e2f1195cc59f88bd73d780067ace0c4c97", size = 2151468, upload-time = "2026-07-24T09:56:55.626Z" }, - { url = "https://files.pythonhosted.org/packages/2a/97/4ff5f9170e4865ec28fb152bc6ebaa8bd1873e695ee98af2103607ba42e8/pydantic_monty-0.0.19-cp312-cp312-manylinux_2_28_ppc64le.whl", hash = "sha256:6fcb4e3a312397432f7c5a0aa04a765670d9f37f4e23f37cb580b3faa2f218b1", size = 2300164, upload-time = "2026-07-24T09:56:57.318Z" }, - { url = "https://files.pythonhosted.org/packages/1a/b1/c9bfb6708bc5d9f6b040db46156c6e3f16de68ab22f07e2de4f25782414b/pydantic_monty-0.0.19-cp312-cp312-manylinux_2_28_s390x.whl", hash = "sha256:955d308171ba0260f75d2dede18c46d9732647f94b8f7fba3076bb539f155dc8", size = 2164984, upload-time = "2026-07-24T09:56:58.66Z" }, - { url = "https://files.pythonhosted.org/packages/35/2d/8c1492216632f53e229cb2886c7fd09613d5990f4856f93f7ae0364da9bc/pydantic_monty-0.0.19-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:3d7e1a6cc1977b24e01b5322888a5ba3f05b112a04a1646ba51f2c34c562a3d9", size = 2357823, upload-time = "2026-07-24T09:57:00.048Z" }, - { url = "https://files.pythonhosted.org/packages/9b/cc/ae4adaaf343de00748fc3598402c1523e48db7094a9c1e0fa26177699440/pydantic_monty-0.0.19-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:4904785a34f71f59aa1eb6dc081bafd09e4caa2f028cd379aa3aa846b8a1c27d", size = 2497387, upload-time = "2026-07-24T09:57:01.686Z" }, - { url = "https://files.pythonhosted.org/packages/ec/09/eaef87ed6cbffed9c720dd3cb850e748b0aab11f5ec1ecffb56250973f7d/pydantic_monty-0.0.19-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:ce15a5326e24cf7045ec9c1456ddb92abd09df91708771c3ae5e72d8e6374c81", size = 2738391, upload-time = "2026-07-24T09:57:03.3Z" }, - { url = "https://files.pythonhosted.org/packages/63/af/58be6fd6ea87e27bd57435013ca1f63d645a0c990c2f23708bcdd048af24/pydantic_monty-0.0.19-cp312-cp312-win32.whl", hash = "sha256:600eb259415e8b2dfef4be38d030c945b3fbb4fb85e727cb97131e7329ab017f", size = 1908274, upload-time = "2026-07-24T09:57:05.023Z" }, - { url = "https://files.pythonhosted.org/packages/20/b7/1cb54e43113cb69c40fb765cfee3be1c222d81153b432131f50508569aee/pydantic_monty-0.0.19-cp312-cp312-win_amd64.whl", hash = "sha256:2c98b1c99994f92ab487a762b107067ab70036f64231e05aa3f6d2b16018688e", size = 2111335, upload-time = "2026-07-24T09:57:06.614Z" }, - { url = "https://files.pythonhosted.org/packages/24/17/0926da051f34ccaa45bf528777dc99e5ea611669cdd7b715be1e086c1fe0/pydantic_monty-0.0.19-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:c01fb1162cf87dbf145b875450eabfde6b35b26f27ed63468398cb4c37732064", size = 2496669, upload-time = "2026-07-24T09:57:07.982Z" }, - { url = "https://files.pythonhosted.org/packages/fd/ba/c38f79e24971b4d2205ee156df6e5c6ccdcc720152f54a956cc26e7488b0/pydantic_monty-0.0.19-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e97dc97de32410003fb5517b72ab47799b4546a3ab48a75e60260cf73734da76", size = 2234599, upload-time = "2026-07-24T09:57:09.458Z" }, - { url = "https://files.pythonhosted.org/packages/5a/89/fb2ed1677ad2c3e2801aa119fdc280d1c64f3480e1015811e0942c9a7d20/pydantic_monty-0.0.19-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:06318e6add780f66259067829ab16062ce7a556dba0884372a881881b4a7c3bf", size = 2305696, upload-time = "2026-07-24T09:57:10.97Z" }, - { url = "https://files.pythonhosted.org/packages/24/ec/e36ff1c2c97f46420d57680199e7ae2e1eb306d2cfda22be2448e53e7484/pydantic_monty-0.0.19-cp313-cp313-manylinux_2_28_armv7l.whl", hash = "sha256:5a765a36141bbeb074d89b424d5488bed4e358c603c72606cc9d81ee44d83de2", size = 2007600, upload-time = "2026-07-24T09:57:12.333Z" }, - { url = "https://files.pythonhosted.org/packages/da/88/b47670d3e28f99dc2f4c2686dc42233843dce1a607f3d3cc8a387f3b9b74/pydantic_monty-0.0.19-cp313-cp313-manylinux_2_28_i686.whl", hash = "sha256:48f5ebc048779c854f993834586ba372df3b65500d1ed7f147023abc29aeb2c4", size = 2151382, upload-time = "2026-07-24T09:57:13.8Z" }, - { url = "https://files.pythonhosted.org/packages/9f/18/559d71fc66a769c22f6b2a8470515aa6e69a141ab617900fe28a806080b7/pydantic_monty-0.0.19-cp313-cp313-manylinux_2_28_ppc64le.whl", hash = "sha256:2eb7502f93d8bb568d255fccca95e3b9daca05bb674b9c5323fcba14d088932c", size = 2302643, upload-time = "2026-07-24T09:57:15.42Z" }, - { url = "https://files.pythonhosted.org/packages/6f/b4/171be42e2ec211bd01fe4be7b6de9ab0c346081edc33830f9f9b781341ab/pydantic_monty-0.0.19-cp313-cp313-manylinux_2_28_s390x.whl", hash = "sha256:608994600a839dd863940ac84810c48d378390a4fea8c77e5145feb84037e610", size = 2169103, upload-time = "2026-07-24T09:57:16.805Z" }, - { url = "https://files.pythonhosted.org/packages/84/ce/8c47253c8f3f528f0fa2d87dde85623dc3f211189a372e2d69cb6ad896b1/pydantic_monty-0.0.19-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:13a652f9dffa6d25ea458fec83108f60f29682caf42cecef91955b5b8cb04365", size = 2358155, upload-time = "2026-07-24T09:57:18.51Z" }, - { url = "https://files.pythonhosted.org/packages/97/4e/239107b83bae3a20e4f5424f6c6f3f88bae1fd1e734603c298167985f8be/pydantic_monty-0.0.19-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:6167c966db7d8d2fe03940f8da596de6dcd184dcf5cf6209f0a1a9af8792550d", size = 2500858, upload-time = "2026-07-24T09:57:19.878Z" }, - { url = "https://files.pythonhosted.org/packages/48/55/02a7059f20e7198e511ac0b88a3957a2c01b8991326359b7c1bb590a09b8/pydantic_monty-0.0.19-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:3102b44307d04f41897ae51d4daeab4fc9b5bf84a78c036781b488876ef998c3", size = 2742212, upload-time = "2026-07-24T09:57:21.296Z" }, - { url = "https://files.pythonhosted.org/packages/a8/28/b632fe0e8eeba3f2b4dbec6ebd4c9b569c20e9597007f2d0fbbf04101307/pydantic_monty-0.0.19-cp313-cp313-win32.whl", hash = "sha256:e43da52776796a894f40533a7e5a322e98d9aaf7d8f6fbb7dc21a0de60a93f41", size = 1908553, upload-time = "2026-07-24T09:57:22.765Z" }, - { url = "https://files.pythonhosted.org/packages/f7/d3/b90872f017871339ceb03e70fa4915ef8682128a476a66adffedfff874d8/pydantic_monty-0.0.19-cp313-cp313-win_amd64.whl", hash = "sha256:fd8c195875f8f44d55bc7d1d53c4e43248b184b3ac803d8131c92d4cc05a1aef", size = 2111345, upload-time = "2026-07-24T09:57:24.351Z" }, - { url = "https://files.pythonhosted.org/packages/eb/11/c2aed55502bfc9620837312f0e2fca7a3d4bc959824a66d81b72144d7256/pydantic_monty-0.0.19-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:2692dc4452937cf2200afd257297e9ac3ccff122b80aa7a69935e8275d684193", size = 2497017, upload-time = "2026-07-24T09:57:25.836Z" }, - { url = "https://files.pythonhosted.org/packages/f5/d0/d4b44a81c71308109cfa642a24b803057615ca1609530c5e66c376780efe/pydantic_monty-0.0.19-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:5a4717f829b35c4bc5d9f6f52a52d19b90729bd494bcb69133bd4c0afcab9c76", size = 2247468, upload-time = "2026-07-24T09:57:27.501Z" }, - { url = "https://files.pythonhosted.org/packages/a9/89/52848dce3acbb1d58df34496db3c4718814cc39f6db01a5025bfdc12b530/pydantic_monty-0.0.19-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:99f7282213ad6ebf7daf149a88d1517e6328afbf6780180512b9de62f30d292b", size = 2306181, upload-time = "2026-07-24T09:57:29.157Z" }, - { url = "https://files.pythonhosted.org/packages/90/05/f7791c79c7be2240c43a9287196ed49034f773e3f4158492f028e486ebc2/pydantic_monty-0.0.19-cp314-cp314-manylinux_2_28_armv7l.whl", hash = "sha256:978788b1c56fa0c49927e0a35151f0a635ef2f8d947d16915541ff864d2112f5", size = 2008738, upload-time = "2026-07-24T09:57:30.423Z" }, - { url = "https://files.pythonhosted.org/packages/93/a1/d714258eeb2583acaab2035834acc54ac6baa65bac456f897d901bc0c03a/pydantic_monty-0.0.19-cp314-cp314-manylinux_2_28_i686.whl", hash = "sha256:8b821cac39deeb1abb2994d5a65f117ce26e55c586d0f570d425ff469ff48e3c", size = 2152275, upload-time = "2026-07-24T09:57:31.725Z" }, - { url = "https://files.pythonhosted.org/packages/2a/e5/cdb1fa761e992a489b07a17adb80c21bf99eabd1286ca62f9e73acda1f4b/pydantic_monty-0.0.19-cp314-cp314-manylinux_2_28_ppc64le.whl", hash = "sha256:b43c4ffa5651f0eca97458dc673d7952064ae5eb5b836d23967a7d41483bb8f4", size = 2303533, upload-time = "2026-07-24T09:57:33.131Z" }, - { url = "https://files.pythonhosted.org/packages/db/9b/e6685cf82521e68e0dcb94e0c97a1a20410b1266fbce7008c0caa3484039/pydantic_monty-0.0.19-cp314-cp314-manylinux_2_28_s390x.whl", hash = "sha256:4054610601358943a3dd740a8d59a6812cc682e94d6903cb72baadae1ef5d2ac", size = 2169585, upload-time = "2026-07-24T09:57:34.511Z" }, - { url = "https://files.pythonhosted.org/packages/df/f4/6f031a628d3de72d95bedbb18292ccf998d9a422aff58eedf37347a0b367/pydantic_monty-0.0.19-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:0b2a34c320968a3cef3d933737e99c932f13c55667315128f366afdaeea2be04", size = 2375171, upload-time = "2026-07-24T09:57:35.907Z" }, - { url = "https://files.pythonhosted.org/packages/0d/d3/4367bccdf2c06a977c0d5ddf816190d570a07199f011034df16d51c84723/pydantic_monty-0.0.19-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:ffaf950f4284bb193a54f18fa4e3e8c225b5b52265813abffe492074e90c65fb", size = 2501475, upload-time = "2026-07-24T09:57:37.872Z" }, - { url = "https://files.pythonhosted.org/packages/0b/0e/7a0e1d5c016848afc9a8605aa4f16e6960d68806e775865b986aacaf8f88/pydantic_monty-0.0.19-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:bb1e5ee6762f9494bfb0f4cc316ad3e9a8c7917e3c984a23eb2e2322d401e5cb", size = 2742547, upload-time = "2026-07-24T09:57:39.284Z" }, - { url = "https://files.pythonhosted.org/packages/42/f0/92f77cf088f1df72a5dbab109c8c0e82f7ddeb18e65835a8dffcf967bb4a/pydantic_monty-0.0.19-cp314-cp314-win32.whl", hash = "sha256:b68ef6503b39f2f014162e3d8e7f48b9722a43ceb7b6d7199dc6d545f4c37b34", size = 1907832, upload-time = "2026-07-24T09:57:40.982Z" }, - { url = "https://files.pythonhosted.org/packages/cc/3e/85e1a914f81659ea25c34916fdef27fcda2b00323dafb76cf2a5321f7c11/pydantic_monty-0.0.19-cp314-cp314-win_amd64.whl", hash = "sha256:14ef37b43c5bf90966ca51bcad0fae892a3c2546cd151fadc039e0a25ca74073", size = 2123187, upload-time = "2026-07-24T09:57:42.352Z" }, + { url = "https://files.pythonhosted.org/packages/7e/f3/4e4975daee5ac4a86e95b08124e304cf7af774df50891b28479bb2c28ede/pydantic_monty_client-0.0.23-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:7f309a03e75c418b34e469e547171059ebd1ebcd11971ed129b6905e2a7c057f", size = 3930126, upload-time = "2026-09-05T19:24:38.269Z" }, + { url = "https://files.pythonhosted.org/packages/6e/23/927d73209b509db9604606d88aaff90f56474bd9bdcfe9f39a873bfc945c/pydantic_monty_client-0.0.23-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0e3efb7e2af60e17c16f3baa49ea3a2e2b90c88610666f294591b9819d3ccc69", size = 3702556, upload-time = "2026-09-05T19:24:39.81Z" }, + { url = "https://files.pythonhosted.org/packages/ff/c9/bfb5847d79e8bdeefbeea531b90f09e2c5f156bc361730975171a7e93146/pydantic_monty_client-0.0.23-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f065e8286378f26d29ceaeb2fb10712f202f45656967706675f93b3a17b99f78", size = 3728660, upload-time = "2026-09-05T19:24:41.129Z" }, + { url = "https://files.pythonhosted.org/packages/13/9c/50af7e5a67876cb7c421401ff5f3885ae4fbd60324fea3c25e5fd31926af/pydantic_monty_client-0.0.23-cp312-cp312-manylinux_2_28_armv7l.whl", hash = "sha256:d3e353536016174f9e0aca289f7d3a81a07ac7574e2a610902152e0ebf5feb5a", size = 3425561, upload-time = "2026-09-05T19:24:42.539Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e3/1d1a53f72576c190ea2f1a6c0fb5eeb722bea6901df8ddfba294ff49b08b/pydantic_monty_client-0.0.23-cp312-cp312-manylinux_2_28_i686.whl", hash = "sha256:ef186bfe5344a84515cd51dc9ee5914084b1d616fad2409328d8edfbfb4d5ac3", size = 3633369, upload-time = "2026-09-05T19:24:44.21Z" }, + { url = "https://files.pythonhosted.org/packages/6f/fb/c5987eeff60e9d736a720ef1c6f3dadee39619f596517d8983a99ce4873a/pydantic_monty_client-0.0.23-cp312-cp312-manylinux_2_28_ppc64le.whl", hash = "sha256:7062196b773cc8b956b0dbece42203eb493270e0707e37447801255ad26f065e", size = 3843321, upload-time = "2026-09-05T19:24:45.55Z" }, + { url = "https://files.pythonhosted.org/packages/19/43/cc1e93c74103a225dc39c30966ff3bec9973be0bbb338b080a55e4653bda/pydantic_monty_client-0.0.23-cp312-cp312-manylinux_2_28_s390x.whl", hash = "sha256:15ddde16a1398c601a1bb791e5c3a796ce83222b79ecb21bd7070e52aa007257", size = 3735663, upload-time = "2026-09-05T19:24:46.866Z" }, + { url = "https://files.pythonhosted.org/packages/42/09/b7103e26fe6103e217c281f6581b918a2455f5a1600fc197b82010953b19/pydantic_monty_client-0.0.23-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b41dbca54254850a5ab68bed4c75992d6eeacd283265f7445c4ae705e29fed5a", size = 4006344, upload-time = "2026-09-05T19:24:48.17Z" }, + { url = "https://files.pythonhosted.org/packages/9f/71/6b21d3d52dcb1cdbeaa593baed170d2454a61bf5aad2692e6e3b3136c5f8/pydantic_monty_client-0.0.23-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:1081ab9038004739ede448249d0b1c3d409fb8641ee9bad9ab5347d2bb415600", size = 3896737, upload-time = "2026-09-05T19:24:50.226Z" }, + { url = "https://files.pythonhosted.org/packages/02/df/6a4ed64e3c21389b7e277a7fc98df1021767f9e27edd487b9c5cb1cc4448/pydantic_monty_client-0.0.23-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:182798245ed7ba503c49c9cf5f313876df9b81f57fcbfdf12efa1aa25be0fcd1", size = 4199695, upload-time = "2026-09-05T19:24:51.844Z" }, + { url = "https://files.pythonhosted.org/packages/98/a7/5f166faeff1af270c1d139c90792fe9a2edf00b3c5e16d5014feb29c8d18/pydantic_monty_client-0.0.23-cp312-cp312-win32.whl", hash = "sha256:4d6354723ba6165d3856eac1439d2abed9d231730069c858219456e866a3bc16", size = 3288615, upload-time = "2026-09-05T19:24:53.371Z" }, + { url = "https://files.pythonhosted.org/packages/99/d9/ff24b9dd6d00b65724962f21f10c5844492d1c8d101d86e601155f6c0425/pydantic_monty_client-0.0.23-cp312-cp312-win_amd64.whl", hash = "sha256:099173c188a063ebecd79ea5a2268d06f0031ec93dbaee03b62f4386e024f5dc", size = 3927504, upload-time = "2026-09-05T19:24:54.748Z" }, + { url = "https://files.pythonhosted.org/packages/cf/75/ca3594de58d97c1f34a5c8c6965a3acc69a73e7e616c8b99b8f6fc221c55/pydantic_monty_client-0.0.23-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:80fc4d09ade86f9d527dd79d8ab6f26e510d687c962ca0d499b391905f6efd9e", size = 3934648, upload-time = "2026-09-05T19:24:56.682Z" }, + { url = "https://files.pythonhosted.org/packages/f4/a2/ee39d943b7eb874a91df0fa0175c8153509322772bc5abace77f5e1763c1/pydantic_monty_client-0.0.23-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:48f6cbb37c7a56dce5e1c4000e286447d5fb51dcf1e1cf5b69185922c43965e6", size = 3703648, upload-time = "2026-09-05T19:24:58.334Z" }, + { url = "https://files.pythonhosted.org/packages/07/2b/67bc41a8a6c178bd840e1e6115a7af3b0289ba7b1f7b4d8f39a23d0abd43/pydantic_monty_client-0.0.23-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:9b71eb5edd68ef9de22064dd5cf28feca283b5aef8f0ff00bb1308deba5e56bd", size = 3730495, upload-time = "2026-09-05T19:25:00.1Z" }, + { url = "https://files.pythonhosted.org/packages/a5/2a/0d94df69fef37fd3d4fbc0330e7f8442eb9f1bcb7f2ece69f90ef0b349b6/pydantic_monty_client-0.0.23-cp313-cp313-manylinux_2_28_armv7l.whl", hash = "sha256:2a0aa01d8c1321b610b1c0633b7ffcf9dcbb7ecaa79e5bd6d355483b23440231", size = 3425587, upload-time = "2026-09-05T19:25:01.593Z" }, + { url = "https://files.pythonhosted.org/packages/d5/4f/61f1f175987a3ad595ee31b883ca1521f7defe3c733b7eb11f228b608a36/pydantic_monty_client-0.0.23-cp313-cp313-manylinux_2_28_i686.whl", hash = "sha256:bfea6ada82fb43308833852980b147edca3a8de48bc2ce4e7cb08c6f1e13144a", size = 3633833, upload-time = "2026-09-05T19:25:03.397Z" }, + { url = "https://files.pythonhosted.org/packages/1b/71/173b79e38316b878534fb268bcb461730301ca2e7d656f71599f3be4cdb8/pydantic_monty_client-0.0.23-cp313-cp313-manylinux_2_28_ppc64le.whl", hash = "sha256:4ca25b02433751eb0735badd7004543f5eb799a7afc7354d0b0caeadc2e754d2", size = 3845772, upload-time = "2026-09-05T19:25:04.92Z" }, + { url = "https://files.pythonhosted.org/packages/d1/78/66707709d9d11222e04efd934486f22b1c44bdf4a62954516cfc746c5845/pydantic_monty_client-0.0.23-cp313-cp313-manylinux_2_28_s390x.whl", hash = "sha256:a4a7b40bb0d48b513a72977ccd9192d585c85bb470447a80a0a128a258d4c2c2", size = 3738982, upload-time = "2026-09-05T19:25:06.322Z" }, + { url = "https://files.pythonhosted.org/packages/91/ac/d3ffac7ea991b490271df273213765408df4dec6943fed99ff36db1f0642/pydantic_monty_client-0.0.23-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6acfe4ffe75aac0b89f13ea1a527f388fc0d09643bca25d7c78f59b4fed5be25", size = 4011071, upload-time = "2026-09-05T19:25:07.746Z" }, + { url = "https://files.pythonhosted.org/packages/fa/26/c6599214e286e55eb7dcc8b43617b74878fe4d55f69b8e0d12abb0576019/pydantic_monty_client-0.0.23-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:538a8edd9afcc3461fe9b744beda644d1daca60fc76423e695795de57d939d4f", size = 3899476, upload-time = "2026-09-05T19:25:09.499Z" }, + { url = "https://files.pythonhosted.org/packages/4b/4d/e3179f513cb4e112d6acd1fafff1217eaf5567ea351f00dc61dd92a7d8ae/pydantic_monty_client-0.0.23-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:ef72181626512a8b7c3c6365d60e416c57bef219ee68f98475f7ec7c59db0db0", size = 4203793, upload-time = "2026-09-05T19:25:11.253Z" }, + { url = "https://files.pythonhosted.org/packages/54/68/2889ff38031ab76eababba29eec92c7d4e9f23a84a58aad306f4609219d8/pydantic_monty_client-0.0.23-cp313-cp313-win32.whl", hash = "sha256:3509dce955db0b7ccbf8a2b458d6562e289295281b7910841e42a1de656a59dc", size = 3288733, upload-time = "2026-09-05T19:25:12.559Z" }, + { url = "https://files.pythonhosted.org/packages/bd/f6/e8344f4b3c4d0b8b37fee2b3d27676ae8ad35c5ef731c6ba48ae4012448c/pydantic_monty_client-0.0.23-cp313-cp313-win_amd64.whl", hash = "sha256:c661a74a80158460d633ac5510ccc5c80f77eb3c6c51a54a71d1d9a6108a6fdb", size = 3931122, upload-time = "2026-09-05T19:25:13.992Z" }, + { url = "https://files.pythonhosted.org/packages/96/73/f4f36e465afd5346737a46b59c50e50ec133b9f05a0ba7f3adab0c5e8ab4/pydantic_monty_client-0.0.23-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:28bbf56a7f48dda2acf1cc11807daa50af8cc94e34e1c57fac5d16f99b3e7392", size = 3935351, upload-time = "2026-09-05T19:25:15.497Z" }, + { url = "https://files.pythonhosted.org/packages/ec/ab/0f6f9e9018c2107c9f0bf036167cadc2d1e7edf8275727c9333311f0f8d5/pydantic_monty_client-0.0.23-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ff363c8a2afb25d6e716f2a67a6b350aaa864a278b0ba60a390231d4b351940b", size = 3704385, upload-time = "2026-09-05T19:25:16.944Z" }, + { url = "https://files.pythonhosted.org/packages/54/71/2a0eded398b705c65a6d56cb8139c9fce4d878e01594570c98574734818a/pydantic_monty_client-0.0.23-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:94b056d3ee44ff3fa39be0790debba41f65fcb7db9d26461f039a01e49349c0a", size = 3730307, upload-time = "2026-09-05T19:25:18.505Z" }, + { url = "https://files.pythonhosted.org/packages/9f/c8/ba9fba363a8c96ef4635cf26f3d6d3715e7663823d166676bdf8f5bc06a3/pydantic_monty_client-0.0.23-cp314-cp314-manylinux_2_28_armv7l.whl", hash = "sha256:f94396bdd59f74918871307fb7f63f4ed852ab770a91a70ed29d159b63823e97", size = 3426378, upload-time = "2026-09-05T19:25:19.958Z" }, + { url = "https://files.pythonhosted.org/packages/72/6f/73d07f0b6f2608e155e121c573b4c5f42fa9dba672f4a5b827045ead3bd9/pydantic_monty_client-0.0.23-cp314-cp314-manylinux_2_28_i686.whl", hash = "sha256:fbc150afe82ab299ceb80a9c308224f3d10ac6a1d98cc9214c24baede8fc42d5", size = 3635275, upload-time = "2026-09-05T19:25:21.652Z" }, + { url = "https://files.pythonhosted.org/packages/aa/e6/6a812e7bce36c48d87d84d0cda46002c86bcda8aa5fdf18667d48174815a/pydantic_monty_client-0.0.23-cp314-cp314-manylinux_2_28_ppc64le.whl", hash = "sha256:a080f9e5427b3223189c456337ec7866cb51c2033592c89f289ef457153800a8", size = 3846543, upload-time = "2026-09-05T19:25:23.174Z" }, + { url = "https://files.pythonhosted.org/packages/2b/06/19788b96be88fe6f58783b7a24225fb3129de7d4807c0837ae50c67172d6/pydantic_monty_client-0.0.23-cp314-cp314-manylinux_2_28_s390x.whl", hash = "sha256:ccae280613d9e15f8344cf82cda25d24dc7b598e191c4549c5f3118f327659b4", size = 3741406, upload-time = "2026-09-05T19:25:24.837Z" }, + { url = "https://files.pythonhosted.org/packages/96/e4/306c0715579eb2ecf63166470535951d11907b162ac8d54cbdcea8eb479a/pydantic_monty_client-0.0.23-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:f41b870c7ce2e0b852ae50355749b47d92452ced6f166d32ee9df5da714356a2", size = 4011087, upload-time = "2026-09-05T19:25:26.483Z" }, + { url = "https://files.pythonhosted.org/packages/7e/9d/07e939e83c4b223a0a5941664c2447dcd24c0d3c11c78721889d02d16443/pydantic_monty_client-0.0.23-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:2d48db758abafcf45cdbfd710f998a6b0af991267587831498739ac11fd7e3b5", size = 3899694, upload-time = "2026-09-05T19:25:27.818Z" }, + { url = "https://files.pythonhosted.org/packages/bf/b7/fdbf223f919b2b8ffdd2df79912b3eec1db274656486534b1ffd84b8dbc1/pydantic_monty_client-0.0.23-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:7fedc8d482e9338e8e97d2d433cf82b10f671d98416b427a52ab034978731e3d", size = 4204183, upload-time = "2026-09-05T19:25:29.279Z" }, + { url = "https://files.pythonhosted.org/packages/f2/6d/b231969ef89209c02b172fddb6182f328b56f89b6c494b238cd923c50766/pydantic_monty_client-0.0.23-cp314-cp314-win32.whl", hash = "sha256:3bc5e57df0e44057b97614149ea749438f3e3cb48c93c9a06f0ffc753824a28d", size = 3289691, upload-time = "2026-09-05T19:25:30.833Z" }, + { url = "https://files.pythonhosted.org/packages/1d/5f/5e681044f43f6e3fd42dffc84548d86c0a15b06486f29194eca54caa1b8e/pydantic_monty_client-0.0.23-cp314-cp314-win_amd64.whl", hash = "sha256:e2b664edc793fda985f7fb6a03fddfe536fd6de4b7dba0765282938a883f4c51", size = 3929684, upload-time = "2026-09-05T19:25:32.378Z" }, ] [[package]] name = "pydantic-monty-runtime" -version = "0.0.19" +version = "0.0.23" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e2/29/e44460fd934584ddecc6b8a8acddbc0605e86ea9d027910acdbf526c8f14/pydantic_monty_runtime-0.0.19.tar.gz", hash = "sha256:717e11349d7234575750cec881ac390961de02a40bd09f875dc5e097705b896b", size = 1322628, upload-time = "2026-07-24T10:00:14.955Z" } +sdist = { url = "https://files.pythonhosted.org/packages/df/1d/c1139f1f82460046d505d7a6871d4d22f9e36cdab967807487f6a70cb51f/pydantic_monty_runtime-0.0.23.tar.gz", hash = "sha256:d181f557cd3d19ee826459dea234a2b440e2b6de093e66384629843959146071", size = 1727575, upload-time = "2026-09-05T19:26:05.353Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/11/a6f4e12982b2232b9036db334fbcfecbacf46b9acaf311f9c3110e431c53/pydantic_monty_runtime-0.0.19-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:be34905548e31237fc683f5a34986489c127727e8f65481e8c87c4ad0b3a4dc2", size = 9449108, upload-time = "2026-07-24T09:58:44.394Z" }, - { url = "https://files.pythonhosted.org/packages/c3/a0/1e720b1bba457c6a1ab4fca20c571ae058238c7df3e1892b5efebad05a8b/pydantic_monty_runtime-0.0.19-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0f402f20672c1e25d4915de7e97023519461b9ffbedacd08c427d40daa77bdd6", size = 9735875, upload-time = "2026-07-24T09:58:46.952Z" }, - { url = "https://files.pythonhosted.org/packages/51/97/4439b312d9463881630dd9d998d2c8fe2b8ef457b451202c71816e25688c/pydantic_monty_runtime-0.0.19-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:9ec0d25dfbe65b9a5dfe77ff86353e95a86774d786a1d496206163f828f072fd", size = 9171496, upload-time = "2026-07-24T09:58:49.504Z" }, - { url = "https://files.pythonhosted.org/packages/1e/67/fa7b99afe1917b89eec3b4b6375f468c76eeea8227dd48361b7f2db90d97/pydantic_monty_runtime-0.0.19-cp312-cp312-manylinux_2_28_armv7l.whl", hash = "sha256:7ecc2f8eeabf6483908db4cc69b7d4c156a95675dcbdda2a7540a22ed904bfb2", size = 9565495, upload-time = "2026-07-24T09:58:51.913Z" }, - { url = "https://files.pythonhosted.org/packages/1b/a6/83a9dceb8d9f5dffd9a082b60590bb61b0ac48dca19aa02847ebbab1ad46/pydantic_monty_runtime-0.0.19-cp312-cp312-manylinux_2_28_i686.whl", hash = "sha256:fc1e508b9dc2cab64e27004d0f4ce44c7e1d255e85bafba390884fa07d696319", size = 10199598, upload-time = "2026-07-24T09:58:54.394Z" }, - { url = "https://files.pythonhosted.org/packages/c6/5e/bafd9dfa9a9a0d26b9882053aaea2280d791225c1e3b33b4b48f23fd5f92/pydantic_monty_runtime-0.0.19-cp312-cp312-manylinux_2_28_ppc64le.whl", hash = "sha256:9e3f6e366f62e54d7bd0fb995f007f42d23c26eee560d57c04d75d5adde3b45c", size = 10355698, upload-time = "2026-07-24T09:58:56.713Z" }, - { url = "https://files.pythonhosted.org/packages/0a/00/ef55cdc9f5ade20475622bb8dba617efce00ef9da2b94b36a76172edadce/pydantic_monty_runtime-0.0.19-cp312-cp312-manylinux_2_28_s390x.whl", hash = "sha256:2ddf4f2d1d063f330bb54b2dc00f6d5bcfbdd3defdc5988b856d350f62160e26", size = 10197859, upload-time = "2026-07-24T09:58:59.158Z" }, - { url = "https://files.pythonhosted.org/packages/a8/67/d1c359658458cf97296fda4359bfc71de8c9f968fb3db68eb7ce00136d43/pydantic_monty_runtime-0.0.19-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:ce2c9550c30a5c94b63dfe578243d0823511feeecb020723f5140f9737505410", size = 10661715, upload-time = "2026-07-24T09:59:01.576Z" }, - { url = "https://files.pythonhosted.org/packages/8f/1f/9841d93f956bfbd0bbbac75edf42bb3b13b5bad7d25b09d9a221b5e54d71/pydantic_monty_runtime-0.0.19-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:05563e178b6be783de088abe4a25cef9cfbc04811f7c1e1873860252e711edac", size = 9143212, upload-time = "2026-07-24T09:59:04.019Z" }, - { url = "https://files.pythonhosted.org/packages/ce/ce/0cd3643cc5051456b7baad6f16d85fb1d05daefd0556e4c82ea8f22cc9a2/pydantic_monty_runtime-0.0.19-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:de74aa4df47147d104bccdac4b39c1f3fd543091be3fd0156f77eeea3e483bc3", size = 9731590, upload-time = "2026-07-24T09:59:06.323Z" }, - { url = "https://files.pythonhosted.org/packages/ab/5a/69f4eabec2538df364242395ab3ef77b30a124a0e4b461231589651a1e97/pydantic_monty_runtime-0.0.19-cp312-cp312-win32.whl", hash = "sha256:5208056d9e23d951768ba4b94df3caf7fd84bfe951f68ec4a1803eb03377bbeb", size = 9227834, upload-time = "2026-07-24T09:59:08.694Z" }, - { url = "https://files.pythonhosted.org/packages/c7/4b/221a21f477aef0c488cbe1467111b0988658bc4a42cfc6b404201bc432af/pydantic_monty_runtime-0.0.19-cp312-cp312-win_amd64.whl", hash = "sha256:e4bba0c6024a3a8bd8c8a8cba25233a19cf686218e97afb4c059aa0c625a4b8b", size = 10941520, upload-time = "2026-07-24T09:59:10.959Z" }, - { url = "https://files.pythonhosted.org/packages/f2/a8/1ba497c35eca33273f2144b8e78d94832cc33aec5f69d8bef56968a61933/pydantic_monty_runtime-0.0.19-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:84dd041652581335503af7c60ee7362a462d946a62e8c4a44a939984034d0d25", size = 9449108, upload-time = "2026-07-24T09:59:13.497Z" }, - { url = "https://files.pythonhosted.org/packages/2d/6e/08c05c9729a35d6c9831bfd57d535bd1257f601d15b62936c27f1c0cfb47/pydantic_monty_runtime-0.0.19-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:263ab33dce2008ca6c83b4013b0280a98a1991d4071c55413d00b539940fe8aa", size = 9735874, upload-time = "2026-07-24T09:59:15.918Z" }, - { url = "https://files.pythonhosted.org/packages/c0/64/63fbdb4c069ceb2af6261fe5923864b5be309799087f16a5dccc96141c12/pydantic_monty_runtime-0.0.19-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:64f56d48097d8a158125534b20e8b22003c5e727082a303907da8e61aae0c7e3", size = 9171498, upload-time = "2026-07-24T09:59:18.799Z" }, - { url = "https://files.pythonhosted.org/packages/1d/cf/82390fe7f0e1100662e6cac7a1c0de716ea4bf851a53aa2b0ef324fc4e3e/pydantic_monty_runtime-0.0.19-cp313-cp313-manylinux_2_28_armv7l.whl", hash = "sha256:a4d80fb4598524cd5039a596f8e6900adb2a5a2da66d012a3734a2b441ad06c3", size = 9565494, upload-time = "2026-07-24T09:59:21.232Z" }, - { url = "https://files.pythonhosted.org/packages/23/37/caa8bf32ba0c0e8ce31ef2773b1a1f60d688e0237cd396e48bbef9f7161f/pydantic_monty_runtime-0.0.19-cp313-cp313-manylinux_2_28_i686.whl", hash = "sha256:1c5cd0b140c765772e0606a7047fbf95cc49d62d0d8b79b4f520dae0e38b3ba7", size = 10199597, upload-time = "2026-07-24T09:59:23.702Z" }, - { url = "https://files.pythonhosted.org/packages/48/25/ae9bb52518b2ce8fe7509d6cad4f4916b4458b748fecb180bef2d78165e9/pydantic_monty_runtime-0.0.19-cp313-cp313-manylinux_2_28_ppc64le.whl", hash = "sha256:7d80975c6e792088285f49a91e26de483ef95e5bffc4939b02d4c5ea1059749b", size = 10355699, upload-time = "2026-07-24T09:59:26.389Z" }, - { url = "https://files.pythonhosted.org/packages/a7/38/0875b1239b8ea57e4ea6de421cd240fc5a88c02e381f88d858f00439b1e9/pydantic_monty_runtime-0.0.19-cp313-cp313-manylinux_2_28_s390x.whl", hash = "sha256:f826cb911f78fe4fddf09fa3f38518484041908d08de927fa7dd134c002a2c77", size = 10197858, upload-time = "2026-07-24T09:59:28.889Z" }, - { url = "https://files.pythonhosted.org/packages/11/d4/9365473fe31e53a6dc4d6fea406d5d8f3f03a9b7d84d1f29a9e6d664c862/pydantic_monty_runtime-0.0.19-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:fdf56c6cd8c6163737d1ce284f9acb00feff77b7e952f041dc281b94336c4fc9", size = 10661715, upload-time = "2026-07-24T09:59:31.498Z" }, - { url = "https://files.pythonhosted.org/packages/20/f9/8d85b8f77d4006a8d80517a0781c49e6ca026f68747f801b63298e64197e/pydantic_monty_runtime-0.0.19-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:a675f165773af0a2e0209dd96014aa9e613115cdc6dad297a7bb96cf87f83315", size = 9143210, upload-time = "2026-07-24T09:59:33.843Z" }, - { url = "https://files.pythonhosted.org/packages/58/53/ded73742ee9fa2160c711141c28ea2ee083f073019e89724049c5e67cd84/pydantic_monty_runtime-0.0.19-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:9e42ab0287c0497d9d5529e8a53dd2f63f9d1f92f6592170baab894ea9956da6", size = 9731590, upload-time = "2026-07-24T09:59:36.642Z" }, - { url = "https://files.pythonhosted.org/packages/4b/06/67be8320dc592b8c4caa8c4c1544c8ddf2760029d46843d078aa5a4cdb14/pydantic_monty_runtime-0.0.19-cp313-cp313-win32.whl", hash = "sha256:1f5ff1b9585e648304096705045fb6bd90d43b561568b0d373265e8d201b1234", size = 9227833, upload-time = "2026-07-24T09:59:39.022Z" }, - { url = "https://files.pythonhosted.org/packages/7b/f4/bab34897974d83640f8773f03d2001142bc13e80e517bfce8c8c4a57157e/pydantic_monty_runtime-0.0.19-cp313-cp313-win_amd64.whl", hash = "sha256:7181a2153ff257fe34109685148167b6e8219d4acfe8f346115b58152fd26aa8", size = 10941519, upload-time = "2026-07-24T09:59:41.538Z" }, - { url = "https://files.pythonhosted.org/packages/63/8d/f46ac4778b2ac64183607bc63ecc783f16ca9981de30eb8ec9aa5e7132cd/pydantic_monty_runtime-0.0.19-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:4cc92139b476469c0d7e5caa147d38c2929de39bf1724cb22bbab80297823963", size = 9449107, upload-time = "2026-07-24T09:59:44.139Z" }, - { url = "https://files.pythonhosted.org/packages/a5/14/34a7bb4630d1bac0d055049568ecc888a87954c176428bde5764a7ff6ed7/pydantic_monty_runtime-0.0.19-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7a97e00bd46305f34a85f2e2e3ac4c92dbd3340e7af694aafb192ff58c4fb40e", size = 9735874, upload-time = "2026-07-24T09:59:46.546Z" }, - { url = "https://files.pythonhosted.org/packages/f1/b7/977de0b258c0858f52b23bd32afbfdf8a8c4614daff5d0b7d5c86332ce6e/pydantic_monty_runtime-0.0.19-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:46ad28b1b3c41113c9e89373da18bcc883a24da371d4eeb9b32d3d194286f1a5", size = 9171497, upload-time = "2026-07-24T09:59:48.768Z" }, - { url = "https://files.pythonhosted.org/packages/39/39/4374200ee4b938fc8b5026056f13bb677e15796bca1323a16210738431ad/pydantic_monty_runtime-0.0.19-cp314-cp314-manylinux_2_28_armv7l.whl", hash = "sha256:dc61844187a32c2f9c2b69846c5d55679eb838b4f7e495b4d03509f89fbf41f9", size = 9565495, upload-time = "2026-07-24T09:59:51.149Z" }, - { url = "https://files.pythonhosted.org/packages/04/55/c4ee4b0a10610359e09cad9d327db19095914a3bf427fb3d8164ec2bcae0/pydantic_monty_runtime-0.0.19-cp314-cp314-manylinux_2_28_i686.whl", hash = "sha256:b23b52a79ee6be0a943e4b8e5d996e6c489a37b23a337838164f22c9e8ed11a7", size = 10199598, upload-time = "2026-07-24T09:59:53.619Z" }, - { url = "https://files.pythonhosted.org/packages/87/62/cc9df084e9f930bbb2873b6dd832b377f76071aec309b1545589d818fd90/pydantic_monty_runtime-0.0.19-cp314-cp314-manylinux_2_28_ppc64le.whl", hash = "sha256:484496817f5f238c42aaea8a1945b545a17d8ef2a21e9e81799d55e481d25485", size = 10355699, upload-time = "2026-07-24T09:59:56.043Z" }, - { url = "https://files.pythonhosted.org/packages/c8/cc/bb13e6f655fcaee340032ad6a3cd1524957d0fa471ddc7e27bdb1f4c240b/pydantic_monty_runtime-0.0.19-cp314-cp314-manylinux_2_28_s390x.whl", hash = "sha256:c3e7cbad58ae9bd3402581faa7d54d719f9c140dc1888f5c9439d45bff528ea1", size = 10197859, upload-time = "2026-07-24T09:59:58.481Z" }, - { url = "https://files.pythonhosted.org/packages/ba/c3/c532715987383668ee835337e1485f51585bc8bf189f033370a05eff17f1/pydantic_monty_runtime-0.0.19-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:5534067dc7ffdae809293da95d0e95b6d8481f4c88aff59385e19f466ba3c0f0", size = 10661715, upload-time = "2026-07-24T10:00:01.163Z" }, - { url = "https://files.pythonhosted.org/packages/ed/6c/9a0ab28a061efc184398a0b4db34e459572bb2316cf766f0d6ce65b475e3/pydantic_monty_runtime-0.0.19-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:c1ae32b06a4456ab223bafa54d29a349066d625d09063c28ba14ee9019f9b7c2", size = 9143211, upload-time = "2026-07-24T10:00:03.856Z" }, - { url = "https://files.pythonhosted.org/packages/23/5f/af5b3e6395834572975d98f4d1a00a57ee8029bf68ca5732347550f32b35/pydantic_monty_runtime-0.0.19-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:497cf8c3f30992f9aafc8707084eaffe2391b7b5dec067d04d5715f9a562c56b", size = 9731591, upload-time = "2026-07-24T10:00:06.351Z" }, - { url = "https://files.pythonhosted.org/packages/1a/98/96797fd269342cfdb49f03c22fd9a05ef91f711089081c4b3861b9c520e2/pydantic_monty_runtime-0.0.19-cp314-cp314-win32.whl", hash = "sha256:942feb948df8edb61ae7ba6ae77dc655e6985be06d72eef886562fe573ba3086", size = 9227832, upload-time = "2026-07-24T10:00:08.708Z" }, - { url = "https://files.pythonhosted.org/packages/39/fc/02d15281c8e00b48df9af8f75a4fe06f3f8f33ef6a910507a45a19f2b61b/pydantic_monty_runtime-0.0.19-cp314-cp314-win_amd64.whl", hash = "sha256:91d93339c70483ed9256b3b15e3375f6597ae65be280f9b89ba9ca0355f95f54", size = 10941519, upload-time = "2026-07-24T10:00:11.226Z" }, + { url = "https://files.pythonhosted.org/packages/00/97/d1575be31396cf662ea8419d427c58279378a2e08b674d436ed4990751ea/pydantic_monty_runtime-0.0.23-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:0c8968491ae44648c11cf075c0f5d798d4aa7f7e4a0d0754add0b74d725c075c", size = 9828049, upload-time = "2026-09-05T19:25:34.446Z" }, + { url = "https://files.pythonhosted.org/packages/15/09/01ae8472d860223618cd61a4656d8ff44519e1bd9e320ee8698dbb206aca/pydantic_monty_runtime-0.0.23-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c60e44cbef9cfdc9bd5e53631650490e66890f57e8c6bea5611dae502d8374d4", size = 11548563, upload-time = "2026-09-05T19:25:37.189Z" }, + { url = "https://files.pythonhosted.org/packages/fe/30/ebf796ec2236b15cb9f5541e3a940536f37de41bd2561989985aa92fb396/pydantic_monty_runtime-0.0.23-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:4bc24603678107b7f5b2b24a5946ca1e103379b2289589212c86959993a052b4", size = 11935977, upload-time = "2026-09-05T19:25:39.862Z" }, + { url = "https://files.pythonhosted.org/packages/c2/40/8ed68e144f80c9d14e0532d631ffd78ee6a0ae8d7d3272ac037d5d703266/pydantic_monty_runtime-0.0.23-py3-none-manylinux_2_28_armv7l.whl", hash = "sha256:7c68422e06c5aae8b52f9835ee949256db0f261301515e4443d6647f249cf789", size = 10180097, upload-time = "2026-09-05T19:25:42.298Z" }, + { url = "https://files.pythonhosted.org/packages/1c/f3/ea3632cebb754b81fa63c01b84c89299bba936c670fa519f3ff363166bcc/pydantic_monty_runtime-0.0.23-py3-none-manylinux_2_28_i686.whl", hash = "sha256:b8a049a554bff6ba7e553c1ae4b6290e554b6a1695bc1d47b1b03db485f7fb91", size = 10515148, upload-time = "2026-09-05T19:25:44.576Z" }, + { url = "https://files.pythonhosted.org/packages/e8/a9/6a43d9d9f01f441ecd6ec6c31d9b8d45b143a37b9368f582b27378e50ec7/pydantic_monty_runtime-0.0.23-py3-none-manylinux_2_28_ppc64le.whl", hash = "sha256:9cea174ed1a56888bb58975feff34b0475504ab3831b3b1dc3189397d5bcccc3", size = 10746650, upload-time = "2026-09-05T19:25:46.823Z" }, + { url = "https://files.pythonhosted.org/packages/18/d0/dc2c2e066ce4f619f34a98b44c9c84fb398095421808c9dd6ac12463aea2/pydantic_monty_runtime-0.0.23-py3-none-manylinux_2_28_s390x.whl", hash = "sha256:28bb59cde9de0cec0d3d5d72029851757249c682a6747251f61e8cdece5995f7", size = 11092334, upload-time = "2026-09-05T19:25:49.314Z" }, + { url = "https://files.pythonhosted.org/packages/3b/82/6af8a19dd357d63e783e9647ef81d47db2f48c13e96558d25484cd0dec91/pydantic_monty_runtime-0.0.23-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:721d2d1455482c929c934412e61ecf56673b88a2c0c049e041295cf543ff1412", size = 12371608, upload-time = "2026-09-05T19:25:51.663Z" }, + { url = "https://files.pythonhosted.org/packages/f7/53/2939bce8c8a688771777ee2301851a712744d809d752c1c8dcb5d6a52574/pydantic_monty_runtime-0.0.23-py3-none-musllinux_1_1_aarch64.whl", hash = "sha256:f0183e589ba18de4fac8a675fe0f93049baa293d5409eb8d39c639dc24af943f", size = 9593312, upload-time = "2026-09-05T19:25:53.8Z" }, + { url = "https://files.pythonhosted.org/packages/e6/f9/180d9acbea18d8050c190280ef591cb52d05e4d0e6d026cc6b7b2613510c/pydantic_monty_runtime-0.0.23-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:cbf56c19712c12072f6607d8e2323941441c9343734dcb07d2d2b4ce2dd98bf3", size = 10133520, upload-time = "2026-09-05T19:25:56.797Z" }, + { url = "https://files.pythonhosted.org/packages/e1/00/dcd69bd9fcc57c7a6ecfcd3b2199f4598853edb14d504f481dcda41b17cd/pydantic_monty_runtime-0.0.23-py3-none-win32.whl", hash = "sha256:854f58a472397f34588a22eb813f2acd65efc0bcb966ef1fe173c32366209d15", size = 9639882, upload-time = "2026-09-05T19:25:58.978Z" }, + { url = "https://files.pythonhosted.org/packages/e6/c0/3f7c604ac93ec47a36a3733d6e7717b7096f0989f669d4133bd3e85f3375/pydantic_monty_runtime-0.0.23-py3-none-win_amd64.whl", hash = "sha256:742375f494e298a4f96933ac3694a0fe34c9cddd9b666669bc193f348bbed0b8", size = 10404629, upload-time = "2026-09-05T19:26:01.143Z" }, ] [[package]] From 95fdb46c3aed30dc998e7e48e20b0dc92e242419 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 7 Sep 2026 12:47:28 +0300 Subject: [PATCH 17/18] Check the sandbox deadline on every host call, one error contract analysis.code_timeout was enforced only in _run_on_loop, the bridge for database-bound reads; metadata.json, the cached JSONL files and in-code search() and list_documents() never looked at the clock, and Monty's watchdog counts compute only. Every host call now checks the deadline before it starts. Host errors keep their message for every caller: the masking added for the MCP server goes, and the sandbox is one path for the capability and the server alike. --- CHANGELOG.md | 6 +- docs/configuration/qa.md | 2 +- docs/mcp.md | 13 +- haiku_rag_slim/haiku/rag/sandbox/sandbox.py | 93 ++++++++------- tests/sandbox/test_sandbox.py | 124 +++++++++----------- tests/test_mcp.py | 21 ++-- 6 files changed, 122 insertions(+), 137 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3070696d..555d9ec8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,9 +50,6 @@ Unknown document, unknown collection, invalid filter, invalid base64 and a failing program carry a message. Anything else is masked (`mask_error_details=True`) and logged server-side. -- A host-side failure inside the analysis sandbox (a document read or an - in-code `search()` raising) reaches the program as - `RuntimeError(" failed: ")`; the traceback is logged. - `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on `search_documents`, `search_documents_by_image` and `execute_code`; `source` on `get_document`; an unknown name is a tool error. `DocumentInfo.source`. @@ -61,6 +58,9 @@ - `toc.json` `item_range` in the analysis sandbox is a line slice into `items.jsonl`, as documented; it held item positions. +- Past `analysis.code_timeout` a sandbox program starts no further host call. + Files served from memory and in-code `search()` / `list_documents()` were + not checked against the deadline. ### Removed diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index 32c4c34b..a7efa866 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -50,7 +50,7 @@ analysis: provider: anthropic name: claude-sonnet-4-20250514 temperature: 0.0 # Default: 0.0 (deterministic for code generation) - code_timeout: 60.0 # Max seconds a call may spend reading documents + code_timeout: 60.0 # Per call: compute stops, no read or search starts past it max_output_chars: 50000 # Truncate output after this many chars max_executions: 15 # Max execute_code calls per question ``` diff --git a/docs/mcp.md b/docs/mcp.md index d1cb4bac..c2afa2ba 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -171,9 +171,10 @@ Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`, reached for: `decimal` and `statistics`. No generator functions, class inheritance or `match` statements, and a file object cannot be iterated. Files are read-only, and there is no network and no filesystem -beyond `/documents`. `analysis.code_timeout` bounds a call, counted separately -for compute and for document reads, and `analysis.max_output_chars` bounds its -output. +beyond `/documents`. `analysis.code_timeout` is the call's budget: compute is +stopped at it, and past it no further host call starts, a file read or an +in-code search alike, though one already running finishes. +`analysis.max_output_chars` bounds the output. ### Filters @@ -192,10 +193,8 @@ title = 'Q3 report' A failure is an MCP error, never an empty result. Expected failures carry a message: a document or section id that matches nothing, a collection the server does not cover, a filter the query engine rejects (with its message), -invalid base64, and a program that fails in `execute_code`. A failure on the -server inside a program, a database read or an in-code search raising, reaches -the program and the client as its exception type only; the traceback goes to -the server log. +invalid base64, and a program that fails in `execute_code`, with the error the +program hit. Anything else reaches the client as `Error calling tool 'name'` and its traceback goes to the server log. diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index 823f0040..512b5b24 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -1,6 +1,5 @@ import asyncio import json -import logging import os from collections.abc import AsyncIterator, Callable, Coroutine from contextlib import asynccontextmanager, suppress @@ -31,21 +30,9 @@ if TYPE_CHECKING: from haiku.rag.client.scope import DatabaseScope -logger = logging.getLogger(__name__) - _MAX_HOST_CALLS = 10_000_000 -def _host_failure(where: str, e: Exception) -> RuntimeError: - """The error a program gets for a failure on the host side of a call. - - The message and traceback go to the log. The program, and through the MCP - server its client, learn the exception type only. - """ - logger.exception("%s failed inside the sandbox", where) - return RuntimeError(f"{where} failed: {type(e).__name__}") - - @dataclass class SandboxResult: """Result of executing code in the sandbox.""" @@ -283,21 +270,47 @@ class Sandbox: loop overruns it by however long the outstanding reads take. Raising from inside the callback answers the worker's suspension, which keeps the session usable — cancelling ``feed_run`` from outside does not, and wedges - the protocol. A failed read reaches the program by type only. + the protocol. """ assert self._loop is not None, ( "VFS reads happen during execute(); the loop must be captured first." ) - if self._deadline is not None and self._loop.time() > self._deadline: + if self._past_deadline(): coro.close() - raise TimeoutError( - "time limit exceeded: no further document reads after " - f"{self._config.analysis.code_timeout}s" - ) - try: - return asyncio.run_coroutine_threadsafe(coro, self._loop).result() - except Exception as e: - raise _host_failure("document read", e) from None + raise self._time_limit() + return asyncio.run_coroutine_threadsafe(coro, self._loop).result() + + def _past_deadline(self) -> bool: + return ( + self._deadline is not None + and self._loop is not None + and self._loop.time() > self._deadline + ) + + def _time_limit(self) -> TimeoutError: + return TimeoutError( + "time limit exceeded: no further document reads or calls after " + f"{self._config.analysis.code_timeout}s" + ) + + def _check_deadline(self) -> None: + """Refuse a host call once the call's time is up. + + Monty's watchdog counts only time the worker spends computing, so every + host call, a file served from memory and an in-code search included, + checks the deadline before it runs. + """ + if self._past_deadline(): + raise self._time_limit() + + def _timed( + self, read: Callable[["PurePosixPath"], str] + ) -> Callable[["PurePosixPath"], str]: + def call(path: "PurePosixPath") -> str: + self._check_deadline() + return read(path) + + return call async def _discard_session(self) -> None: """Drop a session whose worker is gone. @@ -334,6 +347,7 @@ class Sandbox: context = self._context async def search(query: str, limit: int = 10) -> list[dict[str, Any]]: + self._check_deadline() # Picture bytes are deliberately not attached to in-code search # results: the Monty interpreter has no PIL/base64/hashlib, so the # agent's Python can't do anything with them. The driving model @@ -373,6 +387,7 @@ class Sandbox: return out async def list_documents() -> list[dict[str, Any]]: + self._check_deadline() docs, _ = await self._documents() return [ { @@ -387,22 +402,10 @@ class Sandbox: ] return { - "search": self._guarded("search()", search), - "list_documents": self._guarded("list_documents()", list_documents), + "search": search, + "list_documents": list_documents, } - @staticmethod - def _guarded( - where: str, fn: Callable[..., Coroutine[Any, Any, Any]] - ) -> Callable[..., Coroutine[Any, Any, Any]]: - async def call(*args: Any, **kwargs: Any) -> Any: - try: - return await fn(*args, **kwargs) - except Exception as e: - raise _host_failure(where, e) from None - - return call - async def _build_vfs(self) -> OSAccess: """Build the virtual filesystem with document data. @@ -554,7 +557,7 @@ class Sandbox: files.append( CallbackFile( f"{doc_dir}/metadata.json", - read=lambda _path, text=metadata: text, + read=self._timed(lambda _path, text=metadata: text), write=_deny_write, ) ) @@ -575,21 +578,21 @@ class Sandbox: files.append( CallbackFile( f"{doc_dir}/content.txt", - read=_make_content_reader(doc_id), + read=self._timed(_make_content_reader(doc_id)), write=_deny_write, ) ) files.append( CallbackFile( f"{doc_dir}/items.jsonl", - read=_make_items_reader(doc_id), + read=self._timed(_make_items_reader(doc_id)), write=_deny_write, ) ) files.append( CallbackFile( f"{doc_dir}/chunks.jsonl", - read=_make_chunks_reader(doc_id), + read=self._timed(_make_chunks_reader(doc_id)), write=_deny_write, ) ) @@ -600,7 +603,7 @@ class Sandbox: files.append( CallbackFile( f"{doc_dir}/toc.json", - read=_make_toc_reader(doc_id), + read=self._timed(_make_toc_reader(doc_id)), write=_deny_write, ) ) @@ -612,9 +615,9 @@ class Sandbox: Monty spends ``max_duration_secs`` across the session's whole life, and the session is reused so variables persist between calls: the budget - covers the whole run. ``code_timeout`` is enforced per call elsewhere: the read - deadline in ``_run_on_loop`` bounds a call that reads, and the pool's - ``request_timeout`` bounds one that computes. + covers the whole run. ``code_timeout`` is enforced per call elsewhere: past + its deadline no further host call starts (``_check_deadline``), and the + pool's ``request_timeout`` bounds compute. ``max_suspensions`` counts host callbacks per session, document reads included, defaults to 1000 and cannot be disabled. The time budgets are diff --git a/tests/sandbox/test_sandbox.py b/tests/sandbox/test_sandbox.py index 98465329..1d3635dc 100644 --- a/tests/sandbox/test_sandbox.py +++ b/tests/sandbox/test_sandbox.py @@ -1,5 +1,4 @@ import asyncio -import logging import threading from pathlib import Path @@ -334,69 +333,21 @@ class TestSandboxExternalFunctionEdgeCases: assert "external error" in result.stderr @pytest.mark.asyncio - async def test_a_failing_search_reaches_the_program_by_type_only( - self, sandbox, monkeypatch, caplog + async def test_a_failing_search_keeps_its_message_for_the_program( + self, sandbox, monkeypatch ): - """A host-side failure inside search() names its exception type to - the program; the message and traceback go to the log.""" + """A host-side failure inside search() reaches the program with its + message, which the agent reads to repair its code.""" async def boom(self, *args, **kwargs): raise ValueError("failed at /secret/path") monkeypatch.setattr(HaikuRAG, "search", boom) - with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): - result = await sandbox.execute("await search('hello')") + result = await sandbox.execute("await search('hello')") assert not result.success - assert "search() failed: ValueError" in result.stderr - assert "/secret/path" not in result.stderr - assert any( - r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) - for r in caplog.records - ) - - @pytest.mark.asyncio - async def test_a_failing_document_read_reaches_the_program_by_type_only( - self, temp_db_path, monkeypatch, caplog - ): - """A program can catch a failed file read, and what it catches names - the exception type only.""" - from haiku.rag.store.models.document import Document - - async with HaikuRAG(temp_db_path, create=True) as client: - doc = await client.document_repository.create( - Document(content="x", uri="test://read", title="Read") - ) - repository = type(client.document_repository) - - async def boom(self, *args, **kwargs): - raise ValueError("failed at /secret/path") - - monkeypatch.setattr(repository, "get_content", boom) - sb = Sandbox( - db_path=temp_db_path, config=AppConfig(), context=AnalysisContext() - ) - try: - with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): - result = await sb.execute( - "from pathlib import Path\n" - "try:\n" - f" Path('/documents/{doc.id}/content.txt').read_text()\n" - "except Exception as e:\n" - " print('caught:', e)" - ) - finally: - await sb.close() - - assert result.success, result.stderr - assert "caught:" in result.stdout - assert "ValueError" in result.stdout - assert "/secret/path" not in result.stdout - assert any( - r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) - for r in caplog.records - ) + assert "ValueError: failed at /secret/path" in result.stderr class TestSandboxOutputTruncation: @@ -1013,24 +964,59 @@ class TestSandboxReadDeadline: enforces the budget itself, before each read.""" @pytest.mark.asyncio - async def test_a_failed_read_reaches_the_program_by_type_only( - self, sandbox, caplog + async def test_the_deadline_covers_reads_from_memory_and_in_code_calls( + self, temp_db_path, monkeypatch ): - """The bridged read hands the program the exception type, not the - message, and logs the traceback.""" - sandbox._loop = asyncio.get_running_loop() + """Once a call's time is up, a file served from memory and an in-code + listing are refused like a database read. A slow first read spends the + budget; the watchdog does not count time spent waiting on the host.""" + from docling_core.types.doc.document import DoclingDocument + from docling_core.types.doc.labels import DocItemLabel - async def failing_read(): - raise ValueError("failed at /secret/path") + config = AppConfig() + config.analysis.code_timeout = 1.0 + docling = DoclingDocument(name="d") + docling.add_text(label=DocItemLabel.TEXT, text="Foxes and dogs.") + async with HaikuRAG(temp_db_path, create=True) as client: + doc = await client.import_document( + docling, + [ + Chunk( + content="Foxes and dogs.", + embedding=[0.1] * config.embeddings.model.vector_dim, + order=0, + ) + ], + uri="test://deadline-paths", + ) + repository = type(client.document_repository) - with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): - with pytest.raises(RuntimeError, match="document read failed: ValueError"): - await asyncio.to_thread(sandbox._run_on_loop, failing_read()) + async def slow_content(self, *args, **kwargs): + await asyncio.sleep(1.3) + return "body" - assert any( - r.exc_info and "failed at /secret/path" in str(r.exc_info[1]) - for r in caplog.records - ) + monkeypatch.setattr(repository, "get_content", slow_content) + sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext()) + try: + result = await sb.execute( + "from pathlib import Path\n" + f"root = Path('/documents/{doc.id}')\n" + "print(len((root / 'content.txt').read_text()))\n" + "try:\n" + " (root / 'metadata.json').read_text()\n" + " print('static: read')\n" + "except Exception as e:\n" + " print('static:', type(e).__name__)\n" + "await list_documents()\n" + "print('listed')" + ) + finally: + await sb.close() + + assert "static: TimeoutError" in result.stdout + assert "listed" not in result.stdout + assert not result.success + assert "time limit exceeded" in result.stderr @pytest.mark.asyncio async def test_read_after_deadline_raises_without_scheduling(self, sandbox): diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 18d7106d..04e7cb0a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1098,25 +1098,22 @@ class TestMCPErrorContract: assert "base64" in result.content[0].text @pytest.mark.asyncio - async def test_a_host_failure_inside_a_program_names_only_its_type( - self, mcp_db, monkeypatch, caplog + async def test_a_host_failure_inside_a_program_carries_its_message( + self, mcp_db, monkeypatch ): + """One contract for the sandbox: the client reads the same error the + program did, message included.""" + async def boom(self, *args, **kwargs): raise RuntimeError("boom at /secret/path") monkeypatch.setattr(HaikuRAG, "search", boom) - with caplog.at_level(logging.ERROR, logger="haiku.rag.sandbox.sandbox"): - result = await _call( - create_mcp_server(mcp_db), "execute_code", code="await search('x')" - ) + result = await _call( + create_mcp_server(mcp_db), "execute_code", code="await search('x')" + ) assert result.is_error - assert "RuntimeError" in result.content[0].text - assert "/secret/path" not in result.content[0].text - assert any( - r.exc_info and "boom at /secret/path" in str(r.exc_info[1]) - for r in caplog.records - ) + assert "RuntimeError: boom at /secret/path" in result.content[0].text @pytest.mark.asyncio @pytest.mark.parametrize( From 4b1ec096c6125628da7e2a7da88cdc1522fe12e7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 7 Sep 2026 13:14:16 +0300 Subject: [PATCH 18/18] One error contract for the MCP server Every failure reaches the client as an MCP error carrying its message; mask_error_details is set off explicitly, since FastMCP also reads it from the environment. The masking goes, and with it the filter pre-check that ran a count before every filtered call and the UnknownDatabaseError translations that existed only to survive it. Explicit domain errors stay. --- CHANGELOG.md | 6 +- docs/configuration/qa.md | 2 +- docs/mcp.md | 16 +++-- haiku_rag_slim/haiku/rag/mcp.py | 108 ++++++++------------------------ tests/test_mcp.py | 69 ++++---------------- 5 files changed, 50 insertions(+), 151 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 555d9ec8..d858aa1a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,10 +46,8 @@ `ImageContent` blocks, with no structured content. `SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`; `collect_pictures` in `haiku.rag.tools.search`. -- MCP tools raise on failure; an empty result no longer doubles as an error. - Unknown document, unknown collection, invalid filter, invalid base64 and a - failing program carry a message. Anything else is masked - (`mask_error_details=True`) and logged server-side. +- MCP tools raise on failure, with the error's message; an empty result no + longer doubles as an error. - `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on `search_documents`, `search_documents_by_image` and `execute_code`; `source` on `get_document`; an unknown name is a tool error. `DocumentInfo.source`. diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index a7efa866..bee12f8f 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -56,7 +56,7 @@ analysis: ``` - **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`. -- **code_timeout**: Seconds a single `execute_code` call may spend reading documents (default: 60). The sandbox refuses further reads past this point. Code that computes without reading is killed by the worker watchdog at the same limit. `code_timeout * max_executions` is the cumulative ceiling across all calls in one question. +- **code_timeout**: Seconds a single `execute_code` call has (default: 60). Past it the sandbox starts no further host call, a document read or an in-code `search()` / `list_documents()`; one already running finishes. Code that computes without host calls is killed by the worker watchdog at the same limit. `code_timeout * max_executions` is the cumulative ceiling across all calls in one question. - **max_output_chars**: Truncate code output after this many characters (default: 50000) - **max_executions**: Maximum `execute_code` calls per question before the capability is told to answer from what it has (default: 15) diff --git a/docs/mcp.md b/docs/mcp.md index c2afa2ba..35212650 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -31,8 +31,8 @@ The server opens the database read-only. Ingestion goes through the CLI ## 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 `search_documents`, `search_documents_by_image` +them, as `haiku-rag search` does. Results and documents name theirs in +`source`. `sources` on `search_documents`, `search_documents_by_image` and `execute_code` 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 @@ -190,13 +190,11 @@ title = 'Q3 report' ### Errors -A failure is an MCP error, never an empty result. Expected failures carry a -message: a document or section id that matches nothing, a collection the -server does not cover, a filter the query engine rejects (with its message), -invalid base64, and a program that fails in `execute_code`, with the error the -program hit. -Anything else reaches the client as `Error calling tool 'name'` and its -traceback goes to the server log. +A failure is an MCP error carrying its message, never an empty result: a +document or section id that matches nothing, a collection the server does not +cover, a filter the query engine rejects, invalid base64, a program that fails +in `execute_code` with the error it hit, and anything unexpected with its own +message. ### Instructions diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 58a94e9b..72dbf165 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -1,6 +1,5 @@ import asyncio -import logging -import re +import base64 from collections.abc import AsyncIterator from contextlib import AsyncExitStack, asynccontextmanager from importlib import metadata @@ -17,9 +16,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, get_config from haiku.rag.context import build_toc from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint -from haiku.rag.store.exceptions import UnknownDatabaseError from haiku.rag.store.models import Document, SearchResult -from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.schema import DocumentMetaRecord from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode from haiku.rag.tools.search import collect_pictures @@ -28,8 +25,7 @@ if TYPE_CHECKING: from typing import Any from haiku.rag.client.scope import DatabaseScope - -logger = logging.getLogger(__name__) + from haiku.rag.store.models.document_item import DocumentItem _FILTER_COLUMNS = ", ".join(DocumentMetaRecord.model_fields) @@ -55,8 +51,6 @@ def _read_only(title: str) -> ToolAnnotations: def _decode_image(image_base64: str) -> bytes: - import base64 - try: return base64.b64decode(image_base64, validate=True) except ValueError as e: @@ -65,33 +59,6 @@ def _decode_image(image_base64: str) -> bytes: raise ToolError("Invalid base64 image") from e -async def _check_filter( - rag: HaikuRAG, filter: str | None, sources: list[str] | None = None -) -> None: - """Evaluate a filter on its own before the read that would use it. - - A filtered count on one selected database runs the same predicate on the - same table and nothing else, so a ValueError here is the query engine - rejecting the filter; its message names columns and the statement, never - a location. A ValueError raised later in the read stays masked. Only the - selection is touched: every database shares the schema, so one suffices. - """ - if filter is None: - return - selected = await rag.clients_covering(sources) - if not selected: - return - try: - await selected[0].count_documents(filter=filter) - except ValueError as e: - # The engine lists its own columns too, lance internals among them. - reason = re.sub(r"\s*Valid fields are .*", "", str(e), flags=re.DOTALL) - raise ToolError( - f"Invalid filter {filter!r}: {reason.rstrip('. ')}. " - f"Columns: {_FILTER_COLUMNS}." - ) from e - - def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: """What the server is for, naming no tools: the client has every tool's description from the listing.""" @@ -107,7 +74,7 @@ def _instructions(scope: "DatabaseScope", config: AppConfig) -> str: if scope.covers_multiple: lines.append( f"It holds several collections: {', '.join(scope.names)}. Results " - "and citations name theirs in `source`; pass `sources` to use a subset." + "name theirs in `source`; pass `sources` to use a subset." ) if config.prompts.domain_preamble: lines.append(config.prompts.domain_preamble) @@ -119,8 +86,6 @@ def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolRe metadata, then each distinct picture as an image block labelled with its result. No structured content: a client given both shows the model the JSON and drops the text, or shows both.""" - import base64 - total = len(results) text = "\n\n".join( result.format_for_agent( @@ -196,8 +161,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and resolves it, which is its own job. A caller that resolved already passes the - scope, so the configured name survives, which results and citations carry as - ``source``. + scope, so the configured name survives, which results carry as ``source``. """ client: HaikuRAG | None = None stack = AsyncExitStack() @@ -233,14 +197,14 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: finally: client = None - # Masking keeps paths and provider URLs out of an unexpected error's text; - # the traceback goes to the server log. A ToolError reaches the client as is. + # Explicit: the setting is also read from the environment, and the contract + # is that every failure reaches the client with its message. mcp = FastMCP( "haiku-rag", instructions=_instructions(scope, config), version=metadata.version("haiku.rag-slim"), lifespan=lifespan, - mask_error_details=True, + mask_error_details=False, ) @mcp.tool(annotations=_read_only("Search documents")) @@ -272,17 +236,13 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: False for a smaller response. """ rag = await _client() - try: - await _check_filter(rag, filter, sources) - results = await rag.search( - query, - limit=limit, - filter=filter, - include_images=include_images, - sources=sources, - ) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e + results = await rag.search( + query, + limit=limit, + filter=filter, + include_images=include_images, + sources=sources, + ) return _search_result(await rag.expand_context(results), rag.covers_multiple) # Image-as-query tool, only registered when the configured embedder @@ -317,17 +277,13 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: """ raw = _decode_image(image_base64) rag = await _client() - try: - await _check_filter(rag, filter, sources) - results = await rag.search( - raw, - limit=limit, - filter=filter, - include_images=include_images, - sources=sources, - ) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e + results = await rag.search( + raw, + limit=limit, + filter=filter, + include_images=include_images, + sources=sources, + ) return _search_result( await rag.expand_context(results), rag.covers_multiple ) @@ -346,24 +302,18 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: is asked. """ rag = await _client() - try: - document = await rag.get_document_by_id(document_id, source) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e + document = await rag.get_document_by_id(document_id, source) if document is None: raise ToolError(f"No document with id {document_id!r}") return document - async def _items_of(document_id: str, source: str | None) -> list[DocumentItem]: + async def _items_of(document_id: str, source: str | None) -> list["DocumentItem"]: """A document's items in reading order, from the database holding it.""" rag = await _client() - try: - document = await rag.get_document_by_id(document_id, source) - if document is None: - raise ToolError(f"No document with id {document_id!r}") - owner = await rag.reader_for(source or document.source) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e + document = await rag.get_document_by_id(document_id, source) + if document is None: + raise ToolError(f"No document with id {document_id!r}") + owner = await rag.reader_for(source or document.source) assert owner is not None, "a stored document names its database" return await owner.document_item_repository.get_all_items(document_id) @@ -434,7 +384,6 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: offset: How many documents to skip, for paging. """ rag = await _client() - await _check_filter(rag, filter) documents = await rag.list_documents(limit, offset, filter) return [ DocumentInfo( @@ -498,10 +447,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP: scope, config, AnalysisContext(filter=filter, sources=sources), rag=rag ) try: - await _check_filter(rag, filter, sources) result = await sandbox.execute(code) - except UnknownDatabaseError as e: - raise ToolError(str(e)) from e finally: await sandbox.close() if not result.success: diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 04e7cb0a..6252ab8b 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -1,4 +1,3 @@ -import logging import re from pathlib import Path @@ -8,6 +7,7 @@ from fastmcp.exceptions import ToolError from haiku.rag.client import HaikuRAG from haiku.rag.mcp import _covering as _mcp_covering from haiku.rag.mcp import create_mcp_server +from haiku.rag.store.exceptions import UnknownDatabaseError from haiku.rag.store.models import Chunk, Document, SearchResult from haiku.rag.tools.document import DocumentInfo from tests.multi_db.helpers import _config, _seed, _seed_expandable @@ -485,7 +485,7 @@ class TestMCPDocumentNavigation: assert ( await section(document_id=doc.id, section_id="#/texts/0", source="beta") ).title == "Only in beta" - with pytest.raises(ToolError, match="nope"): + with pytest.raises(UnknownDatabaseError, match="nope"): await outline(document_id=doc.id, source="nope") with pytest.raises(ToolError, match=doc.id): await outline(document_id=doc.id, source="alpha") @@ -908,11 +908,10 @@ class TestMCPCoversTheConfiguredSet: async def test_an_unknown_database_is_an_error_not_an_empty_result( self, two_dbs, multimodal_embedder, tool_name, kwargs ): - mcp = _covering_all(two_dbs) - tool = await _get_tool(mcp, tool_name) + result = await _call(_covering_all(two_dbs), tool_name, **kwargs) - with pytest.raises(ToolError, match="nope"): - await tool(**kwargs) + assert result.is_error + assert "nope" in result.content[0].text @pytest.mark.asyncio async def test_a_filtered_search_touches_only_the_selected_databases(self, two_dbs): @@ -1029,9 +1028,8 @@ class TestMCPImageQuery: @pytest.mark.filterwarnings("ignore:Found propagated trace context:RuntimeWarning") class TestMCPErrorContract: - """A failure is an error on the wire, never an empty result. Expected - failures say what went wrong; anything else is masked and logged on the - server.""" + """A failure is an error on the wire carrying its message, never an empty + result.""" @pytest.mark.asyncio async def test_an_unknown_document_is_an_error(self, mcp_db): @@ -1051,37 +1049,13 @@ class TestMCPErrorContract: ("execute_code", {"code": "print(1)"}), ], ) - async def test_an_invalid_filter_is_an_error_naming_the_filter( - self, mcp_db, tool_name, kwargs - ): + async def test_an_invalid_filter_is_an_error(self, mcp_db, tool_name, kwargs): result = await _call( create_mcp_server(mcp_db), tool_name, filter="no_such_column = 1", **kwargs ) assert result.is_error - assert "no_such_column = 1" in result.content[0].text - assert "created_at" in result.content[0].text - assert "_rowid" not in result.content[0].text - - @pytest.mark.asyncio - @pytest.mark.parametrize("filter", [None, "title = 'AI Overview'"]) - async def test_a_value_error_from_the_read_is_not_an_invalid_filter( - self, mcp_db, monkeypatch, filter - ): - """Only the filter check translates ValueError; one raised by the read - itself, with or without a valid filter, stays masked.""" - - async def boom(self, *args, **kw): - raise ValueError("boom at /secret/path") - - monkeypatch.setattr(HaikuRAG, "search", boom) - result = await _call( - create_mcp_server(mcp_db), "search_documents", query="x", filter=filter - ) - - assert result.is_error - assert "filter" not in result.content[0].text - assert "/secret/path" not in result.content[0].text + assert "no_such_column" in result.content[0].text @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1125,34 +1099,17 @@ class TestMCPErrorContract: ("list_documents", "list_documents", {}), ], ) - async def test_an_unexpected_failure_is_masked_and_logged( - self, - mcp_db, - multimodal_embedder, - monkeypatch, - caplog, - client_method, - tool_name, - kwargs, + async def test_an_unexpected_failure_carries_its_message( + self, mcp_db, multimodal_embedder, monkeypatch, client_method, tool_name, kwargs ): async def boom(self, *args, **kw): raise RuntimeError("boom at /secret/path") monkeypatch.setattr(HaikuRAG, client_method, boom) - # fastmcp's logger does not propagate, so listen to it directly. - fastmcp_logger = logging.getLogger("fastmcp") - fastmcp_logger.addHandler(caplog.handler) - try: - result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs) - finally: - fastmcp_logger.removeHandler(caplog.handler) + result = await _call(create_mcp_server(mcp_db), tool_name, **kwargs) assert result.is_error - assert "/secret/path" not in result.content[0].text - assert any( - r.exc_info and "boom at /secret/path" in str(r.exc_info[1]) - for r in caplog.records - ) + assert "boom at /secret/path" in result.content[0].text class TestAgentPlugins: