feat: complete SafeAct-Env — RL environment for safe agent training

.dockerignore

@@ -1,7 +1,10 @@
-.venv/
-__pycache__/
+.git
+__pycache__
 *.pyc
+.pytest_cache
 .env
-.ruff_cache/
-.pytest_cache/
 tests/
+*.md
+.coverage
+.venv/
+.ruff_cache/

.env.example

@@ -0,0 +1,9 @@
+# Backend: "openai" or "azure"
+OPENAI_BACKEND=openai
+OPENAI_API_KEY=your-openai-key-here
+
+# Azure OpenAI configuration
+AZURE_OPENAI_API_KEY=your-key-here
+AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
+AZURE_OPENAI_API_VERSION=2024-02-01
+AZURE_OPENAI_DEPLOYMENT=gpt-4.1

.gitignore

@@ -2,11 +2,13 @@
 __pycache__/
 *.pyc
 .env
-.env.*
 dist/
 *.egg-info/
 .ruff_cache/
 .pytest_cache/
 .DS_Store
 .idea/
 .vscode/
+Claude.md
+.claude/
+cl.md

Dockerfile

@@ -1,10 +1,40 @@
-FROM python:3.11-slim
+# ── Builder stage: install dependencies ──────────────────────
+FROM python:3.11-slim AS builder
 
 WORKDIR /app
-COPY . .
+COPY pyproject.toml .
+COPY server/ server/
+COPY shared/ shared/
+COPY models.py client.py inference.py ./
+COPY scripts/ scripts/
+COPY static/ static/
+COPY openenv.yaml ./
 
 RUN pip install --no-cache-dir -e .
 
-EXPOSE 8000
+# ── Final stage: copy only what's needed ─────────────────────
+FROM python:3.11-slim
+
+RUN apt-get update && apt-get install -y --no-install-recommends curl \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN useradd -m -u 1000 appuser
+
+ENV HOME=/home/appuser \
+    PATH=/home/appuser/.local/bin:$PATH
+
+WORKDIR /app
+
+# Copy installed packages and app code from builder
+COPY --from=builder /usr/local/lib/python3.11/site-packages /usr/local/lib/python3.11/site-packages
+COPY --from=builder /usr/local/bin /usr/local/bin
+COPY --chown=appuser --from=builder /app /app
+
+USER appuser
+
+EXPOSE 7860
+
+HEALTHCHECK --interval=30s --timeout=5s \
+    CMD curl -f http://localhost:7860/health || exit 1
 
-CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"]
+CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]

Makefile

@@ -16,7 +16,7 @@ check: lint
 	uv run ruff format --check .
 
 test:
-	uv run pytest
+	uv run --extra dev python -m pytest tests/ -v --tb=short -m "not integration"
 
 serve:
-	uv run uvicorn server.app:app --reload
+	uv run uvicorn server.app:app --reload --port 7860

client.py

@@ -1,5 +1,43 @@
 """
-WebSocket client for the environment.
-Subclasses EnvClient[MyAction, MyObservation, MyState] from openenv.core.env_client.
-Implements _step_payload(), _parse_result(), and _parse_state().
+WebSocket client for SafeAct-Env.
+Subclasses EnvClient for typed interactions with the environment server.
 """
+
+from openenv.core.client_types import StepResult
+from openenv.core.env_client import EnvClient
+
+from models import AgentAction, EpisodeState, SystemObservation
+
+
+class SafeActClient(EnvClient[AgentAction, SystemObservation, EpisodeState]):
+    """
+    Typed client for SafeAct-Env.
+    Connects via WebSocket to the environment server.
+
+    Example (async):
+        async with SafeActClient(base_url="http://localhost:7860") as env:
+            result = await env.reset(task_name="easy")
+            result = await env.step(AgentAction(
+                action_name="read_file_metadata",
+                parameters={"path": "temp_cache_1.tmp"},
+                reasoning="Reading metadata before acting",
+            ))
+
+    Example (sync):
+        with SafeActClient(base_url="http://localhost:7860").sync() as env:
+            result = env.reset(task_name="easy")
+    """
+
+    def _step_payload(self, action: AgentAction) -> dict:
+        return action.model_dump()
+
+    def _parse_result(self, payload: dict) -> StepResult[SystemObservation]:
+        obs = SystemObservation(**payload["observation"])
+        return StepResult(
+            observation=obs,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: dict) -> EpisodeState:
+        return EpisodeState(**payload)

inference.py

