diff --git a/python/coinbase-agentkit/changelog.d/add-prism-action-provider.feature.md b/python/coinbase-agentkit/changelog.d/add-prism-action-provider.feature.md new file mode 100644 index 000000000..89f10cd40 --- /dev/null +++ b/python/coinbase-agentkit/changelog.d/add-prism-action-provider.feature.md @@ -0,0 +1 @@ +Added a Prism Network action provider for renting real NVIDIA GPUs onchain: `wallet`, `list_gpus`, `lease_and_run`, `run`, and `end_lease`. The provider carries its own funded wallet and settles in USDG, so it composes with any AgentKit wallet provider. diff --git a/python/coinbase-agentkit/coinbase_agentkit/__init__.py b/python/coinbase-agentkit/coinbase_agentkit/__init__.py index e31253873..ed9fdf24d 100644 --- a/python/coinbase-agentkit/coinbase_agentkit/__init__.py +++ b/python/coinbase-agentkit/coinbase_agentkit/__init__.py @@ -18,6 +18,7 @@ morpho_action_provider, nillion_action_provider, onramp_action_provider, + prism_action_provider, pyth_action_provider, ssh_action_provider, superfluid_action_provider, @@ -71,6 +72,7 @@ "morpho_action_provider", "nillion_action_provider", "onramp_action_provider", + "prism_action_provider", "pyth_action_provider", "ssh_action_provider", "superfluid_action_provider", diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/__init__.py b/python/coinbase-agentkit/coinbase_agentkit/action_providers/__init__.py index 68573da62..ab64791c0 100644 --- a/python/coinbase-agentkit/coinbase_agentkit/action_providers/__init__.py +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/__init__.py @@ -26,6 +26,7 @@ from .morpho.morpho_action_provider import MorphoActionProvider, morpho_action_provider from .nillion.nillion_action_provider import NillionActionProvider, nillion_action_provider from .onramp.onramp_action_provider import OnrampActionProvider, onramp_action_provider +from .prism.prism_action_provider import PrismActionProvider, prism_action_provider from .pyth.pyth_action_provider import PythActionProvider, pyth_action_provider from .ssh.ssh_action_provider import SshActionProvider, ssh_action_provider from .superfluid.superfluid_action_provider import ( @@ -54,6 +55,7 @@ "MorphoActionProvider", "NillionActionProvider", "OnrampActionProvider", + "PrismActionProvider", "PythActionProvider", "SshActionProvider", "SuperfluidActionProvider", @@ -75,6 +77,7 @@ "morpho_action_provider", "nillion_action_provider", "onramp_action_provider", + "prism_action_provider", "pyth_action_provider", "ssh_action_provider", "superfluid_action_provider", diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/README.md b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/README.md new file mode 100644 index 000000000..ef397d048 --- /dev/null +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/README.md @@ -0,0 +1,39 @@ +# Prism Action Provider + +This directory contains the `PrismActionProvider` class, which provides actions for renting real NVIDIA GPUs through [Prism Network](https://prismnetwork.tech). + +The provider carries its own funded wallet (`PRISM_AGENT_KEY`) and settles onchain in USDG, so it composes with any AgentKit wallet provider rather than spending the agent's primary wallet. + +## Actions + +- `wallet`: The Prism agent wallet address and its USDG and native balances. +- `list_gpus`: GPUs available to rent right now, with model, VRAM, and price per hour. +- `lease_and_run`: Rent a GPU, run a command, and return the output. Pays onchain in USDG up to `max_usdg`. +- `run`: Run another command on a GPU already leased in this session. +- `end_lease`: Release a leased GPU session. + +## Environment + +- `PRISM_AGENT_KEY` (required): the agent wallet private key, funded with USDG and native gas on Robinhood Chain. +- `PRISM_ESCROW` (optional): the lease-escrow contract address. Defaults to the canonical escrow. + +## Setup + +```python +from coinbase_agentkit import AgentKit, AgentKitConfig +from coinbase_agentkit.action_providers.prism import prism_action_provider + +agent_kit = AgentKit( + AgentKitConfig( + wallet_provider=wallet_provider, + action_providers=[prism_action_provider()], + ) +) +``` + +## Notes + +- `lease_and_run` blocks while the machine provisions, usually one to four minutes. Give the agent a long tool-call timeout. +- Leasing spends real USDG onchain, capped by `max_usdg`. + +For more information on the **Prism Network**, visit [prismnetwork.tech](https://prismnetwork.tech). diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/__init__.py b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/__init__.py new file mode 100644 index 000000000..acf44c168 --- /dev/null +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/__init__.py @@ -0,0 +1,5 @@ +"""Prism Network action provider for renting NVIDIA GPUs.""" + +from .prism_action_provider import PrismActionProvider, prism_action_provider + +__all__ = ["PrismActionProvider", "prism_action_provider"] diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/client.py b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/client.py new file mode 100644 index 000000000..83cc4c229 --- /dev/null +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/client.py @@ -0,0 +1,373 @@ +"""Client for Prism Network: wallet-signature auth, onchain USDG payment, provisioning, SSH. + +The client holds its own wallet, authenticates to the Prism control plane with an +EIP-191 signature, pays into the onchain lease escrow in USDG, waits for the machine +to provision, and runs commands over SSH. +""" + +from __future__ import annotations + +import shutil +import subprocess +import tempfile +import time +from dataclasses import dataclass + +import requests +from eth_account import Account +from eth_account.messages import encode_defunct +from web3 import Web3 + +from .constants import ( + CHAIN_ID, + CONFIRMATIONS, + DEFAULT_API_BASE, + DIGEST_PATTERN, + ERC20_ABI, + ESCROW_ABI, + FETCH_TIMEOUT, + ROBINHOOD_RPC, + USDG_ADDRESS, +) + + +class PrismError(Exception): + """An error returned by the Prism control plane or the chain.""" + + def __init__(self, status: int, code: str, body: object = None) -> None: + """Store the HTTP-like status, machine code, and optional body.""" + super().__init__(f"prism {status}: {code}") + self.status = status + self.code = code + self.body = body + + +@dataclass +class Lease: + """A funded GPU lease and its SSH access material.""" + + lease_id: int + access: dict + key_path: str + key_dir: str + public_key: str + funding_hash: str + quote: dict + + +class PrismClient: + """Headless GPU leasing for a wallet-holding agent.""" + + def __init__( + self, + private_key: str, + escrow: str, + api_base: str = DEFAULT_API_BASE, + rpc_url: str = ROBINHOOD_RPC, + ) -> None: + """Build a client from a wallet key and the lease-escrow address.""" + if not escrow: + raise ValueError("escrow address is required") + self.api_base = api_base.rstrip("/") + self.escrow = Web3.to_checksum_address(escrow) + self.account = Account.from_key(private_key) + self.w3 = Web3(Web3.HTTPProvider(rpc_url)) + self._usdg = self.w3.eth.contract(address=USDG_ADDRESS, abi=ERC20_ABI) + self._escrow = self.w3.eth.contract(address=self.escrow, abi=ESCROW_ABI) + self.session: str | None = None + + @property + def address(self) -> str: + """The agent wallet's checksummed address.""" + return self.account.address + + def authenticate(self) -> dict: + """Prove wallet ownership with a signature and open a session.""" + challenge = self._json("GET", f"/api/agent/challenge?address={self.address}") + signed = self.account.sign_message(encode_defunct(text=challenge["message"])) + sig = signed.signature.hex() + session = self._json( + "POST", + "/api/agent/session", + { + "challenge": challenge["challenge"], + "address": self.address, + "signature": sig if sig.startswith("0x") else "0x" + sig, + }, + ) + self.session = session["session"] + return session + + def offers(self) -> list: + """List GPU offers currently available to lease.""" + return self._proxy("GET", ["offers"]) + + def balances(self) -> dict: + """Return the wallet address and its USDG and native balances.""" + return { + "address": self.address, + "usdg": self._usdg.functions.balanceOf(self.address).call(), + "eth": self.w3.eth.get_balance(self.address), + } + + def quote( + self, + image: str, + duration_seconds: int, + min_vram_mib: int = 16000, + preferred_node_id: str | None = None, + ) -> dict: + """Request a price quote for a lease matching the given constraints.""" + if not isinstance(image, str) or not DIGEST_PATTERN.search(image): + raise PrismError(400, "image_must_be_digest_pinned") + return self._proxy( + "POST", + ["leases", "match"], + { + "request": { + "image": image, + "duration_seconds": duration_seconds, + "min_vram_mib": min_vram_mib, + "preferred_node_id": preferred_node_id, + } + }, + ) + + def confirm(self, quote_id: str, transaction_hash: str, ssh_authorized_key: str) -> dict: + """Confirm a funded quote with its onchain transaction and SSH key.""" + return self._proxy( + "POST", + ["leases", "confirm"], + { + "quote_id": quote_id, + "transaction_hash": transaction_hash, + "ssh_authorized_key": ssh_authorized_key, + }, + ) + + def leases(self) -> list: + """List the wallet's leases.""" + return self._proxy("GET", ["leases"]) + + def access(self, lease_id: int) -> dict: + """Fetch SSH access details for a lease.""" + return self._proxy("GET", ["leases", str(lease_id), "access"]) + + def wait_for_access(self, lease_id: int, timeout: int = 600, interval: int = 10) -> dict: + """Poll until a lease's SSH access is ready or the timeout elapses.""" + deadline = time.time() + timeout + while time.time() < deadline: + status, body = self._proxy("GET", ["leases", str(lease_id), "access"], raw=True) + if status == 200: + return body + if status != 404: + raise PrismError(status, (body or {}).get("error", "access_error")) + time.sleep(interval) + raise PrismError(408, "access_timeout") + + def lease( + self, + image: str, + duration_seconds: int, + min_vram_mib: int = 16000, + preferred_node_id: str | None = None, + max_deposit: int | None = None, + ) -> Lease: + """Quote, fund onchain, confirm, and wait for a provisioned GPU.""" + if not self.session: + self.authenticate() + quote = self.quote(image, duration_seconds, min_vram_mib, preferred_node_id) + if max_deposit is not None and int(quote["maximum_escrow"]) > int(max_deposit): + raise PrismError( + 402, + "cost_exceeds_max", + {"required": quote["maximum_escrow"], "max": str(max_deposit)}, + ) + key = self._generate_ssh_key() + try: + funding = self._fund(quote) + record = self.confirm(quote["quote_id"], funding, key["public_key"]) + lease_id = record["lease_id"] + return Lease( + lease_id, + self.wait_for_access(lease_id), + key["key_path"], + key["dir"], + key["public_key"], + funding, + quote, + ) + except Exception: + shutil.rmtree(key["dir"], ignore_errors=True) + raise + + def run( + self, + lease: Lease, + command: str, + timeout: int = 120, + connect_retries: int = 24, + connect_delay: int = 10, + ) -> dict: + """Run a command on a leased GPU over SSH, retrying while it warms up.""" + a = lease.access + args = [ + "ssh", + "-i", + lease.key_path, + "-p", + str(a["ssh_port"]), + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "BatchMode=yes", + "-o", + "ConnectTimeout=15", + f"{a.get('ssh_user', 'root')}@{a['ssh_host']}", + command, + ] + last = None + for attempt in range(connect_retries + 1): + try: + p = subprocess.run(args, capture_output=True, text=True, timeout=timeout + 20) + res = {"code": p.returncode, "stdout": p.stdout.strip(), "stderr": p.stderr.strip()} + except subprocess.TimeoutExpired: + res = {"code": -1, "stdout": "", "stderr": "timed out"} + if not _is_ssh_warmup(res): + return res + last = res + if attempt < connect_retries: + time.sleep(connect_delay) + return last + + def end_lease(self, lease: Lease) -> None: + """Release local key material. The onchain lease settles at term end.""" + if lease and lease.key_dir: + shutil.rmtree(lease.key_dir, ignore_errors=True) + + def _fund(self, quote: dict) -> str: + deposit = int(quote["maximum_escrow"]) + duration = int(quote["duration_seconds"]) + client_ref = Web3.keccak(text=quote["quote_id"]) + node_id = bytes.fromhex(quote["node_id"].removeprefix("0x")) + allowance = self._usdg.functions.allowance(self.address, self.escrow).call() + if allowance < deposit: + self._send(self._usdg.functions.approve(self.escrow, deposit)) + return self._send( + self._escrow.functions.createLease(node_id, duration, client_ref), + confirmations=CONFIRMATIONS, + ) + + def _send(self, call: object, confirmations: int = 1) -> str: + tx = call.build_transaction( + { + "from": self.address, + "nonce": self.w3.eth.get_transaction_count(self.address), + "chainId": CHAIN_ID, + } + ) + signed = self.account.sign_transaction(tx) + h = self.w3.eth.send_raw_transaction(signed.raw_transaction) + receipt = self.w3.eth.wait_for_transaction_receipt(h) + if receipt.status != 1: + raise PrismError(402, "tx_reverted", {"hash": h.hex()}) + if confirmations > 1: + target = receipt.blockNumber + confirmations - 1 + while self.w3.eth.block_number < target: + time.sleep(2) + return h.hex() + + def _proxy( + self, + method: str, + segments: list, + body: object = None, + raw: bool = False, + reauthed: bool = False, + ): + if not self.session: + self.authenticate() + res = self._request( + f"/api/agent/proxy/{'/'.join(segments)}", + method, + body, + {"authorization": f"Bearer {self.session}"}, + ) + if res.status_code == 401 and not reauthed: + self.session = None + self.authenticate() + return self._proxy(method, segments, body, raw, True) + if raw: + return res.status_code, _safe_json(res) + return self._unwrap(res) + + def _json(self, method: str, path: str, body: object = None): + return self._unwrap(self._request(path, method, body)) + + def _request(self, path: str, method: str, body: object = None, headers: dict | None = None): + try: + return requests.request( + method, + f"{self.api_base}{path}", + json=body, + headers={"accept": "application/json", **(headers or {})}, + timeout=FETCH_TIMEOUT, + ) + except requests.RequestException as e: + raise PrismError(504, "control_plane_unreachable", {"cause": str(e)}) from e + + @staticmethod + def _unwrap(res): + data = _safe_json(res) + if not res.ok: + code = (data or {}).get("error") or (data or {}).get("code") or "request_failed" + raise PrismError(res.status_code, code, data) + return data + + def _generate_ssh_key(self) -> dict: + directory = tempfile.mkdtemp(prefix="prism-ssh-") + try: + key_path = f"{directory}/id_ed25519" + subprocess.run( + [ + "ssh-keygen", + "-t", + "ed25519", + "-N", + "", + "-q", + "-f", + key_path, + "-C", + "prism-agent", + ], + check=True, + capture_output=True, + ) + with open(f"{key_path}.pub") as f: + return {"dir": directory, "key_path": key_path, "public_key": f.read().strip()} + except Exception as e: + shutil.rmtree(directory, ignore_errors=True) + raise PrismError(500, "ssh_keygen_failed", {"cause": str(e)}) from e + + +def _safe_json(res): + try: + return res.json() + except ValueError: + return None + + +def _is_ssh_warmup(res: dict) -> bool: + if res["code"] != 255: + return False + e = res["stderr"] + return ( + e.startswith("ssh: ") + or "\nssh: " in e + or "kex_exchange_identification" in e + or "Connection reset by peer" in e + or ("Permission denied (publickey" in e and res["stdout"] == "") + ) diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/constants.py b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/constants.py new file mode 100644 index 000000000..80e1321fd --- /dev/null +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/constants.py @@ -0,0 +1,66 @@ +"""Constants for the Prism Network action provider.""" + +import re + +from web3 import Web3 + +ROBINHOOD_RPC = "https://rpc.mainnet.chain.robinhood.com" +CHAIN_ID = 4663 +USDG_ADDRESS = Web3.to_checksum_address("0x5fc5360D0400a0Fd4f2af552ADD042D716F1d168") +DEFAULT_ESCROW = Web3.to_checksum_address("0x71Df0eF3bc81022cB3bec0b1a05f52f12bAfcDeD") +USDG_DECIMALS = 6 +CONFIRMATIONS = 12 +FETCH_TIMEOUT = 30 +DEFAULT_API_BASE = "https://prismnetwork.tech" + +# A digest-pinned default image so the boot target can't silently drift. +DEFAULT_IMAGE = ( + "docker.io/ollama/ollama@sha256:" + "a61a8fd395dbb931cc8cb1b5da7a2510746575c87113fdc45b647ee59ef7f808" +) + +DIGEST_PATTERN = re.compile(r"@sha256:[0-9a-f]{64}$") + +ERC20_ABI = [ + { + "name": "approve", + "type": "function", + "stateMutability": "nonpayable", + "inputs": [ + {"name": "spender", "type": "address"}, + {"name": "value", "type": "uint256"}, + ], + "outputs": [{"type": "bool"}], + }, + { + "name": "allowance", + "type": "function", + "stateMutability": "view", + "inputs": [ + {"name": "owner", "type": "address"}, + {"name": "spender", "type": "address"}, + ], + "outputs": [{"type": "uint256"}], + }, + { + "name": "balanceOf", + "type": "function", + "stateMutability": "view", + "inputs": [{"name": "owner", "type": "address"}], + "outputs": [{"type": "uint256"}], + }, +] + +ESCROW_ABI = [ + { + "name": "createLease", + "type": "function", + "stateMutability": "nonpayable", + "inputs": [ + {"name": "nodeId", "type": "bytes32"}, + {"name": "duration", "type": "uint32"}, + {"name": "clientReference", "type": "bytes32"}, + ], + "outputs": [{"type": "uint256"}], + }, +] diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/prism_action_provider.py b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/prism_action_provider.py new file mode 100644 index 000000000..d0d912282 --- /dev/null +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/prism_action_provider.py @@ -0,0 +1,199 @@ +"""Prism Network action provider.""" + +import os +from typing import Any + +from ...network import Network +from ...wallet_providers import WalletProvider +from ..action_decorator import create_action +from ..action_provider import ActionProvider +from .client import Lease, PrismClient +from .constants import DEFAULT_ESCROW, USDG_DECIMALS +from .schemas import ( + EndLeaseSchema, + LeaseAndRunSchema, + ListGpusSchema, + RunSchema, + WalletSchema, +) + + +def _usdg(micros: int) -> str: + return f"{int(micros) / 10**USDG_DECIMALS:.6f} USDG" + + +class PrismActionProvider(ActionProvider[WalletProvider]): + """Action provider for renting real NVIDIA GPUs through Prism Network. + + The provider holds its own funded wallet (``PRISM_AGENT_KEY``) and settles + onchain in USDG, so it composes with any AgentKit wallet provider rather than + spending the agent's primary wallet. + """ + + def __init__(self, client: PrismClient | None = None) -> None: + """Initialize the provider. + + Args: + client: A configured PrismClient. If omitted, one is built from the + ``PRISM_AGENT_KEY`` and ``PRISM_ESCROW`` environment variables. + + """ + super().__init__("prism", []) + if client is None: + key = os.getenv("PRISM_AGENT_KEY") + if not key: + raise ValueError("PRISM_AGENT_KEY is required (or pass client=)") + client = PrismClient(key, os.getenv("PRISM_ESCROW", DEFAULT_ESCROW)) + self.client = client + self._leases: dict[int, Lease] = {} + + @create_action( + name="wallet", + description="Show the Prism agent wallet address and its USDG and native balances.", + schema=WalletSchema, + ) + def wallet(self, args: dict[str, Any]) -> str: + """Return the agent wallet address and balances. + + Args: + args (dict[str, Any]): Input arguments for the action. + + Returns: + str: The wallet address and balances, or an error message. + + """ + try: + b = self.client.balances() + return ( + f"address: {b['address']}\n" + f"usdg: {_usdg(b['usdg'])}\n" + f"eth: {int(b['eth']) / 10**18:.6f}" + ) + except Exception as e: + return f"Error fetching wallet: {e!s}" + + @create_action( + name="list_gpus", + description="List GPUs available to rent right now, with model, VRAM, and price per hour.", + schema=ListGpusSchema, + ) + def list_gpus(self, args: dict[str, Any]) -> str: + """List available GPU offers. + + Args: + args (dict[str, Any]): Input arguments for the action. + + Returns: + str: One line per available GPU, or a message when none are online. + + """ + try: + offers = self.client.offers() + if not offers: + return "No GPUs are online to rent right now." + rows = [] + for o in offers: + gpu = o.get("gpu", {}) + per_hour = int(o.get("rate_per_second", 0)) * 3600 / 10**USDG_DECIMALS + rows.append( + f"{gpu.get('model', 'GPU')} - {gpu.get('vram_mib', '?')} MiB - " + f"${per_hour:.2f}/hr" + ) + return "\n".join(rows) + except Exception as e: + return f"Error listing GPUs: {e!s}" + + @create_action( + name="lease_and_run", + description=( + "Rent a GPU, run one command on it, and return the output. Pays onchain in USDG " + "up to max_usdg. Blocks while the machine provisions, usually one to four minutes." + ), + schema=LeaseAndRunSchema, + ) + def lease_and_run(self, args: dict[str, Any]) -> str: + """Lease a GPU, run a command, and return the output. + + Args: + args (dict[str, Any]): Input arguments for the action. + + Returns: + str: The lease id, funding transaction, and command output. + + """ + try: + p = LeaseAndRunSchema(**args) + lease = self.client.lease( + image=p.image, + duration_seconds=p.duration_seconds, + min_vram_mib=p.min_vram_mib, + max_deposit=int(p.max_usdg * 10**USDG_DECIMALS), + ) + self._leases[lease.lease_id] = lease + res = self.client.run(lease, p.command) + out = res.get("stdout") or res.get("stderr") or "" + return ( + f"lease {lease.lease_id} funded onchain (tx {lease.funding_hash}), " + f"exit {res.get('code')}:\n{out}" + ) + except Exception as e: + return f"Error leasing GPU: {e!s}" + + @create_action( + name="run", + description="Run another command on a GPU already leased in this session.", + schema=RunSchema, + ) + def run(self, args: dict[str, Any]) -> str: + """Run a command on an existing lease. + + Args: + args (dict[str, Any]): Input arguments for the action. + + Returns: + str: The command output, or an error message. + + """ + try: + p = RunSchema(**args) + lease = self._leases.get(p.lease_id) + if lease is None: + return f"No active lease {p.lease_id} in this session." + res = self.client.run(lease, p.command) + return f"exit {res.get('code')}:\n{res.get('stdout') or res.get('stderr') or ''}" + except Exception as e: + return f"Error running command: {e!s}" + + @create_action( + name="end_lease", + description="Release a leased GPU session. The onchain lease settles when its term ends.", + schema=EndLeaseSchema, + ) + def end_lease(self, args: dict[str, Any]) -> str: + """Release a lease's local session. + + Args: + args (dict[str, Any]): Input arguments for the action. + + Returns: + str: A confirmation, or an error message. + + """ + try: + p = EndLeaseSchema(**args) + lease = self._leases.pop(p.lease_id, None) + if lease is None: + return f"No active lease {p.lease_id} in this session." + self.client.end_lease(lease) + return f"Released lease {p.lease_id}." + except Exception as e: + return f"Error ending lease: {e!s}" + + def supports_network(self, network: Network) -> bool: + """Return True: Prism settles on its own network, independent of the agent's.""" + return True + + +def prism_action_provider(client: PrismClient | None = None) -> PrismActionProvider: + """Create a Prism action provider.""" + return PrismActionProvider(client) diff --git a/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/schemas.py b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/schemas.py new file mode 100644 index 000000000..e04a289a9 --- /dev/null +++ b/python/coinbase-agentkit/coinbase_agentkit/action_providers/prism/schemas.py @@ -0,0 +1,36 @@ +"""Schemas for the Prism Network action provider.""" + +from pydantic import BaseModel, Field + +from .constants import DEFAULT_IMAGE + + +class WalletSchema(BaseModel): + """Input schema for the wallet action.""" + + +class ListGpusSchema(BaseModel): + """Input schema for the list_gpus action.""" + + +class LeaseAndRunSchema(BaseModel): + """Input schema for the lease_and_run action.""" + + command: str = Field(..., description="Command to run on the rented GPU") + duration_seconds: int = Field(600, description="How long to hold the lease, in seconds") + min_vram_mib: int = Field(16000, description="Minimum GPU memory in MiB") + image: str = Field(DEFAULT_IMAGE, description="Digest-pinned container image to boot") + max_usdg: float = Field(1.0, description="Hard cap on the USDG this lease may cost") + + +class RunSchema(BaseModel): + """Input schema for the run action.""" + + lease_id: int = Field(..., description="A lease id returned by lease_and_run") + command: str = Field(..., description="Command to run on the leased GPU") + + +class EndLeaseSchema(BaseModel): + """Input schema for the end_lease action.""" + + lease_id: int = Field(..., description="A lease id from this session") diff --git a/python/coinbase-agentkit/tests/action_providers/prism/__init__.py b/python/coinbase-agentkit/tests/action_providers/prism/__init__.py new file mode 100644 index 000000000..85dc45344 --- /dev/null +++ b/python/coinbase-agentkit/tests/action_providers/prism/__init__.py @@ -0,0 +1 @@ +"""Tests for the Prism action provider.""" diff --git a/python/coinbase-agentkit/tests/action_providers/prism/test_prism_action_provider.py b/python/coinbase-agentkit/tests/action_providers/prism/test_prism_action_provider.py new file mode 100644 index 000000000..409183c25 --- /dev/null +++ b/python/coinbase-agentkit/tests/action_providers/prism/test_prism_action_provider.py @@ -0,0 +1,79 @@ +"""Tests for the Prism action provider.""" + +from unittest.mock import MagicMock + +from coinbase_agentkit.action_providers.prism.client import Lease +from coinbase_agentkit.action_providers.prism.prism_action_provider import prism_action_provider + + +def test_wallet_formats_balances(): + """Wallet returns the address and formatted USDG balance.""" + client = MagicMock() + client.balances.return_value = { + "address": "0xEcaaE714912C38fA7e0dAF78afa7C54DbeD11039", + "usdg": 920362, + "eth": 4506706341389206, + } + out = prism_action_provider(client).wallet({}) + assert "0xEcaaE714912C38fA7e0dAF78afa7C54DbeD11039" in out + assert "0.920362 USDG" in out + + +def test_list_gpus_formats_offers(): + """list_gpus renders model, VRAM, and hourly price.""" + client = MagicMock() + client.offers.return_value = [ + {"gpu": {"model": "L40S", "vram_mib": 46068}, "rate_per_second": 222} + ] + out = prism_action_provider(client).list_gpus({}) + assert "L40S" in out + assert "46068 MiB" in out + assert "/hr" in out + + +def test_list_gpus_empty(): + """list_gpus reports when no GPUs are online.""" + client = MagicMock() + client.offers.return_value = [] + assert "No GPUs" in prism_action_provider(client).list_gpus({}) + + +def test_lease_and_run_returns_receipt_and_output(): + """lease_and_run funds a lease, runs the command, and returns both.""" + client = MagicMock() + client.lease.return_value = Lease( + lease_id=12, + access={}, + key_path="/tmp/k", + key_dir="/tmp/k.d", + public_key="ssh-ed25519 AAAA", + funding_hash="0xabc123", + quote={}, + ) + client.run.return_value = {"code": 0, "stdout": "NVIDIA L40S", "stderr": ""} + out = prism_action_provider(client).lease_and_run( + {"command": "nvidia-smi", "duration_seconds": 600, "max_usdg": 0.5} + ) + assert "lease 12" in out + assert "0xabc123" in out + assert "NVIDIA L40S" in out + client.lease.assert_called_once() + + +def test_run_without_active_lease(): + """Run reports when the lease id is not held in this session.""" + client = MagicMock() + out = prism_action_provider(client).run({"lease_id": 999, "command": "ls"}) + assert "No active lease" in out + + +def test_end_lease_without_active_lease(): + """end_lease reports when the lease id is not held in this session.""" + client = MagicMock() + out = prism_action_provider(client).end_lease({"lease_id": 999}) + assert "No active lease" in out + + +def test_supports_network(): + """The provider supports any network since it uses its own wallet.""" + assert prism_action_provider(MagicMock()).supports_network(MagicMock()) is True