62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
"""Provider-agnostic agent tool example centered on vm_run."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from pyro_mcp import Pyro
|
|
|
|
VM_RUN_TOOL: dict[str, Any] = {
|
|
"name": "vm_run",
|
|
"description": "Run one command in an ephemeral Firecracker VM and clean it up.",
|
|
"input_schema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"profile": {"type": "string"},
|
|
"command": {"type": "string"},
|
|
"vcpu_count": {"type": "integer"},
|
|
"mem_mib": {"type": "integer"},
|
|
"timeout_seconds": {"type": "integer", "default": 30},
|
|
"ttl_seconds": {"type": "integer", "default": 600},
|
|
"network": {"type": "boolean", "default": False},
|
|
},
|
|
"required": ["profile", "command", "vcpu_count", "mem_mib"],
|
|
},
|
|
}
|
|
|
|
|
|
def call_vm_run(arguments: dict[str, Any]) -> dict[str, Any]:
|
|
pyro = Pyro()
|
|
return pyro.run_in_vm(
|
|
profile=str(arguments["profile"]),
|
|
command=str(arguments["command"]),
|
|
vcpu_count=int(arguments["vcpu_count"]),
|
|
mem_mib=int(arguments["mem_mib"]),
|
|
timeout_seconds=int(arguments.get("timeout_seconds", 30)),
|
|
ttl_seconds=int(arguments.get("ttl_seconds", 600)),
|
|
network=bool(arguments.get("network", False)),
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
tool_arguments: dict[str, Any] = {
|
|
"profile": "debian-git",
|
|
"command": "git --version",
|
|
"vcpu_count": 1,
|
|
"mem_mib": 1024,
|
|
"timeout_seconds": 30,
|
|
"network": False,
|
|
}
|
|
tool_call: dict[str, Any] = {
|
|
"name": "vm_run",
|
|
"arguments": tool_arguments,
|
|
}
|
|
|
|
print(json.dumps({"tool": VM_RUN_TOOL, "tool_call": tool_call}, indent=2, sort_keys=True))
|
|
result = call_vm_run(tool_arguments)
|
|
print(json.dumps({"tool_result": result}, indent=2, sort_keys=True))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|