mle-kit-mcp 0.0.8__py3-none-any.whl → 0.1.0__py3-none-any.whl

mle_kit_mcp/llm_proxy_source.py

@@ -0,0 +1,51 @@
+import os
+
+import fire  # type: ignore
+import uvicorn
+from openai import AsyncOpenAI
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse
+
+ACCESS_TOKEN = os.getenv("ACCESS_TOKEN", "")
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
+OPENROUTER_BASE_URL = os.getenv("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1")
+
+app = FastAPI()
+client = AsyncOpenAI(base_url=OPENROUTER_BASE_URL, api_key=OPENROUTER_API_KEY)
+
+
+def _check_auth(request: Request) -> None:
+    auth = request.headers.get("authorization", "")
+    if not ACCESS_TOKEN or auth != f"Bearer {ACCESS_TOKEN}":
+        raise HTTPException(status_code=401, detail="Unauthorized")
+    if not OPENROUTER_API_KEY:
+        raise HTTPException(status_code=500, detail="OpenRouter key not configured")
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request) -> JSONResponse:
+    _check_auth(request)
+    payload = await request.json()
+    if isinstance(payload, dict) and payload.get("stream"):
+        payload.pop("stream", None)
+
+    try:
+        resp = await client.chat.completions.create(**payload)
+    except Exception as e:
+        raise HTTPException(status_code=502, detail=str(e))
+    return JSONResponse(resp.model_dump())
+
+
+@app.get("/health")
+async def health() -> JSONResponse:
+    if not OPENROUTER_API_KEY:
+        return JSONResponse({"error": "missing OPENROUTER_API_KEY"}, 500)
+    return JSONResponse({"result": "ok"})
+
+
+def main(host: str = "127.0.0.1", port: int = 8001) -> None:
+    uvicorn.run("llm_proxy:app", host=host, port=port)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

mle_kit_mcp/server.py

@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 from typing import Optional
 
@@ -12,6 +13,10 @@ from .tools.remote_gpu import (
     create_remote_text_editor,
     remote_download,
 )
+from .tools.llm_proxy import (
+    llm_proxy_local,
+    llm_proxy_remote,
+)
 from .files import get_workspace_dir, WorkspaceDirectory
 
 
