nvidia-nemo--speech/examples/voice_agent/server/server.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import asyncio
import copy
import os
import signal
import sys
from contextlib import asynccontextmanager
from datetime import datetime
from typing import Any, Dict

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from loguru import logger
from omegaconf import OmegaConf
from pipecat.audio.vad.silero import SileroVADAnalyzer
from pipecat.frames.frames import EndTaskFrame
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.runner import PipelineRunner
from pipecat.pipeline.task import PipelineParams, PipelineTask
from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext
from pipecat.processors.frameworks.rtvi import RTVIAction, RTVIConfig, RTVIObserverParams, RTVIProcessor
from pipecat.serializers.protobuf import ProtobufFrameSerializer
from websocket_url import build_websocket_url

from nemo.agents.voice_agent.pipecat.processors.frameworks.rtvi import RTVIObserver
from nemo.agents.voice_agent.pipecat.services.nemo.audio_logger import AudioLogger, RTVIAudioLoggerObserver
from nemo.agents.voice_agent.pipecat.services.nemo.diar import NemoDiarService
from nemo.agents.voice_agent.pipecat.services.nemo.llm import get_llm_service_from_config
from nemo.agents.voice_agent.pipecat.services.nemo.stt import NemoSTTService
from nemo.agents.voice_agent.pipecat.services.nemo.tts import get_tts_service_from_config
from nemo.agents.voice_agent.pipecat.services.nemo.turn_taking import NeMoTurnTakingService
from nemo.agents.voice_agent.pipecat.transports.network.websocket_server import (
    WebsocketServerParams,
    WebsocketServerTransport,
)
from nemo.agents.voice_agent.utils.config_manager import ConfigManager
from nemo.agents.voice_agent.utils.tool_calling.basic_tools import tool_get_city_weather
from nemo.agents.voice_agent.utils.tool_calling.mixins import register_direct_tools_to_llm

# Load environment variables
load_dotenv(override=True)


def setup_logging():
    # Configure loguru to output to both console and file
    logger.remove()  # Remove default handler
    logger.add(
        sys.stderr,
        format="<green>{time:YYYY-MM-DD HH:mm:ss.SSSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
        level="DEBUG",
    )

    logger.add("bot_server.log", rotation="1 day", level="DEBUG")


setup_logging()

# Global flag for graceful shutdown
shutdown_event = asyncio.Event()

# Initialize configuration manager
config_manager = ConfigManager(
    server_base_path=os.path.dirname(__file__), server_config_path=os.environ.get("SERVER_CONFIG_PATH", None)
)
server_config = config_manager.get_server_config()

logger.info(f"Server config: {OmegaConf.to_container(server_config, resolve=True)}")

# Access configuration parameters from ConfigManager
SAMPLE_RATE = config_manager.SAMPLE_RATE
RAW_AUDIO_FRAME_LEN_IN_SECS = config_manager.RAW_AUDIO_FRAME_LEN_IN_SECS
SYSTEM_PROMPT = config_manager.SYSTEM_PROMPT
SYSTEM_ROLE = config_manager.SYSTEM_ROLE

# Transport configuration
TRANSPORT_AUDIO_OUT_10MS_CHUNKS = config_manager.TRANSPORT_AUDIO_OUT_10MS_CHUNKS
RECORD_AUDIO_DATA = server_config.transport.get("record_audio_data", False)
AUDIO_LOG_DIR = server_config.transport.get("audio_log_dir", "./audio_logs")
SERVER_HOST = os.getenv("SERVER_HOST", "0.0.0.0")
SERVER_PUBLIC_HOST = os.getenv("SERVER_PUBLIC_HOST", "127.0.0.1")
WEBSOCKET_SCHEME = os.getenv("WEBSOCKET_SCHEME", "ws")
WEBSOCKET_PORT = int(os.getenv("WEBSOCKET_PORT", 8765))
FASTAPI_PORT = int(os.getenv("FASTAPI_PORT", 7860))

# VAD configuration
vad_params = config_manager.get_vad_params()

