diff --git a/src/ifcmcp/ifcmcp/core.py b/src/ifcmcp/ifcmcp/core.py index 2d9a52b57f..4df97cd241 100644 --- a/src/ifcmcp/ifcmcp/core.py +++ b/src/ifcmcp/ifcmcp/core.py @@ -370,10 +370,12 @@ class IfcSession: scale: float = 1.0 / 100.0, png_width: int = 1024, png_height: int = 1024, + output_format: str = "png", ) -> bytes: - """Generate a 2D technical drawing (floor plan, elevation, or section) and return PNG bytes. + """Generate a 2D technical drawing (floor plan, elevation, or section) and return image bytes. - Uses ifcopenshell.draw to produce SVG output which is rasterised to PNG via CairoSVG. + Uses ifcopenshell.draw to produce SVG output which is rasterised to PNG via CairoSVG + when output_format is 'png'. :param selector: ifcopenshell selector to restrict plotted elements (e.g. ``'IfcWall'``). Omit to plot the whole model. @@ -386,12 +388,13 @@ class IfcSession: :param scale: Model-to-paper scale ratio (default 0.01 = 1:100). :param png_width: Raster output width in pixels (default 1024). :param png_height: Raster output height in pixels (default 1024). - :return: PNG image as raw bytes. + :param output_format: ``'svg'`` or ``'png'`` (default ``'png'``). + :return: SVG or PNG bytes depending on output_format. """ model = self._require_model() return plot_mod.plot( model, - output_format="png", + output_format=output_format, selector=selector if selector else None, element_ids=element_ids, view=view, diff --git a/src/ifcmcp/ifcmcp/server.py b/src/ifcmcp/ifcmcp/server.py index 13c43d55a6..95a51f6036 100644 --- a/src/ifcmcp/ifcmcp/server.py +++ b/src/ifcmcp/ifcmcp/server.py @@ -150,11 +150,14 @@ def build_server() -> Any: scale: float = 1.0 / 100.0, png_width: int = 1024, png_height: int = 1024, + output_path: str = "", ) -> list[ImageContent]: """Generate a 2D technical drawing of the loaded IFC model. - Returns an inline image (floor plan, elevation, or section) that the - LLM can inspect to understand the 2D layout of the model. + Returns an inline PNG image (floor plan, elevation, or section) that the + LLM can inspect to understand the 2D layout of the model. If + ``output_path`` is provided the drawing is also saved to disk — as SVG + when the path ends in ``.svg``, otherwise as PNG. :param selector: ifcopenshell selector to restrict plotted elements (e.g. ``'IfcWall'``). Omit to plot the whole model. @@ -167,6 +170,7 @@ def build_server() -> Any: :param scale: Model-to-paper scale ratio (default 0.01 = 1:100). :param png_width: Raster output width in pixels (default 1024). :param png_height: Raster output height in pixels (default 1024). + :param output_path: Optional file path to save the drawing to disk. """ png_bytes = session.ifc_plot( selector=selector, @@ -177,7 +181,24 @@ def build_server() -> Any: scale=scale, png_width=png_width, png_height=png_height, + output_format="png", ) + if output_path: + if output_path.endswith(".svg"): + svg_bytes = session.ifc_plot( + selector=selector, + element_ids=element_ids, + view=view, + width_mm=width_mm, + height_mm=height_mm, + scale=scale, + output_format="svg", + ) + with open(output_path, "wb") as f: + f.write(svg_bytes) + else: + with open(output_path, "wb") as f: + f.write(png_bytes) return [ImageContent(type="image", data=base64.b64encode(png_bytes).decode(), mimeType="image/png")] @server.tool(structured_output=False) @@ -185,11 +206,13 @@ def build_server() -> Any: selector: str = "", element_ids: list[int] | None = None, view: str = "iso", + output_path: str = "", ) -> list[ImageContent]: """Render the loaded IFC model to a PNG image. Returns an inline image the LLM can inspect to understand the spatial - layout of the model or a specific element in context. + layout of the model or a specific element in context. If + ``output_path`` is provided the PNG is also saved to that file path. :param selector: ifcopenshell selector to restrict rendered elements (e.g. ``'IfcWall'``, ``'IfcBuildingStorey[Name="0"]'``). @@ -198,8 +221,12 @@ def build_server() -> Any: are rendered in translucent grey so the subject stands out. :param view: Camera angle — ``iso`` (default), ``top``, ``south``, ``north``, ``east``, or ``west``. + :param output_path: Optional file path to save the PNG to disk. """ png_bytes = session.ifc_render(selector=selector, element_ids=element_ids, view=view) + if output_path: + with open(output_path, "wb") as f: + f.write(png_bytes) return [ImageContent(type="image", data=base64.b64encode(png_bytes).decode(), mimeType="image/png")] return server \ No newline at end of file diff --git a/src/ifcmcp/tests/test_server.py b/src/ifcmcp/tests/test_server.py index a45731631b..996be3b024 100644 --- a/src/ifcmcp/tests/test_server.py +++ b/src/ifcmcp/tests/test_server.py @@ -1,4 +1,8 @@ # This file was generated with the assistance of an AI coding tool. +from unittest.mock import patch + +import pytest + from ifcmcp.server import build_server @@ -25,3 +29,54 @@ class TestServerRegistration: ] for name in expected: assert name in tools, f"Tool {name} not registered" + + +@pytest.fixture +def tool_fns(): + """Return a dict of tool name → raw function from a freshly built server.""" + server = build_server() + return {t.name: t.fn for t in server._tool_manager.list_tools()} + + +PNG_FAKE = b"\x89PNG\r\n\x1a\nFAKE" +SVG_FAKE = b"FAKE" + + +class TestRenderOutputPath: + def test_no_output_path_no_file_written(self, tool_fns, tmp_path): + with patch("ifcmcp.core.IfcSession.ifc_render", return_value=PNG_FAKE): + tool_fns["ifc_render"](selector="", element_ids=None, view="iso", output_path="") + assert list(tmp_path.iterdir()) == [] + + def test_png_output_path_writes_file(self, tool_fns, tmp_path): + out = str(tmp_path / "render.png") + with patch("ifcmcp.core.IfcSession.ifc_render", return_value=PNG_FAKE): + tool_fns["ifc_render"](selector="", element_ids=None, view="iso", output_path=out) + assert open(out, "rb").read() == PNG_FAKE + + +class TestPlotOutputPath: + def test_no_output_path_no_file_written(self, tool_fns, tmp_path): + with patch("ifcmcp.core.IfcSession.ifc_plot", return_value=PNG_FAKE): + tool_fns["ifc_plot"](selector="", element_ids=None, view="floorplan", + width_mm=297.0, height_mm=420.0, scale=0.01, + png_width=1024, png_height=1024, output_path="") + assert list(tmp_path.iterdir()) == [] + + def test_png_output_path_writes_png(self, tool_fns, tmp_path): + out = str(tmp_path / "plot.png") + with patch("ifcmcp.core.IfcSession.ifc_plot", return_value=PNG_FAKE): + tool_fns["ifc_plot"](selector="", element_ids=None, view="floorplan", + width_mm=297.0, height_mm=420.0, scale=0.01, + png_width=1024, png_height=1024, output_path=out) + assert open(out, "rb").read() == PNG_FAKE + + def test_svg_output_path_writes_svg(self, tool_fns, tmp_path): + out = str(tmp_path / "plot.svg") + # ifc_plot is called twice: once with "png" for the inline image, + # once with "svg" for the file. + with patch("ifcmcp.core.IfcSession.ifc_plot", side_effect=[PNG_FAKE, SVG_FAKE]): + tool_fns["ifc_plot"](selector="", element_ids=None, view="floorplan", + width_mm=297.0, height_mm=420.0, scale=0.01, + png_width=1024, png_height=1024, output_path=out) + assert open(out, "rb").read() == SVG_FAKE diff --git a/src/ifcmcp/tests/test_session.py b/src/ifcmcp/tests/test_session.py index 0e933a8eb4..4fa59120c0 100644 --- a/src/ifcmcp/tests/test_session.py +++ b/src/ifcmcp/tests/test_session.py @@ -1,4 +1,6 @@ # This file was generated with the assistance of an AI coding tool. +from unittest.mock import patch + import ifcopenshell import pytest @@ -42,3 +44,19 @@ class TestSave: def test_save_no_path_no_original(self, loaded_session): with pytest.raises(IfcSessionError, match="No path specified"): loaded_session.ifc_save() + + +class TestIfcPlotOutputFormat: + """ifc_plot should pass output_format through to the underlying plot function.""" + + def test_default_output_format_is_png(self, loaded_session): + with patch("ifcmcp.core.plot_mod.plot", return_value=b"PNG_FAKE") as mock_plot: + loaded_session.ifc_plot() + mock_plot.assert_called_once() + assert mock_plot.call_args.kwargs["output_format"] == "png" + + def test_svg_output_format(self, loaded_session): + with patch("ifcmcp.core.plot_mod.plot", return_value=b"SVG_FAKE") as mock_plot: + result = loaded_session.ifc_plot(output_format="svg") + assert result == b"SVG_FAKE" + assert mock_plot.call_args.kwargs["output_format"] == "svg"