import argparse import asyncio import logging import json import sys from pathlib import Path from urllib.parse import urlencode import aiohttp import websockets def build_ws_url(base_url: str, api_key: str, username: str) -> str: normalized = base_url.rstrip("/") if normalized.startswith("https://"): ws_base = "wss://" + normalized.removeprefix("https://").removeprefix("wss://") elif normalized.startswith("http://"): ws_base = "ws://" + normalized.removeprefix("http://").removeprefix("wss://") else: ws_base = f"ws://{normalized}" query = urlencode( { "api_key": api_key, "username": username, "ct": "live", } ) return f"{ws_base}/api/v1/live/ws?{query}" def build_headers(api_key: str) -> dict[str, str]: return {"X-API-Key": api_key} def create_logger(log_file: str | None) -> logging.Logger: logger = logging.getLogger("live_upload_flow") if logger.handlers: return logger logger.setLevel(logging.INFO) formatter = logging.Formatter( "%(asctime)s | %(levelname)s | %(name)s | %(message)s" ) console_handler = logging.StreamHandler(sys.stdout) console_handler.setLevel(logging.INFO) console_handler.setFormatter(formatter) logger.addHandler(console_handler) if log_file: file_handler = logging.FileHandler(log_file, encoding="utf-8") file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) logger.propagate = False return logger async def upload_file(session: aiohttp.ClientSession, base_url: str, api_key: str, file_path: Path) -> str: logger = logging.getLogger("live_upload_flow") form = aiohttp.FormData() with file_path.open("rb") as file_handle: form.add_field( "file", file_handle, filename=file_path.name, content_type="application/octet-stream", ) async with session.post( f"{base_url.rstrip('/')}/api/v1/file", data=form, headers=build_headers(api_key), ) as resp: payload = await resp.json() logger.info( "[UPLOAD] status=%s, attachment_id=%s", payload.get("status"), payload.get("data", {}).get("attachment_id"), ) if payload.get("status") != "ok": raise RuntimeError(f"Upload failed: {payload}") attachment_id = payload["data"]["attachment_id"] logger.info("[UPLOAD] attachment_id=%s", attachment_id) return attachment_id async def get_file(session: aiohttp.ClientSession, base_url: str, api_key: str, attachment_id: str) -> bytes: logger = logging.getLogger("live_upload_flow") url = f"{base_url.rstrip('/')}/api/v1/file?{urlencode({'attachment_id': attachment_id})}" async with session.get(url, headers=build_headers(api_key)) as resp: logger.info("[GET] status=%s, content_type=%s", resp.status, resp.headers.get("Content-Type")) if resp.status != 200: text = await resp.text() raise RuntimeError(f"Failed to fetch attachment: {resp.status} {text}") return await resp.read() async def run_live_check( base_url: str, api_key: str, username: str, attachment_id: str | None, text: str, ) -> None: logger = logging.getLogger("live_upload_flow") ws_url = build_ws_url(base_url, api_key, username) message_parts = [] if attachment_id: message_parts.append({"type": "file", "attachment_id": attachment_id}) if text: message_parts.append({"type": "plain", "text": text}) if not message_parts: raise ValueError("text is empty and no attachment_id provided") message = { "t": "text_input", "message": message_parts, } async with websockets.connect(ws_url) as websocket: logger.info("[WS] connected: %s", ws_url) logger.info("[WS] send: %s", json.dumps(message, ensure_ascii=False)) await websocket.send(json.dumps(message)) try: while True: raw = await asyncio.wait_for(websocket.recv(), timeout=90) data = json.loads(raw) logger.info("[WS] recv: %s", json.dumps(data, ensure_ascii=False)) if data.get("t") == "end": break except asyncio.TimeoutError: logger.warning("[WS] timeout reached, stop collecting messages") async def main() -> None: parser = argparse.ArgumentParser(description="Upload file and test live mode input path.") parser.add_argument("--base-url", default="http://localhost:6185", help="Server base URL") parser.add_argument("--api-key", required=True, help="OpenAPI key") parser.add_argument("--username", default="alice", help="OpenAPI username") parser.add_argument( "--file", type=Path, required=False, default=None, help="Local file to upload", ) parser.add_argument( "--attachment-id", help="Existing attachment_id to use directly", default=None, ) parser.add_argument( "--text", default="Please analyze the uploaded file.", help="Additional text for live message", ) parser.add_argument( "--log-file", default=str(Path.cwd() / "test.log"), help="Write logs to this file in addition to terminal output", ) parser.add_argument( "--skip-download-check", action="store_true", help="Skip GET attachment content verification", ) args = parser.parse_args() create_logger(args.log_file) logger = logging.getLogger("live_upload_flow") if args.file is not None: if not args.file.exists(): raise FileNotFoundError(f"file not found: {args.file}") if not args.file.is_file(): raise ValueError(f"not a regular file: {args.file}") timeout = aiohttp.ClientTimeout(total=30) async with aiohttp.ClientSession(timeout=timeout) as session: attachment_id = args.attachment_id if args.file is not None: attachment_id = await upload_file( session, args.base_url, args.api_key, args.file, ) if not args.skip_download_check: content = await get_file( session, args.base_url, args.api_key, attachment_id ) logger.info("[GET] attachment size=%s bytes", len(content)) elif attachment_id: logger.info("[INPUT] using attachment_id=%s", attachment_id) if args.file is None and not attachment_id: logger.info("[INPUT] no attachment_id provided, using text message only") await run_live_check( args.base_url, args.api_key, args.username, attachment_id, args.text, ) if __name__ == "__main__": create_logger(None) try: asyncio.run(main()) except Exception as e: logging.getLogger("live_upload_flow").error("error: %s", e, exc_info=True) raise SystemExit(1)