AmazonAIConvertPredictionsProvider.ts

import { AbstractConvertPredictionsProvider } from "../types/Providers/AbstractConvertPredictionsProvider";
import * as Translate from 'aws-sdk/clients/translate';
import * as TextToSpeech from 'aws-sdk/clients/polly';
import {
    TranslateTextInput, TextToSpeechInput,
    SpeechToTextInput, TranslateTextOutput, TextToSpeechOutput,
    SpeechToTextOutput, isBytesSource} from "../types";
import { Credentials, ConsoleLogger as Logger, Signer } from '@aws-amplify/core';
import { EventStreamMarshaller, MessageHeaderValue } from '@aws-sdk/eventstream-marshaller';
import { fromUtf8, toUtf8 } from '@aws-sdk/util-utf8-node';

const logger = new Logger('AmazonAIConvertPredictionsProvider');
const eventBuilder = new EventStreamMarshaller(toUtf8, fromUtf8);

export default class AmazonAIConvertPredictionsProvider extends AbstractConvertPredictionsProvider {

    private translate: Translate;
    private textToSpeech: TextToSpeech;
    constructor() {
        super();
    }

    getProviderName() {
        return "AmazonAIConvertPredictionsProvider";
    }

    protected translateText(input: TranslateTextInput): Promise<TranslateTextOutput> {
        logger.debug("Starting translation");
        return new Promise(async (res, rej) => {
            const { translateText: {
                defaults: { sourceLanguage = "", targetLanguage = "" } = {},
                region = "" } = {}
            } = this._config;

            if (!region) {
                return rej("region not configured for transcription");
            }

            const credentials = await Credentials.get();
            if (!credentials) { return rej('No credentials'); }
            const sourceLanguageCode = input.translateText.source.language || sourceLanguage;
            const targetLanguageCode = input.translateText.targetLanguage || targetLanguage;
            if (!sourceLanguageCode || !targetLanguageCode) {
                return rej("Please provide both source and target language");
            }

            this.translate = new Translate({ region, credentials });
            this.translate.translateText({
                SourceLanguageCode: sourceLanguageCode,
                TargetLanguageCode: targetLanguageCode,
                Text: input.translateText.source.text
                // tslint:disable-next-line: align
            }, (err, data) => {
                logger.debug({ err, data });
                if (err) {
                    return rej(err);
                } else {
                    return res({ text: data.TranslatedText, language: data.TargetLanguageCode } as TranslateTextOutput);
                }
            });
        });
    }

    protected convertTextToSpeech(input: TextToSpeechInput): Promise<TextToSpeechOutput> {
        return new Promise(async (res, rej) => {
            const credentials = await Credentials.get();
            if (!credentials) { return rej('No credentials'); }
            const { speechGenerator: {
                defaults: { VoiceId = "" } = {},
                region = "" } = {}
            } = this._config;

            if(!input.textToSpeech.source) {
                return rej("Source needs to be provided in the input");
            }
            const voiceId = input.textToSpeech.voiceId || VoiceId;

            if(!region) {
                return rej("Region was undefined. Did you enable speech generator using amplify CLI?");
            }

            if(!voiceId) {
                return rej("VoiceId was undefined.");
            }

            this.textToSpeech = new TextToSpeech({ region, credentials });
            this.textToSpeech.synthesizeSpeech({
                OutputFormat: 'mp3',
                Text: input.textToSpeech.source.text,
                VoiceId: voiceId,
                TextType: 'text',
                SampleRate: '24000'
                // tslint:disable-next-line: align
            }, (err, data) => {
                if (err) {
                    rej(err);
                } else {
                    const blob = new Blob([data.AudioStream as ArrayBuffer], { type: data.ContentType });
                    const url = URL.createObjectURL(blob);
                    res({
                        speech: { url },
                        audioStream: (data.AudioStream as any).buffer,
                        text: input.textToSpeech.source.text
                    } as TextToSpeechOutput);
                }
            });
        });
    }

    protected convertSpeechToText(input: SpeechToTextInput): Promise<SpeechToTextOutput> {
        return new Promise(async (res, rej) => {
            try {
                logger.debug('starting transcription..');
                const credentials = await Credentials.get();
                if (!credentials) { return rej('No credentials'); }
                const { transcription: {
                    defaults: { language: languageCode = "" } = {},
                    region = "" } = {}
                } = this._config;
                if (!region) {
                    return rej("region not configured for transcription");
                }
                if (!languageCode) {
                    return rej("languageCode not configured or provided for transcription");
                }
                const { transcription: { source, language = languageCode } } = input;

                if (isBytesSource(source)) {
                    const connection
                        = await this.openConnectionWithTranscribe({ credentials, region, languageCode: language });

                    try {
                        const fullText = await this.sendDataToTranscribe({ connection, raw: source.bytes });
                        return res({
                            transcription: {
                                fullText,
                            }
                        });
                    } catch (err) {
                        rej(err);
                    }
                }

                return rej("Source types other than byte source are not supported.");
            } catch (err) {
                return rej(err.name + ': ' + err.message);
            }
        });
    }

