mle-kit-mcp 0.0.2__tar.gz → 0.0.4__tar.gz

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mle-kit-mcp
-Version: 0.0.2
+Version: 0.0.4
 Summary: MCP server that provides different tools for MLE
 Author-email: Ilya Gusev <phoenixilya@gmail.com>
 Project-URL: Homepage, https://github.com/IlyaGusev/mle_kit_mcp

mle_kit_mcp-0.0.4/mle_kit_mcp/files.py

@@ -0,0 +1,25 @@
+import os
+from typing import Optional
+from pathlib import Path
+
+DIR_PATH = Path(__file__).parent
+ROOT_PATH = DIR_PATH.parent
+DEFAULT_WORKSPACE_DIR_PATH: Path = DIR_PATH / "workdir"
+
+
+class WorkspaceDirectory:
+    workspace_dir: Optional[Path] = None
+
+    @classmethod
+    def get_dir(cls) -> Path:
+        if cls.workspace_dir is None:
+            return Path(os.getenv("WORKSPACE_DIR", DEFAULT_WORKSPACE_DIR_PATH))
+        return cls.workspace_dir
+
+    @classmethod
+    def set_dir(cls, workspace_dir: Path) -> None:
+        cls.workspace_dir = workspace_dir
+
+
+def get_workspace_dir() -> Path:
+    return WorkspaceDirectory.get_dir()

mle_kit_mcp-0.0.4/mle_kit_mcp/server.py

