AgapHost/ai/pipecat/test_pipeline.py

"""
End-to-end pipeline test:
1. Call /connect to get a room + token
2. Join the LiveKit room as a Python client
3. Publish TTS audio (pre-generated) as microphone input
4. Capture bot's audio response and save to file
"""
import asyncio
import wave
import struct
import httpx
import numpy as np
from livekit import rtc

PIPECAT_URL = "http://localhost:8882"
TTS_URL = "http://host.docker.internal:8881"
OUTPUT_FILE = "/tmp/bot_response.wav"
SAMPLE_RATE = 48000
NUM_CHANNELS = 1


async def generate_tts_pcm(text: str) -> bytes:
    """Get WAV audio from Silero TTS, return raw PCM int16 bytes."""
    async with httpx.AsyncClient(timeout=30) as c:
        r = await c.post(f"{TTS_URL}/v1/audio/speech", json={
            "input": text, "voice": "onyx", "response_format": "wav"
        })
        r.raise_for_status()
    # Skip WAV header (44 bytes) to get raw PCM
    return r.content[44:]


async def main():
    # Step 1 — create room
    print("[test] Creating room...")
    async with httpx.AsyncClient() as c:
        r = await c.post(f"{PIPECAT_URL}/connect")
        r.raise_for_status()
        creds = r.json()
    print(f"[test] Room: {creds['room']}  URL: {creds['url']}")

    # Step 2 — generate test audio
    test_phrase = "Привет! Как тебя зовут?"
    print(f"[test] Generating TTS for: {test_phrase!r}")
    pcm_bytes = await generate_tts_pcm(test_phrase)
    print(f"[test] TTS PCM: {len(pcm_bytes)} bytes (~{len(pcm_bytes)//(SAMPLE_RATE*2):.1f}s)")

    # Step 3 — join room
    room = rtc.Room()
    received_frames: list[bytes] = []

    @room.on("track_subscribed")
    def on_track(track, pub, participant):
        if track.kind == rtc.TrackKind.KIND_AUDIO and participant.identity == "pipecat-bot":
            print(f"[test] Subscribed to bot audio track")
            audio_stream = rtc.AudioStream(track, sample_rate=SAMPLE_RATE, num_channels=NUM_CHANNELS)
            asyncio.ensure_future(_collect_audio(audio_stream, received_frames))

    ws_url = creds["url"].replace("https://", "wss://").replace("http://", "ws://")
    # Connect internally via host.docker.internal
    internal_url = "ws://host.docker.internal:7880"
    print(f"[test] Connecting to LiveKit at {internal_url}...")
    await room.connect(internal_url, creds["token"])
    print(f"[test] Connected. Waiting for bot to join...")

    # Wait for bot participant
    for _ in range(20):
        if any(p.identity == "pipecat-bot" for p in room.remote_participants.values()):
            break
        await asyncio.sleep(0.5)
    print(f"[test] Participants: {[p.identity for p in room.remote_participants.values()]}")

    # Step 4 — publish audio as microphone
    print("[test] Publishing audio track...")
    source = rtc.AudioSource(SAMPLE_RATE, NUM_CHANNELS)
    local_track = rtc.LocalAudioTrack.create_audio_track("microphone", source)
    opts = rtc.TrackPublishOptions(source=rtc.TrackSource.SOURCE_MICROPHONE)
    await room.local_participant.publish_track(local_track, opts)

    # Send PCM in 20ms chunks
    chunk_samples = SAMPLE_RATE * 20 // 1000  # 960 samples per chunk
    chunk_bytes = chunk_samples * 2  # int16
    print(f"[test] Sending {len(pcm_bytes) // chunk_bytes} audio chunks...")
    for i in range(0, len(pcm_bytes), chunk_bytes):
        chunk = pcm_bytes[i:i + chunk_bytes]
        if len(chunk) < chunk_bytes:
            chunk = chunk + b'\x00' * (chunk_bytes - len(chunk))
        samples = np.frombuffer(chunk, dtype=np.int16)
        frame = rtc.AudioFrame(
            data=samples.tobytes(),
            sample_rate=SAMPLE_RATE,
            num_channels=NUM_CHANNELS,
            samples_per_channel=chunk_samples,
        )
        await source.capture_frame(frame)
        await asyncio.sleep(0.02)

    print("[test] Audio sent. Waiting for bot response (up to 30s)...")
    await asyncio.sleep(30)

    await room.disconnect()

    # Step 5 — save response
    if received_frames:
        total = b"".join(received_frames)
        print(f"[test] Received {len(total)} bytes of bot audio ({len(total)//(SAMPLE_RATE*2):.1f}s)")
        with wave.open(OUTPUT_FILE, "wb") as wf:
            wf.setnchannels(NUM_CHANNELS)
            wf.setsampwidth(2)
            wf.setframerate(SAMPLE_RATE)
            wf.writeframes(total)
        print(f"[test] Saved to {OUTPUT_FILE}")
    else:
        print("[test] No audio received from bot!")


async def _collect_audio(stream: rtc.AudioStream, buf: list):
    async for event in stream:
        buf.append(bytes(event.frame.data))


asyncio.run(main())