scripts/oakbench/classify.py 148 lines · 5.4 KB · python
1
"""Command-intent classification shared by transcript analysis and metrics."""
2
3
from __future__ import annotations
4
5
import re
6
import shlex
7
8
UNKNOWN_COMMAND_PATTERN = re.compile(
9
    r"(unknown (sub)?command|unrecognized (option|subcommand|argument)|no such (sub)?command"
10
    r"|is not a git command|unexpected argument|invalid option|unknown option)",
11
    re.IGNORECASE,
12
)
13
14
15
def unwrap_shell_command(command: str) -> str:
16
    try:
17
        parts = shlex.split(command)
18
    except ValueError:
19
        return command
20
    if len(parts) >= 3 and parts[0] in {"/bin/zsh", "zsh", "/bin/bash", "bash", "sh"} and parts[1] in {"-c", "-lc"}:
21
        return parts[2]
22
    return command
23
24
25
# Read-only VCS interrogation commands. A successful info command immediately
26
# re-asked in the same family ("status, then status with more flags") means
27
# the first answer did not carry enough information to act on — the cost of
28
# compact output shows up here, not in the bytes of the first call.
29
VCS_INFO_FAMILIES = ("status", "diff", "log", "show", "blame", "branch")
30
# Harnesses and agents routinely invoke VCS binaries by absolute path
31
# (/usr/bin/git, ../oak/target/release/oak); classification must not depend
32
# on how the binary was addressed.
33
_VCS_SUBCOMMAND = re.compile(r"(^|[;&|]\s*|\s)(\S*/)?(git|oak)\s+([\w][\w-]*)")
34
35
36
def _shell_segment_end(command: str, start: int) -> int:
37
    quote: str | None = None
38
    escaped = False
39
    for index in range(start, len(command)):
40
        char = command[index]
41
        if escaped:
42
            escaped = False
43
            continue
44
        if char == "\\":
45
            escaped = True
46
            continue
47
        if quote is not None:
48
            if char == quote:
49
                quote = None
50
            continue
51
        if char in {"'", '"'}:
52
            quote = char
53
            continue
54
        if char in {";", "\n", "|"}:
55
            return index
56
        if char == "&":
57
            return index
58
    return len(command)
59
60
61
def _quoted_offsets(command: str) -> set[int]:
62
    quoted: set[int] = set()
63
    quote: str | None = None
64
    escaped = False
65
    for index, char in enumerate(command):
66
        if escaped:
67
            if quote is not None:
68
                quoted.add(index)
69
            escaped = False
70
            continue
71
        if char == "\\":
72
            if quote is not None:
73
                quoted.add(index)
74
            escaped = True
75
            continue
76
        if quote is not None:
77
            quoted.add(index)
78
            if char == quote:
79
                quote = None
80
            continue
81
        if char in {"'", '"'}:
82
            quote = char
83
            quoted.add(index)
84
    return quoted
85
86
87
def vcs_invocations(command: str) -> list[tuple[str, str]]:
88
    """VCS invocation shell segment and subcommand for every VCS call in a chain."""
89
    stripped = unwrap_shell_command(command).strip()
90
    quoted = _quoted_offsets(stripped)
91
    invocations: list[tuple[str, str]] = []
92
    for match in _VCS_SUBCOMMAND.finditer(stripped):
93
        binary_start = match.start(2) if match.group(2) is not None else match.start(3)
94
        if binary_start in quoted:
95
            continue
96
        segment_start = match.start(2) if match.group(2) is not None else match.start(3)
97
        segment = stripped[segment_start:_shell_segment_end(stripped, segment_start)].strip()
98
        invocations.append((segment, match.group(4)))
99
    return invocations
100
101
102
def vcs_subcommand(command: str) -> str | None:
103
    """The VCS subcommand of an invocation (any path form), or None if not a VCS call.
104
105
    The single classification authority: every consumer that needs to know
106
    "which VCS operation was this" goes through here, so transcript
107
    sub-counts can never disagree with classify_command about what is a VCS
108
    call (the reviewer-found vcs_commands_total=0 vs status_commands=1 rows).
109
    """
110
    invocations = vcs_invocations(command)
111
    return invocations[0][1] if invocations else None
112
113
114
def vcs_subcommands(command: str) -> list[str]:
115
    """All VCS subcommands in a shell command chain.
116
117
    `git add . && git commit ...` is one terminal tool call but two VCS
118
    invocations; snapshot metrics need to see the commit rather than stopping
119
    at the first `add`.
120
    """
121
    return [subcommand for _, subcommand in vcs_invocations(command)]
122
123
124
def vcs_info_family(command: str) -> str | None:
125
    """The info-command family of a VCS invocation, or None if not read-only info."""
126
    subcommand = vcs_subcommand(command)
127
    return subcommand if subcommand in VCS_INFO_FAMILIES else None
128
129
130
def classify_command(command: str) -> set[str]:
131
    kinds: set[str] = {"terminal"}
132
    stripped = unwrap_shell_command(command).strip()
133
    if vcs_subcommands(stripped):
134
        kinds.add("vcs")
135
    # validate[_.]\w* (not the old validate_\b) so validator scripts like
136
    # validate_configs.py classify as test commands; \b cannot follow "_".
137
    if re.search(r"\b(unittest|pytest|cargo test|npm test|validate[_.]\w*|test)\b", stripped):
138
        kinds.add("test")
139
    if re.search(r"(^|[;&|]\s*)(\S*/)?(rg|grep|find)\b", stripped):
140
        kinds.add("search")
141
    if re.search(r"(^|[;&|]\s*)(\S*/)?(cat|sed|head|tail|less)\b|(^|[;&|]\s*)python3? - <<", stripped):
142
        kinds.add("read")
143
    if re.search(r"\b(write_text|replace\(|apply_patch|python3? - <<|perl -pi|ruby -pi)\b", stripped):
144
        kinds.add("edit")
145
        kinds.add("file_write")
146
    if re.search(r"(\s--help\b|^(\S*/)?(git|oak)\s+help\b|^man\s|^tldr\s|^help\b)", stripped):
147
        kinds.add("help")
148
    return kinds