Grok Voice

from __future__ import annotations
import asyncio
import base64
import json
import logging
import os
from typing import Any, Optional
import websockets
from dotenv import load_dotenv
from websockets.asyncio.client import ClientConnection
from websockets.exceptions import ConnectionClosed, InvalidStatus
from agentduet import (
    BufferFullError,
    Call,
    CallAudioConfig,
    CallClosedError,
    IncomingCallNotification,
    SessionManager,
    SessionManagerConfig,
    new_session_id,
)
load_dotenv()
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(name)s %(levelname)s %(message)s",
)
logger = logging.getLogger(__name__)
AGENTDUET_API_KEY = os.environ["AGENTDUET_API_KEY"]
AGENTDUET_CONNECTOR_UUID = os.environ["AGENTDUET_CONNECTOR_UUID"]
XAI_API_KEY = os.environ["XAI_API_KEY"]
GROK_MODEL = "grok-voice-think-fast-1.0"
GROK_REALTIME_URL = f"wss://api.x.ai/v1/realtime?model={GROK_MODEL}"
GROK_SAMPLE_RATE = 24000
GROK_VOICE = os.environ.get("GROK_VOICE", "leo").lower()
AGENT_NAME = "Grok"
SYSTEM_PROMPT = (
    f"Your name is {AGENT_NAME}. You are a witty, helpful voice assistant on a phone call. "
    "Greet the caller briefly, keep answers short and conversational, and ask clarifying "
    "questions when needed. When the caller wants to hang up or says goodbye, say a brief "
    "goodbye and call the hang_up tool."
)
HANG_UP_TOOL = {
    "type": "function",
    "name": "hang_up",
    "description": (
        "End the phone call. Use when the caller asks to hang up, "
        "end the call, or says goodbye and wants to leave."
    ),
    "parameters": {"type": "object", "properties": {}, "required": []},
}
class GrokLiveIntegration:
    """Bidirectional audio bridge: AgentDuet Call ↔ xAI Grok Realtime."""
    def __init__(self, call: Call, grok_ws: ClientConnection):
        self._call = call
        self._grok_ws = grok_ws
        self._send_to_grok_task: Optional[asyncio.Task] = None
        self._recv_from_grok_task: Optional[asyncio.Task] = None
        self._playback_task: Optional[asyncio.Task] = None
        self._terminated = False
        # Playback must not block the Grok event loop, or speech_started
        # arrives too late for clear_send_audio_buffer() to help.
        self._playback_gen = 0
        self._audio_queue: asyncio.Queue[tuple[int, bytes] | None] = asyncio.Queue()
        self._active_response_id: Optional[str] = None
        self._cancelled_response_ids: set[str] = set()
    async def _on_hangup(self, _evt: Any) -> None:
        logger.info("Call %s hung up", self._call.id)
        self._terminated = True
        try:
            await self._grok_ws.close()
        except Exception:
            logger.exception("Error closing Grok WebSocket")
        for task in (
            self._send_to_grok_task,
            self._recv_from_grok_task,
            self._playback_task,
        ):
            if task and not task.done():
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
    async def run(self) -> None:
        self._call.on_hangup(self._on_hangup)
        await self._configure_session()
        await self._greet_caller()
        self._playback_task = asyncio.create_task(self._playback_worker())
        self._send_to_grok_task = asyncio.create_task(self._stream_to_grok())
        self._recv_from_grok_task = asyncio.create_task(self._receive_from_grok())
        await asyncio.gather(
            self._send_to_grok_task,
            self._recv_from_grok_task,
            return_exceptions=True,
        )
        if self._playback_task and not self._playback_task.done():
            self._audio_queue.put_nowait(None)
            try:
                await self._playback_task
            except asyncio.CancelledError:
                pass
    async def _configure_session(self) -> None:
        await self._grok_ws.send(
            json.dumps(
                {
                    "type": "session.update",
                    "session": {
                        "voice": GROK_VOICE,
                        "instructions": SYSTEM_PROMPT,
                        "reasoning": {"effort": "none"},
                        "turn_detection": {
                            "type": "server_vad",
                            "threshold": 0.5,
                            "silence_duration_ms": 300,
                            "prefix_padding_ms": 200,
                        },
                        "tools": [HANG_UP_TOOL],
                        "audio": {
                            "input": {
                                "format": {
                                    "type": "audio/pcm",
                                    "rate": GROK_SAMPLE_RATE,
                                },
                            },
                            "output": {
                                "format": {
                                    "type": "audio/pcm",
                                    "rate": GROK_SAMPLE_RATE,
                                },
                            },
                        },
                    },
                }
            )
        )
    async def _greet_caller(self) -> None:
        await self._grok_ws.send(
            json.dumps(
                {
                    "type": "response.create",
                    "response": {
                        "instructions": (
                            f"Greet the caller warmly. Introduce yourself as {AGENT_NAME} "
                            "and ask how you can help today."
                        ),
                    },
                }
            )
        )
    async def _stream_to_grok(self) -> None:
        try:
            async for chunk in self._call.caller.audio_stream():
                if self._terminated:
                    break
                await self._grok_ws.send(
                    json.dumps(
                        {
                            "type": "input_audio_buffer.append",
                            "audio": base64.b64encode(chunk).decode("ascii"),
                        }
                    )
                )
        except (CallClosedError, ConnectionClosed):
            pass
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("Error streaming caller audio to Grok")
            raise
    async def _playback_worker(self) -> None:
        while True:
            item = await self._audio_queue.get()
            if item is None:
                break
            gen, audio = item
            if gen != self._playback_gen:
                continue
            try:
                await self._call.send_audio(audio)
            except BufferFullError:
                logger.warning("Outgoing buffer full - dropping chunk")
            except CallClosedError:
                break
            except asyncio.CancelledError:
                raise
    async def _flush_playback(self) -> None:
        self._playback_gen += 1
        self._active_response_id = None
        while not self._audio_queue.empty():
            try:
                self._audio_queue.get_nowait()
            except asyncio.QueueEmpty:
                break
        try:
            await self._call.clear_send_audio_buffer()
        except CallClosedError:
            return
    async def _agent_hang_up(self) -> None:
        logger.info("Agent hanging up call %s", self._call.id)
        try:
            for _ in range(40):  # up to ~4s for goodbye audio to drain
                if await self._call.get_send_audio_buffer_size() == 0:
                    break
                await asyncio.sleep(0.1)
        except CallClosedError:
            return
        result = await self._call.close()
        if not result:
            logger.error(
                "Hang up failed for %s: %s (%s)",
                self._call.id,
                result.error_message,
                result.error_code,
            )
    async def _receive_from_grok(self) -> None:
        try:
            async for raw_event in self._grok_ws:
                if self._terminated:
                    break
                event = json.loads(raw_event)
                etype = event.get("type")
                if etype == "response.created":
                    self._active_response_id = event.get("response", {}).get("id")
                elif etype == "input_audio_buffer.speech_started":
                    try:
                        buf_size = await self._call.get_send_audio_buffer_size()
                    except CallClosedError:
                        buf_size = 0
                    should_flush = (
                        self._active_response_id is not None
                        or buf_size > 0
                        or self._audio_queue.qsize() > 0
                    )
                    if should_flush:
                        if self._active_response_id is not None:
                            self._cancelled_response_ids.add(self._active_response_id)
                        logger.info("Caller interrupted - stopping playback")
                        await self._flush_playback()
                elif etype == "response.done":
                    response_id = event.get("response", {}).get("id")
                    if response_id:
                        self._cancelled_response_ids.discard(response_id)
                    if response_id == self._active_response_id:
                        self._active_response_id = None
                elif etype in ("response.output_audio.delta", "response.audio.delta"):
                    response_id = event.get("response_id")
                    if response_id and response_id in self._cancelled_response_ids:
                        continue
                    if (
                        response_id
                        and self._active_response_id
                        and response_id != self._active_response_id
                    ):
                        continue
                    audio = base64.b64decode(event["delta"])
                    self._audio_queue.put_nowait((self._playback_gen, audio))
                elif etype == "response.function_call_arguments.done":
                    if event.get("name") != "hang_up":
                        logger.warning("Unknown tool: %s", event.get("name"))
                        continue
                    await self._grok_ws.send(
                        json.dumps(
                            {
                                "type": "conversation.item.create",
                                "item": {
                                    "type": "function_call_output",
                                    "call_id": event["call_id"],
                                    "output": json.dumps({"status": "hanging_up"}),
                                },
                            }
                        )
                    )
                    await self._agent_hang_up()
                    return
                elif etype == "error":
                    logger.error("Grok error event: %s", event)
        except (CallClosedError, ConnectionClosed):
            pass
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("Error receiving from Grok")
            raise
        finally:
            self._audio_queue.put_nowait(None)
