streaming-json.py

#!/usr/local/bin/python3
import datetime
import itertools
import json
import logging
import os
import time
import uuid
from argparse import ArgumentParser
from datetime import datetime

import grpc
# import sounddevice # required in mac m4
import pyaudio
import requests
from six.moves import queue

import speech_pb2
import speech_pb2_grpc
from webex_ci_sso import sso_access

parser = ArgumentParser(description='Audio Client')
parser.add_argument('--speech_hints', type=str, default="")
parser.add_argument('--expected_language_codes', type=str, default="")
parser.add_argument('--language_code', type=str, default="en")
parser.add_argument('--sample_rate', type=int, default=16000)
parser.add_argument('--encoding', type=str, default="LINEAR16")
parser.add_argument('--enable_automatic_punctuation', type=bool, default=True)
parser.add_argument('--enable_multilang_detection', type=str, default=True)
parser.add_argument('--enable_training_deprecated', type=bool, default=False)
parser.add_argument('--final-only', type=bool, default=False),
parser.add_argument('--file', type=str, default="")
parser.add_argument('--longrunning', type=bool, default=False)
parser.add_argument('--url', type=str, default="")
parser.add_argument('--callback_url', type=str, default="")
parser.add_argument('--engine', type=str, default="k2")
parser.add_argument('--enable_phonetic_output', type=bool, default=False)
parser.add_argument('--server', type=str, default="asr-proxy.uscentral1-0.cint.vcra.co:443")
parser.add_argument('--config_key_name', type=str, default="")
parser.add_argument('--recording_upload', type=bool, default=False)
parser.add_argument('--csi_file', type=str, default="")
parser.add_argument('--csi_file_url', type=str, default="")
parser.add_argument('--meeting_id', type=str, default="")
parser.add_argument('--get_transcript', type=bool, default=False)
parser.add_argument('--cleanup', type=str, default="yes")
parser.add_argument('--json_output', type=str, default="no")


args = parser.parse_args()

supportedAudioFormats = {
    "mp3": "audio/mp3",
    "mpeg": "audio/mpeg",
    "ogg": "audio/ogg",
    "wav": "audio/wav",
    "flac": "audio/flac",
}

_LOGGER = logging.getLogger(__name__)
_LOGGER.setLevel(logging.INFO)

_SERVER_ADDR_TEMPLATE = args.server

#ci related:
USER_ID=os.getenv("CI_USER_ID","voicea-asr-sample-client-bts")
USER_PASSWD=os.getenv("CI_USER_PASSWD", "UGMP.stgh.38.AGXI.vwym.37.XHKL.rwkp.2578") # Vault path: cint/asrproxy/sample_python_client
USER_REALM=os.getenv("CI_ORG_ID","a93dde14-65b7-4a59-81dd-28962a8473e3")
BROKER_HOST=os.getenv("BORKER_HOST", "idbrokerbts.webex.com")
CLIENT_ID=os.getenv("CI_CLIENT_ID", "C6ac166d70e8ea8a59e0a341c7690ddc0e3de836bc1c702d31c211501358b594b")
CLIENT_SECRET=os.getenv("CI_CLIENT_SECRET", "19d82f38891800bcd9bc636d986b6a0d037ec0fc7aae72161e60c7eb0c67e930")
SCOPE=os.getenv("CI_SCOPE", "Identity:SCIM Identity:Organization Identity:Config webex-voice-assistant:asr")
__ci_token__ = "" 
__token_expire__ = None
__sso__ = sso_access(orgId=USER_REALM, machineAccount=USER_ID, machineAccountPW=USER_PASSWD, oauth_host=BROKER_HOST)
 