    public static serializeDataFromTranscribe(message) {
        let decodedMessage = "";
        const transcribeMessage = eventBuilder.unmarshall(Buffer.from(message.data));
        const transcribeMessageJson = JSON.parse(String.fromCharCode.apply(String, transcribeMessage.body));
        if (transcribeMessage.headers[":message-type"].value === "exception") {
            logger.debug('exception', JSON.stringify(transcribeMessageJson.Message, null, 2));
            throw new Error(transcribeMessageJson.Message);
        }
        else if (transcribeMessage.headers[":message-type"].value === "event") {
            if (transcribeMessageJson.Transcript.Results.length > 0) {
                if (transcribeMessageJson.Transcript.Results[0].Alternatives.length > 0) {
                    if (transcribeMessageJson.Transcript.Results[0].Alternatives[0].Transcript.length > 0) {
                        if (transcribeMessageJson.Transcript.Results[0].IsPartial === false) {
                            decodedMessage = 
                                transcribeMessageJson.Transcript.Results[0].Alternatives[0].Transcript + "\n";
                            logger.debug({ decodedMessage });
                        } else {
                            logger.debug({
                                transcript: transcribeMessageJson.Transcript.Results[0].Alternatives[0]
                            });
                        }
                    }
                }
            }
        }
        return decodedMessage;
    }

    private sendDataToTranscribe({ connection, raw }): Promise<string> {
        return new Promise((res, rej) => {
            let fullText = "";
            connection.onmessage = (message) => {
                try {
                    const decodedMessage = AmazonAIConvertPredictionsProvider.serializeDataFromTranscribe(message);
                    if (decodedMessage) {
                        fullText += decodedMessage + " ";
                    }
                } catch (err) {
                    logger.debug(err);
                    rej(err.message);
                }
            };

            connection.onerror = (errorEvent) => {
                logger.debug({ errorEvent });
                rej('failed to transcribe, network error');
            };

            connection.onclose = (closeEvent) => {
                logger.debug({ closeEvent });
                return res(fullText.trim());
            };

            logger.debug({ raw });

            if (Array.isArray(raw)) {
                for (let i = 0; i < raw.length - 1023; i += 1024) {
                    const data = raw.slice(i, i + 1024);
                    this.sendEncodedDataToTranscribe(connection, data);
                }
            }

            // sending end frame
            const endFrameEventMessage = this.getAudioEventMessage(Buffer.from([]));
            const endFrameBinary = eventBuilder.marshall(endFrameEventMessage);
            connection.send(endFrameBinary);
        });
    }

    private sendEncodedDataToTranscribe(connection, data) {
        const downsampledBuffer = this.downsampleBuffer({ buffer: data });
        const pcmEncodedBuffer = this.pcmEncode(downsampledBuffer);
        const audioEventMessage = this.getAudioEventMessage(Buffer.from(pcmEncodedBuffer));
        const binary = eventBuilder.marshall(audioEventMessage);
        connection.send(binary);
    }

    private getAudioEventMessage(buffer) {
        const audioEventMessage = {
            body: buffer as Uint8Array,
            headers: {
                ':message-type': {
                    type: 'string',
                    value: 'event'
                } as MessageHeaderValue,
                ':event-type': {
                    type: 'string',
                    value: 'AudioEvent'
                } as MessageHeaderValue
            },
        };

        return audioEventMessage;
    }

    private pcmEncode(input) {
        let offset = 0;
        const buffer = new ArrayBuffer(input.length * 2);
        const view = new DataView(buffer);
        for (let i = 0; i < input.length; i++ , offset += 2) {
            const s = Math.max(-1, Math.min(1, input[i]));
            view.setInt16(offset, s < 0 ? s * 0x8000 : s * 0x7FFF, true);
        }
        return buffer;
    }

    private inputSampleRate = 44100;
    private outputSampleRate = 16000;

    private downsampleBuffer({ buffer }) {

        if (this.outputSampleRate === this.inputSampleRate) {
            return buffer;
        }

        const sampleRateRatio = this.inputSampleRate / this.outputSampleRate;
        const newLength = Math.round(buffer.length / sampleRateRatio);
        const result = new Float32Array(newLength);
        let offsetResult = 0;
        let offsetBuffer = 0;
        while (offsetResult < result.length) {
            const nextOffsetBuffer = Math.round((offsetResult + 1) * sampleRateRatio);
            let accum = 0,
                count = 0;
            for (let i = offsetBuffer; i < nextOffsetBuffer && i < buffer.length; i++) {
                accum += buffer[i];
                count++;
            }
            result[offsetResult] = accum / count;
            offsetResult++;
            offsetBuffer = nextOffsetBuffer;
        }

        return result;
    }

    private openConnectionWithTranscribe({ credentials: userCredentials, region, languageCode }): Promise<WebSocket> {
        return new Promise(async (res, rej) => {
            const { accessKeyId: access_key,
                secretAccessKey: secret_key,
                sessionToken: session_token
            } = userCredentials;

            const credentials = {
                access_key,
                secret_key,
                session_token
            };

            const signedUrl = this.generateTranscribeUrl({ credentials, region, languageCode });

            logger.debug('connecting...');
            const connection = new WebSocket(signedUrl);

            connection.binaryType = "arraybuffer";
            connection.onopen = () => {
                logger.debug('connected');
                res(connection);
            };
            
        });
    }

    private generateTranscribeUrl({ credentials, region, languageCode }): string {
        const url = [`wss://transcribestreaming.${region}.amazonaws.com:8443`,
            '/stream-transcription-websocket?',
            `media-encoding=pcm&`,
            `sample-rate=16000&`,
        `language-code=${languageCode}`]
            .join('');

        const signedUrl = Signer.signUrl(url, credentials, { region, service: 'transcribe' }, 300);

        return signedUrl;
    }

}