Spaces:

trialdesignbench
/

tdb-intake

Running on CPU Upgrade

Turn Run agent into a real agentic loop with a workspace and tools

74239a8 3 months ago

9.37 kB

"""Sandboxed workspace tools for the agent.
The agent gets a directory containing the SAP text and must write its answers
there as files. Every path is confined to the workspace root — traversal
outside it is rejected.
NOTE: `run_r` executes model-generated code. It is intended to run inside an
ephemeral, isolated container (HF Space) with a timeout.
"""

from __future__ import annotations

import re
import shutil
import subprocess
from pathlib import Path
from typing import Any, Dict, List

MAX_READ_LINES = 400
MAX_GREP_MATCHES = 60
MAX_WRITE_BYTES = 2_000_000
R_TIMEOUT_SEC = 120


class Workspace:
    """Filesystem operations confined to `root`."""

    def __init__(self, root: Path):
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)

    # -- path safety ------------------------------------------------------
    def _resolve(self, rel: str) -> Path:
        p = (self.root / (rel or "").lstrip("/")).resolve()
        if p != self.root and self.root not in p.parents:
            raise ValueError(f"path escapes the workspace: {rel}")
        return p

    def rel(self, p: Path) -> str:
        return str(p.relative_to(self.root))

    # -- tools ------------------------------------------------------------
    def list_files(self) -> str:
        out = []
        for p in sorted(self.root.rglob("*")):
            if p.is_file():
                out.append(f"{self.rel(p)}  ({p.stat().st_size:,} bytes)")
        return "\n".join(out) or "(workspace is empty)"

    def read_file(self, path: str, offset: int = 1, limit: int = MAX_READ_LINES) -> str:
        p = self._resolve(path)
        if not p.is_file():
            return f"ERROR: no such file: {path}"
        limit = max(1, min(int(limit or MAX_READ_LINES), MAX_READ_LINES))
        offset = max(1, int(offset or 1))
        lines = p.read_text(encoding="utf-8", errors="replace").splitlines()
        total = len(lines)
        chunk = lines[offset - 1 : offset - 1 + limit]
        if not chunk:
            return f"(no lines at offset {offset}; file has {total} lines)"
        body = "\n".join(f"{offset + i}\t{ln}" for i, ln in enumerate(chunk))
        end = offset + len(chunk) - 1
        more = (
            f"\n\n[showing lines {offset}-{end} of {total}. "
            f"Use offset={end + 1} to continue.]"
            if end < total
            else f"\n\n[end of file, {total} lines]"
        )
        return body + more

    def grep(self, pattern: str, path: str = "", max_matches: int = MAX_GREP_MATCHES) -> str:
        max_matches = max(1, min(int(max_matches or MAX_GREP_MATCHES), MAX_GREP_MATCHES))
        try:
            rx = re.compile(pattern, re.IGNORECASE)
        except re.error as e:
            return f"ERROR: bad regex: {e}"
        targets: List[Path] = []
        if path:
            p = self._resolve(path)
            if not p.is_file():
                return f"ERROR: no such file: {path}"
            targets = [p]
        else:
            targets = [p for p in sorted(self.root.rglob("*")) if p.is_file()]

        hits: List[str] = []
        for p in targets:
            try:
                lines = p.read_text(encoding="utf-8", errors="replace").splitlines()
            except Exception:
                continue
            for i, ln in enumerate(lines, start=1):
                if rx.search(ln):
                    hits.append(f"{self.rel(p)}:{i}\t{ln.strip()[:300]}")
                    if len(hits) >= max_matches:
                        return "\n".join(hits) + f"\n\n[stopped at {max_matches} matches]"
        return "\n".join(hits) or f"(no matches for /{pattern}/)"

    def write_file(self, path: str, content: str) -> str:
        p = self._resolve(path)
        data = (content or "").encode("utf-8")
        if len(data) > MAX_WRITE_BYTES:
            return f"ERROR: content too large ({len(data):,} bytes)"
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(content or "", encoding="utf-8")
        return f"wrote {self.rel(p)} ({len(data):,} bytes)"

    def run_r(self, path: str) -> str:
        p = self._resolve(path)
        if not p.is_file():
            return f"ERROR: no such file: {path}"
        exe = shutil.which("Rscript")
        if not exe:
            return (
                "ERROR: Rscript is not available in this environment, so the R "
                "script could not be executed. Continue without running it."
            )
        try:
            proc = subprocess.run(
                [exe, "--vanilla", str(p)],
                capture_output=True,
                text=True,
                timeout=R_TIMEOUT_SEC,
                cwd=str(self.root),
            )
        except subprocess.TimeoutExpired:
            return f"ERROR: Rscript timed out after {R_TIMEOUT_SEC}s"
        parts = [f"exit code: {proc.returncode}"]
        if proc.stdout.strip():
            parts.append("--- stdout ---\n" + proc.stdout.strip()[:6000])
        if proc.stderr.strip():
            parts.append("--- stderr ---\n" + proc.stderr.strip()[:4000])
        return "\n".join(parts)

    # -- dispatch ---------------------------------------------------------
    def call(self, name: str, args: Dict[str, Any]) -> str:
        try:
            if name == "list_files":
                return self.list_files()
            if name == "read_file":
                return self.read_file(
                    args.get("path", ""), args.get("offset", 1), args.get("limit", MAX_READ_LINES)
                )
            if name == "grep":
                return self.grep(
                    args.get("pattern", ""), args.get("path", ""), args.get("max_matches", MAX_GREP_MATCHES)
                )
            if name == "write_file":
                return self.write_file(args.get("path", ""), args.get("content", ""))
            if name == "run_r":
                return self.run_r(args.get("path", "output.R"))
            return f"ERROR: unknown tool {name}"
        except Exception as e:
            return f"ERROR: {type(e).__name__}: {e}"


