Use cases

"""A data analyst agent: notebook-style analysis over local files.

The model explores a CSV across several turns, keeps variables between them,
writes a report, and tries to read something it should not. The host allows
reading ./data and writing ./out, for 20 seconds per step.
"""

import csv
import random
import tempfile
from pathlib import Path

import pynox

# What a model writes, turn by turn, for "Which region grew the most? Write a report."
TURNS = [
    "import csv\n"
    "with open(DATA) as f:\n"
    "    rows = list(csv.DictReader(f))\n"
    "print(len(rows), 'rows;', 'columns:', ', '.join(rows[0]))",
    "from collections import defaultdict\n"
    "totals = defaultdict(lambda: [0.0, 0.0])\n"
    "for r in rows:\n"
    "    totals[r['region']][r['quarter'] == 'Q2'] += float(r['revenue'])\n"
    "growth = {k: (q2 - q1) / q1 for k, (q1, q2) in totals.items()}\n"
    "best = max(growth, key=growth.get)\n"
    "print(f'{best} grew {growth[best]:.1%}')",
    "with open(OUT, 'w') as f:\n"
    "    f.write('# Revenue growth Q1 to Q2\\n\\n')\n"
    "    for k, g in sorted(growth.items(), key=lambda kv: -kv[1]):\n"
    "        f.write(f'- {k}: {g:+.1%}\\n')\n"
    "print('report written')",
    # The model gets curious about the project's secrets "to check the data source".
    "print(open(ENV).read())",
]


