scripts/generate_traffic.py

"""Replay sampled task inputs through a (deployed) contimp-app to generate traffic.

    uv run python scripts/generate_traffic.py --target https://... --passcode X \
        --task pr-area --n 25

Traces land in LangFuse tagged `synthetic` with user `traffic-bot`, so they can
be segmented from human traffic.
"""

import argparse
import concurrent.futures
import json
import urllib.request
import uuid


def call(target: str, passcode: str, path: str, payload: dict | None = None) -> dict:
    req = urllib.request.Request(
        f"{target.rstrip('/')}{path}",
        data=json.dumps(payload).encode() if payload is not None else b"",
        headers={"Content-Type": "application/json", "X-Contimp-Passcode": passcode},
        method="POST",
    )
    return json.loads(urllib.request.urlopen(req, timeout=300).read())


def one_run(target: str, passcode: str, task: str, session_id: str) -> str:
    sample = call(target, passcode, f"/api/tasks/{task}/sample")
    result = call(target, passcode, f"/api/tasks/{task}/run", {
        "text": sample["text"],
        "input_id": sample["input_id"],
        "user": "traffic-bot",
        "session_id": session_id,
        "source": "synthetic",
    })
    return f"{sample['input_id']}: {result['scores']}"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--target", default="http://localhost:7860")
    parser.add_argument("--passcode", default="")
    parser.add_argument("--task", required=True)
    parser.add_argument("--n", type=int, default=10)
    parser.add_argument("--concurrency", type=int, default=3)
    args = parser.parse_args()

    session_id = f"traffic-{uuid.uuid4().hex[:8]}"
    ok = failed = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.concurrency) as pool:
        futures = [
            pool.submit(one_run, args.target, args.passcode, args.task, session_id)
            for _ in range(args.n)
        ]
        for future in concurrent.futures.as_completed(futures):
            try:
                print(" ", future.result())
                ok += 1
            except Exception as e:  # noqa: BLE001 - keep the batch going
                print("  run failed:", e)
                failed += 1
    print(f"done: {ok} ok, {failed} failed (session {session_id})")


if __name__ == "__main__":
    main()