补充设备语音归档与联调脚本
This commit is contained in:
stu2not
348
talkingq-url/scripts/simulate_device_voice_exchange.py
Normal file
348
talkingq-url/scripts/simulate_device_voice_exchange.py
Normal file
@@ -0,0 +1,348 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import struct
|
||||
import sys
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import pymysql
|
||||
import websockets
|
||||
from pymysql.cursors import DictCursor
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
if str(PROJECT_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
|
||||
from config import settings
|
||||
|
||||
|
||||
DEFAULT_WS_URL = "ws://127.0.0.1:8080/ws"
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeviceContext:
|
||||
device_id: str
|
||||
serial_number: str
|
||||
child_id: int
|
||||
card_uuid: str
|
||||
inbox: asyncio.Queue[str] = field(default_factory=asyncio.Queue)
|
||||
websocket: websockets.WebSocketClientProtocol | None = None
|
||||
receiver_task: asyncio.Task | None = None
|
||||
|
||||
|
||||
def build_runtime_args() -> argparse.Namespace:
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
audio_dir = repo_root / "tmp_voice_sim"
|
||||
|
||||
parser = argparse.ArgumentParser(description="Simulate real websocket voice exchange between two devices")
|
||||
parser.add_argument("--ws-url", default=DEFAULT_WS_URL)
|
||||
parser.add_argument("--device-a", default="TalkingQ_device001")
|
||||
parser.add_argument("--device-b", default="TalkingQ_device002")
|
||||
parser.add_argument("--a1-audio", default=str(audio_dir / "device001_to_device002_1.mp3"))
|
||||
parser.add_argument("--b1-audio", default=str(audio_dir / "device002_to_device001_1.mp3"))
|
||||
parser.add_argument("--a2-audio", default=str(audio_dir / "device001_to_device002_2.mp3"))
|
||||
parser.add_argument("--b2-audio", default=str(audio_dir / "device002_to_device001_2.mp3"))
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def connect_db():
|
||||
return pymysql.connect(
|
||||
host=settings.db_host,
|
||||
port=settings.db_port,
|
||||
user=settings.db_user,
|
||||
password=settings.db_password,
|
||||
database=settings.db_name,
|
||||
charset="utf8mb4",
|
||||
cursorclass=DictCursor,
|
||||
autocommit=True,
|
||||
)
|
||||
|
||||
|
||||
def load_device_contexts(device_ids: list[str]) -> dict[str, DeviceContext]:
|
||||
placeholders = ", ".join(["%s"] * len(device_ids))
|
||||
sql = f"""
|
||||
SELECT
|
||||
da.device_id,
|
||||
da.serial_number,
|
||||
db.child_id,
|
||||
c.card_uuid
|
||||
FROM device_auth AS da
|
||||
LEFT JOIN device_bindings AS db
|
||||
ON db.device_id = da.device_id
|
||||
AND db.status = 1
|
||||
LEFT JOIN cards AS c
|
||||
ON c.device_id = da.device_id
|
||||
AND c.status = 1
|
||||
WHERE da.device_id IN ({placeholders})
|
||||
AND da.is_active = 1
|
||||
"""
|
||||
|
||||
contexts: dict[str, DeviceContext] = {}
|
||||
with connect_db() as connection:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(sql, device_ids)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
for row in rows:
|
||||
if row["child_id"] is None:
|
||||
raise RuntimeError(f"device {row['device_id']} is not bound to any child")
|
||||
if not row["card_uuid"]:
|
||||
raise RuntimeError(f"device {row['device_id']} does not have an active card")
|
||||
contexts[row["device_id"]] = DeviceContext(
|
||||
device_id=str(row["device_id"]),
|
||||
serial_number=str(row["serial_number"]),
|
||||
child_id=int(row["child_id"]),
|
||||
card_uuid=str(row["card_uuid"]),
|
||||
)
|
||||
|
||||
missing = [device_id for device_id in device_ids if device_id not in contexts]
|
||||
if missing:
|
||||
raise RuntimeError(f"device context not found: {', '.join(missing)}")
|
||||
return contexts
|
||||
|
||||
|
||||
def fetch_conversation_snapshot(child_a_id: int, child_b_id: int) -> tuple[int | None, int, int]:
|
||||
pair_key = f"{min(child_a_id, child_b_id)}:{max(child_a_id, child_b_id)}"
|
||||
with connect_db() as connection:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT id, message_count
|
||||
FROM im_conversations
|
||||
WHERE conversation_type = 1
|
||||
AND pair_key = %s
|
||||
LIMIT 1
|
||||
""",
|
||||
(pair_key,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if not row:
|
||||
return None, 0, 0
|
||||
|
||||
conversation_id = int(row["id"])
|
||||
message_count = int(row["message_count"])
|
||||
cursor.execute(
|
||||
"SELECT COALESCE(MAX(id), 0) AS max_message_id FROM im_messages WHERE conversation_id = %s",
|
||||
(conversation_id,),
|
||||
)
|
||||
max_row = cursor.fetchone() or {"max_message_id": 0}
|
||||
return conversation_id, message_count, int(max_row["max_message_id"] or 0)
|
||||
|
||||
|
||||
def fetch_new_messages(conversation_id: int, min_message_id: int) -> list[dict]:
|
||||
with connect_db() as connection:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT
|
||||
id,
|
||||
conversation_id,
|
||||
seq,
|
||||
sender_type,
|
||||
sender_id,
|
||||
receiver_type,
|
||||
receiver_id,
|
||||
content_type,
|
||||
media_file_key,
|
||||
client_msg_id,
|
||||
created_at
|
||||
FROM im_messages
|
||||
WHERE conversation_id = %s
|
||||
AND id > %s
|
||||
ORDER BY id ASC
|
||||
""",
|
||||
(conversation_id, min_message_id),
|
||||
)
|
||||
return cursor.fetchall()
|
||||
|
||||
|
||||
async def receiver_loop(ctx: DeviceContext) -> None:
|
||||
assert ctx.websocket is not None
|
||||
async for message in ctx.websocket:
|
||||
if isinstance(message, str):
|
||||
print(f"[{ctx.device_id}] <- {message}", flush=True)
|
||||
await ctx.inbox.put(message)
|
||||
else:
|
||||
print(f"[{ctx.device_id}] <- <binary:{len(message)}>", flush=True)
|
||||
|
||||
|
||||
async def wait_for_text(
|
||||
ctx: DeviceContext,
|
||||
predicate: Callable[[str], bool],
|
||||
*,
|
||||
timeout: float,
|
||||
description: str,
|
||||
) -> str:
|
||||
loop = asyncio.get_running_loop()
|
||||
deadline = loop.time() + timeout
|
||||
while True:
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
raise TimeoutError(f"{ctx.device_id} timed out waiting for {description}")
|
||||
message = await asyncio.wait_for(ctx.inbox.get(), timeout=remaining)
|
||||
if predicate(message):
|
||||
return message
|
||||
|
||||
|
||||
def create_packet(device_id: str, session_id: str, sequence_number: int, packet_type: int, audio_data: bytes) -> bytes:
|
||||
if len(session_id) != 32:
|
||||
raise ValueError("session_id must be exactly 32 ascii chars")
|
||||
return (
|
||||
device_id.encode("ascii")
|
||||
+ b"\x00"
|
||||
+ session_id.encode("ascii")
|
||||
+ b"\x00"
|
||||
+ struct.pack("<I", sequence_number)
|
||||
+ struct.pack("<B", packet_type)
|
||||
+ struct.pack("<I", len(audio_data))
|
||||
+ audio_data
|
||||
)
|
||||
|
||||
|
||||
async def connect_device(ctx: DeviceContext, ws_url: str) -> None:
|
||||
websocket = await websockets.connect(ws_url, ping_interval=None, max_size=None)
|
||||
ctx.websocket = websocket
|
||||
ctx.receiver_task = asyncio.create_task(receiver_loop(ctx))
|
||||
await websocket.send(json.dumps({"device_id": ctx.device_id, "serial_number": ctx.serial_number}, ensure_ascii=False))
|
||||
await wait_for_text(
|
||||
ctx,
|
||||
lambda text: json.loads(text).get("status") == "authenticated",
|
||||
timeout=10,
|
||||
description="authentication response",
|
||||
)
|
||||
print(f"[{ctx.device_id}] authenticated", flush=True)
|
||||
|
||||
|
||||
async def send_voice_message(
|
||||
sender: DeviceContext,
|
||||
target: DeviceContext,
|
||||
audio_path: Path,
|
||||
*,
|
||||
label: str,
|
||||
) -> None:
|
||||
if sender.websocket is None:
|
||||
raise RuntimeError(f"{sender.device_id} is not connected")
|
||||
if not audio_path.exists():
|
||||
raise FileNotFoundError(f"audio file not found: {audio_path}")
|
||||
|
||||
audio_bytes = audio_path.read_bytes()
|
||||
session_id = uuid.uuid4().hex
|
||||
|
||||
print(f"[{label}] register target card {target.card_uuid}", flush=True)
|
||||
await sender.websocket.send(f"REGISTER_TARGET_DEVICE:{target.card_uuid}")
|
||||
await wait_for_text(
|
||||
sender,
|
||||
lambda text: text.startswith("TARGET_DEVICE_REGISTERED_URL:"),
|
||||
timeout=5,
|
||||
description="target registration response",
|
||||
)
|
||||
|
||||
await sender.websocket.send(create_packet(sender.device_id, session_id, 0, 1, b""))
|
||||
print(f"[{label}] start session {session_id}", flush=True)
|
||||
|
||||
chunk_size = 2048
|
||||
sequence_number = 1
|
||||
for offset in range(0, len(audio_bytes), chunk_size):
|
||||
chunk = audio_bytes[offset : offset + chunk_size]
|
||||
await sender.websocket.send(create_packet(sender.device_id, session_id, sequence_number, 4, chunk))
|
||||
sequence_number += 1
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
await sender.websocket.send(create_packet(sender.device_id, session_id, sequence_number, 2, b""))
|
||||
print(f"[{label}] finish session {session_id}, bytes={len(audio_bytes)}", flush=True)
|
||||
|
||||
await wait_for_text(
|
||||
sender,
|
||||
lambda text: text.startswith("PROMPT_SOUND_URL:"),
|
||||
timeout=20,
|
||||
description="message stored response",
|
||||
)
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
|
||||
async def close_device(ctx: DeviceContext) -> None:
|
||||
if ctx.websocket is not None:
|
||||
await ctx.websocket.close()
|
||||
if ctx.receiver_task is not None:
|
||||
try:
|
||||
await asyncio.wait_for(ctx.receiver_task, timeout=2)
|
||||
except Exception:
|
||||
ctx.receiver_task.cancel()
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
args = build_runtime_args()
|
||||
audio_paths = {
|
||||
"a1": Path(args.a1_audio),
|
||||
"b1": Path(args.b1_audio),
|
||||
"a2": Path(args.a2_audio),
|
||||
"b2": Path(args.b2_audio),
|
||||
}
|
||||
|
||||
contexts = load_device_contexts([args.device_a, args.device_b])
|
||||
device_a = contexts[args.device_a]
|
||||
device_b = contexts[args.device_b]
|
||||
|
||||
conversation_id, message_count_before, max_message_id_before = fetch_conversation_snapshot(
|
||||
device_a.child_id,
|
||||
device_b.child_id,
|
||||
)
|
||||
print(
|
||||
f"[snapshot-before] conversation_id={conversation_id} message_count={message_count_before} max_message_id={max_message_id_before}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
await connect_device(device_a, args.ws_url)
|
||||
await connect_device(device_b, args.ws_url)
|
||||
|
||||
try:
|
||||
await send_voice_message(device_a, device_b, audio_paths["a1"], label="A->B #1")
|
||||
await send_voice_message(device_b, device_a, audio_paths["b1"], label="B->A #1")
|
||||
await send_voice_message(device_a, device_b, audio_paths["a2"], label="A->B #2")
|
||||
await send_voice_message(device_b, device_a, audio_paths["b2"], label="B->A #2")
|
||||
finally:
|
||||
await close_device(device_a)
|
||||
await close_device(device_b)
|
||||
|
||||
await asyncio.sleep(2)
|
||||
|
||||
conversation_id_after, message_count_after, max_message_id_after = fetch_conversation_snapshot(
|
||||
device_a.child_id,
|
||||
device_b.child_id,
|
||||
)
|
||||
print(
|
||||
f"[snapshot-after] conversation_id={conversation_id_after} message_count={message_count_after} max_message_id={max_message_id_after}",
|
||||
flush=True,
|
||||
)
|
||||
if conversation_id_after is None:
|
||||
raise RuntimeError("child-peer conversation was not created")
|
||||
|
||||
new_messages = fetch_new_messages(conversation_id_after, max_message_id_before)
|
||||
print(f"[new-messages] count={len(new_messages)}", flush=True)
|
||||
for row in new_messages:
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"id": int(row["id"]),
|
||||
"seq": int(row["seq"]),
|
||||
"sender_id": row["sender_id"],
|
||||
"receiver_id": row["receiver_id"],
|
||||
"content_type": int(row["content_type"]),
|
||||
"media_file_key": row["media_file_key"],
|
||||
"client_msg_id": row["client_msg_id"],
|
||||
"created_at": row["created_at"].isoformat(sep=" "),
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user