class MicrophoneStream(object):
    """Opens a recording stream as a generator yielding the audio chunks."""
    def __init__(self, rate, chunk):
        self._rate = rate
        self._chunk = chunk

        # Create a thread-safe buffer of audio data
        self._buff = queue.Queue()
        self.closed = True

    def __enter__(self):
        self._audio_interface = pyaudio.PyAudio()
        self._audio_stream = self._audio_interface.open(
            format=pyaudio.paInt16,
            # The API currently only supports 1-channel (mono) audio
            # https://goo.gl/z757pE
            channels=1, rate=self._rate,
            input=True, frames_per_buffer=self._chunk,
            # Run the audio stream asynchronously to fill the buffer object.
            # This is necessary so that the input device's buffer doesn't
            # overflow while the calling thread makes network requests, etc.
            stream_callback=self._fill_buffer,
        )

        self.closed = False

        return self

    def __exit__(self, type, value, traceback):
        self._audio_stream.stop_stream()
        self._audio_stream.close()
        self.closed = True
        # Signal the generator to terminate so that the client's
        # streaming_recognize method will not block the process termination.
        self._buff.put(None)
        self._audio_interface.terminate()

    def _fill_buffer(self, in_data, frame_count, time_info, status_flags):
        """Continuously collect data from the audio stream, into the buffer."""
        self._buff.put(in_data)
        return None, pyaudio.paContinue

    def generator(self):
        while not self.closed:
            # Use a blocking get() to ensure there's at least one chunk of
            # data, and stop iteration if the chunk is None, indicating the
            # end of the audio stream.
            chunk = self._buff.get()
            if chunk is None:
                return
            data = [chunk]

            # Now consume whatever other data's still buffered.
            while True:
                try:
                    chunk = self._buff.get(block=False)
                    if chunk is None:
                        return
                    data.append(chunk)
                except queue.Empty:
                    break

            yield b''.join(data)


class FileStream(object):
    def __init__(self, file_name, rate, chunk_size):
        self.file_name = file_name
        self.rate = rate
        self.chunk_size = chunk_size
        self.file = open(self.file_name, 'rb')

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.file.close()

    def generator(self):
        chunks = chunk_array(self.file.read(), self.chunk_size)
        for chunk in chunks:
            yield chunk


def make_channel():
    """Creates a secure channel with auth credentials from the environment."""
    # Grab application default credentials from the environment
    creds = grpc.ssl_channel_credentials()
    #options = (('grpc.ssl_target_name_override', 'nutcracker.grpc.staging.vcra.co',),)
    max_message_length = 16 * 1024 * 1024  # Example: 16 MB
    options=[
        ('grpc.max_send_message_length', max_message_length),
        ('grpc.max_receive_message_length', max_message_length)
    ]
    #secure
    if _SERVER_ADDR_TEMPLATE.endswith("443"):
        return grpc.secure_channel(_SERVER_ADDR_TEMPLATE, creds, options=options)
    else:
        return grpc.insecure_channel(_SERVER_ADDR_TEMPLATE, options=options)


def getAccessToken():
    global __ci_token__
    global __token_expire__
    if 1==1:
        if __ci_token__ == "" or \
           (__token_expire__ - datetime.now()).total_seconds() < 5:
            __ci_token__, __token_expire__ = __sso__.getAccessToken(CLIENT_ID, CLIENT_SECRET, scope=SCOPE)
        
        # print ("token", __ci_token__)
        return __ci_token__
    return None

def chunk_array(arr, length):
    return [arr[0+i:length+i] for i in range(0, len(arr), length)]


def word_time_to_seconds(word_time):
    return word_time.seconds + (word_time.nanos / 1_000_000_000)


def print_utterances(transcripts):
    for tr in transcripts:
        _time_from = (
            word_time_to_seconds(tr.start_time)
            if tr.start_time
            else None
        )
        _time_to = (
            word_time_to_seconds(tr.end_time)
            if tr.end_time
            else None
        )
        _speaker_id = tr.speaker_ids[0] if len(tr.speaker_ids) > 0 else None
        _speaker_label = tr.speaker_labels[0] if len(tr.speaker_labels) > 0 else None
        
        meta = ""
        if _speaker_id:
            meta += f"[speaker_id:{_speaker_id}]"
        if _speaker_label:
            meta += f"[speaker_label:{_speaker_label}]"
        if tr.language_code:
            meta += f"[language:{tr.language_code}]"

        time_range = ""
        if _time_from and _time_to:
            time_range = f"[{_time_from}-{_time_to}] "

        text = tr.alternatives[0].transcript
        
        if meta:
            print(f"{meta}")
        print(f"{time_range}{text}")
        print("=" * 80)


