mle-kit-mcp 0.0.1__py3-none-any.whl

mle_kit_mcp/__init__.py

@@ -0,0 +1,7 @@
+import fire  # type: ignore
+
+from .server import run
+
+
+def main() -> None:
+    fire.Fire(run)

mle_kit_mcp/__main__.py

@@ -0,0 +1,5 @@
+from . import main
+
+
+if __name__ == "__main__":
+    main()

mle_kit_mcp/files.py

@@ -0,0 +1,12 @@
+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/server.py

@@ -0,0 +1,36 @@
+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)

mle_kit_mcp/tools/__init__.py

@@ -0,0 +1,14 @@
+from .bash import bash
+from .text_editor import text_editor
+from .remote_gpu import (
+    remote_bash,
+    remote_download,
+)
+
+
+__all__ = [
+    "bash",
+    "text_editor",
+    "remote_bash",
+    "remote_download",
+]

mle_kit_mcp/tools/bash.py

@@ -0,0 +1,81 @@
+import docker  # type: ignore
+import atexit
+import signal
+from typing import Optional, Any
+
+from mle_kit_mcp.files import WORKSPACE_DIR_PATH
+
+
+_container = None
+_client = None
+
+BASE_IMAGE = "python:3.9-slim"
+DOCKER_WORKSPACE_DIR_PATH = "/workdir"
+
+
+def cleanup_container(signum: Optional[Any] = None, frame: Optional[Any] = None) -> None:
+    global _container
+    if _container:
+        try:
+            _container.remove(force=True)
+            _container = None
+        except Exception:
+            pass
+    if signum == signal.SIGINT:
+        raise KeyboardInterrupt()
+
+
+atexit.register(cleanup_container)
+signal.signal(signal.SIGINT, cleanup_container)
+signal.signal(signal.SIGTERM, cleanup_container)
+
+
+def bash(command: str) -> str:
+    """
+    Run commands in a bash shell.
+    When invoking this tool, the contents of the "command" parameter does NOT need to be XML-escaped.
+    You don't have access to the internet via this tool.
+    You do have access to a mirror of common linux and python packages via apt and pip.
+    State is persistent across command calls and discussions with the user.
+    To inspect a particular line range of a file, e.g. lines 10-25, try 'sed -n 10,25p /path/to/the/file'.
+    Please avoid commands that may produce a very large amount of output.
+    Please run long lived commands in the background, e.g. 'sleep 10 &' or start a server in the background.
+
+    Args:
+        command: The bash command to run.
+    """
+
+    global _container, _client
+
+    if not _client:
+        _client = docker.from_env()
+
+    if not _container:
+        try:
+            _container = _client.containers.get("bash_runner")
+        except docker.errors.NotFound:
+            _container = _client.containers.run(
+                BASE_IMAGE,
+                "tail -f /dev/null",
+                detach=True,
+                remove=True,
+                name="bash_runner",
+                tty=True,
+                stdin_open=True,
+                volumes={
+                    WORKSPACE_DIR_PATH: {
+                        "bind": DOCKER_WORKSPACE_DIR_PATH,
+                        "mode": "rw",
+                    }
+                },
+                working_dir=DOCKER_WORKSPACE_DIR_PATH,
+            )
+
+    result = _container.exec_run(
+        ["bash", "-c", command],
+        workdir=DOCKER_WORKSPACE_DIR_PATH,
+        stdout=True,
+        stderr=True,
+    )
+    output: str = result.output.decode("utf-8").strip()
+    return output

mle_kit_mcp/tools/remote_gpu.py

