AIvoices/server-deno/models/grok.ts

import { Buffer } from "node:buffer";
import type { RawData } from "npm:@types/ws";
import { WebSocket } from "npm:ws";
import { addConversation, getDeviceInfo } from "../supabase.ts";
import { createOpusPacketizer, isDev, xaiApiKey, defaultGrokVoice } from "../utils.ts";

const XAI_REALTIME_URL = "wss://api.x.ai/v1/realtime";

export const connectToGrok = async ({
    ws,
    payload,
    connectionPcmFile,
    firstMessage,
    systemPrompt,
    closeHandler,
}: ProviderArgs) => {
    const { user, supabase } = payload;

    if (!xaiApiKey) {
        throw new Error("XAI_API_KEY is not set");
    }

    const voice = user.personality?.oai_voice ?? defaultGrokVoice;

    const opus = createOpusPacketizer((packet) => ws.send(packet));

    const grokWs = new WebSocket(XAI_REALTIME_URL, {
        headers: {
            Authorization: `Bearer ${xaiApiKey}`,
            "Content-Type": "application/json",
        },
    });

    let isConnected = false;
    const messageQueue: RawData[] = [];

    let createdSent = false;
    let outputTranscript = "";

    const sendResponseCreated = async () => {
        try {
            const device = await getDeviceInfo(supabase, user.user_id);
            opus.reset();
            ws.send(
                JSON.stringify({
                    type: "server",
                    msg: "RESPONSE.CREATED",
                    volume_control: device?.volume ?? 100,
                }),
            );
        } catch {
            ws.send(JSON.stringify({ type: "server", msg: "RESPONSE.CREATED" }));
        }
    };

    const sendFirstMessage = () => {
        if (!firstMessage) return;
        grokWs.send(
            JSON.stringify({
                type: "conversation.item.create",
                item: {
                    type: "message",
                    role: "user",
                    content: [{ type: "input_text", text: firstMessage }],
                },
            }),
        );
        grokWs.send(JSON.stringify({ type: "response.create" }));
    };

    grokWs.on("open", () => {
        isConnected = true;

        grokWs.send(
            JSON.stringify({
                type: "session.update",
                session: {
                    voice,
                    instructions: systemPrompt,
                    turn_detection: { type: "server_vad" },
                    audio: {
                        input: { format: { type: "audio/pcm", rate: 16000 } },
                        output: { format: { type: "audio/pcm", rate: 24000 } },
                    },
                },
            }),
        );

        sendFirstMessage();

        while (messageQueue.length > 0) {
            const queuedMessage = messageQueue.shift();
            if (queuedMessage) {
                messageHandler(queuedMessage, true);
            }
        }
    });

    grokWs.on("message", async (data: Buffer) => {
        let event: any;
        try {
            event = JSON.parse(data.toString("utf-8"));
        } catch {
            return;
        }

        try {
            switch (event.type) {
                case "response.created":
                    if (!createdSent) {
                        await sendResponseCreated();
                        createdSent = true;
                    }
                    break;

                case "response.output_audio_transcript.delta":
                    if (typeof event.delta === "string") {
                        outputTranscript += event.delta;
                    }
                    break;

                case "response.output_audio.delta":
                    if (typeof event.delta === "string") {
                        const pcmChunk = Buffer.from(event.delta, "base64");
                        // Use Opus packetizer to encode and send audio
                        opus.push(pcmChunk);
                    }
                    break;

                case "conversation.item.input_audio_transcription.completed":
                    if (typeof event.transcript === "string" && event.transcript.length > 0) {
                        await addConversation(supabase, "user", event.transcript, user);
                    }
                    break;

                case "input_audio_buffer.committed":
                    ws.send(JSON.stringify({ type: "server", msg: "AUDIO.COMMITTED" }));
                    break;

                case "response.done":
                    // Flush any remaining audio
                    opus.flush(true);

                    if (outputTranscript) {
                        await addConversation(supabase, "assistant", outputTranscript, user);
                        outputTranscript = "";
                    }
                    ws.send(JSON.stringify({ type: "server", msg: "RESPONSE.COMPLETE" }));
                    createdSent = false;
                    break;

                case "error":
                    ws.send(JSON.stringify({ type: "server", msg: "RESPONSE.ERROR" }));
                    createdSent = false;
                    break;
            }
        } catch (err) {
            console.error("Error processing Grok event:", err);
            ws.send(JSON.stringify({ type: "server", msg: "RESPONSE.ERROR" }));
            createdSent = false;
        }
    });

    grokWs.on("close", () => {
        ws.close();
    });

    grokWs.on("error", (error: any) => {
        console.error("Grok WebSocket error:", error);
        ws.send(JSON.stringify({ type: "server", msg: "RESPONSE.ERROR" }));
    });

    const messageHandler = async (data: RawData, isBinary: boolean) => {
        if (isBinary) {
            const base64Data = (data as Buffer).toString("base64");
            grokWs.send(JSON.stringify({ type: "input_audio_buffer.append", audio: base64Data }));

            if (isDev && connectionPcmFile) {
                await connectionPcmFile.write(data as Buffer);
            }
            return;
        }

        let message: any;
        try {
            message = JSON.parse((data as Buffer).toString("utf-8"));
        } catch {
            return;
        }

        if (message?.type !== "instruction") return;

        if (message.msg === "end_of_speech") {
            grokWs.send(JSON.stringify({ type: "input_audio_buffer.commit" }));
            grokWs.send(JSON.stringify({ type: "response.create" }));
            grokWs.send(JSON.stringify({ type: "input_audio_buffer.clear" }));
        } else if (message.msg === "INTERRUPT") {
            grokWs.send(JSON.stringify({ type: "input_audio_buffer.clear" }));
        }
    };

    ws.on("message", (data: RawData, isBinary: boolean) => {
        if (!isConnected) {
            messageQueue.push(data);
        } else {
            messageHandler(data, isBinary);
        }
    });

    ws.on("error", (error: any) => {
        console.error("ESP32 WebSocket error:", error);
        grokWs.close();
    });

    ws.on("close", async (code: number, reason: string) => {
        console.log(`ESP32 WebSocket closed with code ${code}, reason: ${reason}`);
        await closeHandler();
        opus.close();
        grokWs.close();
        if (isDev && connectionPcmFile) {
            connectionPcmFile.close();
        }
    });

    return new Promise<void>((resolve, reject) => {
        const timeout = setTimeout(() => reject(new Error("Grok connection timeout")), 10000);
        grokWs.on("open", () => {
            clearTimeout(timeout);
            resolve();
        });
        grokWs.on("error", (error: any) => {
            clearTimeout(timeout);
            reject(error);
        });
    });
};