def make_data(root: Path) -> Path:
    rng = random.Random(7)
    data = root / "data"
    data.mkdir()
    with open(data / "sales.csv", "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["region", "quarter", "revenue"])
        for region, trend in [("north", 1.05), ("south", 1.31), ("east", 0.92), ("west", 1.12)]:
            for _ in range(50):
                base = rng.uniform(800, 1200)
                writer.writerow([region, "Q1", f"{base:.2f}"])
                writer.writerow([region, "Q2", f"{base * trend:.2f}"])
    (root / "out").mkdir()
    (root / ".env").write_text("WAREHOUSE_PASSWORD=hunter2\n")
    return data / "sales.csv"


def run(root: Path) -> list[tuple[str, str]]:
    source = make_data(root)
    report = root / "out" / "report.md"
    notebook = pynox.tool(
        session=True, allow_read=[root / "data"], allow_write=[root / "out"], timeout="20s"
    )
    notebook(f"DATA = {str(source)!r}\nOUT = {str(report)!r}\nENV = {str(root / '.env')!r}")
    return [(code, notebook(code)) for code in TURNS]


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as temporary:
        for code, answer in run(Path(temporary)):
            print(f">>> model code\n{code}\n<<< tool result\n{answer}\n")
        print((Path(temporary) / "out" / "report.md").read_text())

"""A support agent that calls your backend through host functions.

The model's code chains lookups and a refund in one run. The functions run in
your process, where the database and the business rules live; the sandboxed
code can call them but cannot open the database itself.
"""

import json
import sqlite3
import tempfile
from pathlib import Path

import pynox


def make_backend(root: Path):
    database = root / "orders.db"
    with sqlite3.connect(database) as db:
        db.execute("create table orders (id text, customer text, total real, status text)")
        db.executemany(
            "insert into orders values (?, ?, ?, ?)",
            [
                ("A-100", "ana@example.com", 42.0, "delivered"),
                ("A-101", "ana@example.com", 310.0, "damaged"),
                ("B-200", "bo@example.com", 18.5, "delivered"),
            ],
        )
    refunds = []

    def find_orders(customer: str) -> list[dict]:
        """List a customer's orders with id, total and status."""
        with sqlite3.connect(database) as db:
            rows = db.execute(
                "select id, total, status from orders where customer = ?", (customer,)
            ).fetchall()
        return [{"id": i, "total": t, "status": s} for i, t, s in rows]

    def issue_refund(order_id: str, amount: float) -> str:
        """Refund an order. Amounts above 200 need a human."""
        if amount > 200:
            raise PermissionError("refunds above 200 need approval; ticket opened")
        refunds.append((order_id, amount))
        return f"refunded {amount:.2f} on {order_id}"

    return database, refunds, {"find_orders": find_orders, "issue_refund": issue_refund}


# What a model writes for "Ana says an order arrived damaged. Sort it out."
MODEL_CODE = """
from pynox_host import find_orders, issue_refund, HostError

orders = find_orders("ana@example.com")
damaged = [o for o in orders if o["status"] == "damaged"]
for order in damaged:
    try:
        print(issue_refund(order["id"], order["total"]))
    except HostError as error:
        print(f"{order['id']}: {error}")
print(f"checked {len(orders)} orders, {len(damaged)} damaged")
"""


def run(root: Path) -> dict:
    database, refunds, functions = make_backend(root)
    run_python = pynox.tool(functions=functions)
    answer = run_python(MODEL_CODE)
    # The model then tries to skip the rules and edit the database directly.
    bypass = run_python(
        f"import sqlite3\nsqlite3.connect({str(database)!r}).execute(\"update orders set status='refunded'\")"
    )
    return {"answer": answer, "bypass": bypass, "refunds": refunds, "doc": run_python.__doc__}


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as temporary:
        print(json.dumps(run(Path(temporary)), indent=2))

"""A coding agent that runs the project's tests after editing it.

`pynox -m unittest` makes the project directory readable and nothing writable,
so a test, or code the model just wrote, cannot delete files or read secrets
outside the project, and cannot write inside it either.
"""

import os
import subprocess
import tempfile
from pathlib import Path

PROJECT = {
    "calc.py": "def add(a, b):\n    return a + b\n",
    "test_calc.py": (
        "import os, shutil, unittest\n"
        "from calc import add\n\n"
        "class CalcTest(unittest.TestCase):\n"
        "    def test_add(self):\n"
        "        self.assertEqual(add(2, 3), 5)\n\n"
        "    def test_cleanup_gone_wrong(self):\n"
        "        # A model 'cleaning up fixtures' with the wrong path.\n"
        "        shutil.rmtree(os.path.expanduser('~/important-notes'))\n"
    ),
}


def run(root: Path, binary: str) -> dict:
    project = root / "project"
    project.mkdir()
    for name, source in PROJECT.items():
        (project / name).write_text(source)
    notes = root / "home" / "important-notes"
    notes.mkdir(parents=True)
    (notes / "plan.txt").write_text("do not delete")
    env = {**os.environ, "HOME": str(root / "home")}
    result = subprocess.run(
        [binary, "--allow-env=HOME", "-m", "unittest", "-q"],
        cwd=project,
        env=env,
        capture_output=True,
        text=True,
        timeout=60,
    )
    return {
        "exit": result.returncode,
        "stderr": result.stderr,
        "notes_survived": (notes / "plan.txt").exists(),
    }


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as temporary:
        outcome = run(Path(temporary), os.environ.get("PYNOX_BIN", "pynox"))
        print(outcome["stderr"])
        print("exit", outcome["exit"], "| notes survived:", outcome["notes_survived"])

"""An agent that calls an external API with a key it never sees in its prompt.

The host forwards one environment variable and turns the network on for this
profile only. Without the grant, the same code cannot open a socket.
A local HTTP server stands in for the API.
"""

import json
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

import pynox


class WeatherAPI(BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        authorized = self.headers.get("Authorization") == "Bearer demo-key"
        body = json.dumps(
            {"city": "Lisbon", "temp_c": 21} if authorized else {"error": "unauthorized"}
        )
        self.send_response(200 if authorized else 401)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body.encode())


def model_code(port: int) -> str:
    # What a model writes for "What's the weather in Lisbon?"
    return (
        "import json, os, urllib.request\n"
        f"request = urllib.request.Request('http://127.0.0.1:{port}/weather?city=Lisbon',\n"
        "    headers={'Authorization': 'Bearer ' + os.environ['WEATHER_KEY']})\n"
        "print(json.load(urllib.request.urlopen(request, timeout=5)))\n"
    )


def run() -> dict:
    server = ThreadingHTTPServer(("127.0.0.1", 0), WeatherAPI)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    try:
        code = model_code(server.server_address[1])
        key = {"WEATHER_KEY": "demo-key"}
        offline = pynox.tool(env=key)(code)  # same key, no network grant
        online = pynox.tool(allow_network=True, env=key)(code)
    finally:
        server.shutdown()
    return {"offline": offline, "online": online}


if __name__ == "__main__":
    print(json.dumps(run(), indent=2))