normalize.py

"""
Pre-process a trajectory that alternates
    {"thought": "...", "action": "<bash command>"}
    {"observation": "<stdout / stderr text>"}

The only network identifier that can appear is an IPv4 address,
optionally as part of   ssh user@IP   commands.

Output (per step):
    {
        "cwd"        : "<ROOT>/…",
        "stdout_prev": "...",
        "stderr_prev": "...",
        "cmd"        : "<canonicalised command>"
    }
"""

from __future__ import annotations
import json
import os
import re
from pathlib import Path
from typing import Dict, Generator, List

# ──────────────────────────────────────────────
# CONFIG
# ──────────────────────────────────────────────
MAX_STDOUT = 10000
MAX_STDERR = 10000

REPO_ROOT_PLACEHOLDER = "<ROOT>"
HOME_PLACEHOLDER      = "<HOME>"
TMP_PLACEHOLDER       = "<TMP>"
SHA_PLACEHOLDER       = "<SHA>"
DATE_PLACEHOLDER      = "<DATE>"
IP_PLACEHOLDER        = "<IP>"
SSH_PLACEHOLDER       = "ssh <USER>@<IP>"

# absolute path of the repo that was used when recording the demos
# set to "" to disable
DEMO_REPO_ABS_PATH = str(Path("/your/demo/repo").resolve())  

# ──────────────────────────────────────────────
# REGEXES
# ──────────────────────────────────────────────
RE_HOME  = re.compile(r"/home/[^/\s]+")
RE_TMP   = re.compile(r"/tmp/[^/\s]+")
RE_SHA   = re.compile(r"\b[0-9a-f]{7,40}\b")
RE_DATE  = re.compile(r"\b\d{4}-\d{2}-\d{2}\b")

# IPv4  (simplified: 1‒3 digits each byte)
RE_IPV4  = re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b")

# ssh user@IP
RE_SSH   = re.compile(r"\bssh\s+([^\s@]+)@((?:\d{1,3}\.){3}\d{1,3})")

# ──────────────────────────────────────────────
def canonicalise(text: str) -> str:
    """Replace host-specific and machine-specific artefacts with placeholders."""
    if not text:
        return text

    text = text.replace("\r\n", "\n")

    if DEMO_REPO_ABS_PATH:
        text = text.replace(DEMO_REPO_ABS_PATH, REPO_ROOT_PLACEHOLDER)

    text = RE_HOME.sub(HOME_PLACEHOLDER, text)
    text = RE_TMP.sub(TMP_PLACEHOLDER, text)
    text = RE_SHA.sub(SHA_PLACEHOLDER, text)
    text = RE_DATE.sub(DATE_PLACEHOLDER, text)

    # ssh user@IP   → ssh <USER>@<IP>
    text = RE_SSH.sub(SSH_PLACEHOLDER, text)

    # plain IP addresses
    text = RE_IPV4.sub(IP_PLACEHOLDER, text)

    return text


def norm_cd(current: str, arg: str) -> str:
    """Cheap string-only resolution for a cd command, no real IO."""
    arg = arg.strip()
    if not arg or arg == "-":
        return current
    if arg.startswith("~"):
        arg = HOME_PLACEHOLDER + arg[1:]
    p = Path(arg) if Path(arg).is_absolute() else Path(current) / arg
    try:
        return str(p.resolve())
    except Exception:
        return str(p)


def trim_tail(s: str, limit: int) -> str:
    return s if len(s) <= limit else s[-limit:]


# ──────────────────────────────────────────────
def iter_steps(traj: List[Dict[str, str]]):
    """Yield cleaned steps one by one."""
    cwd = REPO_ROOT_PLACEHOLDER
    stdout_prev, stderr_prev = "", ""

    it = iter(traj)
    for action_rec, obs_rec in zip(it, it):
        cmd         = action_rec.get("action", "")
        raw_out     = obs_rec.get("observation", "")

        # rudimentary stdout / stderr split; change if you have a delimiter
        if "\nSTDERR:\n" in raw_out:
            stdout_now, _, stderr_now = raw_out.partition("\nSTDERR:\n")
        else:
            stdout_now, stderr_now = raw_out, ""

        yield {
            "cwd"        : canonicalise(cwd),
            "stdout_prev": canonicalise(trim_tail(stdout_prev, MAX_STDOUT)),
            "stderr_prev": canonicalise(trim_tail(stderr_prev, MAX_STDERR)),
            "cmd"        : canonicalise(cmd.strip()),
        }

        # update synthetic state for next step
        if cmd.lstrip().startswith("cd "):
            cwd = norm_cd(cwd, cmd.lstrip()[3:])
        stdout_prev, stderr_prev = stdout_now, stderr_now


def process_trajectory(
    traj: List[Dict[str, str]],
    out_path: str | os.PathLike | None = None,
) -> List[Dict[str, str]]:
    processed = list(iter_steps(traj))
    if out_path:
        with open(out_path, "w", encoding="utf-8") as fh:
            for rec in processed:
                fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
    return processed


# ──────────────────────────────────────────────
if __name__ == "__main__":
    demo_traj = [
        {"thought": "connect", "action": "ssh alice@172.18.0.3"},
        {"observation": "Welcome to Ubuntu\n"},
        {"thought": "run tests", "action": "pytest -q"},
        {"observation": "F..\n=== FAILURES ===\nE AssertionError: ..."},
    ]
    steps = process_trajectory(demo_traj)
    print(json.dumps(steps, indent=2))