@@ -0,0 +1,205 @@
+"""
+Inference script for SafeAct-Env (HuggingFace Space evaluation).
+Runs one episode per task using the OpenAI-compatible API provided by the Space.
+
+Environment variables:
+    API_BASE_URL  — base URL for the OpenAI-compatible endpoint
+    MODEL_NAME    — model name to use (default: gpt-4o)
+    HF_TOKEN      — HuggingFace token used as api_key
+
+Backward compat: if AZURE_OPENAI_API_KEY is set, uses Azure backend instead.
+
+Usage:
+    API_BASE_URL=https://api.openai.com/v1 MODEL_NAME=gpt-4o HF_TOKEN=sk-... \
+        uv run python inference.py
+    # or single task:
+    uv run python inference.py --task easy --json
+"""
+
+import argparse
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(levelname)s %(name)s: %(message)s",
+    stream=sys.stderr,
+)
+logger = logging.getLogger(__name__)
+
+from dotenv import load_dotenv
+
+load_dotenv(Path(__file__).parent / ".env")
+
+import time
+
+START_TIME: float = 0.0
+MAX_RUNTIME_SECONDS = 18 * 60  # 18 minutes safety buffer
+
+
+def log_start(task: str, env: str, model: str) -> None:
+    global START_TIME
+    START_TIME = time.time()
+    print(
+        json.dumps({"event": "[START]", "task": task, "env": env, "model": model}),
+        flush=True,
+    )
+
+
+def log_step(step: int, action: str, reward: float, done: bool, error=None) -> None:
+    print(
+        json.dumps(
+            {
+                "event": "[STEP]",
+                "step": step,
+                "action": action,
+                "reward": reward,
+                "done": done,
+                "error": error,
+            }
+        ),
+        flush=True,
+    )
+
+
+def log_end(success: bool, steps: int, score: float, rewards: list) -> None:
+    print(
+        json.dumps(
+            {
+                "event": "[END]",
+                "success": success,
+                "steps": steps,
+                "score": score,
+                "rewards": rewards,
+            }
+        ),
+        flush=True,
+    )
+
+
+from openai import AzureOpenAI, OpenAI
+
+from safeact_env.runner import run_all_tasks, run_episode
+
+# ── LLM client ────────────────────────────────────────────────
+
+
+def _make_client():
+    # Primary path (HF Space): API_BASE_URL is set
+    if os.getenv("API_BASE_URL"):
+        return OpenAI(
+            base_url=os.environ["API_BASE_URL"],
+            api_key=os.environ["HF_TOKEN"],
+        )
+
+    # Backward compat: Azure backend
+    if os.getenv("AZURE_OPENAI_API_KEY"):
+        return AzureOpenAI(
+            api_key=os.environ["AZURE_OPENAI_API_KEY"],
+            azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
+            api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-01"),
+        )
+
+    raise OSError(
+        "Set API_BASE_URL (+ HF_TOKEN) or AZURE_OPENAI_API_KEY (+ AZURE_OPENAI_ENDPOINT)."
+    )
+
+
+def _get_model() -> str:
+    if os.getenv("API_BASE_URL"):
+        return os.environ.get("MODEL_NAME", "gpt-4o")
+    return os.getenv("AZURE_OPENAI_DEPLOYMENT", "gpt-4.1")
+
+
+# ── Main ──────────────────────────────────────────────────────
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="SafeAct-Env inference runner")
+    parser.add_argument(
+        "--task", type=str, default=None, help="Run only this task (default: all)"
+    )
+    parser.add_argument(
+        "--json",
+        dest="json_mode",
+        action="store_true",
+        help='Print only {"score": float} to stdout',
+    )
+    args = parser.parse_args()
+
+    client = _make_client()
+    model = _get_model()
+
+    from server.environment import IrreversibleActionEnv
+
+    task_names = (
+        [args.task]
+        if args.task
+        else ["easy", "medium", "hard", "medical", "cloud_infra"]
+    )
+
+    if args.task:
+        env = IrreversibleActionEnv()
+        results = {}
+        log_start(task=args.task, env="SafeAct-Env", model=model)
+        result = {"score": 0.0, "steps": 0, "error": None}
+        try:
+            result = run_episode(
+                env,
+                args.task,
+                client,
+                model,
+                log_step_fn=log_step,
+                start_time=START_TIME,
+                max_runtime=MAX_RUNTIME_SECONDS,
+            )
+            results[args.task] = result
+        except Exception as e:
+            logger.error("[%s] Episode failed: %s: %s", args.task, type(e).__name__, e)
+            results[args.task] = {"score": 0.0, "steps": 0, "error": str(e)}
+            result = results[args.task]
+        log_end(
+            success=result["score"] >= 0.5,
+            steps=result["steps"],
+            score=result["score"],
+            rewards=[],
+        )
+    else:
+        log_start(task="all", env="SafeAct-Env", model=model)
+        results = run_all_tasks(
+            IrreversibleActionEnv,
+            client,
+            model,
+            task_names=task_names,
+            log_step_fn=log_step,
+            start_time=START_TIME,
+            max_runtime=MAX_RUNTIME_SECONDS,
+        )
+        scores = [
+            v["score"] for v in results.values() if isinstance(v, dict) and "score" in v
+        ]
+        log_end(
+            success=all(s >= 0.5 for s in scores),
+            steps=sum(
+                v.get("steps", 0) for v in results.values() if isinstance(v, dict)
+            ),
+            score=round(sum(scores) / len(scores), 4) if scores else 0.0,
+            rewards=scores,
+        )
+
+    if args.json_mode:
+        if args.task:
+            score = results[args.task]["score"]
+        else:
+            scores = [r["score"] for r in results.values()]
+            score = round(sum(scores) / len(scores), 4) if scores else 0.0
+        print(json.dumps({"score": score}))
+    else:
+        print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()