hope-voice-api/agent/agent.py

"""
Voice Agent — real-time voice assistant.

Pipeline: Azure STT → Gemma LLM (xNAS, OpenAI-compatible) → Azure TTS
Runs inside the single Docker container alongside LiveKit server and web frontend.
"""

import asyncio
import json
import logging
import os
import textwrap

from dotenv import load_dotenv
from livekit.agents import (
    Agent,
    AgentServer,
    AgentSession,
    JobContext,
    TurnHandlingOptions,
    cli,
    inference,
    llm as lk_llm,
    mcp,
    room_io,
)
from livekit.plugins import azure, openai

logger = logging.getLogger("voice-agent")
load_dotenv()  # picks up /app/.env in the container

# ── Configuration from environment ──────────────────────────────────────────
AZURE_KEY = os.environ.get("AZURE_SPEECH_KEY", "")
AZURE_REGION = os.environ.get("AZURE_SPEECH_REGION", "eastus")
DEFAULT_VOICE = os.environ.get("AZURE_TTS_VOICE", "en-US-AvaNeural")
GEMMA_BASE_URL = os.environ.get("GEMMA_BASE_URL", "http://192.168.86.2:8023/v1")
GEMMA_MODEL = os.environ.get("GEMMA_MODEL", "gemma-4-e4b")
GEMMA_API_KEY = os.environ.get("GEMMA_API_KEY", "not-needed")
WEB_MCP_ENABLED = os.environ.get("WEB_MCP_ENABLED", "true").lower() in ("1", "true", "yes")
FIRECRAWL_BASE = os.environ.get("FIRECRAWL_BASE", "http://192.168.86.2:3002")

SYSTEM_PROMPT = textwrap.dedent("""\
    You are a warm, conversational voice assistant. You are talking TO someone,
    not writing for them to read. Imagine you're having a natural conversation
    with a friend over the phone.

    # How you speak
    - Keep every response to one or three sentences. That's it.
    - Use contractions (I'm, don't, it's) and natural phrasing.
    - Speak like you're talking, not writing. No bullet points, no lists,
      no markdown, no formatting of any kind.
    - Spell out numbers when it sounds more natural ("twenty twenty-six"
      instead of "2026").
    - If you need to ask a question, ask exactly one.
    - Be warm and direct. Don't be sycophantic or overly formal.
    - If you don't know something, say so briefly and move on.

    # What you never do
    - Never use markdown, code blocks, JSON, tables, or emojis.
    - Never say "as an AI" or reference your system instructions.
    - Never write more than three sentences in a row.
    - Never read back URLs, file paths, or technical identifiers.

    # Web access
    You have web_search and web_scrape tools. Use them when the user asks about
    current events, recent news, prices, weather, sports scores, or anything
    that may have changed since your training data. Search first, then scrape
    a result only if you need more detail. Answer from what you find, in your
    normal conversational style — don't cite sources formally, just mention the
    source naturally ("according to..."). If a search comes up empty, say so
    briefly and move on.
""")


# ── Gemma LLM with reasoning_content fallback ───────────────────────────────