def recording_upload(stub, audio_file_name, csi_file_name, language_code):
    tracking_id = 'asr_client_python_'+str(uuid.uuid4())
    print("Tracking ID will be " + str(tracking_id))
    metadata = (('authorization', 'Bearer ' + getAccessToken()),
                ('trackingid', str(tracking_id)),
                ('configkeyname', str(args.config_key_name)))
    audio_file_size = os.path.getsize(audio_file_name)
    csi_file_size = 0
    if csi_file_name:
        csi_file_size = os.path.getsize(csi_file_name)
    audio_file_extension = audio_file_name.split(".")[-1]
    audio_file_format = supportedAudioFormats[audio_file_extension]
    print("audio_file_size " + str(audio_file_size))
    print("csi_file_size " + str(csi_file_size))
    recording_upload_request = speech_pb2.RecordingUploadRequest(recording_size=audio_file_size, csi_size=csi_file_size,
                                                                 recording_type=audio_file_format, callback_url="",
                                                                 language_code=language_code)
    response = stub.RecordingUpload(recording_upload_request, metadata=metadata)
    if len(response.meeting_id) == 0:
        return ""
    meeting_id = response.meeting_id
    start_time = datetime.now()
    print("Meeting ID is " + str(meeting_id))
    print("csi upload URL " + str(response.csi_upload_url))
    print("audio upload URL " + str(response.audio_upload_url))

    audio_file = open(audio_file_name, 'rb')
    audio_request_headers = {"Content-Type": "application/octet-stream", "Content-Length": str(audio_file_size)}
    audio_file_upload_response = requests.put(response.audio_upload_url, data=audio_file, headers=audio_request_headers)
    print("audio upload response " + str(audio_file_upload_response.status_code))

    if csi_file_size > 0:
        csi_files = open(csi_file_name, 'rb')
        csi_request_headers = {"Content-Type": "application/json", "Content-Length": str(csi_file_size)}
        csi_file_upload_response = requests.put(response.csi_upload_url, data=csi_files, headers=csi_request_headers)
        print("csi upload response " + str(csi_file_upload_response.status_code))

    upload_completed_request = speech_pb2.UploadCompletedRequest(meeting_id=meeting_id)
    stub.UploadCompleted(upload_completed_request, metadata=metadata)

    done = False
    while not done:
        (done, transcripts, words, jsonOutput) = getLongRunningRecognizeOperation(stub, meeting_id)
        current = datetime.now()
        elapsed = current - start_time
        print("time elapsed", elapsed)
        time.sleep(5)

    if args.json_output == "yes" and jsonOutput is not None:
        with open(f'{meeting_id}.json', 'w+') as fd:
            json.dump(jsonOutput, fd)
        print(f"JSON output written to {meeting_id}.json")
        print(jsonOutput)
    else:
        print("Final Utterances:")
        print_utterances(transcripts)

    print("args: %s" % args)
    if args.cleanup == "yes":
        print("meeting delete: %s" % meeting_id)
        deleteLongRunningRecognizeOperation(stub, meeting_id)

    #deleteLongRunningRecognizeOperation(stub, meeting_id)
    return meeting_id

