diff --git a/src/vtk_mcp/config.py b/src/vtk_mcp/config.py index 91ee40e..322f439 100644 --- a/src/vtk_mcp/config.py +++ b/src/vtk_mcp/config.py @@ -23,6 +23,11 @@ class Settings(BaseSettings): # Layer 3 — validation enable_validation: bool = True + # DSL translation + translate_model: str = "anthropic/claude-haiku-4-5" + translate_base_url: Optional[str] = None # e.g. http://localhost:11434 for Ollama + translate_api_key: Optional[str] = None # set to "ollama" for Ollama + # Transport transport: str = "stdio" http_host: str = "0.0.0.0" diff --git a/src/vtk_mcp/server.py b/src/vtk_mcp/server.py index 5019938..78ef7d6 100644 --- a/src/vtk_mcp/server.py +++ b/src/vtk_mcp/server.py @@ -250,6 +250,48 @@ def vtk_validate_import(import_statement: str) -> dict: return _f(import_statement, _ctx()) +# ── DSL translation ──────────────────────────────────────────────────────── + + +@mcp.tool() +def translate_prompt_to_dsl( + query: str, + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, +) -> str: + """Translate a natural language VTK request into the pipeline DSL. + + The DSL encodes the full pipeline as a structured specification + (sources, filters, auxiliary objects, render settings) with explicit + class slugs and parameter names, which produces more accurate code + generation than plain natural language. + + Args: + query: Natural language description of the desired VTK visualization. + model: LiteLLM model identifier. Overrides VTK_MCP_TRANSLATE_MODEL when set. + Use e.g. ``ollama/llama3`` for a local Ollama model. + base_url: OpenAI-compatible base URL. Overrides VTK_MCP_TRANSLATE_BASE_URL. + Example: ``http://localhost:11434`` for Ollama. + api_key: API key for the endpoint. Overrides VTK_MCP_TRANSLATE_API_KEY. + Pass ``"ollama"`` for Ollama (requires non-empty key). + + Returns: + A VTK pipeline DSL string ready to be used as a vtk-prompt input. + """ + from .tools.dsl import translate_prompt_to_dsl as _f + + return _f(query, _ctx(), model=model, base_url=base_url, api_key=api_key) + + +@mcp.tool() +def is_dsl_prompt(text: str) -> bool: + """Return True if *text* is already in the VTK pipeline DSL format.""" + from .tools.dsl import is_dsl_prompt as _f + + return _f(text) + + # ── Meta ─────────────────────────────────────────────────────────────────── diff --git a/src/vtk_mcp/tools/dsl.py b/src/vtk_mcp/tools/dsl.py new file mode 100644 index 0000000..fcb8dad --- /dev/null +++ b/src/vtk_mcp/tools/dsl.py @@ -0,0 +1,52 @@ +"""DSL translation tool — delegates to vtk-validate's dsl subpackage.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ..composition import VTKMCPContext + + +def translate_prompt_to_dsl( + query: str, + ctx: "VTKMCPContext", + model: str | None = None, + base_url: str | None = None, + api_key: str | None = None, +) -> str: + """Translate a natural language VTK request into the pipeline DSL. + + Uses vtk-validate's DSL translator with the loaded api_index for class + context. All parameters fall back to VTK_MCP_TRANSLATE_* env vars when + not provided. + + Args: + query: Natural language description (e.g. "create a warped sine surface"). + model: LiteLLM model identifier. Overrides VTK_MCP_TRANSLATE_MODEL when set. + base_url: OpenAI-compatible base URL (e.g. ``http://localhost:11434`` for Ollama). + Overrides VTK_MCP_TRANSLATE_BASE_URL when set. + api_key: API key for the endpoint. Overrides VTK_MCP_TRANSLATE_API_KEY when set. + + Returns: + A VTK pipeline DSL string ready for code generation. + """ + try: + from vtk_validate.dsl import translate_to_dsl + except ImportError as e: + return f"Error: vtk-validate[translate] not installed — {e}" + + return translate_to_dsl( + query=query, + api_index=ctx.api_index, + model=model or ctx.settings.translate_model, + base_url=base_url or ctx.settings.translate_base_url, + api_key=api_key or ctx.settings.translate_api_key, + ) + + +def is_dsl_prompt(text: str) -> bool: + """Return True if *text* is already in the VTK pipeline DSL format.""" + from vtk_validate.dsl import is_dsl + + return is_dsl(text) diff --git a/tests/test_dsl_tools.py b/tests/test_dsl_tools.py new file mode 100644 index 0000000..b3f6b79 --- /dev/null +++ b/tests/test_dsl_tools.py @@ -0,0 +1,103 @@ +"""Unit tests for the DSL tools in vtk_mcp.tools.dsl.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +pytestmark = pytest.mark.unit + + +def _make_ctx( + translate_model: str = "test-model", + translate_base_url: str | None = None, + translate_api_key: str | None = None, +) -> MagicMock: + ctx = MagicMock() + ctx.settings.translate_model = translate_model + ctx.settings.translate_base_url = translate_base_url + ctx.settings.translate_api_key = translate_api_key + ctx.api_index = MagicMock() + ctx.api_index.vtk_version = "9.6.1" + return ctx + + +class TestIsDslPrompt: + def test_dsl_text_returns_true(self): + from vtk_mcp.tools.dsl import is_dsl_prompt + + assert is_dsl_prompt("create plane_source called src with x_resolution 60") is True + + def test_natural_language_returns_false(self): + from vtk_mcp.tools.dsl import is_dsl_prompt + + assert is_dsl_prompt("make a sphere with a colormap") is False + + def test_empty_string_returns_false(self): + from vtk_mcp.tools.dsl import is_dsl_prompt + + assert is_dsl_prompt("") is False + + +class TestTranslatePromptToDsl: + def test_delegates_to_vtk_validate(self): + from vtk_mcp.tools.dsl import translate_prompt_to_dsl + + ctx = _make_ctx(translate_model="haiku") + expected_dsl = "create plane_source called src with x_resolution 10" + + with patch("vtk_validate.dsl.translate_to_dsl", return_value=expected_dsl) as mock_fn: + result = translate_prompt_to_dsl("make a plane", ctx) + + mock_fn.assert_called_once_with( + query="make a plane", + api_index=ctx.api_index, + model="haiku", + base_url=None, + api_key=None, + ) + assert result == expected_dsl + + def test_model_override_takes_precedence(self): + from vtk_mcp.tools.dsl import translate_prompt_to_dsl + + ctx = _make_ctx(translate_model="default-model") + with patch("vtk_validate.dsl.translate_to_dsl", return_value="render render with background [0,0,0]") as mock_fn: + translate_prompt_to_dsl("make a scene", ctx, model="override-model") + + mock_fn.assert_called_once_with( + query="make a scene", + api_index=ctx.api_index, + model="override-model", + base_url=None, + api_key=None, + ) + + def test_no_override_uses_settings_model(self): + from vtk_mcp.tools.dsl import translate_prompt_to_dsl + + ctx = _make_ctx(translate_model="settings-model") + with patch("vtk_validate.dsl.translate_to_dsl", return_value="render render with background [0,0,0]") as mock_fn: + translate_prompt_to_dsl("make a scene", ctx) + + assert mock_fn.call_args.kwargs["model"] == "settings-model" + + def test_missing_vtk_validate_returns_error(self): + from vtk_mcp.tools.dsl import translate_prompt_to_dsl + + ctx = _make_ctx() + with patch.dict("sys.modules", {"vtk_validate.dsl": None}): + result = translate_prompt_to_dsl("make a sphere", ctx) + + assert result.startswith("Error:") + + def test_returns_string(self): + from vtk_mcp.tools.dsl import translate_prompt_to_dsl + + ctx = _make_ctx() + with patch("vtk_validate.dsl.translate_to_dsl", return_value="render render with background [0,0,0]"): + result = translate_prompt_to_dsl("render a black background", ctx) + + assert isinstance(result, str) + assert len(result) > 0