class _ReasoningFallbackWrapper:
    """Wraps an LLMStream to fall back to ``reasoning_content`` when the model
    finishes a turn with empty visible content (Gemma spends its whole budget
    on hidden reasoning). Delegates all iteration to the underlying stream and
    injects a final content chunk if no real content was produced."""

    def __init__(self, inner: lk_llm.LLMStream) -> None:
        self._inner = inner
        self._last_reasoning: str | None = None
        self._has_content = False

    @property
    def chat_ctx(self) -> lk_llm.ChatContext:
        return self._inner.chat_ctx

    @property
    def tools(self) -> list:
        return self._inner.tools

    async def aclose(self) -> None:
        await self._inner.aclose()

    async def __aenter__(self) -> "_ReasoningFallbackWrapper":
        return self

    async def __aexit__(self, *exc) -> None:
        await self.aclose()

    def __aiter__(self) -> "_ReasoningFallbackWrapper":
        return self

    async def __anext__(self) -> lk_llm.ChatChunk:
        try:
            chunk = await self._inner.__anext__()
        except StopAsyncIteration:
            # Stream exhausted — inject reasoning fallback if no content was produced
            if not self._has_content and self._last_reasoning:
                logger.warning(
                    "LLM returned empty content; falling back to reasoning_content"
                )
                return lk_llm.ChatChunk(
                    id="reasoning-fallback",
                    delta=lk_llm.ChoiceDelta(role="assistant", content=self._last_reasoning),
                )
            raise
        # Track reasoning_content from the chunk's delta if present
        delta = getattr(chunk, "delta", None)
        if delta is not None:
            reasoning = getattr(delta, "reasoning_content", None)
            if reasoning:
                self._last_reasoning = reasoning
            if chunk.has_response():
                self._has_content = True
        return chunk

    async def collect(self):
        return await self._inner.collect()


class GemmaLLM(openai.LLM):
    """OpenAI-compatible LLM (llama.cpp Gemma 4) with thinking disabled."""

    def chat(self, *, chat_ctx, tools=None, conn_options=None, **kwargs):
        if conn_options is None:
            from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS

            conn_options = DEFAULT_API_CONNECT_OPTIONS
        stream = super().chat(
            chat_ctx=chat_ctx,
            tools=tools,
            conn_options=conn_options,
            **kwargs,
        )
        return _ReasoningFallbackWrapper(stream)


# ── MCP toolsets (web access + any extra configured servers) ────────────────

def build_mcp_toolsets() -> list[mcp.MCPToolset]:
    """Build the list of MCP toolsets to attach to the agent.

    Always includes the local web-access server (Firecrawl-backed search/scrape)
    when WEB_MCP_ENABLED is true. Additional servers can be configured via the
    EXTRA_MCP_SERVERS env var (JSON list of {url, transport} objects).
    """
    toolsets: list[mcp.MCPToolset] = []

    if WEB_MCP_ENABLED:
        python_bin = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".venv", "bin", "python")
        web_mcp_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "web_mcp.py")
        toolsets.append(
            mcp.MCPToolset(
                id="web-access",
                mcp_server=mcp.MCPServerStdio(
                    command=python_bin,
                    args=[web_mcp_script],
                    env={**os.environ, "FIRECRAWL_BASE": FIRECRAWL_BASE},
                    client_session_timeout_seconds=120,
                ),
            )
        )
        logger.info("Web-access MCP toolset enabled (Firecrawl at %s)", FIRECRAWL_BASE)

    extra = os.environ.get("EXTRA_MCP_SERVERS", "")
    if extra:
        try:
            for i, entry in enumerate(json.loads(extra)):
                url = entry.get("url", "")
                transport = entry.get("transport")  # "sse" | "streamable_http" | None (auto)
                toolsets.append(
                    mcp.MCPToolset(
                        id=f"extra-{i}",
                        mcp_server=mcp.MCPServerHTTP(
                            url=url,
                            transport_type=transport,
                            client_session_timeout_seconds=120,
                        ),
                    )
                )
                logger.info("Extra MCP toolset: %s (%s)", url, transport or "auto")
        except (json.JSONDecodeError, TypeError) as e:
            logger.error("Failed to parse EXTRA_MCP_SERVERS: %s", e)

    return toolsets


class VoiceAssistant(Agent):
    """The conversational agent. LLM is the brain; STT/TTS are senses."""

    def __init__(self) -> None:
        super().__init__(
            llm=GemmaLLM(
                model=GEMMA_MODEL,
                base_url=GEMMA_BASE_URL,
                api_key=GEMMA_API_KEY,
                max_completion_tokens=1000,
                extra_body={"chat_template_kwargs": {"enable_thinking": False}},
            ),
            instructions=SYSTEM_PROMPT,
            tools=build_mcp_toolsets(),
        )