# ---------------------------------------------------------------------------
# Tool schemas
# ---------------------------------------------------------------------------

_TOOLS: List[Dict[str, Any]] = [
    {
        "name": "list_files",
        "description": "List all files in the workspace with their sizes.",
        "properties": {},
        "required": [],
    },
    {
        "name": "read_file",
        "description": (
            "Read a slice of a text file with line numbers. The SAP is large, so "
            "read it in chunks (or grep first to find the right lines)."
        ),
        "properties": {
            "path": {"type": "string", "description": "Workspace-relative path, e.g. sap.txt"},
            "offset": {"type": "integer", "description": "1-based first line to read (default 1)"},
            "limit": {"type": "integer", "description": f"Lines to read (max {MAX_READ_LINES})"},
        },
        "required": ["path"],
    },
    {
        "name": "grep",
        "description": (
            "Case-insensitive regex search across the workspace (or one file). "
            "Returns path:line and the matching text — use it to locate sections, "
            "then read_file around those line numbers."
        ),
        "properties": {
            "pattern": {"type": "string", "description": "Regular expression"},
            "path": {"type": "string", "description": "Optional single file to search"},
            "max_matches": {"type": "integer", "description": f"Max matches (default {MAX_GREP_MATCHES})"},
        },
        "required": ["pattern"],
    },
    {
        "name": "write_file",
        "description": "Create or overwrite a file in the workspace with the given content.",
        "properties": {
            "path": {"type": "string", "description": "Workspace-relative path, e.g. output.json"},
            "content": {"type": "string", "description": "Full file content"},
        },
        "required": ["path", "content"],
    },
    {
        "name": "run_r",
        "description": (
            "Execute an R script with Rscript and return its exit code, stdout and "
            "stderr. Use it to verify output.R runs and prints the values you "
            "reported; fix the script and re-run if it errors or disagrees."
        ),
        "properties": {
            "path": {"type": "string", "description": "Script path (default output.R)"},
        },
        "required": [],
    },
]


def anthropic_tools() -> List[Dict[str, Any]]:
    return [
        {
            "name": t["name"],
            "description": t["description"],
            "input_schema": {
                "type": "object",
                "properties": t["properties"],
                "required": t["required"],
            },
        }
        for t in _TOOLS
    ]


