Evil hack to allow for runtime specification of resource mimetype

Author: gleachkr

src/pdf2sqlite/mcp_server/resources.py

@@ -3,10 +3,14 @@
 import base64
 import io
 from dataclasses import dataclass
+from types import MethodType
+from typing import Any, Awaitable, Callable, Mapping, cast
 
 from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.fastmcp.resources.types import FunctionResource
 from mcp.server.fastmcp.utilities.types import Image
 from mcp.types import BlobResourceContents, EmbeddedResource
+from pydantic import AnyUrl
 from pypdf import PdfReader, PdfWriter
 
 from .config import ServerConfig
@@ -71,7 +75,7 @@ async def make_embedded_pdf(self, uri: str, data: bytes) -> EmbeddedResource:
         return EmbeddedResource(
             type="resource",
             resource=BlobResourceContents(
-                uri=uri,
+                uri=cast(AnyUrl, uri),
                 mimeType="application/pdf",
                 blob=encoded,
                 _meta={"size": len(data)},
@@ -85,6 +89,44 @@ def as_image(self, data: bytes, mime_type: str | None) -> Image:
         return Image(data=data, format=subtype)
 
 
+def _patch_dynamic_blob_template(
+    server: FastMCP,
+    uri_template: str,
+    *,
+    default_mime: str,
+    loader: Callable[
+        [Mapping[str, object], Context | None],
+        Awaitable[tuple[bytes, str | None]],
+    ],
+) -> None:
+    template = server._resource_manager._templates.get(uri_template)
+    if template is None:
+        raise RuntimeError(
+            f"Resource template {uri_template} is not registered"
+        )
+
+    async def create_resource(
+        self: Any,
+        uri: str,
+        params: Mapping[str, object],
+        context: Context | None = None,
+    ) -> FunctionResource:
+        payload, mime_type = await loader(params, context)
+        return FunctionResource(
+            uri=uri,  # type: ignore[arg-type]
+            name=self.name,
+            title=self.title,
+            description=self.description,
+            mime_type=mime_type or default_mime,
+            icons=self.icons,
+            annotations=self.annotations,
+            fn=lambda payload=payload: payload,
+        )
+
+    template.mime_type = default_mime
+    object.__setattr__(template, "create_resource", MethodType(create_resource, template))
+
+
 def register_resources(server: FastMCP, service: ResourceService) -> None:
     @server.resource(
         "pdf2sqlite://pdf/{pdf_id}",
@@ -127,59 +169,127 @@ async def read_figure(figure_id: int, ctx: Context | None = None) -> bytes:  # n
     async def read_table_image(table_id: int, ctx: Context | None = None) -> bytes:  # noqa: ARG001
         return await service.load_table_image_blob(TableImageResource(table_id))
 
+    async def _figure_loader(
+        params: Mapping[str, object],
+        ctx: Context | None,
+    ) -> tuple[bytes, str | None]:  # noqa: ARG001
+        figure_id = _require_int(params.get("figure_id"), "figure.figure_id")
+        return await service.load_figure_blob(FigureResource(figure_id))
+
+    async def _table_loader(
+        params: Mapping[str, object],
+        ctx: Context | None,
+    ) -> tuple[bytes, str | None]:  # noqa: ARG001
+        table_id = _require_int(params.get("table_id"), "table.table_id")
+        data = await service.load_table_image_blob(
+            TableImageResource(table_id)
+        )
+        return data, "image/jpeg"
 
-def build_page_payload(page: dict[str, int | str | None]) -> dict[str, object]:
-    pdf_id = int(page["pdf_id"])
-    page_number = int(page["page_number"])
+    _patch_dynamic_blob_template(
+        server,
+        "pdf2sqlite://figure/{figure_id}",
+        default_mime="application/octet-stream",
+        loader=_figure_loader,
+    )
+    _patch_dynamic_blob_template(
+        server,
+        "pdf2sqlite://table-image/{table_id}",
+        default_mime="image/jpeg",
+        loader=_table_loader,
+    )
+
+
+def _require_int(value: object | None, label: str) -> int:
+    if isinstance(value, bool):
+        return int(value)
+    if isinstance(value, int):
+        return value
+    if isinstance(value, str):
+        stripped = value.strip()
+        if not stripped:
+            raise ValueError(f"{label} must not be empty")
+        try:
+            return int(stripped)
+        except ValueError as exc:
+            msg = f"{label} must be an integer-compatible value, got {value!r}"
+            raise ValueError(msg) from exc
+    if value is None:
+        raise ValueError(f"{label} is required")
+    raise TypeError(f"{label} must be int-like, got {type(value).__name__}")
+
+
+def _optional_int(value: object | None, label: str) -> int | None:
+    if value is None:
+        return None
+    return _require_int(value, label)
+
+
+def build_page_payload(page: Mapping[str, object]) -> dict[str, object]:
+    pdf_id = _require_int(page.get("pdf_id"), "page.pdf_id")
+    page_number = _require_int(page.get("page_number"), "page.page_number")
     resource_uri = build_pdf_page_uri(pdf_id, page_number)
     return {
-        "page_id": int(page["id"]),
+        "page_id": _require_int(page.get("id"), "page.id"),
         "pdf_id": pdf_id,
         "page_number": page_number,
         "gist": page.get("gist"),
-        "text_length": page.get("text_length"),
-        "data_bytes": page.get("data_bytes"),
+        "text_length": _optional_int(page.get("text_length"), "page.text_length"),
+        "data_bytes": _optional_int(page.get("data_bytes"), "page.data_bytes"),
         "resource": resource_uri,
     }
 
 
-def build_pdf_payload(pdf: dict[str, object]) -> dict[str, object]:
-    pdf_id = int(pdf["id"])
+def build_pdf_payload(pdf: Mapping[str, object]) -> dict[str, object]:
+    pdf_id = _require_int(pdf.get("id"), "pdf.id")
     return {
         "pdf_id": pdf_id,
         "title": pdf.get("title"),
         "description": pdf.get("description"),
-        "page_count": pdf.get("page_count"),
+        "page_count": _optional_int(pdf.get("page_count"), "pdf.page_count"),
         "resource": build_pdf_uri(pdf_id),
     }
 
 
-def build_figure_payload(figure: dict[str, object]) -> dict[str, object]:
-    figure_id = int(figure["id"])
+def build_figure_payload(figure: Mapping[str, object]) -> dict[str, object]:
+    figure_id = _require_int(figure.get("id"), "figure.id")
     return {
         "figure_id": figure_id,
         "description": figure.get("description"),
         "mime_type": figure.get("mime_type"),
-        "data_bytes": figure.get("data_bytes"),
+        "data_bytes": _optional_int(figure.get("data_bytes"), "figure.data_bytes"),
         "resource": build_figure_uri(figure_id),
     }
 
 
-def build_table_payload(table: dict[str, object]) -> dict[str, object]:
-    table_id = int(table["id"])
+def build_table_payload(table: Mapping[str, object]) -> dict[str, object]:
+    table_id = _require_int(table.get("id"), "table.id")
     payload: dict[str, object] = {
         "table_id": table_id,
         "description": table.get("description"),
         "caption_above": table.get("caption_above"),
         "caption_below": table.get("caption_below"),
-        "data_bytes": table.get("data_bytes"),
+        "text": table.get("text"),
         "resource": build_table_image_uri(table_id),
     }
-    if "pdf_id" in table and "page_number" in table:
-        payload["pdf_id"] = table.get("pdf_id")
-        payload["page_number"] = table.get("page_number")
-        payload["page_resource"] = build_pdf_page_uri(
-            int(table["pdf_id"]),
-            int(table["page_number"]),
-        )
+    text_length = _optional_int(table.get("text_length"), "table.text_length")
+    if text_length is not None:
+        payload["text_length"] = text_length
+    data_bytes = _optional_int(table.get("data_bytes"), "table.data_bytes")
+    if data_bytes is not None:
+        payload["data_bytes"] = data_bytes
+    xmin = _optional_int(table.get("xmin"), "table.xmin")
+    if xmin is not None:
+        payload["xmin"] = xmin
+    ymin = _optional_int(table.get("ymin"), "table.ymin")
+    if ymin is not None:
+        payload["ymin"] = ymin
+    pdf_id = _optional_int(table.get("pdf_id"), "table.pdf_id")
+    page_number = _optional_int(table.get("page_number"), "table.page_number")
+    if pdf_id is not None:
+        payload["pdf_id"] = pdf_id
+    if page_number is not None:
+        payload["page_number"] = page_number
+    if pdf_id is not None and page_number is not None:
+        payload["page_resource"] = build_pdf_page_uri(pdf_id, page_number)
     return payload