133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""repl_drive.py - 用 pty 驱动 ligent 的 REPL, 跑多轮对话并打印全部输出.
|
|
|
|
交互式行为 (多轮上下文、/命令、行编辑)没法用管道测 —— 管道模式会把
|
|
整段 stdin 当成一个任务. 所以这里开一个伪终端, 逐条发送提示词.
|
|
|
|
用法：
|
|
repl_drive.py --bin ./ligent --base-url URL --prompt "第一轮" --prompt "第二轮"
|
|
"""
|
|
import argparse
|
|
import os
|
|
import pty
|
|
import re
|
|
import select
|
|
import sys
|
|
import time
|
|
|
|
|
|
def drive(binary, args, prompts, wait_per_prompt=2.0, quiet_after=1.5,
|
|
raw_keys=None):
|
|
pid, fd = pty.fork()
|
|
if pid == 0:
|
|
os.execv(binary, [binary] + args)
|
|
os._exit(127)
|
|
|
|
out = bytearray()
|
|
# 行编辑器每次回到提示符都会打开括号粘贴, 所以这个序列就是"轮到用户输入了"
|
|
# 的信号. 用它做同步点, 比"等输出安静下来"可靠也快得多.
|
|
PROMPT_READY = b"\x1b[?2004h"
|
|
ready_seen = 0
|
|
tail = b""
|
|
|
|
def note(chunk):
|
|
# 标记可能被 read 切断, 所以要连着上一块的尾巴一起找
|
|
nonlocal tail, ready_seen
|
|
buf = tail + chunk
|
|
ready_seen += buf.count(PROMPT_READY)
|
|
tail = buf[-(len(PROMPT_READY) - 1):]
|
|
|
|
def drain(timeout, idle=0.5, until_ready=False):
|
|
"""等到"回到提示符" (until_ready), 或者输出彻底安静下来.
|
|
|
|
必须按标记等, 不能只按安静等：等待动画每 250ms 画一帧, pty 永远
|
|
不会安静超过 idle, 于是每次都要白等满 timeout —— 而 mock 其实几
|
|
毫秒就答复了. 这是整个 e2e 从 126 秒掉下来的关键.
|
|
"""
|
|
want = ready_seen + 1 if until_ready else None
|
|
deadline = time.time() + timeout
|
|
last = time.time()
|
|
while time.time() < deadline:
|
|
r, _, _ = select.select([fd], [], [], 0.02)
|
|
if r:
|
|
try:
|
|
chunk = os.read(fd, 65536)
|
|
except OSError:
|
|
return False
|
|
if not chunk:
|
|
return False
|
|
out.extend(chunk)
|
|
note(chunk)
|
|
last = time.time()
|
|
if want is not None and ready_seen >= want:
|
|
return True
|
|
elif time.time() - last > idle:
|
|
return True
|
|
return True
|
|
|
|
drain(quiet_after, until_ready=True) # 等横幅画完、第一次回到提示符
|
|
if raw_keys is not None:
|
|
# 逐字节发送原始按键序列 (用于测方向键等编辑行为)
|
|
for chunk, pause in raw_keys:
|
|
os.write(fd, chunk if isinstance(chunk, bytes) else chunk.encode("utf-8"))
|
|
drain(pause)
|
|
os.write(fd, b"/exit\n")
|
|
drain(0.8)
|
|
try:
|
|
os.close(fd)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
os.waitpid(pid, 0)
|
|
except ChildProcessError:
|
|
pass
|
|
return re.sub(r"\x1b\[[0-9;]*m", "", out.decode("utf-8", "replace"))
|
|
|
|
for p in prompts:
|
|
os.write(fd, (p + "\n").encode("utf-8"))
|
|
if not drain(wait_per_prompt, until_ready=True):
|
|
break
|
|
os.write(fd, b"/exit\n")
|
|
drain(0.8)
|
|
|
|
try:
|
|
os.close(fd)
|
|
except OSError:
|
|
pass
|
|
try:
|
|
os.waitpid(pid, 0)
|
|
except ChildProcessError:
|
|
pass
|
|
|
|
text = out.decode("utf-8", "replace")
|
|
return re.sub(r"\x1b\[[0-9;]*m", "", text) # 去掉 ANSI 样式
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--bin", default="./ligent")
|
|
ap.add_argument("--base-url", required=True)
|
|
ap.add_argument("--api-key", default="test")
|
|
ap.add_argument("--prompt", action="append", default=[])
|
|
ap.add_argument("--extra", action="append", default=[])
|
|
ap.add_argument("--wait", type=float, default=2.0)
|
|
ap.add_argument("--keys", action="append", default=[],
|
|
help="原始按键序列 (支持 \\x1b 转义), 按顺序发送")
|
|
a = ap.parse_args()
|
|
|
|
# --no-history：被驱动的实例不该往用户真实的 ligent.history 里写测试垃圾.
|
|
# 需要历史的调用方用 --extra=--history PATH (排在后面, 会覆盖这一条).
|
|
args = ["--api-key", a.api_key, "--base-url", a.base_url, "--no-color",
|
|
"--no-confirm", "--no-auto-save", "--no-history"] + a.extra
|
|
if a.keys:
|
|
raw = [(k.encode("utf-8").decode("unicode_escape").encode("latin-1"), 0.25)
|
|
for k in a.keys]
|
|
sys.stdout.write(drive(a.bin, args, [], a.wait, raw_keys=raw))
|
|
else:
|
|
sys.stdout.write(drive(a.bin, args, a.prompt, a.wait))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|