agent-lab/scripts/ws_smoke.py

import argparse
import asyncio
import json
from copy import deepcopy
from typing import Any

from agent_lab.application.contracts import DEFAULT_EXTRA_BODY


DEFAULT_URL = "ws://127.0.0.1:8000/ws/debug"
DEFAULT_MESSAGE = "Debug this agent handoff path."
DEFAULT_MODEL = ""


def build_payload(
    message: str,
    model: str,
    handoff_note_enabled: bool,
) -> dict[str, Any]:
    return {
        "user_message": message,
        "system_prompts": [
            "You are a chat agent running inside Agent Lab. "
            "Stream a concise answer and call handoff_note when an event handoff helps."
        ],
        "pre_messages": [],
        "chat_agent": {
            "model": model,
            "temperature": 0.2,
            "max_tokens": 1024,
            "system_prompt": "",
            "extra_body": deepcopy(DEFAULT_EXTRA_BODY),
        },
        "event_agent": {
            "model": model,
            "temperature": 0.2,
            "max_tokens": 1024,
            "system_prompt": "",
            "extra_body": deepcopy(DEFAULT_EXTRA_BODY),
            "enabled_tools": ["handoff_note"] if handoff_note_enabled else [],
            "max_event_loops": 1,
        },
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Send one Agent Lab WebSocket debug request and print streamed JSON."
    )
    parser.add_argument("--url", default=DEFAULT_URL, help=f"WebSocket URL. Default: {DEFAULT_URL}")
    parser.add_argument("--message", default=DEFAULT_MESSAGE, help="User message to send.")
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help="Chat model name. Defaults to the backend configured model.",
    )

    tool_group = parser.add_mutually_exclusive_group()
    tool_group.add_argument(
        "--enable-handoff-note",
        dest="handoff_note_enabled",
        action="store_true",
        default=True,
        help="Enable the handoff_note tool. This is the default.",
    )
    tool_group.add_argument(
        "--disable-handoff-note",
        dest="handoff_note_enabled",
        action="store_false",
        help="Disable the handoff_note tool.",
    )
    return parser


async def run_smoke(url: str, payload: dict[str, Any]) -> None:
    import websockets

    async with websockets.connect(url) as websocket:
        await websocket.send(_compact_json(payload))
        async for raw_message in websocket:
            try:
                message = json.loads(raw_message)
            except json.JSONDecodeError:
                print(raw_message, flush=True)
                continue

            print(_compact_json(message), flush=True)
            if message.get("type") in {"done", "error"}:
                break


def _compact_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def main() -> int:
    args = build_parser().parse_args()
    payload = build_payload(
        message=args.message,
        model=args.model,
        handoff_note_enabled=args.handoff_note_enabled,
    )
    asyncio.run(run_smoke(args.url, payload))
    return 0


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