def longRunningRecognize(stub, url, fileName, callbackUrl, expectedLanguageCodes, csi_file_url=""):
    trackingId = 'asr_client_python_'+str(uuid.uuid4())
    print("Tracking ID will be "+ str(trackingId))
    if fileName:
        with open(fileName, 'rb') as f:
            data = f.read()
        RATE = args.sample_rate
        metadata = (('authorization', 'Bearer ' + getAccessToken()),
                    ('trackingid', str(trackingId)),
                    ('configkeyname', str(args.config_key_name)))
        request_config = speech_pb2.RecognitionConfig(sample_rate_hertz=RATE,max_alternatives=1)
        request_config.enable_multilang_detection = args.enable_multilang_detection
        request_config.language_code = args.language_code
        if(len(expectedLanguageCodes)>0):
            request_config.enable_automatic_language_detection = True
            request_config.expected_language_codes.extend(expectedLanguageCodes)
        request_config.enable_training_deprecated = args.enable_training_deprecated        
        request_config.enable_phonetic_output = args.enable_phonetic_output
        # only LINEAR16 is supported for now when sending bytes directly
        request_config.encoding = speech_pb2.RecognitionConfig.LINEAR16 
        print(request_config)
        request_audio = speech_pb2.RecognitionAudio(content=data)
        
        recognize_request = speech_pb2.LongRunningRecognizeRequest(
                config=request_config, audio=request_audio, callback_url=callbackUrl, csi_file_url=csi_file_url)
        
    else:
        RATE = args.sample_rate
        metadata = (('authorization', 'Bearer ' + getAccessToken()),
                    ('trackingid', str(trackingId)),
                    ('configkeyname', str(args.config_key_name)))
        request_config = speech_pb2.RecognitionConfig(sample_rate_hertz=RATE, max_alternatives=1)
        request_config.enable_automatic_punctuation = args.enable_automatic_punctuation
        request_config.enable_multilang_detection = args.enable_multilang_detection
        request_config.enable_training_deprecated = args.enable_training_deprecated
        request_config.language_code = args.language_code
        if len(expectedLanguageCodes) > 0:
            request_config.enable_automatic_language_detection = True
            request_config.expected_language_codes.extend(expectedLanguageCodes)

        request_audio = speech_pb2.RecognitionAudio(uri=url)
        recognize_request = speech_pb2.LongRunningRecognizeRequest(
                config=request_config, audio=request_audio, callback_url=callbackUrl, csi_file_url=csi_file_url)
        print(recognize_request)
            
        #in case it is url, we can post a message directly for sharder to pickup on server side
    response = stub.LongRunningRecognize(recognize_request, metadata=metadata)
    if len(response.name) > 0:
        print("Job Id is "+ str(response.name))
        return response.name
    return ""


def getLongRunningRecognizeOperation(stub, meetingId):
    transcripts = []
    try:
        trackingId = 'asr_client_python_'+str(uuid.uuid4())
        metadata = (('authorization', 'Bearer ' + getAccessToken()),
                    ('trackingid', str(trackingId)))
        operationRequest = speech_pb2.LongRunningRecognizeOperationRequest(name=meetingId)
        operation = stub.GetLongRunningRecognizeOperationV2(operationRequest, metadata=metadata)
        if operation.done:
            if operation.error is not None and operation.error.code > 0:
                print("Response error:",operation.error)
                return (True, None, None, None)

            # Ucomment below line to see the full response
            # print(operation.response) 
            jsonOutput = {}
            jsonTranscripts = []
            jsonOutput["interactionId"] = meetingId
            jsonOutput["languageCode"] = "en-US"
            jsonOutput["responseContents"] = jsonTranscripts

            for r in operation.response.results:
                words = []
                for w in r.alternatives[0].words:
                    start_time = w.start_time
                    end_time = w.end_time
                    words.append({
                        "word": w.word,
                        "start_time": { "seconds": start_time.seconds, "nanos": start_time.nanos },
                        "end_time": { "seconds": end_time.seconds, "nanos": end_time.nanos },
                    })
                # r.append(words)
                transcripts.append(r)
                alt={}
                alt['words'] = words
                alt['transcript'] = r.alternatives[0].transcript
                alt['confidence'] = r.alternatives[0].confidence
                responseResults = { "recognitionResult": 
                                   { "alternatives": [alt],
                                     "language_code": r.language_code ,
                                     "role": "CALLER" if r.channel_tag==1 else "AGENT"
                                   } } 
                
                jsonTranscripts.append(responseResults)
                
            return True, transcripts, words, jsonOutput
        else:
            print("Transcript is not ready yet, current progress: %i, meetingId=%s" % (operation.metadata.progress_percent, meetingId))
            return False, None, None, None
    except Exception as ex:
        print("failed at", datetime.now())
        print("failed ex %s" % ex)

def deleteLongRunningRecognizeOperation(stub, meetingId):
    try:
        trackingId = 'asr_client_python_'+str(uuid.uuid4())
        metadata = (('authorization', 'Bearer ' + getAccessToken()),
                    ('trackingid', str(trackingId))) 
        operationRequest = speech_pb2.LongRunningRecognizeOperationRequest(name=meetingId)
        stub.DeleteLongRunningRecognizeOperation(operationRequest, metadata=metadata)
    except Exception as ex:
        print ("failed to delete the meeting %s" % ex)

