ResearchHarness/agent_base/trace_utils.py at main · InternScience/ResearchHarness · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import argparse
import datetime
from pathlib import Path
from typing import Any, Callable, Optional
from uuid import uuid4

from agent_base.utils import append_jsonl, safe_jsonable


TRACE_FIELD_NAMES = [
    "run_id",
    "event_index",
    "turn_index",
    "timestamp",
    "model_name",
    "workspace_root",
    "role",
    "text",
    "tool_call_ids",
    "tool_names",
    "tool_arguments",
    "finish_reason",
    "termination",
    "error",
    "image_paths",
    "capture_type",
    "payload",
]


class FlatTraceWriter:
    def __init__(
        self,
        *,
        trace_dir: Optional[str | Path],
        model_name: str,
        workspace_root: str | Path,
        on_event: Optional[Callable[[dict[str, Any]], None]] = None,
    ):
        self.model_name = model_name
        self.workspace_root = str(workspace_root)
        self.on_event = on_event
        self.run_id = uuid4().hex
        self.path = resolve_trace_path(trace_dir, run_id=self.run_id) if trace_dir else None
        self.event_index = 0

    def append(
        self,
        *,
        role: str,
        text: str = "",
        turn_index: int = 0,
        tool_call_ids: Optional[list[str]] = None,
        tool_names: Optional[list[str]] = None,
        tool_arguments: Optional[list[Any]] = None,
        finish_reason: Optional[str] = None,
        termination: Optional[str] = None,
        error: Optional[str] = None,
        image_paths: Optional[list[str]] = None,
        capture_type: str = "",
        payload: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        self.event_index += 1
        row = {
            "run_id": self.run_id,
            "event_index": self.event_index,
            "turn_index": turn_index,
            "timestamp": datetime.datetime.now().astimezone().isoformat(timespec="seconds"),
            "model_name": self.model_name,
            "workspace_root": self.workspace_root,
            "role": role,
            "text": text,
            "tool_call_ids": tool_call_ids or [],
            "tool_names": tool_names or [],
            "tool_arguments": safe_jsonable(tool_arguments or []),
            "finish_reason": finish_reason or "",
            "termination": termination or "",
            "error": error or "",
            "image_paths": image_paths or [],
            "capture_type": capture_type or "",
            "payload": safe_jsonable(payload or {}),
        }
        if self.path is not None:
            append_jsonl(self.path, row)
        if self.on_event is not None:
            self.on_event(row)
        return row


def resolve_trace_path(
    trace_dir: str | Path,
    *,
    run_id: str,
    prefix: str = "trace",
    suffix: str = ".jsonl",
) -> Path:
    directory = Path(trace_dir)
    timestamp = datetime.datetime.now().astimezone().strftime("%Y%m%d_%H%M%S")
    short_run_id = run_id[:12]
    filename = f"{prefix}_{timestamp}_{short_run_id}{suffix}"
    return directory / filename


def main(argv: Optional[list[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Inspect the flat trace field order used by the agent.")
    parser.parse_args(argv)
    print("\n".join(TRACE_FIELD_NAMES))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())