def openai_tools() -> List[Dict[str, Any]]:
    return [
        {
            "type": "function",
            "function": {
                "name": t["name"],
                "description": t["description"],
                "parameters": {
                    "type": "object",
                    "properties": t["properties"],
                    "required": t["required"],
                },
            },
        }
        for t in _TOOLS
    ]

	"""Sandboxed workspace tools for the agent.

	The agent gets a directory containing the SAP text and must write its answers
	there as files. Every path is confined to the workspace root — traversal
	outside it is rejected.

	NOTE: `run_r` executes model-generated code. It is intended to run inside an
	ephemeral, isolated container (HF Space) with a timeout.
	"""

	from __future__ import annotations

	import re
	import shutil
	import subprocess
	from pathlib import Path
	from typing import Any, Dict, List

	MAX_READ_LINES = 400
	MAX_GREP_MATCHES = 60
	MAX_WRITE_BYTES = 2_000_000
	R_TIMEOUT_SEC = 120


	class Workspace:
	"""Filesystem operations confined to `root`."""

	def __init__(self, root: Path):
	self.root = Path(root).resolve()
	self.root.mkdir(parents=True, exist_ok=True)

	# -- path safety ------------------------------------------------------
	def _resolve(self, rel: str) -> Path:
	p = (self.root / (rel or "").lstrip("/")).resolve()
	if p != self.root and self.root not in p.parents:
	raise ValueError(f"path escapes the workspace: {rel}")
	return p

	def rel(self, p: Path) -> str:
	return str(p.relative_to(self.root))

	# -- tools ------------------------------------------------------------
	def list_files(self) -> str:
	out = []
	for p in sorted(self.root.rglob("*")):
	if p.is_file():
	out.append(f"{self.rel(p)} ({p.stat().st_size:,} bytes)")
	return "\n".join(out) or "(workspace is empty)"

	def read_file(self, path: str, offset: int = 1, limit: int = MAX_READ_LINES) -> str:
	p = self._resolve(path)
	if not p.is_file():
	return f"ERROR: no such file: {path}"
	limit = max(1, min(int(limit or MAX_READ_LINES), MAX_READ_LINES))
	offset = max(1, int(offset or 1))
	lines = p.read_text(encoding="utf-8", errors="replace").splitlines()
	total = len(lines)
	chunk = lines[offset - 1 : offset - 1 + limit]
	if not chunk:
	return f"(no lines at offset {offset}; file has {total} lines)"
	body = "\n".join(f"{offset + i}\t{ln}" for i, ln in enumerate(chunk))
	end = offset + len(chunk) - 1
	more = (
	f"\n\n[showing lines {offset}-{end} of {total}. "
	f"Use offset={end + 1} to continue.]"
	if end < total
	else f"\n\n[end of file, {total} lines]"
	)
	return body + more

	def grep(self, pattern: str, path: str = "", max_matches: int = MAX_GREP_MATCHES) -> str:
	max_matches = max(1, min(int(max_matches or MAX_GREP_MATCHES), MAX_GREP_MATCHES))
	try:
	rx = re.compile(pattern, re.IGNORECASE)
	except re.error as e:
	return f"ERROR: bad regex: {e}"
	targets: List[Path] = []
	if path:
	p = self._resolve(path)
	if not p.is_file():
	return f"ERROR: no such file: {path}"
	targets = [p]
	else:
	targets = [p for p in sorted(self.root.rglob("*")) if p.is_file()]

	hits: List[str] = []
	for p in targets:
	try:
	lines = p.read_text(encoding="utf-8", errors="replace").splitlines()
	except Exception:
	continue
	for i, ln in enumerate(lines, start=1):
	if rx.search(ln):
	hits.append(f"{self.rel(p)}:{i}\t{ln.strip()[:300]}")
	if len(hits) >= max_matches:
	return "\n".join(hits) + f"\n\n[stopped at {max_matches} matches]"
	return "\n".join(hits) or f"(no matches for /{pattern}/)"

	def write_file(self, path: str, content: str) -> str:
	p = self._resolve(path)
	data = (content or "").encode("utf-8")
	if len(data) > MAX_WRITE_BYTES:
	return f"ERROR: content too large ({len(data):,} bytes)"
	p.parent.mkdir(parents=True, exist_ok=True)
	p.write_text(content or "", encoding="utf-8")
	return f"wrote {self.rel(p)} ({len(data):,} bytes)"

	def run_r(self, path: str) -> str:
	p = self._resolve(path)
	if not p.is_file():
	return f"ERROR: no such file: {path}"
	exe = shutil.which("Rscript")
	if not exe:
	return (
	"ERROR: Rscript is not available in this environment, so the R "
	"script could not be executed. Continue without running it."
	)
	try:
	proc = subprocess.run(
	[exe, "--vanilla", str(p)],
	capture_output=True,
	text=True,
	timeout=R_TIMEOUT_SEC,
	cwd=str(self.root),
	)
	except subprocess.TimeoutExpired:
	return f"ERROR: Rscript timed out after {R_TIMEOUT_SEC}s"
	parts = [f"exit code: {proc.returncode}"]
	if proc.stdout.strip():
	parts.append("--- stdout ---\n" + proc.stdout.strip()[:6000])
	if proc.stderr.strip():
	parts.append("--- stderr ---\n" + proc.stderr.strip()[:4000])
	return "\n".join(parts)

	# -- dispatch ---------------------------------------------------------
	def call(self, name: str, args: Dict[str, Any]) -> str:
	try:
	if name == "list_files":
	return self.list_files()
	if name == "read_file":
	return self.read_file(
	args.get("path", ""), args.get("offset", 1), args.get("limit", MAX_READ_LINES)
	)
	if name == "grep":
	return self.grep(
	args.get("pattern", ""), args.get("path", ""), args.get("max_matches", MAX_GREP_MATCHES)
	)
	if name == "write_file":
	return self.write_file(args.get("path", ""), args.get("content", ""))
	if name == "run_r":
	return self.run_r(args.get("path", "output.R"))
	return f"ERROR: unknown tool {name}"
	except Exception as e:
	return f"ERROR: {type(e).__name__}: {e}"


	# ---------------------------------------------------------------------------
	# Tool schemas
	# ---------------------------------------------------------------------------

	_TOOLS: List[Dict[str, Any]] = [
	{
	"name": "list_files",
	"description": "List all files in the workspace with their sizes.",
	"properties": {},
	"required": [],
	},
	{
	"name": "read_file",
	"description": (
	"Read a slice of a text file with line numbers. The SAP is large, so "
	"read it in chunks (or grep first to find the right lines)."
	),
	"properties": {
	"path": {"type": "string", "description": "Workspace-relative path, e.g. sap.txt"},
	"offset": {"type": "integer", "description": "1-based first line to read (default 1)"},
	"limit": {"type": "integer", "description": f"Lines to read (max {MAX_READ_LINES})"},
	},
	"required": ["path"],
	},
	{
	"name": "grep",
	"description": (
	"Case-insensitive regex search across the workspace (or one file). "
	"Returns path:line and the matching text — use it to locate sections, "
	"then read_file around those line numbers."
	),
	"properties": {
	"pattern": {"type": "string", "description": "Regular expression"},
	"path": {"type": "string", "description": "Optional single file to search"},
	"max_matches": {"type": "integer", "description": f"Max matches (default {MAX_GREP_MATCHES})"},
	},
	"required": ["pattern"],
	},
	{
	"name": "write_file",
	"description": "Create or overwrite a file in the workspace with the given content.",
	"properties": {
	"path": {"type": "string", "description": "Workspace-relative path, e.g. output.json"},
	"content": {"type": "string", "description": "Full file content"},
	},
	"required": ["path", "content"],
	},
	{
	"name": "run_r",
	"description": (
	"Execute an R script with Rscript and return its exit code, stdout and "
	"stderr. Use it to verify output.R runs and prints the values you "
	"reported; fix the script and re-run if it errors or disagrees."
	),
	"properties": {
	"path": {"type": "string", "description": "Script path (default output.R)"},
	},
	"required": [],
	},
	]


	def anthropic_tools() -> List[Dict[str, Any]]:
	return [
	{
	"name": t["name"],
	"description": t["description"],
	"input_schema": {
	"type": "object",
	"properties": t["properties"],
	"required": t["required"],
	},
	}
	for t in _TOOLS
	]


	def openai_tools() -> List[Dict[str, Any]]:
	return [
	{
	"type": "function",
	"function": {
	"name": t["name"],
	"description": t["description"],
	"parameters": {
	"type": "object",
	"properties": t["properties"],
	"required": t["required"],
	},
	},
	}
	for t in _TOOLS
	]