def main():
    speech_hints = []
    if args.speech_hints != "":
       speech_hints = args.speech_hints.split(",")
    
    expected_language_codes = []
    if args.expected_language_codes != "":
       expected_language_codes = args.expected_language_codes.split(",")

    # In practice, stream should be a generator yielding chunks of audio data.
    stream = []
    RATE = args.sample_rate
    CHUNK = int(RATE / 10)  # 100ms
    channel = make_channel()
    stub = speech_pb2_grpc.SpeechStub(channel)
    s = MicrophoneStream(RATE, CHUNK)
    starttime = datetime.now()
    print("=== start time", starttime) 
    if args.longrunning:
        long_running_recognize(stub, expected_language_codes)
    elif args.recording_upload:
        recording_upload(stub, args.file, args.csi_file, args.language_code)
    elif args.file != "":
        s = FileStream(args.file, RATE, CHUNK)
        stream_recognize(s, speech_hints, RATE, CHUNK, channel, stub)
    elif args.get_transcript:
        getLongRunningRecognizeOperation(stub, args.meeting_id)
    else:
        stream_recognize(s, speech_hints, RATE, CHUNK, channel, stub)

def long_running_recognize(stub, expected_language_codes):
    file = args.file
    url = args.url
    if (file == None or len(file)==0) and (url == None or len(url)==0):
        url = "https://www2.cs.uic.edu/~i101/SoundFiles/preamble.wav"

    callback = args.callback_url
    starttime = datetime.now()
    csi_file_url = args.csi_file_url
    

    meetingId = longRunningRecognize(stub, url=url, fileName=file, callbackUrl=callback,
                                     expectedLanguageCodes=expected_language_codes,
                                     csi_file_url=csi_file_url)
    done = False
    words = []
    transcripts = ""
    while not done:
        current = datetime.now()
        (done, transcripts, words, jsonOutput) =  getLongRunningRecognizeOperation(stub, meetingId)
        print("time-elapsed", (current-starttime))
        time.sleep(5)
    if args.json_output == "yes" and jsonOutput is not None:
        with open(f'{meetingId}.json', 'w+') as fd:
            json.dump(jsonOutput, fd)
        print(f"JSON output written to {meetingId}.json")
        print(jsonOutput)
    else:
        print("Final Utterances:")
        print_utterances(transcripts)

    with open(f'{file}.words.json', 'w+') as fd:
        json.dump(words, fd)

    print("args: %s" % args)
    if args.cleanup == "yes":
        print("meeting delete: %s" % meetingId)
        deleteLongRunningRecognizeOperation(stub, meetingId)

def stream_recognize(s, speech_hints, RATE, CHUNK, channel, stub):
    if s is not None:
        with s as stream:
            audio_generator = stream.generator()
            request_config = speech_pb2.RecognitionConfig(sample_rate_hertz=RATE)
            if len(speech_hints) > 0:
                request_config.speech_contexts.append(speech_pb2.SpeechContext(phrases=speech_hints))
            if args.language_code != "":
                request_config.language_code = args.language_code
            if args.encoding == "MULAW":
                request_config.encoding = speech_pb2.RecognitionConfig.MULAW
            else:
                request_config.encoding = speech_pb2.RecognitionConfig.LINEAR16
            request_config.enable_automatic_punctuation = args.enable_automatic_punctuation
            request_config.enable_training_deprecated = args.enable_training_deprecated        
            print(request_config)

            streaming_request_config = speech_pb2.StreamingRecognitionConfig(config=request_config,
                                                                            single_utterance=False,
                                                                            interim_results=False)
            config_request = (speech_pb2.StreamingRecognizeRequest(streaming_config=streaming_request_config) for i in range(1))
            requests = (speech_pb2.StreamingRecognizeRequest(audio_content=chunk) for chunk in audio_generator)
            trackingId = 'asr_client_python_'+str(uuid.uuid4())
            print("Tracking ID will be "+ str(trackingId))
            print("engine:", args.engine)

            metadata = (('authorization', 'Bearer ' + getAccessToken()),
                        ('trackingid', str(trackingId)), 
                        ('engine', args.engine))
            recognize_responses = stub.StreamingRecognize(itertools.chain(config_request, requests), metadata=metadata)

            for response in recognize_responses:
                results = response.results
                if not args.final_only or results[0].is_final:
                    print(response)


if __name__ == '__main__':
    main()