Ai.Audio.Continuous.js

const Security = require("Security");
const Docker = require("Docker");
const WebSocket = require("ws");

const Ai = {
    Audio: {
        Continuous: (seedPrompt, settings, saveChunk, getPrompt) => {
            Security.String(seedPrompt)
            Security.Check(seedPrompt.length > 0, `Prompt is too short, please provide more context`);
            Security.Object("Ai.Audio.Continuous")(settings)
            Security.Callback(saveChunk)
            Security.Callback(getPrompt)

            const serviceName = `ai-continuous`;
            const serverPort = 3003;
            const serverUri = `ws://localhost:${serverPort}`;
            delete settings._is

            const isStarted = Docker.Container.Exists(serviceName, true);
            const connHandler = () => {
                console.log("Attempting to connect...")
                const ws = new WebSocket(serverUri);
                ws.on('open', function () {
                    const payload = JSON.stringify({
                        started: true,
                        prompt: seedPrompt,
                        settings
                    });
                    console.log("Handshake complete")
                    ws.send(payload);
                });

                ws.on('message', function (data) {
                    const {success, chunk, progress} = JSON.parse(data);
                    if (success !== undefined) {
                        console.log("Server finished")
                    } else if (chunk !== undefined) {
                        saveChunk(chunk)
                    } else if (progress !== undefined) {
                        const newPrompt = getPrompt()
                        saveChunk(progress)
                        console.log("sending new chunk - ", {
                            finished: newPrompt==="",
                            prompt: newPrompt
                        })
                        ws.send(JSON.stringify({
                            finished: newPrompt==="",
                            prompt: newPrompt
                        }))
                    } else {
                        console.log(data)
                        throw new Error(`Unrecognized server response`);
                    }
                });
                ws.on('error', function (error) {
                    setTimeout(connHandler, 5000)
                });
            }
            if (!isStarted) {
                console.log("Starting server...")
                const isGpu = settings.device.toLowerCase() === "cuda"
                let containerId = Docker.Service.AI.AUDIO_STREAM(serviceName, serverPort, isGpu);
                console.log("Started audiocraft server with ID: " + containerId)
                setTimeout(connHandler, 10000);
            }
            else {
                console.log("Ai server already running")
                connHandler()
            }
        }
    }
}

module.exports = Ai.Audio.Continuous