diff --git a/autofit/mcp/__init__.py b/autofit/mcp/__init__.py new file mode 100644 index 000000000..62056c242 --- /dev/null +++ b/autofit/mcp/__init__.py @@ -0,0 +1,38 @@ +""" +The read-only results-inspector MCP core, distributed as the optional +``autofit[mcp]`` extra. + +It exposes PyAutoFit output directories to chat harnesses that cannot execute +code (Claude Desktop, ChatGPT) as a small set of read-only tools — list fits +ranked by evidence, read model/posterior/result summaries, view result images +inline. It is deliberately *glue, not code*: every tool is one existing public +aggregator call plus serialization; composing models and running searches stay +Python-first. + +Layout: + +- ``tools`` — the plain aggregator-wrapper functions (no ``mcp`` dependency; the + test suite exercises them directly). +- ``server.core_server`` — registers those functions on a FastMCP stdio server. +- ``bootstrap`` — the launch-time config-pin + stdout protections, kept free of + any autofit import so a launcher can call them *before* importing the + autofit-backed modules. + +Run standalone with ``python -m autofit.mcp`` (needs the extra: +``pip install autofit[mcp]``). Downstream servers — e.g. the PyAutoLens +results-inspector, which adds a lens image/FITS layer — call ``core_server()`` +and register their own tools on top. +""" + + +def __getattr__(name): + # Lazy re-exports (PEP 562): ``autofit.mcp.core_server`` / ``autofit.mcp.png`` + # resolve without importing ``autofit.mcp.server`` (and therefore autofit and + # the ``mcp`` package) at ``import autofit.mcp`` time — keeping this package + # import-light so a launcher can pin config before the heavy import. + if name in ("core_server", "png", "_png"): + from autofit.mcp import server + + return server.core_server if name == "core_server" else server._png + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + diff --git a/autofit/mcp/__main__.py b/autofit/mcp/__main__.py new file mode 100644 index 000000000..55ea8672f --- /dev/null +++ b/autofit/mcp/__main__.py @@ -0,0 +1,28 @@ +""" +Standalone runner for the ``autofit[mcp]`` extra: ``python -m autofit.mcp``. + +Serves the seven read-only core tools over stdio. Downstream servers that add +their own tools (e.g. the PyAutoLens lens layer) provide their own launcher and +call :func:`autofit.mcp.server.core_server` directly. + +The ordering here is load-bearing: force JAX onto CPU and pin config to autofit's +own bundled ``config/`` (so the server does not depend on the launch directory) +*before* importing the autofit-backed server, and guard that import so no +import-time stdout reaches the JSON-RPC channel. +""" + +import contextlib +import os +import sys +from pathlib import Path + +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +from autofit.mcp.bootstrap import pin_config + +pin_config(Path(__file__).resolve().parents[1] / "config") + +with contextlib.redirect_stdout(sys.stderr): + from autofit.mcp.server import core_server + +core_server().run() diff --git a/autofit/mcp/bootstrap.py b/autofit/mcp/bootstrap.py new file mode 100644 index 000000000..f15ab4b30 --- /dev/null +++ b/autofit/mcp/bootstrap.py @@ -0,0 +1,53 @@ +""" +Launch-time hardening for the MCP stdio server. + +This module is deliberately kept free of any ``autofit`` import so a launcher can +call it *before* importing the autofit-backed tool modules — autofit reads its +configuration at import time, and jax's backend probe logs to stdout during that +import, both of which must be handled first. It imports only the standard library +and ``autonerves``. +""" + +import logging +import sys +from pathlib import Path +from typing import Union + + +def pin_config(config_path: Union[str, Path]) -> None: + """ + Pin ``autonerves`` config to ``config_path`` instead of letting it default to + ``os.getcwd()/config``. + + autofit reads config at import time; if the server is launched from a foreign + working directory — e.g. a chat client spawning ``wsl.exe`` in a Windows + folder — the default would scan unrelated files and crash on import (a + ``desktop.ini`` trips configparser interpolation). Call this before importing + ``autofit.mcp.server`` (or anything else that pulls in autofit). + """ + from autonerves import conf + + config_path = Path(config_path) + conf.instance = conf.Config( + str(config_path), output_path=str(config_path.parent / "output") + ) + + +def route_logging_to_stderr() -> None: + """ + Rebind every stdout-bound logging handler to stderr. + + stdout carries the JSON-RPC channel, but autofit's logging config (loaded on + import) attaches stdout stream handlers — one stray log line corrupts the + protocol. Call this after autofit has been imported. + """ + loggers = [logging.getLogger()] + [ + logging.getLogger(name) for name in logging.root.manager.loggerDict + ] + for logger in loggers: + for handler in getattr(logger, "handlers", []): + if ( + isinstance(handler, logging.StreamHandler) + and getattr(handler, "stream", None) is sys.stdout + ): + handler.setStream(sys.stderr) diff --git a/autofit/mcp/server.py b/autofit/mcp/server.py new file mode 100644 index 000000000..940c22d89 --- /dev/null +++ b/autofit/mcp/server.py @@ -0,0 +1,113 @@ +""" +The read-only results-inspector MCP stdio server core. + +``core_server()`` returns a FastMCP with the seven read-only tools registered +over ``autofit.mcp.tools``. Callers — the assistants, or ``python -m autofit.mcp`` +— register any extra tools on the returned object and call ``.run()``. + +Importing this module imports autofit. A launcher must therefore, *before* +importing it: pin config with ``autofit.mcp.bootstrap.pin_config`` (so config +does not depend on the working directory), set ``JAX_PLATFORMS=cpu`` (skips the +jax backend probe, which logs to stdout), and guard the import with +``contextlib.redirect_stdout(sys.stderr)`` — stdout carries the JSON-RPC channel. +``__main__`` does exactly this for the standalone case. +""" + +import io + +from mcp.server.fastmcp import FastMCP, Image + +from autofit.mcp import tools +from autofit.mcp.bootstrap import route_logging_to_stderr + + +def _png(image) -> Image: + buffer = io.BytesIO() + image.save(buffer, format="PNG") + return Image(data=buffer.getvalue(), format="png") + + +def core_server(name: str = "pyauto-results-inspector") -> FastMCP: + """ + Build a FastMCP server with the seven read-only results-inspector tools + registered. Downstream servers add their own tools on the returned object. + """ + route_logging_to_stderr() + + mcp = FastMCP(name) + + @mcp.tool() + def list_searches( + directory: str, + sort_by: str = "log_evidence", + limit: int = 20, + completed_only: bool = False, + ) -> list: + """ + List every model-fit found under `directory` (searched recursively): one + row per fit with its name, unique tag, output directory, completion state, + log evidence, maximum log likelihood and free-parameter count. + + Rows are sorted by `sort_by` (descending; fits without that value last) — + use "log_evidence" for nested samplers or "max_log_likelihood" generally — + and truncated to `limit` (pass 0 for all). The returned `directory` of a + row is what the other tools take as their `directory` argument. + """ + return tools.list_searches( + directory, sort_by=sort_by, limit=limit, completed_only=completed_only + ) + + @mcp.tool() + def get_model(directory: str) -> dict: + """ + The model that was fitted in one search-output directory: a human-readable + `info` block (component classes, priors) and the full model as a dict. + """ + return tools.get_model(directory) + + @mcp.tool() + def get_result_summary(directory: str) -> str: + """ + The `model.results` text for one search-output directory: the fit's own + summary of the maximum-likelihood model and (when the search produces + them) parameter estimates with errors. + """ + return tools.get_result_summary(directory) + + @mcp.tool() + def get_samples_summary(directory: str) -> dict: + """ + Posterior summary for one search-output directory: log evidence (None for + MLE/MCMC searches without one), maximum log likelihood, the model's + parameter paths, and the maximum-likelihood and median-PDF parameter + vectors (`median_pdf_parameters` is None for MLE searches, which have no + PDF). + """ + return tools.get_samples_summary(directory) + + @mcp.tool() + def get_search_info(directory: str) -> dict: + """ + The non-linear search used in one search-output directory: name, unique + tag, completion state, and the search's serialized settings. + """ + return tools.get_search_info(directory) + + @mcp.tool() + def list_images(directory: str) -> list: + """ + Names of the visualization images (`image/*.png`) available in one + search-output directory — pass a name (without `.png`) to `fetch_image`. + """ + return tools.list_images(directory) + + @mcp.tool() + def fetch_image(directory: str, name: str = "subplot_fit") -> Image: + """ + One visualization image from a search-output directory (e.g. + "subplot_fit"), returned inline so it renders directly in chat. Use + `list_images` to see what is available. + """ + return _png(tools.fetch_image(directory, name=name)) + + return mcp diff --git a/autofit/mcp/tools.py b/autofit/mcp/tools.py new file mode 100644 index 000000000..21f0e2376 --- /dev/null +++ b/autofit/mcp/tools.py @@ -0,0 +1,147 @@ +""" +Read-only results-inspector tools over PyAutoFit output directories. + +Each function is a thin wrapper over an existing public PyAutoFit aggregator +API (`autofit.aggregator.Aggregator`, `af.SearchOutput`): argument parsing, one call, and +JSON-friendly serialization — nothing more. Any behaviour beyond that belongs +in PyAutoFit itself, not here ("glue, not code"). + +This module deliberately has no MCP dependency: `server.py` registers these +functions as MCP tools, and the test suite exercises them without the `mcp` +package installed. +""" + +import contextlib +import json +import sys +from pathlib import Path + +from autonerves.dictable import to_dict + +import autofit as af + +# The directory-backed aggregator — af.Aggregator is the database-backed one, +# and the alias also keeps the API audit from resolving it there. +from autofit.aggregator import Aggregator as DirectoryAggregator + + +@contextlib.contextmanager +def _stdout_to_stderr(): + """ + An MCP stdio server must keep stdout clean — it carries the JSON-RPC + channel — but the directory aggregator prints progress to stdout, + so every autofit call runs with stdout redirected to stderr. + """ + with contextlib.redirect_stdout(sys.stderr): + yield + + +def _float_or_none(value): + try: + return None if value is None else float(value) + except (TypeError, ValueError): + return None + + +def _search_row(search) -> dict: + summary = search.samples_summary + max_lh_sample = getattr(summary, "max_log_likelihood_sample", None) + return dict( + name=search.name, + unique_tag=search.unique_tag, + directory=str(search.directory), + is_complete=search.is_complete, + log_evidence=_float_or_none(getattr(summary, "log_evidence", None)), + max_log_likelihood=_float_or_none( + getattr(max_lh_sample, "log_likelihood", None) + ), + model_free_parameters=getattr(search.model, "prior_count", None), + ) + + +def list_searches( + directory: str, + sort_by: str = "log_evidence", + limit: int = 20, + completed_only: bool = False, +) -> list: + with _stdout_to_stderr(): + aggregator = DirectoryAggregator.from_directory( + directory, completed_only=completed_only + ) + rows = [_search_row(search) for search in aggregator] + rows.sort( + key=lambda row: (row.get(sort_by) is not None, row.get(sort_by)), + reverse=True, + ) + return rows[:limit] if limit else rows + + +def get_model(directory: str) -> dict: + with _stdout_to_stderr(): + model = af.SearchOutput(Path(directory)).model + return dict(info=model.info, model=to_dict(model)) + + +def get_result_summary(directory: str) -> str: + with _stdout_to_stderr(): + return af.SearchOutput(Path(directory)).model_results + + +def get_samples_summary(directory: str) -> dict: + with _stdout_to_stderr(): + search_output = af.SearchOutput(Path(directory)) + summary = search_output.samples_summary + if summary is None: + raise FileNotFoundError( + f"No samples summary found under {directory}/files/ — " + "is this a completed search output directory?" + ) + model = search_output.model + return dict( + log_evidence=_float_or_none(summary.log_evidence), + max_log_likelihood=_float_or_none( + summary.max_log_likelihood_sample.log_likelihood + ), + parameter_paths=[".".join(path) for path in model.paths], + max_log_likelihood_parameters=[ + float(value) + for value in summary.max_log_likelihood_sample.parameter_lists_for_model( + model + ) + ], + # MLE searches (e.g. LBFGS) have no PDF, so no median-PDF sample. + median_pdf_parameters=None + if summary.median_pdf_sample is None + else [ + float(value) + for value in summary.median_pdf_sample.parameter_lists_for_model( + model + ) + ], + ) + + +def get_search_info(directory: str) -> dict: + search_json = Path(directory) / "files" / "search.json" + with _stdout_to_stderr(): + search_output = af.SearchOutput(Path(directory)) + return dict( + name=search_output.name, + unique_tag=search_output.unique_tag, + is_complete=search_output.is_complete, + search=json.loads(search_json.read_text()) + if search_json.exists() + else None, + ) + + +def list_images(directory: str) -> list: + # Visualization outputs live under /image/ (SearchOutput.image + # reads there, despite its docstring saying `files/`). + return sorted(path.name for path in (Path(directory) / "image").glob("*.png")) + + +def fetch_image(directory: str, name: str = "subplot_fit"): + with _stdout_to_stderr(): + return af.SearchOutput(Path(directory)).image(name) diff --git a/pyproject.toml b/pyproject.toml index d55d03e09..0f1aba552 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ local_scheme = "no-local-version" [project.optional-dependencies] jax = ["autonerves[jax]", "optax>=0.2.5"] +mcp = ["mcp"] optional = [ "autofit[jax]", "astropy>=5.0", diff --git a/test_autofit/mcp/__init__.py b/test_autofit/mcp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autofit/mcp/test_mcp_tools.py b/test_autofit/mcp/test_mcp_tools.py new file mode 100644 index 000000000..5f1e1dc5a --- /dev/null +++ b/test_autofit/mcp/test_mcp_tools.py @@ -0,0 +1,130 @@ +""" +Tests for the results-inspector MCP tool functions (`autofit/mcp/tools.py`, the +``autofit[mcp]`` extra). + +The fixture runs a real (tiny) `af.LBFGS` fit of the `af.ex.Gaussian` toy so the +output directory always matches the current on-disk format — a frozen fixture +directory would silently drift. LBFGS is an MLE search, so `log_evidence` is +legitimately absent; tests assert that shape rather than skipping it. +""" + +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image + +import autofit as af +from autonerves import conf + +from autofit.mcp import tools + +# The shipped config (autofit/config), not test_autofit/config: it keeps search +# output files on disk (remove_files: false), whereas the test config zips them +# away — and the read-only tools operate on real, unzipped output directories, +# exactly as an end user's `pip install autofit` config leaves them. +CONFIG = Path(__file__).resolve().parents[2] / "autofit" / "config" + + +@pytest.fixture(scope="module") +def output_root(tmp_path_factory): + root = tmp_path_factory.mktemp("output") + conf.instance.push(new_path=str(CONFIG), output_path=str(root)) + + xvalues = np.arange(100.0) + gaussian = af.ex.Gaussian(centre=50.0, normalization=25.0, sigma=10.0) + data = gaussian.model_data_from(xvalues=xvalues) + noise_map = np.full(fill_value=1.0, shape=data.shape) + analysis = af.ex.Analysis(data=data, noise_map=noise_map) + + for name in ("fit_0", "fit_1"): + search = af.LBFGS(name=name, path_prefix="mcp_fixture") + search.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + return root + + +@pytest.fixture(scope="module") +def fit_directory(output_root): + directory = sorted( + marker.parent for marker in output_root.rglob(".completed") + )[0] + (directory / "image").mkdir(exist_ok=True) + Image.new("RGB", (32, 16), color=(200, 40, 40)).save( + directory / "image" / "subplot_fit.png" + ) + return directory + + +def test_list_searches(output_root, fit_directory): + rows = tools.list_searches(str(output_root), sort_by="max_log_likelihood") + + assert len(rows) == 2 + assert {row["name"] for row in rows} == {"fit_0", "fit_1"} + for row in rows: + assert row["is_complete"] is True + assert isinstance(row["max_log_likelihood"], float) + assert row["log_evidence"] is None + assert row["model_free_parameters"] == 3 + + +def test_list_searches_limit(output_root): + assert len(tools.list_searches(str(output_root), limit=1)) == 1 + + +def test_list_searches_sorts_none_last(output_root): + rows = tools.list_searches(str(output_root), sort_by="log_evidence") + assert len(rows) == 2 + + +def test_get_model(fit_directory): + result = tools.get_model(str(fit_directory)) + + assert "Gaussian" in result["info"] + assert result["model"]["class_path"].endswith("Gaussian") + + +def test_get_result_summary(fit_directory): + text = tools.get_result_summary(str(fit_directory)) + + assert "Maximum Log Likelihood" in text + + +def test_get_samples_summary(fit_directory): + summary = tools.get_samples_summary(str(fit_directory)) + + assert summary["log_evidence"] is None + assert isinstance(summary["max_log_likelihood"], float) + assert summary["parameter_paths"] == [ + "centre", + "normalization", + "sigma", + ] + assert len(summary["max_log_likelihood_parameters"]) == 3 + assert summary["median_pdf_parameters"] is None + assert summary["max_log_likelihood_parameters"][0] == pytest.approx( + 50.0, abs=1.0 + ) + + +def test_get_search_info(fit_directory): + info = tools.get_search_info(str(fit_directory)) + + assert info["name"] == "fit_0" + assert info["is_complete"] is True + assert info["search"] is not None + + +def test_list_images(fit_directory): + assert "subplot_fit.png" in tools.list_images(str(fit_directory)) + + +def test_fetch_image(fit_directory): + image = tools.fetch_image(str(fit_directory), name="subplot_fit") + + assert image.size == (32, 16) + + +def test_fetch_image_missing(fit_directory): + with pytest.raises(Exception): + tools.fetch_image(str(fit_directory), name="no_such_image")