@@ -30,6 +35,9 @@ def run(host: str = "0.0.0.0", port: int = 5050, workspace: Optional[str] = None
     server.add_tool(remote_bash)
     server.add_tool(remote_text_editor)
     server.add_tool(remote_download)
+    if os.getenv("OPENROUTER_API_KEY"):
+        server.add_tool(llm_proxy_local)
+        server.add_tool(llm_proxy_remote)
 
     http_app = server.streamable_http_app()
 

mle_kit_mcp/tools/__init__.py

@@ -4,6 +4,10 @@ from .remote_gpu import (
     remote_bash,
     remote_download,
 )
+from .llm_proxy import (
+    llm_proxy_local,
+    llm_proxy_remote,
+)
 
 
 __all__ = [
@@ -11,4 +15,6 @@ __all__ = [
     "text_editor",
     "remote_bash",
     "remote_download",
+    "llm_proxy_local",
+    "llm_proxy_remote",
 ]

mle_kit_mcp/tools/bash.py

@@ -1,8 +1,11 @@
-import docker  # type: ignore
 import atexit
 import signal
 from typing import Optional, Any
 
+from docker import from_env as docker_from_env  # type: ignore
+from docker import DockerClient
+from docker.models.containers import Container  # type: ignore
+
 from mle_kit_mcp.files import get_workspace_dir
 
 
@@ -13,6 +16,40 @@ BASE_IMAGE = "python:3.12-slim"
 DOCKER_WORKSPACE_DIR_PATH = "/workdir"
 
 
+def get_docker_client() -> DockerClient:
+    global _client
+    if not _client:
+        _client = docker_from_env()
+    return _client
+
+
+def create_container() -> Container:
+    client = get_docker_client()
+    container = client.containers.run(
+        BASE_IMAGE,
+        "tail -f /dev/null",
+        detach=True,
+        remove=True,
+        tty=True,
+        stdin_open=True,
+        volumes={
+            get_workspace_dir(): {
+                "bind": DOCKER_WORKSPACE_DIR_PATH,
+                "mode": "rw",
+            }
+        },
+        working_dir=DOCKER_WORKSPACE_DIR_PATH,
+    )
+    return container
+
+
+def get_container() -> Container:
+    global _container
+    if not _container:
+        _container = create_container()
+    return _container
+
+
 def cleanup_container(signum: Optional[Any] = None, frame: Optional[Any] = None) -> None:
     global _container
     if _container:
@@ -45,33 +82,8 @@ def bash(command: str) -> str:
         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={
-                    get_workspace_dir(): {
-                        "bind": DOCKER_WORKSPACE_DIR_PATH,
-                        "mode": "rw",
-                    }
-                },
-                working_dir=DOCKER_WORKSPACE_DIR_PATH,
-            )
-
-    result = _container.exec_run(
+    container = get_container()
+    result = container.exec_run(
         ["bash", "-c", command],
         workdir=DOCKER_WORKSPACE_DIR_PATH,
         stdout=True,

mle_kit_mcp/tools/llm_proxy.py

@@ -0,0 +1,148 @@
+import os
+import json
+import time
+import random
+import secrets
+from pathlib import Path
+from typing import Optional
+
+from dotenv import load_dotenv
+
+from mle_kit_mcp.tools.bash import get_container
+from mle_kit_mcp.files import get_workspace_dir
+from mle_kit_mcp.utils import find_free_port
+
+from mle_kit_mcp.tools.remote_gpu import (
+    get_instance as _remote_get_instance,
+    run_command as _remote_run_command,
+    send_rsync as _remote_send_rsync,
+)
+
+INPUT_SCRIPT_FILE_NAME = "llm_proxy_source.py"
+OUTPUT_SCRIPT_FILE_NAME = "llm_proxy.py"
+DEPENDENCIES = "fastapi uvicorn httpx openai fire"
+START_TIMEOUT = 30
+
+
+def _write_proxy_script(script_path: Path) -> None:
+    source_script_path = Path(__file__).parent.parent / INPUT_SCRIPT_FILE_NAME
+    script = source_script_path.read_text()
+    script_path.write_text(script)
+
+
+def llm_proxy_local(port: Optional[int] = None) -> str:
+    """
+    Start a lightweight OpenRouter proxy inside the same Docker container used by the "bash" tool.
+
+    Returns a JSON string with fields: url, token, scope.
+    The url is reachable from inside the "bash" container as localhost.
+    Use the token in the Authorization header: "Bearer <token>" when calling the proxy.
+    It runs a standard OpenAI compatible server, so you can use it with any OpenAI compatible client.
+    You can use all models available on OpenRouter, for instance:
+    - openai/gpt-5-mini
+    - google/gemini-2.5-pro
+    - anthropic/claude-sonnet-4
+
+    Args:
+        port: Optional fixed port to bind inside the container. Random if omitted.
+    """
+
+    load_dotenv()
+    api_key = os.getenv("OPENROUTER_API_KEY")
+    assert api_key, "Set OPENROUTER_API_KEY in the environment before starting the proxy."
+
+    _write_proxy_script(get_workspace_dir() / OUTPUT_SCRIPT_FILE_NAME)
+
+    container = get_container()
+    dependencies_cmd = f"python -m pip install --quiet --no-input {DEPENDENCIES}"
+    container.exec_run(["bash", "-lc", dependencies_cmd])
+
+    chosen_port = port or find_free_port()
+    token = secrets.token_urlsafe(24)
+    launch_cmd = (
+        f"OPENROUTER_API_KEY='{api_key}' ACCESS_TOKEN='{token}' "
+        f"nohup python {OUTPUT_SCRIPT_FILE_NAME} "
+        f"--host 127.0.0.1 --port {chosen_port} "
+        f"> llm_proxy.log 2>&1 "
+        f"& echo $! > llm_proxy.pid"
+    )
+    container.exec_run(["bash", "-lc", launch_cmd])
+
+    health_cmd = f'import httpx; print(httpx.get("http://127.0.0.1:{chosen_port}/health").json())'
+    start_time = time.time()
+    while time.time() - start_time < START_TIMEOUT:
+        result = container.exec_run(["python", "-c", health_cmd])
+        if result.exit_code == 0 and "ok" in result.output.decode("utf-8").strip():
+            break
+        time.sleep(1)
+    else:
+        raise Exception("Failed to start the proxy")
+
+    return json.dumps(
+        {
+            "url": f"http://127.0.0.1:{chosen_port}/v1/chat/completions",
+            "token": token,
+            "scope": "bash-container",
+        }
+    )
+
+
+def llm_proxy_remote(port: Optional[int] = None) -> str:
+    """
+    Start a lightweight OpenRouter proxy on the remote GPU machine.
+
+    Returns a JSON string with fields: url, token, scope.
+    The url is reachable from inside the remote machine as localhost.
+    Use the token in the Authorization header: "Bearer <token>" when calling the proxy.
+    It runs a standard OpenAI compatible server, so you can use it with any OpenAI compatible client.
+    You can use all models available on OpenRouter, for instance:
+    - openai/gpt-5-mini
+    - google/gemini-2.5-pro
+    - anthropic/claude-sonnet-4
+
+    Args:
+        port: Optional fixed port to bind on the remote. Random if omitted.
+    """
+
+    load_dotenv()
+    api_key = os.getenv("OPENROUTER_API_KEY")
+    assert api_key, "Set OPENROUTER_API_KEY in the environment before starting the proxy."
+
+    instance = _remote_get_instance()
+    script_path = get_workspace_dir() / OUTPUT_SCRIPT_FILE_NAME
+    _write_proxy_script(script_path)
+    _remote_send_rsync(instance, f"{script_path}", "/root")
+
+    chosen_port = port or random.randint(5000, 6000)
+    token = secrets.token_urlsafe(24)
+    dependencies_cmd = f"python3 -m pip install -q --no-input {DEPENDENCIES}"
+    _remote_run_command(instance, dependencies_cmd, timeout=300, raise_exc=True)
+
+    launch_cmd = (
+        f"OPENROUTER_API_KEY='{api_key}' ACCESS_TOKEN='{token}' "
+        f"nohup python {OUTPUT_SCRIPT_FILE_NAME} "
+        f"--host 127.0.0.1 --port {chosen_port} "
+        f"> openrouter_proxy.log 2>&1 "
+        f"& echo $! > openrouter_proxy.pid"
+    )
+    _remote_run_command(instance, launch_cmd, timeout=60, raise_exc=True)
+
+    health_cmd = f'import httpx; print(httpx.get("http://127.0.0.1:{chosen_port}/health").json())'
+    start_time = time.time()
+    while time.time() - start_time < START_TIMEOUT:
+        result = _remote_run_command(
+            instance, f"python -c '{health_cmd}'", timeout=10, raise_exc=False
+        )
+        if result.returncode == 0 and "ok" in result.stdout.strip():
+            break
+        time.sleep(1)
+    else:
+        raise Exception("Failed to start the proxy")
+
+    return json.dumps(
+        {
+            "url": f"http://127.0.0.1:{chosen_port}/v1/chat/completions",
+            "token": token,
+            "scope": "remote-gpu",
+        }
+    )

mle_kit_mcp/tools/remote_gpu.py

@@ -36,7 +36,28 @@ _sdk: Optional[VastAI] = None
 _instance_info: Optional[InstanceInfo] = None
 
 
-def cleanup_machine(signum: Optional[Any] = None, frame: Optional[Any] = None) -> None:
+def get_sdk() -> VastAI:
+    global _sdk
+    if not _sdk:
+        _sdk = VastAI(api_key=os.getenv("VAST_AI_KEY"))
+    return _sdk
+
+
+def get_instance() -> InstanceInfo:
+    load_dotenv()
+    signal.alarm(GLOBAL_TIMEOUT)
+    global _instance_info
+    if not _instance_info:
+        _instance_info = launch_instance(get_sdk(), DEFAULT_GPU_TYPE)
+
+    if _instance_info:
+        send_scripts()
+
+    assert _instance_info, "Failed to connect to a remote instance! Try again"
+    return _instance_info
+
+
+def cleanup_instance(signum: Optional[Any] = None, frame: Optional[Any] = None) -> None:
     global _instance_info
     signal.alarm(0)
     if _instance_info and _sdk:
@@ -50,10 +71,10 @@ def cleanup_machine(signum: Optional[Any] = None, frame: Optional[Any] = None) -
         raise KeyboardInterrupt()
 
 
-atexit.register(cleanup_machine)
-signal.signal(signal.SIGINT, cleanup_machine)
-signal.signal(signal.SIGTERM, cleanup_machine)
-signal.signal(signal.SIGALRM, cleanup_machine)
+atexit.register(cleanup_instance)
+signal.signal(signal.SIGINT, cleanup_instance)
+signal.signal(signal.SIGTERM, cleanup_instance)
+signal.signal(signal.SIGALRM, cleanup_instance)
 
 
 def wait_for_instance(vast_sdk: VastAI, instance_id: str, max_wait_time: int = 300) -> bool:
@@ -182,7 +203,7 @@ def launch_instance(vast_sdk: VastAI, gpu_name: str) -> Optional[InstanceInfo]:
 
     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)
+        instance = vast_sdk.create_instance(id=offer_id, image=BASE_IMAGE, disk=300.0)
         if not instance["success"]:
             continue
         instance_id = instance["new_contract"]
@@ -264,25 +285,6 @@ def send_scripts() -> None:
             send_rsync(_instance_info, f"{get_workspace_dir()}/{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: int = 60) -> str:
     """
     Run commands in a bash shell on a remote machine with GPU cards.
@@ -300,10 +302,10 @@ def remote_bash(command: str, timeout: int = 60) -> str:
         timeout: Timeout for the command execution. 60 seconds by default. Set a higher value for heavy jobs.
     """
 
-    init_all()
-    assert _instance_info
+    instance = get_instance()
+    assert instance
     assert timeout
-    result = run_command(_instance_info, command, timeout=timeout, raise_exc=False)
+    result = run_command(instance, 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,9 +318,9 @@ def remote_download(file_path: str) -> str:
     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"{get_workspace_dir()}")
+    instance = get_instance()
+    assert instance
+    recieve_rsync(instance, f"/root/{file_path}", f"{get_workspace_dir()}")
     return f"File '{file_path}' downloaded!"
 
 
@@ -327,8 +329,7 @@ def create_remote_text_editor(
 ) -> Callable[..., str]:
     @functools.wraps(text_editor_func)
     def wrapper(*args: Any, **kwargs: Any) -> str:
-        init_all()
-        assert _instance_info
+        instance = get_instance()
 
         args_dict = {k: v for k, v in kwargs.items()}
         if args:
@@ -337,12 +338,12 @@ def create_remote_text_editor(
         command = args_dict["command"]
 
         if command != "write":
-            recieve_rsync(_instance_info, f"/root/{path}", f"{get_workspace_dir()}")
+            recieve_rsync(instance, f"/root/{path}", f"{get_workspace_dir()}")
 
         result: str = text_editor_func(*args, **kwargs)
 
         if command != "view":
-            send_rsync(_instance_info, f"{get_workspace_dir()}/{path}", "/root")
+            send_rsync(instance, f"{get_workspace_dir()}/{path}", "/root")
 
         return result
 

mle_kit_mcp/utils.py

@@ -1,3 +1,5 @@
+import random
+import socket
 from typing import Optional
 
 
@@ -48,3 +50,16 @@ def truncate_content(
     prefix = content[:half_length]
     suffix = content[-half_length:]
     return prefix + disclaimer + suffix
+
+
+def find_free_port() -> Optional[int]:
+    ports = list(range(5000, 6001))
+    random.shuffle(ports)
+    for port in ports:
+        try:
+            with socket.socket() as s:
+                s.bind(("", port))
+                return port
+        except OSError:
+            continue
+    return None

{mle_kit_mcp-0.0.8.dist-info → mle_kit_mcp-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mle-kit-mcp
-Version: 0.0.8
+Version: 0.1.0
 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.1.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+mle_kit_mcp/__init__.py,sha256=2Ru2I5u4cE7DrkkAsibDUEF1K6sYtqppb9VyFrRoQKI,94
+mle_kit_mcp/__main__.py,sha256=rcmsOtJd3SA82exjrcGBuxuptcoxF8AXI7jNjiVq2BY,59
+mle_kit_mcp/files.py,sha256=ux53kWw7hBAcOmS9qNI4gpQX8XcQPT2LICC--S5-TGI,635
+mle_kit_mcp/llm_proxy_source.py,sha256=zDfNDvG1nAkft-irj0KhdB-1z-UUtKweRJYG0fCveVE,1637
+mle_kit_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mle_kit_mcp/server.py,sha256=W4YJ3m1-NKheJ5QlGggfMQyXPghzKgg5zXqb-TLYH1U,1271
+mle_kit_mcp/utils.py,sha256=iHNcEZZzPD37bEYE18SzJ3WUjLP3Ym-kc91SwcW1vlI,1984
+mle_kit_mcp/tools/__init__.py,sha256=r2fIg2mZ6zaeq0CzEKCEdeUTjV0pcA9NZaaOfBNVTnE,332
+mle_kit_mcp/tools/bash.py,sha256=9YmwRCR2HHTJ6Jhnn-fHuLRkDjsHo9m26tpdXoGCFtQ,2608
+mle_kit_mcp/tools/llm_proxy.py,sha256=9tm3k1R5y9o6Z4nkZXwILkYt6Ludc_TnQSwRw0KcNHA,5312
+mle_kit_mcp/tools/remote_gpu.py,sha256=Fv8SJC8bE3Oo8JHOKXoYpISXFMzUMFjEcCQOTyW7V1I,11544
+mle_kit_mcp/tools/text_editor.py,sha256=EJ832vgdFdycrLOf3eOYMXnlMgB2_iH_u4ANvZVrMsE,9139
+mle_kit_mcp-0.1.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mle_kit_mcp-0.1.0.dist-info/METADATA,sha256=aNQJCc6kZW4IaOvKi7KrxPJ0NqQgnFCmQqVk1uJ3yTg,1011
+mle_kit_mcp-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+mle_kit_mcp-0.1.0.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
+mle_kit_mcp-0.1.0.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
+mle_kit_mcp-0.1.0.dist-info/RECORD,,

mle_kit_mcp-0.0.8.dist-info/RECORD

@@ -1,16 +0,0 @@
-mle_kit_mcp/__init__.py,sha256=2Ru2I5u4cE7DrkkAsibDUEF1K6sYtqppb9VyFrRoQKI,94
-mle_kit_mcp/__main__.py,sha256=rcmsOtJd3SA82exjrcGBuxuptcoxF8AXI7jNjiVq2BY,59
-mle_kit_mcp/files.py,sha256=ux53kWw7hBAcOmS9qNI4gpQX8XcQPT2LICC--S5-TGI,635
-mle_kit_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mle_kit_mcp/server.py,sha256=E9cXdKwVBASBzKyrZGHCTvj6BIMN-EbVSQZDFNg0YnE,1062
-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=10hYo_hIBvQh8wbjssixb89SXG05H6nHtndShr6V31E,2461
-mle_kit_mcp/tools/remote_gpu.py,sha256=NHdg52BuTL8fBFXHgv0s3CDpsWs4I1WykfTdik3c6F4,11482
-mle_kit_mcp/tools/text_editor.py,sha256=EJ832vgdFdycrLOf3eOYMXnlMgB2_iH_u4ANvZVrMsE,9139
-mle_kit_mcp-0.0.8.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mle_kit_mcp-0.0.8.dist-info/METADATA,sha256=uynrL08zcDIQCyEiVhmRA6KJ-Umk1c2vHUrT2qJnCRw,1011
-mle_kit_mcp-0.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-mle_kit_mcp-0.0.8.dist-info/entry_points.txt,sha256=-iHSUVPN49jkBj1ySpc-P0rVF5-IPHw-KWNayNIiEsk,49
-mle_kit_mcp-0.0.8.dist-info/top_level.txt,sha256=XeBtCq_CnVI0gh0Z_daZOLmGl5XPlkA8RgHaj5s5VQY,12
-mle_kit_mcp-0.0.8.dist-info/RECORD,,