Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 84 additions & 0 deletions sandboxes/pydantic-ai-code-agent/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
# Pydantic AI Code Agent with Agent Sandboxes

This example shows how to use Agent Sandboxes as tools in a lightweight Python agent built with [Pydantic AI](https://ai.pydantic.dev/).

Use this example when you want an agent framework to run code, execute shell commands, and read or write files in an isolated prokube Sandbox instead of on the user's machine.

## What This Example Shows

- Claiming a sandbox from a WarmPool.
- Exposing sandbox operations as Pydantic AI tools.
- Running Python code in a stateful sandbox kernel.
- Running shell commands inside the sandbox.
- Reading and writing files under `/workspace`.
- Cleaning up the sandbox after the agent run.

This example focuses on agent-framework integration. For a step-by-step SDK walkthrough, see `sandboxes/sdk-quickstart`.

## Files

| File | Purpose |
|---|---|
| `agent.py` | Pydantic AI agent with Sandbox-backed tools. |
| `requirements.txt` | Python dependencies for the example. |

## Prerequisites

- A prokube workspace with the Sandbox module enabled.
- A ready WarmPool, for example `sandbox-sdk-quickstart`.
- Access to an OpenAI-compatible model supported by Pydantic AI.

In a managed Lab, this repository is usually available at `~/examples`. Run the example from its directory:

```bash
cd ~/examples/sandboxes/pydantic-ai-code-agent
pip install -r requirements.txt
```

## Configuration

For the Sandbox SDK in a managed Lab:

```bash
export PROKUBE_API_URL="http://agentgateway-proxy.agentgateway-system.svc.cluster.local"
export PROKUBE_WORKSPACE="<workspace>"
export PROKUBE_USER_ID="<user-or-workspace>"
export SANDBOX_POOL="sandbox-sdk-quickstart"
```

For external access, use an API key instead of `PROKUBE_USER_ID`:

```bash
export PROKUBE_API_URL="https://<cluster-domain>/pkui"
export PROKUBE_WORKSPACE="<workspace>"
export PROKUBE_API_KEY="<api-key>"
export SANDBOX_POOL="sandbox-sdk-quickstart"
```

For the agent model:

```bash
export OPENAI_API_KEY="<model-api-key>"
export OPENAI_BASE_URL="<optional-openai-compatible-base-url>"
export PYDANTIC_AI_MODEL="openai:gpt-4o-mini"
```

Do not store real API keys in source code, notebooks, screenshots, tickets, or chat messages.

## Run

```bash
python agent.py
```

You can pass a custom task:

```bash
python agent.py "Create a CSV with five rows, analyze it, and write a short markdown report."
```

The default task asks the agent to create `/workspace/scores.csv`, analyze it, and write `/workspace/report.md`.

## Safety Notes

The LLM can decide to call tools that run code or shell commands. Use a sandbox image and workspace policy appropriate for the data and dependencies you expose. Keep cleanup explicit, and do not pass secrets to the agent unless the task requires them.
127 changes: 127 additions & 0 deletions sandboxes/pydantic-ai-code-agent/agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
"""Pydantic AI agent that uses a prokube Sandbox as its execution toolset."""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from pathlib import Path

from prokube.sandbox import Sandbox
from pydantic_ai import Agent, RunContext


DEFAULT_MODEL = "openai:gpt-4o-mini"
DEFAULT_POOL = "sandbox-sdk-quickstart"
IN_CLUSTER_AGENT_GATEWAY = "http://agentgateway-proxy.agentgateway-system.svc.cluster.local"
SERVICE_ACCOUNT_NAMESPACE = Path(
"/var/run/secrets/kubernetes.io/serviceaccount/namespace"
)


@dataclass
class SandboxDeps:
"""Dependencies available to Pydantic AI tools."""

sandbox: Sandbox


def configure_managed_lab_defaults() -> None:
"""Set SDK defaults for managed Labs without overriding explicit values."""

if SERVICE_ACCOUNT_NAMESPACE.exists():
workspace = SERVICE_ACCOUNT_NAMESPACE.read_text().strip()
os.environ.setdefault("PROKUBE_API_URL", IN_CLUSTER_AGENT_GATEWAY)
os.environ.setdefault("PROKUBE_WORKSPACE", workspace)
os.environ.setdefault("PROKUBE_USER_ID", workspace)


def build_agent() -> Agent[SandboxDeps, str]:
"""Create the Pydantic AI agent and register sandbox-backed tools."""

agent = Agent(
os.environ.get("PYDANTIC_AI_MODEL", DEFAULT_MODEL),
deps_type=SandboxDeps,
system_prompt=(
"You are a coding agent. Use the sandbox tools to run Python, run "
"shell commands, and read or write files under /workspace. Keep "
"outputs concise. Do not print secrets or environment variables "
"unless the user explicitly asks for a safe diagnostic."
),
)

@agent.tool
def run_python(ctx: RunContext[SandboxDeps], code: str) -> str:
"""Run Python code in the stateful sandbox kernel."""

result = ctx.deps.sandbox.run_code(code)
return _format_result(result.stdout, result.stderr)

@agent.tool
def run_shell(ctx: RunContext[SandboxDeps], command: str) -> str:
"""Run a shell command inside the sandbox."""

result = ctx.deps.sandbox.commands.run(command)
output = _format_result(result.stdout, result.stderr)
return f"exit_code={result.exit_code}\n{output}".strip()

@agent.tool
def write_file(ctx: RunContext[SandboxDeps], path: str, content: str) -> str:
"""Write a text file inside the sandbox."""

_require_workspace_path(path)
ctx.deps.sandbox.files.write(path, content)
return f"Wrote {path}"

@agent.tool
def read_file(ctx: RunContext[SandboxDeps], path: str) -> str:
"""Read a text file from the sandbox."""

_require_workspace_path(path)
content = ctx.deps.sandbox.files.read(path)
if isinstance(content, bytes):
return content.decode("utf-8", errors="replace")
return content

@agent.tool
def list_files(ctx: RunContext[SandboxDeps], path: str = "/workspace") -> str:
"""List files in a sandbox directory."""

_require_workspace_path(path)
files = ctx.deps.sandbox.files.list(path)
return "\n".join(str(file_info) for file_info in files) or "No files found."

return agent


def _format_result(stdout: str | None, stderr: str | None) -> str:
parts = []
if stdout:
parts.append(f"stdout:\n{stdout}")
if stderr:
parts.append(f"stderr:\n{stderr}")
return "\n\n".join(parts) or "No output."


def _require_workspace_path(path: str) -> None:
if not path.startswith("/workspace/") and path != "/workspace":
raise ValueError("This example only allows file access under /workspace")


def main() -> None:
configure_managed_lab_defaults()

pool = os.environ.get("SANDBOX_POOL", DEFAULT_POOL)
task = " ".join(sys.argv[1:]) or (
"Create /workspace/scores.csv with five rows of sample data, analyze it "
"with Python, and write /workspace/report.md with the findings."
)
agent = build_agent()

with Sandbox.from_pool(pool) as sandbox:
result = agent.run_sync(task, deps=SandboxDeps(sandbox=sandbox))
print(result.output)


if __name__ == "__main__":
main()
2 changes: 2 additions & 0 deletions sandboxes/pydantic-ai-code-agent/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
git+https://github.com/prokube/prokube-sdk.git@v0.1.3
pydantic-ai>=2.0.0