# STT configuration
STT_MODEL = config_manager.STT_MODEL
STT_DEVICE = config_manager.STT_DEVICE
stt_params = config_manager.get_stt_params()

# Diarization configuration
DIAR_MODEL = config_manager.DIAR_MODEL
USE_DIAR = config_manager.USE_DIAR
diar_params = config_manager.get_diar_params()

# Turn taking configuration
TURN_TAKING_BACKCHANNEL_PHRASES_PATH = config_manager.TURN_TAKING_BACKCHANNEL_PHRASES_PATH
TURN_TAKING_MAX_BUFFER_SIZE = config_manager.TURN_TAKING_MAX_BUFFER_SIZE
TURN_TAKING_BOT_STOP_DELAY = config_manager.TURN_TAKING_BOT_STOP_DELAY

# TTS configuration
TTS_TYPE = config_manager.server_config.tts.type


def signal_handler(signum, frame):
    """Handle shutdown signals gracefully"""
    logger.info(f"Received signal {signum}, initiating graceful shutdown...")
    shutdown_event.set()


async def run_bot_websocket_server(host: str = "0.0.0.0", port: int = None):
    """
    NO-TIMEOUT CONFIGURATION:
    - session_timeout=None: Disables WebSocket session timeout
    - idle_timeout=None: Disables pipeline idle timeout
    - asyncio.wait_for(timeout=None): No timeout on pipeline runner
    - Server will run indefinitely until manually stopped (Ctrl+C)
    """
    if port is None:
        port = WEBSOCKET_PORT
    logger.info(f"Starting websocket server on {host}:{port}")
    logger.info(f"Server configured to run indefinitely with no timeouts, use Ctrl+C to quit.")

    # Set up signal handlers for graceful shutdown
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    logger.info("Initializing WebSocket server transport...")
    logger.info("Server configured to run indefinitely with no timeouts")
    # Initialize AudioLogger if recording is enabled
    audio_logger = None
    if RECORD_AUDIO_DATA:
        session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        audio_logger = AudioLogger(
            log_dir=AUDIO_LOG_DIR,
            session_id=session_id,
            enabled=True,
        )
        logger.info(f"AudioLogger initialized for session: {session_id} at {AUDIO_LOG_DIR}")

    vad_analyzer = SileroVADAnalyzer(
        sample_rate=SAMPLE_RATE,
        params=vad_params,
    )
    logger.info("VAD analyzer initialized")

    ws_transport = WebsocketServerTransport(
        params=WebsocketServerParams(
            serializer=ProtobufFrameSerializer(),
            audio_in_enabled=True,
            audio_out_enabled=True,
            add_wav_header=False,
            vad_analyzer=vad_analyzer,
            session_timeout=None,  # Disable session timeout
            audio_in_sample_rate=SAMPLE_RATE,
            can_create_user_frames=False,
            audio_out_10ms_chunks=TRANSPORT_AUDIO_OUT_10MS_CHUNKS,
        ),
        host=host,
        port=port,
    )

    logger.info("Initializing STT service...")

    stt = NemoSTTService(
        model=STT_MODEL,
        device=STT_DEVICE,
        params=stt_params,
        sample_rate=SAMPLE_RATE,
        audio_passthrough=True,
        backend="legacy",
        decoder_type="rnnt",
        audio_logger=audio_logger,
    )
    logger.info("STT service initialized")

    if USE_DIAR:
        diar = NemoDiarService(
            model=DIAR_MODEL,
            device=STT_DEVICE,
            params=diar_params,
            sample_rate=SAMPLE_RATE,
            backend="legacy",
            enabled=USE_DIAR,
        )
        logger.info("Diarization service initialized")
    else:
        diar = None

    turn_taking = NeMoTurnTakingService(
        use_diar=USE_DIAR,
        max_buffer_size=TURN_TAKING_MAX_BUFFER_SIZE,
        bot_stop_delay=TURN_TAKING_BOT_STOP_DELAY,
        backchannel_phrases=TURN_TAKING_BACKCHANNEL_PHRASES_PATH,
        audio_logger=audio_logger,
    )
    logger.info("Turn taking service initialized")

    if TTS_TYPE == "nemo":
        tts = get_tts_service_from_config(config_manager.server_config.tts, audio_logger)
    else:
        raise ValueError(f"Invalid TTS type: {TTS_TYPE}")

    logger.info("TTS service initialized")

    # Setup logging again to avoid logger from being overwritten during setting up the pipeline components
    setup_logging()

    # Put LLM in the end of model initialization to reduce the chance of running out of HBM memory
    logger.info("Initializing LLM service...")
    llm = get_llm_service_from_config(server_config.llm)
    logger.info("LLM service initialized")

    messages = [
        {
            "role": SYSTEM_ROLE,
            "content": SYSTEM_PROMPT,
        }
    ]
    inject_dummy_user_message = server_config.llm.get("inject_dummy_user_message", False)
    if inject_dummy_user_message:
        messages.append(
            {
                "role": "user",
                "content": "Hello, who are you?",
            }
        )
    context = OpenAILLMContext(messages=messages)

    if server_config.llm.get("enable_tool_calling", False):
        logger.info("Tools calling for LLM is enabled by config, registering tools...")
        register_direct_tools_to_llm(llm=llm, context=context, tool_mixins=[tts], tools=[tool_get_city_weather])
    else:
        logger.info("Tools calling for LLM is disabled by config, skipping tool registration.")

    original_messages = copy.deepcopy(context.get_messages())
    original_context = copy.deepcopy(context)
    original_context.set_llm_adapter(llm.get_llm_adapter())

    context_aggregator = llm.create_context_aggregator(context)
    user_context_aggregator = context_aggregator.user()
    assistant_context_aggregator = context_aggregator.assistant()

    # RTVI events for Pipecat client UI
    rtvi = RTVIProcessor(config=RTVIConfig(config=[]))

    # Add reset action to RTVI processor
    async def reset_context_handler(rtvi_processor: RTVIProcessor, service: str, arguments: dict[str, Any]) -> bool:
        """Reset both user and assistant context aggregators"""
        logger.info("Resetting conversation context...")
        try:
            user_context_aggregator.reset()
            assistant_context_aggregator.reset()
            user_context_aggregator.set_messages(copy.deepcopy(original_messages))
            assistant_context_aggregator.set_messages(copy.deepcopy(original_messages))
            tts.reset()
            if diar is not None:
                diar.reset()
            turn_taking.reset()
            logger.info("Conversation context reset successfully")
            return True
        except Exception as e:
            logger.error(f"Error resetting context: {e}")
            return False

    reset_action = RTVIAction(
        service="context",
        action="reset",
        result="bool",
        arguments=[],
        handler=reset_context_handler,
    )
    rtvi.register_action(reset_action)

    logger.info("Setting up pipeline...")

    pipeline = [
        ws_transport.input(),
        rtvi,
        stt,
    ]

    if USE_DIAR:
        pipeline.append(diar)

    pipeline.extend(
        [turn_taking, user_context_aggregator, llm, tts, ws_transport.output(), assistant_context_aggregator]
    )

    pipeline = Pipeline(pipeline)

    rtvi_params = RTVIObserverParams(bot_llm_enabled=False)
    task = PipelineTask(
        pipeline,
        params=PipelineParams(
            allow_interruptions=True,
            enable_metrics=False,
            enable_usage_metrics=False,
            send_initial_empty_metrics=True,
            report_only_initial_ttfb=True,
            idle_timeout=None,  # Disable idle timeout
        ),
        observers=[
            RTVIObserver(rtvi, params=rtvi_params),
            RTVIAudioLoggerObserver(audio_logger=audio_logger),
        ],
        idle_timeout_secs=None,
        cancel_on_idle_timeout=False,
    )

    # Track task state
    task_running = True

    # Setup logging again to avoid logger from being overwritten during setting up the pipeline components
    setup_logging()

    @rtvi.event_handler("on_client_ready")
    async def on_client_ready(rtvi: RTVIProcessor):
        logger.info("Pipecat client ready.")
        await rtvi.set_bot_ready()
        # Kick off the conversation.
        try:
            await task.queue_frames([user_context_aggregator.get_context_frame()])
        except Exception as e:
            logger.error(f"Error queuing context frame: {e}")

    @ws_transport.event_handler("on_client_connected")
    async def on_client_connected(transport, client):
        logger.info(f"Pipecat Client connected from {client.remote_address}")
        # Reset RTVI state for new connection
        rtvi._client_ready = False
        rtvi._bot_ready = False

    @ws_transport.event_handler("on_client_disconnected")
    async def on_client_disconnected(transport, client):
        logger.info(f"Pipecat Client disconnected from {client.remote_address}")
        # Finalize audio logger session if enabled
        if audio_logger:
            audio_logger.finalize_session()
            logger.info("Audio logger session finalized")
        # Don't cancel the task immediately - let it handle the disconnection gracefully
        # The task will continue running and can accept new connections
        # Only send an EndTaskFrame to clean up the current session
        if task_running:
            try:
                await task.queue_frames([EndTaskFrame()])
            except Exception as e:
                # Don't log warnings for normal connection closures
                if "ConnectionClosedOK" not in str(e) and "1005" not in str(e):
                    logger.warning(f"Error sending EndTaskFrame: {e}")
                else:
                    logger.info(f"Normal connection closure: {e}")

    @ws_transport.event_handler("on_session_timeout")
    async def on_session_timeout(transport, client):
        logger.info(f"Session timeout for {client.remote_address}")
        # Don't cancel the task - keep server running indefinitely
        logger.info("Session timeout occurred but keeping server running")
        # Note: With session_timeout=None, this handler should never be called

    logger.info("Starting pipeline runner...")

    try:
        runner = PipelineRunner()
        # Run the task until shutdown is requested
        await asyncio.wait_for(runner.run(task), timeout=None)  # No timeout - run indefinitely
    except asyncio.TimeoutError:
        logger.info("Pipeline runner timeout (should not happen with no timeout)")
    except Exception as e:
        logger.error(f"Pipeline runner error: {e}")
        task_running = False
    finally:
        # Finalize audio logger on shutdown
        if audio_logger:
            audio_logger.finalize_session()
            logger.info("Audio logger session finalized on shutdown")
        logger.info("Pipeline runner stopped")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Handles FastAPI startup and shutdown."""
    yield  # Run app


# Initialize FastAPI app with lifespan manager
app = FastAPI(lifespan=lifespan)

# Configure CORS to allow requests from any origin
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    print("WebSocket connection accepted")
    try:
        # TODO: [heh] Implement FastAPI websocket endpoint
        # await run_bot_fastapi_server(websocket)
        raise NotImplementedError("FastAPI websocket endpoint is not implemented")
    except Exception as e:
        print(f"Exception in run_bot: {e}")


@app.post("/connect")
async def bot_connect() -> Dict[Any, Any]:
    print("Received /connect request")
    ws_url = build_websocket_url(SERVER_PUBLIC_HOST, WEBSOCKET_PORT, WEBSOCKET_SCHEME)
    print(f"Returning WebSocket URL: {ws_url}")
    return {"ws_url": ws_url}


async def main():
    """Main function to run both websocket server and FastAPI server concurrently."""
    logger.info(f"Starting servers - WebSocket on port {WEBSOCKET_PORT}, FastAPI on port {FASTAPI_PORT}")
    tasks = []
    try:
        # Start websocket server
        tasks.append(run_bot_websocket_server(host=SERVER_HOST, port=WEBSOCKET_PORT))

        # Start FastAPI server
        config = uvicorn.Config(app, host=SERVER_HOST, port=FASTAPI_PORT)
        server = uvicorn.Server(config)
        tasks.append(server.serve())

        await asyncio.gather(*tasks)
    except asyncio.CancelledError:
        logger.info("Tasks cancelled (probably due to shutdown).")


if __name__ == "__main__":
    asyncio.run(main())