95 lines
2.6 KiB
Python
95 lines
2.6 KiB
Python
"""Thin LangChain wrapper over the public vm_run contract.
|
|
|
|
Requirements:
|
|
- `pip install langchain-core` or `uv add langchain-core`
|
|
|
|
This example keeps the framework integration intentionally narrow: a single tool
|
|
that delegates straight to `Pyro.run_in_vm(...)`.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any, Callable, TypeVar, cast
|
|
|
|
from pyro_mcp import Pyro
|
|
from pyro_mcp.vm_manager import (
|
|
DEFAULT_ALLOW_HOST_COMPAT,
|
|
DEFAULT_MEM_MIB,
|
|
DEFAULT_TIMEOUT_SECONDS,
|
|
DEFAULT_TTL_SECONDS,
|
|
DEFAULT_VCPU_COUNT,
|
|
)
|
|
|
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
|
|
|
|
def run_vm_run_tool(
|
|
*,
|
|
environment: str,
|
|
command: str,
|
|
vcpu_count: int = DEFAULT_VCPU_COUNT,
|
|
mem_mib: int = DEFAULT_MEM_MIB,
|
|
timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
|
|
ttl_seconds: int = DEFAULT_TTL_SECONDS,
|
|
network: bool = False,
|
|
allow_host_compat: bool = DEFAULT_ALLOW_HOST_COMPAT,
|
|
) -> str:
|
|
pyro = Pyro()
|
|
result = pyro.run_in_vm(
|
|
environment=environment,
|
|
command=command,
|
|
vcpu_count=vcpu_count,
|
|
mem_mib=mem_mib,
|
|
timeout_seconds=timeout_seconds,
|
|
ttl_seconds=ttl_seconds,
|
|
network=network,
|
|
allow_host_compat=allow_host_compat,
|
|
)
|
|
return json.dumps(result, sort_keys=True)
|
|
|
|
|
|
def build_langchain_vm_run_tool() -> Any:
|
|
try:
|
|
from langchain_core.tools import tool # type: ignore[import-not-found]
|
|
except ImportError as exc: # pragma: no cover - dependency guard
|
|
raise RuntimeError(
|
|
"langchain-core is required for this example. Install it with "
|
|
"`uv add langchain-core` or `pip install langchain-core`."
|
|
) from exc
|
|
|
|
decorator = cast(Callable[[F], F], tool("vm_run"))
|
|
|
|
@decorator
|
|
def vm_run(
|
|
environment: str,
|
|
command: str,
|
|
vcpu_count: int = DEFAULT_VCPU_COUNT,
|
|
mem_mib: int = DEFAULT_MEM_MIB,
|
|
timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
|
|
ttl_seconds: int = DEFAULT_TTL_SECONDS,
|
|
network: bool = False,
|
|
allow_host_compat: bool = DEFAULT_ALLOW_HOST_COMPAT,
|
|
) -> str:
|
|
"""Run one command in an ephemeral Firecracker VM and clean it up."""
|
|
return run_vm_run_tool(
|
|
environment=environment,
|
|
command=command,
|
|
vcpu_count=vcpu_count,
|
|
mem_mib=mem_mib,
|
|
timeout_seconds=timeout_seconds,
|
|
ttl_seconds=ttl_seconds,
|
|
network=network,
|
|
allow_host_compat=allow_host_compat,
|
|
)
|
|
|
|
return vm_run
|
|
|
|
|
|
def main() -> None:
|
|
tool = build_langchain_vm_run_tool()
|
|
print(tool)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|