ResearchHarness/tests/test_console_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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#!/usr/bin/env python3

import contextlib
import io
import json
import os
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from agent_base.console_utils import ConsoleEventPrinter
from agent_base.console_utils import _display_width
from test_support import TEST_RUNS_DIR


@dataclass
class ConsoleUtilsResult:
    status: str
    detail: str
    output_preview: str


def main() -> int:
    old_force_color = os.environ.get("FORCE_COLOR")
    old_no_color = os.environ.get("NO_COLOR")
    old_columns = os.environ.get("COLUMNS")
    os.environ["FORCE_COLOR"] = "1"
    os.environ.pop("NO_COLOR", None)
    os.environ["COLUMNS"] = "110"
    try:
        buffer = io.StringIO()
        with contextlib.redirect_stdout(buffer):
            printer = ConsoleEventPrinter(
                model_name="demo-model",
                workspace_root=TEST_RUNS_DIR / "console_utils" / "demo",
                prompt="Summarize **this** task.",
            )
            printer.print_header()
            printer.handle_event(
                {
                    "role": "assistant",
                    "turn_index": 1,
                    "text": "",
                    "tool_names": ["Read"],
                    "tool_arguments": [{"path": "demo.txt"}],
                    "finish_reason": "tool_calls",
                }
            )
            printer.handle_event(
                {
                    "role": "tool",
                    "turn_index": 1,
                    "text": "demo output",
                    "tool_names": ["Read"],
                }
            )
            printer.handle_event(
                {
                    "role": "runtime",
                    "turn_index": 2,
                    "text": "done",
                    "error": "sample error",
                }
            )
        output = buffer.getvalue()

        cjk_buffer = io.StringIO()
        os.environ.pop("FORCE_COLOR", None)
        os.environ["NO_COLOR"] = "1"
        with contextlib.redirect_stdout(cjk_buffer):
            cjk_printer = ConsoleEventPrinter(
                model_name="gpt-5.4",
                workspace_root=TEST_RUNS_DIR / "console_utils" / "cjk",
                prompt="Who proposed the transformer architecture, and in what year was the paper published?",
            )
            cjk_printer.print_header()
            cjk_printer.handle_event(
                {
                    "role": "assistant",
                    "turn_index": 1,
                    "text": (
                        "Transformer 架构是由 Ashish Vaswani、Noam Shazeer、Niki Parmar、Jakob Uszkoreit、"
                        "Llion Jones、Aidan N. Gomez、Lukasz Kaiser 和 Illia Polosukhin 在论文"
                        "《Attention Is All You Need》中提出的，发表于 2017 年。"
                    ),
                }
            )
        cjk_output = cjk_buffer.getvalue()
    finally:
        if old_force_color is None:
            os.environ.pop("FORCE_COLOR", None)
        else:
            os.environ["FORCE_COLOR"] = old_force_color
        if old_no_color is None:
            os.environ.pop("NO_COLOR", None)
        else:
            os.environ["NO_COLOR"] = old_no_color
        if old_columns is None:
            os.environ.pop("COLUMNS", None)
        else:
            os.environ["COLUMNS"] = old_columns

    boxed_cjk_lines = [line for line in cjk_output.splitlines() if line.startswith(("+", "|"))]
    cjk_line_widths = {_display_width(line) for line in boxed_cjk_lines}

    checks = [
        "+ ResearchHarness CLI" in output,
        "+ Assistant | round 1" in output,
        "+ Read Result | round 1" in output,
        "+ Runtime | round 2" in output,
        "Assistant Tool Calls:" in output,
        "demo output" in output,
        "Runtime Error: sample error" in output,
        "== Round" not in output,
        "[Assistant]" not in output,
        "\033[" in output,
        "Transformer 架构是由" in cjk_output,
        len(cjk_line_widths) == 1,
    ]
    result = ConsoleUtilsResult(
        status="PASS" if all(checks) else "FAIL",
        detail="CLI console formatter uses boxed colored steps.",
        output_preview=(output + "\n" + cjk_output)[:1600],
    )
    print(json.dumps(asdict(result), ensure_ascii=False, indent=2))
    return 0 if result.status == "PASS" else 1


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