async def bridge_call_to_grok(call: Call) -> None:
    try:
        async with websockets.connect(
            GROK_REALTIME_URL,
            additional_headers={"Authorization": f"Bearer {XAI_API_KEY}"},
            open_timeout=15,
        ) as grok_ws:
            result = await call.answer()
            if not result:
                logger.error(
                    "Answer failed for call %s: %s (%s)",
                    call.id,
                    result.error_message,
                    result.error_code,
                )
                return
            await GrokLiveIntegration(call, grok_ws).run()
    except InvalidStatus as e:
        body = getattr(e.response, "body", b"") or b""
        detail = body.decode("utf-8", errors="replace") if body else str(e)
        logger.error(
            "Grok WebSocket rejected (HTTP %s). Check XAI_API_KEY - %s",
            e.response.status_code,
            detail,
        )
    except Exception:
        logger.exception("Failed during Grok bridge for call %s", call.id)
async def main() -> None:
    config = SessionManagerConfig.create(
        api_key=AGENTDUET_API_KEY,
        connector_uuid=AGENTDUET_CONNECTOR_UUID,
        call_audio=CallAudioConfig(
            sample_rate=GROK_SAMPLE_RATE,
            buffer_size=1024 * 1024,
        ),
    )
    async with SessionManager(config) as sm:
        logger.info("Connected. Waiting for calls...")
        @sm.on_incoming_call
        async def on_call(noti: IncomingCallNotification) -> None:
            logger.info("Incoming call %s from %s", noti.call_id, noti.participant)
            session = await sm.open_session(new_session_id(), noti.subscriber)
            call = await session.process_call(noti)
            try:
                await bridge_call_to_grok(call)
            finally:
                await call.close()
        await sm.run_forever()
if __name__ == "__main__":
    asyncio.run(main())