@@ -0,0 +1,354 @@
+import os
+import time
+import subprocess
+import atexit
+import signal
+import inspect
+import functools
+from pathlib import Path
+from typing import List, Optional, Any, Callable
+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
+
+BASE_IMAGE = "phoenix120/holosophos_mle"
+DEFAULT_GPU_TYPE = "RTX_3090"
+GLOBAL_TIMEOUT = 43200
+VAST_AI_GREETING = """Welcome to vast.ai. If authentication fails, try again after a few seconds, and double check your ssh key.
+Have fun!"""
+
+
+@dataclass
+class InstanceInfo:
+    instance_id: int
+    ip: str = ""
+    port: int = 0
+    username: str = ""
+    ssh_key_path: str = ""
+    gpu_name: str = ""
+    start_time: int = 0
+
+
+_sdk: Optional[VastAI] = None
+_instance_info: Optional[InstanceInfo] = None
+
+
+def cleanup_machine(signum: Optional[Any] = None, frame: Optional[Any] = None) -> None:
+    print("Cleaning up...")
+    global _instance_info
+    signal.alarm(0)
+    if _instance_info and _sdk:
+        try:
+            _sdk.destroy_instance(id=_instance_info.instance_id)
+            _instance_info = None
+        except Exception:
+            pass
+    if signum == signal.SIGINT:
+        raise KeyboardInterrupt()
+
+
+atexit.register(cleanup_machine)
+signal.signal(signal.SIGINT, cleanup_machine)
+signal.signal(signal.SIGTERM, cleanup_machine)
+signal.signal(signal.SIGALRM, cleanup_machine)
+
+
+def wait_for_instance(vast_sdk: VastAI, instance_id: str, max_wait_time: int = 300) -> bool:
+    print(f"Waiting for instance {instance_id} to be ready...")
+    start_wait = int(time.time())
+    instance_ready = False
+    while time.time() - start_wait < max_wait_time:
+        instance_details = vast_sdk.show_instance(id=instance_id)
+        if (
+            isinstance(instance_details, dict)
+            and instance_details.get("actual_status") == "running"
+        ):
+            instance_ready = True
+            print(f"Instance {instance_id} is running and ready.")
+            break
+        print(f"Instance {instance_id} not ready yet. Waiting...")
+        time.sleep(15)
+    return instance_ready
+
+
+def get_offers(vast_sdk: VastAI, gpu_name: str) -> List[int]:
+    params = [
+        f"gpu_name={gpu_name}",
+        "cuda_vers>=12.1",
+        "num_gpus=1",
+        "reliability > 0.99",
+        "inet_up > 400",
+        "inet_down > 400",
+        "verified=True",
+        "rentable=True",
+        "disk_space > 100",
+    ]
+    query = "  ".join(params)
+    order = "score-"
+    offers = vast_sdk.search_offers(query=query, order=order)
+    assert offers
+    return [int(o["id"]) for o in offers]
+
+
+def run_command(
+    instance: InstanceInfo, command: str, timeout: int = 60
+) -> subprocess.CompletedProcess[str]:
+    cmd = [
+        "ssh",
+        "-i",
+        instance.ssh_key_path,
+        "-p",
+        str(instance.port),
+        "-o",
+        "StrictHostKeyChecking=no",
+        "-o",
+        "ConnectTimeout=10",
+        "-o",
+        "ServerAliveInterval=30",
+        "-o",
+        "ServerAliveCountMax=3",
+        "-o",
+        "TCPKeepAlive=yes",
+        f"{instance.username}@{instance.ip}",
+        command,
+    ]
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+        if result.returncode != 0:
+            raise Exception(
+                f"Error running command: {command}; "
+                f"Output: {result.stdout}; "
+                f"Error: {result.stderr}"
+            )
+    except subprocess.TimeoutExpired:
+        raise Exception(
+            f"Command timed out after {timeout} seconds: {command}; "
+            f"Host: {instance.username}@{instance.ip}:{instance.port}"
+        )
+    return result
+
+
+def recieve_rsync(
+    info: InstanceInfo, remote_path: str, local_path: str
+) -> subprocess.CompletedProcess[str]:
+    rsync_cmd = [
+        "rsync",
+        "-avz",
+        "-e",
+        f"ssh -i {info.ssh_key_path} -p {info.port} -o StrictHostKeyChecking=no",
+        f"{info.username}@{info.ip}:{remote_path}",
+        local_path,
+    ]
+
+    result = subprocess.run(rsync_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        error_output = (
+            f"Error syncing directory: {remote_path} to {local_path}. Error: {result.stderr}"
+        )
+        raise Exception(error_output)
+    return result
+
+
+def send_rsync(
+    info: InstanceInfo, local_path: str, remote_path: str
+) -> subprocess.CompletedProcess[str]:
+    rsync_cmd = [
+        "rsync",
+        "-avz",
+        "-e",
+        f"ssh -i {info.ssh_key_path} -p {info.port} -o StrictHostKeyChecking=no",
+        local_path,
+        f"{info.username}@{info.ip}:{remote_path}",
+    ]
+
+    result = subprocess.run(rsync_cmd, capture_output=True, text=True)
+    if result.returncode != 0:
+        error_output = (
+            f"Error syncing directory: {local_path} to {remote_path}. Error: {result.stderr}"
+        )
+        raise Exception(error_output)
+    return result
+
+
+def launch_instance(vast_sdk: VastAI, gpu_name: str) -> Optional[InstanceInfo]:
+    print(f"Selecting instance with {gpu_name}...")
+    offer_ids = get_offers(vast_sdk, gpu_name)
+
+    instance_id = None
+    for offer_id in offer_ids:
+        print(f"Launching offer {offer_id}...")
+        instance = vast_sdk.create_instance(id=offer_id, image=BASE_IMAGE, disk=50.0)
+        if not instance["success"]:
+            continue
+        instance_id = instance["new_contract"]
+        assert instance_id
+        global _instance_info
+        _instance_info = InstanceInfo(instance_id=instance_id)
+        print(f"Instance launched successfully. ID: {instance_id}")
+        is_ready = wait_for_instance(vast_sdk, instance_id)
+        if not is_ready:
+            print(f"Destroying instance {instance_id}...")
+            vast_sdk.destroy_instance(id=instance_id)
+            continue
+
+        print("Attaching SSH key...")
+        ssh_key_path = Path("~/.ssh/id_rsa").expanduser()
+        if not ssh_key_path.exists():
+            print(f"Generating SSH key at {ssh_key_path}...")
+            os.makedirs(ssh_key_path.parent, exist_ok=True)
+            subprocess.run(
+                [
+                    "ssh-keygen",
+                    "-t",
+                    "rsa",
+                    "-b",
+                    "4096",
+                    "-f",
+                    str(ssh_key_path),
+                    "-N",
+                    "",
+                ]
+            )
+
+        public_key = Path(f"{ssh_key_path}.pub").read_text().strip()
+        vast_sdk.attach_ssh(instance_id=instance_id, ssh_key=public_key)
+        instance_details = vast_sdk.show_instance(id=instance_id)
+
+        info = InstanceInfo(
+            instance_id=instance_details.get("id"),
+            ip=instance_details.get("ssh_host"),
+            port=instance_details.get("ssh_port"),
+            username="root",
+            ssh_key_path=str(ssh_key_path),
+            gpu_name=instance_details.get("gpu_name"),
+            start_time=int(time.time()),
+        )
+
+        print(info)
+        print(f"Checking SSH connection to {info.ip}:{info.port}...")
+        max_attempts = 10
+        is_okay = False
+        for attempt in range(max_attempts):
+            try:
+                result = run_command(info, "echo 'SSH connection successful'")
+            except Exception as e:
+                print(f"Waiting for SSH... {e}\n(Attempt {attempt+1}/{max_attempts})")
+                time.sleep(30)
+                continue
+            if "SSH connection successful" in result.stdout:
+                print("SSH connection established successfully!")
+                is_okay = True
+                break
+            print(f"Waiting for SSH... (Attempt {attempt+1}/{max_attempts})")
+            time.sleep(30)
+
+        if not is_okay:
+            print(f"Destroying instance {instance_id}...")
+            vast_sdk.destroy_instance(id=instance_id)
+            continue
+
+        break
+
+    return info
+
+
+def send_scripts() -> None:
+    assert _instance_info
+    for name in os.listdir(WORKSPACE_DIR_PATH):
+        if name.endswith(".py"):
+            send_rsync(_instance_info, f"{WORKSPACE_DIR_PATH}/{name}", "/root")
+
+
+def init_all() -> None:
+    global _sdk, _instance_info
+
+    load_dotenv()
+
+    if not _sdk:
+        _sdk = VastAI(api_key=os.getenv("VAST_AI_KEY"))
+    assert _sdk
+
+    signal.alarm(GLOBAL_TIMEOUT)
+    if not _instance_info:
+        _instance_info = launch_instance(_sdk, DEFAULT_GPU_TYPE)
+
+    if _instance_info:
+        send_scripts()
+
+    assert _instance_info, "Failed to connect to a remote instance! Try again"
+
+
+def remote_bash(command: str, timeout: Optional[int] = 60) -> str:
+    """
+    Run commands in a bash shell on a remote machine with GPU cards.
+    When invoking this tool, the contents of the "command" parameter does NOT need to be XML-escaped.
+    You don't have access to the internet via this tool.
+    You do have access to a mirror of common linux and python packages via apt and pip.
+    State is persistent across command calls and discussions with the user.
+    To inspect a particular line range of a file, e.g. lines 10-25, try 'sed -n 10,25p /path/to/the/file'.
+    Please avoid commands that may produce a very large amount of output.
+    Do not run commands in the background.
+    You can use python3.
+
+    Args:
+        command: The bash command to run.
+        timeout: Timeout for the command execution. 60 seconds by default. Set a higher value for heavy jobs.
+    """
+
+    init_all()
+    assert _instance_info
+    assert timeout
+    result = run_command(_instance_info, command, timeout=timeout)
+    output = ("STDOUT: " + result.stdout + "\n") if result.stdout else ""
+    output += ("STDERR: " + result.stderr) if result.stderr else ""
+    return output.replace(VAST_AI_GREETING, "")
+
+
+def remote_download(file_path: str) -> str:
+    """
+    Copies a file from a remote machine to the local work directory.
+    Use it to download final artefacts of the runs.
+    Args:
+        file_path: Path to the file on a remote machine.
+    """
+    init_all()
+    assert _instance_info
+    recieve_rsync(_instance_info, f"/root/{file_path}", f"{WORKSPACE_DIR_PATH}")
+    return f"File '{file_path}' downloaded!"
+
+
+def create_remote_text_editor(
+    text_editor_func: Callable[..., str],
+) -> Callable[..., str]:
+    @functools.wraps(text_editor_func)
+    def wrapper(*args: Any, **kwargs: Any) -> str:
+        init_all()
+        assert _instance_info
+
+        args_dict = {k: v for k, v in kwargs.items()}
+        if args:
+            args_dict.update(dict(zip(("command", "path"), args)))
+        path = args_dict["path"]
+        command = args_dict["command"]
+
+        if command != "write":
+            recieve_rsync(_instance_info, f"/root/{path}", f"{WORKSPACE_DIR_PATH}")
+
+        result: str = text_editor_func(*args, **kwargs)
+
+        if command != "view":
+            send_rsync(_instance_info, f"{WORKSPACE_DIR_PATH}/{path}", "/root")
+
+        return result
+
+    orig_sig = inspect.signature(text_editor_func)
+    wrapper.__signature__ = orig_sig  # type: ignore
+    if text_editor_func.__doc__:
+        orig_doc = text_editor_func.__doc__
+        new_doc = orig_doc.replace("text_editor", "remote_text_editor")
+        wrapper.__doc__ = "Executes on a remote machine with GPU.\n" + new_doc
+        wrapper.__name__ = "remote_text_editor"
+    return wrapper

mle_kit_mcp/tools/text_editor.py

@@ -0,0 +1,198 @@
+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.utils import truncate_content
+
+WRITE_MAX_OUTPUT_LENGTH = 500
+READ_MAX_OUTPUT_LENGTH = 3000
+
+# Global state for undo operations
+FILE_HISTORY: Dict[str, List[List[str]]] = defaultdict(list)
+
+
+def _save_file_state(path: Path, content: List[str]) -> None:
+    FILE_HISTORY[str(path.resolve())].append(content.copy())
+
+
+def _write(path: Path, file_text: str, overwrite: bool) -> str:
+    if not overwrite:
+        assert (
+            not path.exists()
+        ), f"Cannot write file, path already exists: {path}. Pass overwrite=True"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    lines = []
+    if path.exists():
+        lines = path.read_text().splitlines(True)
+    _save_file_state(path, lines)
+    path.write_text(file_text)
+    return f"Write was successful, the content of the '{path.name}' has changed!"
+
+
+def _append(path: Path, new_str: str) -> str:
+    assert path.exists(), "You can 'append' only to existing files"
+    content = path.open().read()
+    _save_file_state(path, content.splitlines(True))
+    new_content = "\n".join((content, new_str))
+    path.write_text(new_content)
+    return truncate_content(new_content, WRITE_MAX_OUTPUT_LENGTH, suffix_only=True)
+
+
+def _insert(path: Path, insert_line: int, new_str: str) -> str:
+    assert path.is_file(), f"File not found: {path}"
+    lines = path.open().readlines()
+    assert 0 <= insert_line <= len(lines), f"Invalid insert_line: {insert_line}"
+    _save_file_state(path, lines)
+    lines.insert(insert_line, new_str if new_str.endswith("\n") else new_str + "\n")
+    new_content = "".join(lines)
+    path.write_text(new_content)
+    return truncate_content(new_content, WRITE_MAX_OUTPUT_LENGTH, target_line=insert_line)
+
+
+def _str_replace(path: Path, old_str: str, new_str: str) -> str:
+    assert path.is_file(), f"File not found: {path}"
+    content = path.open().read()
+    count = content.count(old_str)
+    assert count != 0, "old_str not found in file"
+    assert count == 1, "old_str is not unique in file"
+    target_line = content[: content.find(old_str) + len(old_str)].count("\n")
+    _save_file_state(path, content.splitlines(True))
+    new_content = content.replace(old_str, new_str)
+    path.write_text(new_content)
+    return truncate_content(new_content, WRITE_MAX_OUTPUT_LENGTH, target_line=target_line)
+
+
+def _undo_edit(path: Path) -> str:
+    text_path = str(path.resolve())
+    assert text_path in FILE_HISTORY, f"No edit history available for: {text_path}"
+    assert FILE_HISTORY[text_path], f"No edit history available for: {text_path}"
+    previous_state = FILE_HISTORY[text_path].pop()
+    new_content = "".join(previous_state)
+    path.write_text(new_content)
+    return truncate_content(new_content, WRITE_MAX_OUTPUT_LENGTH)
+
+
+def _view(
+    path: Path,
+    view_start_line: Optional[int] = None,
+    view_end_line: Optional[int] = None,
+    show_lines: bool = False,
+) -> str:
+    assert path.exists(), f"Path does not exist: {path}"
+    if not path.is_file():
+        output = []
+        for level1 in path.iterdir():
+            if level1.name.startswith("."):
+                continue
+            output.append(str(level1.relative_to(path)))
+            if level1.is_dir():
+                for level2 in level1.iterdir():
+                    if level2.name.startswith("."):
+                        continue
+                    output.append(f"  {level2.relative_to(path)}")
+        return "\n".join(output)
+
+    lines = path.open().readlines()
+    enum_start_line = 1
+    if view_start_line or view_end_line:
+        if not view_start_line:
+            view_start_line = 1
+        if not view_end_line:
+            view_end_line = len(lines)
+        view_end_line = view_end_line if view_end_line <= len(lines) else len(lines)
+        view_end_line = view_end_line if view_end_line != -1 else len(lines)
+        assert view_start_line >= 1, "Line numbers must start at 1"
+        assert view_end_line >= 1, "Line numbers must start at 1"
+        assert (
+            view_start_line <= view_end_line
+        ), "Incorrect view parameters, start is higher than end"
+        lines = lines[view_start_line - 1 : view_end_line]
+        enum_start_line = view_start_line
+    output = []
+    for i, line in enumerate(lines, enum_start_line):
+        prefix = f"{i:6d}\t"
+        current_line = line
+        if show_lines:
+            current_line = prefix + current_line
+        output.append(current_line)
+    return truncate_content("".join(output), READ_MAX_OUTPUT_LENGTH)
+
+
+def text_editor(
+    command: str,
+    path: str,
+    new_str: Optional[str] = None,
+    file_text: Optional[str] = None,
+    old_str: Optional[str] = None,
+    insert_line: Optional[int] = None,
+    view_start_line: Optional[int] = None,
+    view_end_line: Optional[int] = None,
+    show_lines: Optional[bool] = False,
+) -> str:
+    """
+    Custom editing tool for viewing, creating and editing files.
+    State is persistent across command calls and discussions with the user.
+    If `path` is a file and `show_lines` is True, `view` displays the result of applying `cat -n`.
+    If `path` is a file and `show_lines` is False, `view` displays the result of applying `cat`.
+    If `path` is a directory, `view` lists non-hidden files and directories up to 2 levels deep.
+    The `write` command cannot be used if the specified `path` already exists as a file.
+    If a `command` generates a long output, it will be truncated and marked with `<response clipped>`.
+    The `undo_edit` command will revert the last edit made to the file at `path`.
+    Always write arguments with keys, do not rely on positions.
+    Be careful with escaping strings and line breaks in Python code.
+
+    Notes for using the `text_editor` command:
+    - The `old_str` parameter should match EXACTLY one or more consecutive lines from the original file.
+    - Be mindful of whitespaces!
+    - If the `old_str` parameter is not unique in the file, the replacement will not be performed.
+    - Make sure to include enough context in `old_str` to make it unique
+    - The `new_str` parameter should contain the edited lines that should replace the `old_str`
+
+    Examples:
+        Write a file with "Hello world!": text_editor("write", "file.txt", file_text="Hello world!")
+        View a file with enumerated lines: text_editor("view", "file.txt")
+        View first three lines of a file: text_editor("view", "file.txt", view_start_line=1, view_end_line=3)
+        View lines from 5 to the end of a file: text_editor("view", "file.txt", view_start_line=5)
+        Replace "line1" with "line2": text_editor("str_replace", "file.txt", old_str="line", new_str="line2")
+        Insert "line2" after line 1: text_editor("insert", "file.txt", insert_line=1, new_str="line2")
+        Append "line2" to the file: text_editor("append", "file.txt", new_str="line2")
+
+    Args:
+        command: The commands to run. Allowed options are: `view`, `write`, `str_replace`, `insert`, `undo_edit`.
+        path: Path to file or directory inside current work directory. Should not be absolute.
+        view_start_line: Optional for view command, use to view specific lines
+        view_end_line: Optional for view command, use to view specific lines
+        file_text: Required for `write` command, with the content of the file to be writed.
+        insert_line: Required for `insert` command. `new_str` will be inserted AFTER the line `insert_line` of `path`.
+        new_str: Required for `str_replace`, `insert` and `append`.
+        old_str: Required for `str_replace` containing the string in `path` to replace.
+        show_lines: Optional for view command. If True, the command will also output line numbers.
+    """
+    assert not path.startswith(
+        "/"
+    ), "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
+
+    if command == "view":
+        show_lines = show_lines if show_lines is not None else False
+        return _view(path_obj, view_start_line, view_end_line, show_lines)
+    if command == "write":
+        assert file_text is not None, "'file_text' is required for 'write' command"
+        return _write(path_obj, file_text, overwrite=True)
+    if command == "append":
+        assert new_str is not None, "'new_str' is required for 'append' command"
+        return _append(path_obj, new_str)
+    if command == "insert":
+        assert insert_line is not None, "'insert_line' is required for 'insert' command"
+        assert new_str is not None, "'new_str' is required for 'insert' command"
+        return _insert(path_obj, insert_line, new_str)
+    if command == "str_replace":
+        assert old_str is not None, "'old_str' is required for 'str_replace' command"
+        assert new_str is not None, "'new_str' is required for 'str_replace' command"
+        return _str_replace(path_obj, old_str, new_str)
+    if command == "undo_edit":
+        return _undo_edit(path_obj)
+    assert False, f"Not a valid command! List of commands: {valid_commands}"

mle_kit_mcp/utils.py

@@ -0,0 +1,50 @@
+from typing import Optional
+
+
+def truncate_content(
+    content: str,
+    max_length: int,
+    prefix_only: bool = False,
+    suffix_only: bool = False,
+    target_line: Optional[int] = None,
+) -> str:
+    assert int(prefix_only) + int(suffix_only) + int(target_line is not None) <= 1
+    disclaimer = (
+        f"\n\n..._This content has been truncated to stay below {max_length} characters_...\n\n"
+    )
+    half_length = max_length // 2
+    if len(content) <= max_length:
+        return content
+
+    if prefix_only:
+        prefix = content[:max_length]
+        return prefix + disclaimer
+
+    elif suffix_only:
+        suffix = content[-max_length:]
+        return disclaimer + suffix
+
+    elif target_line:
+        line_start_pos = 0
+        next_pos = content.find("\n") + 1
+        line_end_pos = next_pos
+        for _ in range(target_line):
+            next_pos = content.find("\n", next_pos) + 1
+            line_start_pos = line_end_pos
+            line_end_pos = next_pos
+        assert line_end_pos >= line_start_pos
+        length = line_end_pos - line_start_pos
+        remaining_length = max(0, max_length - length)
+        half_length = remaining_length // 2
+        start = max(0, line_start_pos - half_length)
+        end = min(len(content), line_end_pos + half_length)
+        final_content = content[start:end]
+        if start == 0:
+            return final_content + disclaimer
+        elif end == len(content):
+            return disclaimer + final_content
+        return disclaimer + content[start:end] + disclaimer
+
+    prefix = content[:half_length]
+    suffix = content[-half_length:]
+    return prefix + disclaimer + suffix

mle_kit_mcp-0.0.1.dist-info/METADATA

@@ -0,0 +1,43 @@
+Metadata-Version: 2.4
+Name: mle-kit-mcp
+Version: 0.0.1
+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
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mcp>=1.9.2
+Requires-Dist: black==25.1.0
+Requires-Dist: mypy==1.16.0
+Requires-Dist: flake8==7.2.0
+Requires-Dist: fire>=0.7.0
+Requires-Dist: docker>=7.1.0
+Requires-Dist: vastai-sdk==0.1.16
+Requires-Dist: pytest>=8.4.1
+Dynamic: license-file
+
+# MLE KIT MCP
+
+A collection of MCP tools related to the experiments on remote gpu:
+- Bash and remote bash
+- Text editor and remote text editor 
+- Remote download
+
+
+Run the mcp server:
+```
+uv run python -m mle_kit_mcp --host 127.0.0.1 --port 5056 --workspace workdir
+```
+
+Claude Desktop config:
+```
+{
+  "mcpServers": {
+      "mle_kit": {"url": "http://127.0.0.1:5056/mcp", "transport": "streamable-http"}
+  }
+}
+```

mle_kit_mcp-0.0.1.dist-info/RECORD

@@ -0,0 +1,16 @@
+mle_kit_mcp/__init__.py,sha256=2Ru2I5u4cE7DrkkAsibDUEF1K6sYtqppb9VyFrRoQKI,94
+mle_kit_mcp/__main__.py,sha256=rcmsOtJd3SA82exjrcGBuxuptcoxF8AXI7jNjiVq2BY,59
+mle_kit_mcp/files.py,sha256=loZnpQ7ZLxjX8NgUq-hrvhwbQwmGw1F6gZ3xD3UQI9k,286
+mle_kit_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mle_kit_mcp/server.py,sha256=l81KAfyZ_G_0Suy9381qV40I_UAZtfT0FKim3_89pSI,886
+mle_kit_mcp/utils.py,sha256=wd7wSyddHRHOYdxmXw8uoAOBxVZOL2_vjNomss07inc,1654
+mle_kit_mcp/tools/__init__.py,sha256=0aLl0gD-JteSvOs2PgVhbv0Wnh6fodFySgQWQvoI1xI,215
+mle_kit_mcp/tools/bash.py,sha256=QcPKQBjD6-LHvNGOAhyWXHuzkM70GATN0Vp85iykSXY,2460
+mle_kit_mcp/tools/remote_gpu.py,sha256=Wh02B-hFUFdRRtZ0J4feLw9P8Hq0AW7xwZ3Q27LcdW4,11370
+mle_kit_mcp/tools/text_editor.py,sha256=zwxyHvf8DdJwuUa7HOWfSbhaSTveaE_WMwsOdDnbq1w,9228
+mle_kit_mcp-0.0.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mle_kit_mcp-0.0.1.dist-info/METADATA,sha256=SIyqk-Ay6o2v7SCaSZuX0UbVPU7pB80PP7JQ5xH1rT0,1126
+mle_kit_mcp-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mle_kit_mcp-0.0.1.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
+mle_kit_mcp-0.0.1.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
+mle_kit_mcp-0.0.1.dist-info/RECORD,,

mle_kit_mcp-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

mle_kit_mcp-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+mle_kit_mcp = mle_kit_mcp:main

mle_kit_mcp-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

mle_kit_mcp-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+mle_kit_mcp