gradio_webrtc_r1.py

import gradio as gr
from gradio_webrtc import WebRTC, ReplyOnPause, AdditionalOutputs
import requests
from pyht import Client as PyHtClient, TTSOptions
import dataclasses
import os
import numpy as np
from deepgram import DeepgramClient, PrerecordedOptions, FileSource
import io
from pydub import AudioSegment
from dotenv import load_dotenv

load_dotenv()

account_sid = os.environ.get("TWILIO_ACCOUNT_SID")
auth_token = os.environ.get("TWILIO_AUTH_TOKEN")

if account_sid and auth_token:
    from twilio.rest import Client
    client = Client(account_sid, auth_token)
    token = client.tokens.create()
    rtc_configuration = {
        "iceServers": token.ice_servers,
        "iceTransportPolicy": "relay",
    }
else:
    rtc_configuration = None

@dataclasses.dataclass
class Clients:
    deepseek_key: str
    play_ht: PyHtClient
    deepgram: DeepgramClient

child_tts_options = TTSOptions(voice="s3://play-fal/nv_peppa2_0af1bf1e-eb94-47ef-8149-77b18e18e12b/voices/speaker/manifest.json",
                              sample_rate=24000)

tts_options = TTSOptions(voice="s3://voice-cloning-zero-shot/8cec70ca-f404-4b69-a1a1-708a3ccc77e7/ralph-ineson-instant-2/manifest.json",
                        sample_rate=24000)

def aggregate_chunks(chunks_iterator):
    leftover = b''
    for chunk in chunks_iterator:
        current_bytes = leftover + chunk
        n_complete_samples = len(current_bytes) // 2
        bytes_to_process = n_complete_samples * 2
        to_process = current_bytes[:bytes_to_process]
        leftover = current_bytes[bytes_to_process:]
        if to_process:
            audio_array = np.frombuffer(to_process, dtype=np.int16).reshape(1, -1)
            yield audio_array

def audio_to_bytes(audio: tuple[int, np.ndarray]) -> bytes:
    audio_buffer = io.BytesIO()
    segment = AudioSegment(
        audio[1].tobytes(),
        frame_rate=audio[0],
        sample_width=audio[1].dtype.itemsize,
        channels=1,
    )
    segment.export(audio_buffer, format="mp3")
    return audio_buffer.getvalue()

def query_deepseek(api_key: str, user_content: str, conversation_history: list) -> str:
    url = "https://api.deepseek.com/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    messages.extend(conversation_history)
    messages.append({"role": "user", "content": user_content})
    
    data = {
        "model": "deepseek-reasoner",
        "messages": messages,
        "stream": False
    }
    try:
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()
        result = response.json()
        return result["choices"][0]["message"]["content"], result["choices"][0]["message"]["reasoning_content"]
    except requests.exceptions.RequestException as e:
        raise gr.Error(f"DeepSeek API error: {str(e)}")

def set_api_key(deepseek_key, play_ht_username, play_ht_key, deepgram_api_key):
    try:
        play_ht_client = PyHtClient(user_id=play_ht_username, api_key=play_ht_key)
        deepgram_client = DeepgramClient(api_key=deepgram_api_key)
    except Exception as e:
        raise gr.Error(f"Invalid API keys: {str(e)}")
    gr.Info("Successfully set API keys.", duration=3)
    return Clients(deepseek_key=deepseek_key, 
                  play_ht=play_ht_client,
                  deepgram=deepgram_client), gr.skip()

def response(audio: tuple[int, np.ndarray], conversation_llm_format: list[dict],
            chatbot: list[dict], client_state: Clients):
    if not client_state:
        raise gr.Error("Please set your API keys first.")
    
    # Convert audio to bytes for Deepgram
    audio_bytes = audio_to_bytes(audio)
    
    # Configure Deepgram options
    options = PrerecordedOptions(
        model="nova-2",
        smart_format=True,
    )
    
    # Get transcription from Deepgram
    try:
        payload = {"buffer": audio_bytes}
        result = client_state.deepgram.listen.rest.v("1").transcribe_file(payload, options)
        prompt = result.results.channels[0].alternatives[0].transcript
    except Exception as e:
        raise gr.Error(f"Deepgram transcription error: {str(e)}")
    
    conversation_llm_format.append({"role": "user", "content": prompt})
    
    response_text, response_reasoning_text = query_deepseek(
        client_state.deepseek_key,
        prompt,
        conversation_llm_format[:-1]
    )
    
    conversation_llm_format.append({"role": "assistant", "content": response_text})
    chatbot.append({"role": "user", "content": prompt})
    chatbot.append({"role": "assistant", "content": "**Think**\n" + response_reasoning_text})
    chatbot.append({"role": "assistant", "content": "**Response**\n" + response_text})
    
    yield AdditionalOutputs(conversation_llm_format, chatbot)
    iterator = client_state.play_ht.tts(response_reasoning_text, options=child_tts_options, voice_engine='PlayDialog-http')
    for chunk in aggregate_chunks(iterator):
        audio_array = np.frombuffer(chunk, dtype=np.int16).reshape(1, -1)
        yield (24000, audio_array, "mono")
    
    # Add 0.5 second silence
    silence_duration = int(24000 * 0.5)  # 0.5 seconds at 24kHz
    silence = np.zeros((1, silence_duration), dtype=np.int16)
    yield (24000, silence, "mono")

    iterator = client_state.play_ht.tts(response_text, options=tts_options, voice_engine="PlayDialog-http")
    for chunk in aggregate_chunks(iterator):
        audio_array = np.frombuffer(chunk, dtype=np.int16).reshape(1, -1)
        yield (24000, audio_array, "mono")

with gr.Blocks() as demo:
    with gr.Group():
        with gr.Row():
            chatbot = gr.Chatbot(label="Conversation", type="messages")
        with gr.Row(equal_height=True):
            with gr.Column(scale=1):
                with gr.Row():
                    deepseek_key = gr.Textbox(
                        type="password",
                        value=os.getenv("DEEPSEEK_API_KEY"),
                        label="Enter your DeepSeek API Key"
                    )
                    play_ht_username = gr.Textbox(
                        type="password",
                        value=os.getenv("PLAY_HT_USER_ID"),
                        label="Enter your PlayHt Username"
                    )
                    play_ht_key = gr.Textbox(
                        type="password",
                        value=os.getenv("PLAY_HT_API_KEY"),
                        label="Enter your PlayHt API Key"
                    )
                    deepgram_key = gr.Textbox(
                        type="password",
                        value=os.getenv("DEEPGRAM_API_KEY"),
                        label="Enter your Deepgram API Key"
                    )
                with gr.Row():
                    set_key_button = gr.Button("Set Keys", variant="primary")
            with gr.Column(scale=5):
                audio = WebRTC(
                    modality="audio",
                    mode="send-receive",
                    label="Audio Stream",
                    rtc_configuration=rtc_configuration
                )

    client_state = gr.State(None)
    conversation_llm_format = gr.State([])

    set_key_button.click(
        set_api_key,
        inputs=[deepseek_key, play_ht_username, play_ht_key, deepgram_key],
        outputs=[client_state, set_key_button]
    )
    audio.stream(
        ReplyOnPause(response),
        inputs=[audio, conversation_llm_format, chatbot, client_state],
        outputs=[audio]
    )
    audio.on_additional_outputs(
        lambda l, g: (l, g),
        outputs=[conversation_llm_format, chatbot]
    )

if __name__ == "__main__":
    demo.launch(share=True)