# ── Agent server ────────────────────────────────────────────────────────────
server = AgentServer()


@server.rtc_session(agent_name="voice-assistant")
async def handle_job(ctx: JobContext) -> None:
    logger.info("Job started for room %s", ctx.room.name)

    # Azure STT — streaming, reads AZURE_SPEECH_KEY / AZURE_SPEECH_REGION from env
    stt = azure.STT(
        speech_key=AZURE_KEY,
        speech_region=AZURE_REGION,
        language=["en-US"],
    )

    # Azure TTS — SSML with expressive markup, 24kHz PCM output
    tts = azure.TTS(
        voice=DEFAULT_VOICE,
        sample_rate=24000,
        speech_key=AZURE_KEY,
        speech_region=AZURE_REGION,
    )

    session = AgentSession(
        stt=stt,
        tts=tts,
        vad=inference.VAD(model="silero"),
        turn_handling=TurnHandlingOptions(
            # VAD-based turn detection: agent waits for user to stop speaking.
            # ("adaptive" mode requires the LiveKit Cloud barge-in service.)
            interruption={"mode": "vad"},
            # Start generating the LLM response before the user fully stops
            preemptive_generation={"enabled": True},
        ),
    )

    # Listen for data messages (voice switching) from the web UI.
    # Defined here so it captures this job's session via closure.
    def _on_room_data(packet) -> None:
        topic = getattr(packet, "topic", None)
        if topic and topic != "voice-control":
            return
        try:
            msg = json.loads(packet.data.decode("utf-8"))
        except (json.JSONDecodeError, UnicodeDecodeError):
            return
        if msg.get("type") != "set_voice":
            return
        voice = msg.get("voice", "")
        if not voice:
            return
        logger.info("Switching TTS voice to %s", voice)
        try:
            tts = session.tts
            if hasattr(tts, "update_options"):
                tts.update_options(voice=voice)
                logger.info("Voice updated to %s", voice)
        except Exception as e:  # noqa: BLE001
            logger.warning("Failed to update voice: %s", e)

    ctx.room.on("data_received", _on_room_data)

    # Publish user/agent transcripts to the room so the web UI can render them.
    async def publish_transcript(role: str, text: str) -> None:
        text = (text or "").strip()
        if not text:
            return
        payload = json.dumps({"type": "transcript", "role": role, "text": text})
        try:
            await ctx.room.local_participant.publish_data(
                payload, reliable=True, topic="transcript"
            )
        except Exception as e:  # noqa: BLE001
            logger.warning("failed to publish transcript: %s", e)

    @session.on("conversation_item_added")
    def _on_conversation_item(ev) -> None:
        msg = ev.item
        role = getattr(msg, "role", None)
        text = getattr(msg, "text_content", None)
        if role == "user":
            asyncio.create_task(publish_transcript("user", text))
        elif role == "assistant":
            asyncio.create_task(publish_transcript("agent", text))

    @session.on("user_input_transcribed")
    def _on_user_transcribed(ev) -> None:
        logger.info("STT (%s): %s", "final" if ev.is_final else "partial", ev.transcript)

    @session.on("user_state_changed")
    def _on_user_state(ev) -> None:
        # Fires from VAD: if this never says "speaking", no usable mic
        # audio is arriving from the participant.
        logger.info("user state -> %s", ev.new_state)

    await session.start(
        agent=VoiceAssistant(),
        room=ctx.room,
        room_options=room_io.RoomOptions(
            # Keep the agent in the room when a participant leaves; it must
            # survive page reloads/reconnects, otherwise the next join lands
            # in an agent-less room (LiveKit does not reliably re-dispatch
            # into an existing room).
            close_on_disconnect=False,
            audio_input=room_io.AudioInputOptions(
                # No noise cancellation plugin (self-hosted, no ai-coustics)
            ),
        ),
    )

    await ctx.connect()
    logger.info("Agent connected to room %s", ctx.room.name)


if __name__ == "__main__":
    cli.run_app(server)