@@ -0,0 +1,40 @@
+from pathlib import Path
+from typing import Optional
+
+import fire  # type: ignore
+import uvicorn
+from mcp.server.fastmcp import FastMCP
+
+from .tools.bash import bash
+from .tools.text_editor import text_editor
+from .tools.remote_gpu import (
+    remote_bash,
+    create_remote_text_editor,
+    remote_download,
+)
+from .files import get_workspace_dir, WorkspaceDirectory
+
+
+def run(host: str = "0.0.0.0", port: int = 5050, workspace: Optional[str] = None) -> None:
+    if workspace:
+        WorkspaceDirectory.set_dir(Path(workspace))
+    workspace_path = get_workspace_dir()
+    workspace_path.mkdir(parents=True, exist_ok=True)
+
+    server = FastMCP("MLE kit MCP", stateless_http=True)
+
+    remote_text_editor = create_remote_text_editor(text_editor)
+
+    server.add_tool(bash)
+    server.add_tool(text_editor)
+    server.add_tool(remote_bash)
+    server.add_tool(remote_text_editor)
+    server.add_tool(remote_download)
+
+    http_app = server.streamable_http_app()
+
+    uvicorn.run(http_app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    fire.Fire(run)

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/mle_kit_mcp/tools/bash.py

@@ -3,7 +3,7 @@ import atexit
 import signal
 from typing import Optional, Any
 
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 
 
 _container = None
@@ -63,7 +63,7 @@ def bash(command: str) -> str:
                 tty=True,
                 stdin_open=True,
                 volumes={
-                    WORKSPACE_DIR_PATH: {
+                    get_workspace_dir(): {
                         "bind": DOCKER_WORKSPACE_DIR_PATH,
                         "mode": "rw",
                     }

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/mle_kit_mcp/tools/remote_gpu.py

@@ -12,7 +12,7 @@ from dataclasses import dataclass
 from dotenv import load_dotenv
 from vastai_sdk import VastAI  # type: ignore
 
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 
 BASE_IMAGE = "phoenix120/holosophos_mle"
 DEFAULT_GPU_TYPE = "RTX_3090"
@@ -94,7 +94,7 @@ def get_offers(vast_sdk: VastAI, gpu_name: str) -> List[int]:
 
 
 def run_command(
-    instance: InstanceInfo, command: str, timeout: int = 60
+    instance: InstanceInfo, command: str, timeout: int = 60, raise_exc: bool = True
 ) -> subprocess.CompletedProcess[str]:
     cmd = [
         "ssh",
@@ -117,7 +117,7 @@ def run_command(
     ]
     try:
         result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
-        if result.returncode != 0:
+        if result.returncode != 0 and raise_exc:
             raise Exception(
                 f"Error running command: {command}; "
                 f"Output: {result.stdout}; "
@@ -257,9 +257,9 @@ def launch_instance(vast_sdk: VastAI, gpu_name: str) -> Optional[InstanceInfo]:
 
 def send_scripts() -> None:
     assert _instance_info
-    for name in os.listdir(WORKSPACE_DIR_PATH):
+    for name in os.listdir(get_workspace_dir()):
         if name.endswith(".py"):
-            send_rsync(_instance_info, f"{WORKSPACE_DIR_PATH}/{name}", "/root")
+            send_rsync(_instance_info, f"{get_workspace_dir()}/{name}", "/root")
 
 
 def init_all() -> None:
@@ -301,7 +301,7 @@ def remote_bash(command: str, timeout: Optional[int] = 60) -> str:
     init_all()
     assert _instance_info
     assert timeout
-    result = run_command(_instance_info, command, timeout=timeout)
+    result = run_command(_instance_info, command, timeout=timeout, raise_exc=False)
     output = ("STDOUT: " + result.stdout + "\n") if result.stdout else ""
     output += ("STDERR: " + result.stderr) if result.stderr else ""
     return output.replace(VAST_AI_GREETING, "")
@@ -316,7 +316,7 @@ def remote_download(file_path: str) -> str:
     """
     init_all()
     assert _instance_info
-    recieve_rsync(_instance_info, f"/root/{file_path}", f"{WORKSPACE_DIR_PATH}")
+    recieve_rsync(_instance_info, f"/root/{file_path}", f"{get_workspace_dir()}")
     return f"File '{file_path}' downloaded!"
 
 
@@ -335,12 +335,12 @@ def create_remote_text_editor(
         command = args_dict["command"]
 
         if command != "write":
-            recieve_rsync(_instance_info, f"/root/{path}", f"{WORKSPACE_DIR_PATH}")
+            recieve_rsync(_instance_info, f"/root/{path}", f"{get_workspace_dir()}")
 
         result: str = text_editor_func(*args, **kwargs)
 
         if command != "view":
-            send_rsync(_instance_info, f"{WORKSPACE_DIR_PATH}/{path}", "/root")
+            send_rsync(_instance_info, f"{get_workspace_dir()}/{path}", "/root")
 
         return result
 

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/mle_kit_mcp/tools/text_editor.py

@@ -2,7 +2,7 @@ from collections import defaultdict
 from typing import Dict, List, Optional
 from pathlib import Path
 
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 from mle_kit_mcp.utils import truncate_content
 
 WRITE_MAX_OUTPUT_LENGTH = 500
@@ -174,7 +174,7 @@ def text_editor(
     ), "Absolute path is not supported, only relative to the work directory"
     valid_commands = ("view", "write", "str_replace", "insert", "undo_edit", "append")
 
-    path_obj = WORKSPACE_DIR_PATH / path
+    path_obj = get_workspace_dir() / path
 
     if command == "view":
         show_lines = show_lines if show_lines is not None else False

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/mle_kit_mcp.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mle-kit-mcp
-Version: 0.0.2
+Version: 0.0.4
 Summary: MCP server that provides different tools for MLE
 Author-email: Ilya Gusev <phoenixilya@gmail.com>
 Project-URL: Homepage, https://github.com/IlyaGusev/mle_kit_mcp

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mle-kit-mcp"
-version = "0.0.2"
+version = "0.0.4"
 description = "MCP server that provides different tools for MLE"
 readme = "README.md"
 authors = [

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/tests/test_bash.py

@@ -1,7 +1,7 @@
 import os
 
 from mle_kit_mcp.tools import bash
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 
 
 def test_bash() -> None:
@@ -12,7 +12,7 @@ def test_bash() -> None:
     assert result == "/workdir"
 
     result = bash("touch dummy")
-    assert os.path.exists(WORKSPACE_DIR_PATH / "dummy")
+    assert os.path.exists(get_workspace_dir() / "dummy")
 
     result = bash("fddafad")
-    assert "fddafad: command not found" in result
+    assert "fddafad: command not found" in result

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/tests/test_text_editor.py

@@ -4,7 +4,7 @@ import os
 import pytest
 
 from mle_kit_mcp.tools import text_editor
-from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+from mle_kit_mcp.files import get_workspace_dir
 
 DOCUMENT1 = """
 The dominant sequence transduction models are based on complex recurrent or convolutional
@@ -23,11 +23,11 @@ English constituency parsing both with large and limited training data.
 
 
 def test_text_editor_view() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         f.write(DOCUMENT1)
         f.flush()
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
 
         result = text_editor("view", name, show_lines=True)
         cmd_result = os.popen(f"cat -n {str(test_file.resolve())}").read()
@@ -52,9 +52,9 @@ def test_text_editor_view() -> None:
 
 
 def test_text_editor_write() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.unlink(missing_ok=True)
         assert not test_file.exists()
 
@@ -64,9 +64,9 @@ def test_text_editor_write() -> None:
 
 
 def test_text_editor_insert() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         text_editor("insert", name, insert_line=0, new_str="Hello there!")
@@ -77,9 +77,9 @@ def test_text_editor_insert() -> None:
 
 
 def test_text_editor_str_replace() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         text_editor("str_replace", name, old_str="41.8", new_str="41.9")
@@ -88,9 +88,9 @@ def test_text_editor_str_replace() -> None:
 
 
 def test_text_editor_undo_edit() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         text_editor("str_replace", name, old_str="41.8", new_str="41.9")
@@ -114,9 +114,9 @@ def test_text_editor_view_directory() -> None:
 
 
 def test_text_editor_view_invalid_range() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         with pytest.raises(AssertionError):
@@ -127,17 +127,17 @@ def test_text_editor_view_invalid_range() -> None:
 
 
 def test_text_editor_write_existing_file() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
         text_editor("write", name, file_text="New content")
 
 
 def test_text_editor_missing_required_params() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         with pytest.raises(AssertionError):
@@ -151,9 +151,9 @@ def test_text_editor_missing_required_params() -> None:
 
 
 def test_text_editor_undo_multiple() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         text_editor("str_replace", name, old_str="41.8", new_str="41.9")
@@ -168,9 +168,9 @@ def test_text_editor_undo_multiple() -> None:
 
 
 def test_text_editor_str_replace_no_match() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         with pytest.raises(AssertionError):
@@ -178,9 +178,9 @@ def test_text_editor_str_replace_no_match() -> None:
 
 
 def test_text_editor_str_replace_multiple_matches() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text("Hello\nHello\nHello")
 
         with pytest.raises(AssertionError):
@@ -188,9 +188,9 @@ def test_text_editor_str_replace_multiple_matches() -> None:
 
 
 def test_text_editor_large_file_handling() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
 
         large_content = ""
         for i in range(1000):
@@ -207,12 +207,12 @@ def test_text_editor_large_file_handling() -> None:
 
 
 def test_text_editor_append() -> None:
-    with tempfile.NamedTemporaryFile(dir=WORKSPACE_DIR_PATH, mode="w+") as f:
+    with tempfile.NamedTemporaryFile(dir=get_workspace_dir(), mode="w+") as f:
         name = os.path.basename(f.name)
-        test_file = WORKSPACE_DIR_PATH / name
+        test_file = get_workspace_dir() / name
         test_file.write_text(DOCUMENT1)
 
         text_editor("append", name, new_str="New line")
 
         assert DOCUMENT1.strip() in test_file.read_text().strip()
-        assert "New line" in test_file.read_text().strip()
+        assert "New line" in test_file.read_text().strip()

{mle_kit_mcp-0.0.2 → mle_kit_mcp-0.0.4}/tests/test_truncate_context.py

@@ -32,4 +32,4 @@ def test_truncate_content_suffix() -> None:
     result = truncate_content(DOCUMENT, max_length=40, suffix_only=True)
     parts = result.split("\n\n")
     assert DOCUMENT.endswith(parts[2])
-    assert len(parts[2]) == 40
+    assert len(parts[2]) == 40

mle_kit_mcp-0.0.2/mle_kit_mcp/files.py

@@ -1,12 +0,0 @@
-from pathlib import Path
-
-DIR_PATH = Path(__file__).parent
-ROOT_PATH = DIR_PATH.parent
-
-WORKSPACE_DIR_PATH: Path = ROOT_PATH / "workdir"
-
-
-def set_workspace_dir(path: Path) -> None:
-    global WORKSPACE_DIR_PATH
-    path.mkdir(parents=True, exist_ok=True)
-    WORKSPACE_DIR_PATH = path

mle_kit_mcp-0.0.2/mle_kit_mcp/server.py

@@ -1,36 +0,0 @@
-import fire  # type: ignore
-from pathlib import Path
-import uvicorn
-from mcp.server.fastmcp import FastMCP
-
-from .files import set_workspace_dir
-from .tools.bash import bash
-from .tools.text_editor import text_editor
-from .tools.remote_gpu import (
-    remote_bash,
-    create_remote_text_editor,
-    remote_download,
-)
-
-
-server = FastMCP("MLE kit MCP", stateless_http=True)
-
-remote_text_editor = create_remote_text_editor(text_editor)
-
-server.add_tool(bash)
-server.add_tool(text_editor)
-server.add_tool(remote_bash)
-server.add_tool(remote_text_editor)
-server.add_tool(remote_download)
-
-http_app = server.streamable_http_app()
-
-
-def run(host: str = "0.0.0.0", port: int = 5050, workspace: str = "workdir") -> None:
-    workspace_path = Path(workspace)
-    set_workspace_dir(workspace_path)
-    uvicorn.run(http_app, host=host, port=port)
-
-
-if __name__ == "__main__":
-    fire.Fire(run)