a bgT@sDdZddlZddlZddlZddlZddlZddlZddlZddlm Z ddl m Z m Z m Z mZmZmZmZmZmZmZmZmZddlmZmZddlmZddlmZmZddlm Z m!Z!m"Z"m#Z#m$Z$m%Z%e rddl&Z'ddl(m)m*m+Z,e-e.Z/d Z0d Z1d Z2d d ddZ3GdddZ4e4Z5e4Z6Gddde7e Z8Gddde!Z9Gddde!Z:GdddZ;ee7ede7ddd Z?edee7ddfd!d"d#Z@eeAe4fZBe7ZCGd$d%d%ZDeDZEe7ZFGd&d'd'e9e:eeEeFfZGee7eefZHeAZIGd(d)d)e9e:eeHeIfZJeGZKeJZLeDZMdS)*z*A common module for NVIDIA Riva Runnables.N)Enum) TYPE_CHECKINGAnyAsyncGenerator AsyncIteratorDict GeneratorIteratorListOptionalTupleUnioncast) AnyMessage BaseMessage) PromptValue)RunnableConfigRunnableSerializable) AnyHttpUrl BaseModelField parse_obj_asroot_validator validator?i) .!?¡¿z riva.clientreturnc CsBz ddl}Wn.ty:}ztd|WYd}~n d}~00|jS)z5Import the riva client and raise an error on failure.rNziCould not import the NVIDIA Riva client library. Please install it with `pip install nvidia-riva-client`.) riva.client ImportErrorclient)rivaerrr(w/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/utilities/nvidia_riva.py_import_riva_client1s r*c@seZdZdZdS) SentinelTzAn empty Sentinel type.N)__name__ __module__ __qualname____doc__r(r(r(r)r+>sr+c@sNeZdZdZdZdZdZdZdZdZ e e ddd d Z e d d d dZdS)RivaAudioEncodingaAn enum of the possible choices for Riva audio encoding. The list of types exposed by the Riva GRPC Protobuf files can be found with the following commands: ```python import riva.client print(riva.client.AudioEncoding.keys()) # noqa: T201 ``` ALAWENCODING_UNSPECIFIEDFLAC LINEAR_PCMMULAWOGGOPUS) format_coder"c CsRz|j|j|jd|WStyL}ztd||WYd}~n d}~00dS)zReturn the audio encoding specified by the format code in the wave file. ref: https://mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html )z>The following wave file format code is not supported by Riva: N)r4r1r5KeyErrorNotImplementedError)clsr7r'r(r(r)from_wave_format_codeXsz'RivaAudioEncoding.from_wave_format_codezriva.client.AudioEncodingr!cCst}t|j|S)z-Returns the Riva API object for the encoding.)r*getattrZ AudioEncodingself riva_clientr(r(r)riva_pb2fszRivaAudioEncoding.riva_pb2N)r,r-r.r/r1r2r3r4r5r6 classmethodintr>propertyrCr(r(r(r)r0Fs  r0c@seZdZUdZeeddddgdZeeefe d<eddd Z e ee d <e d d d dZ eddddeeedddZdS) RivaAuthMixinzBConfiguration for the authentication to a Riva service connection.zhttp://localhost:50051z1The full URL where the Riva service can be found.z"https://user@pass:riva.example.com) descriptionZexamplesurlNz@A full path to the file where Riva's public ssl key can be read.rHssl_certzriva.client.Authr!cCsBt}tt|j}|jdk}t|jdd}|j|j||dS)z!Return a riva client auth object.https/)rKuse_ssluri) r*rrrIschemestrsplitZAuthrK)rArBrIrOZ url_no_schemer(r(r)authzs  zRivaAuthMixin.authT)preZ allow_reusevalr"cCs$t|trtttt|Stt|S)z:Do some initial conversations for the URL before checking.) isinstancerRrrr)r=rWr(r(r) _validate_urls zRivaAuthMixin._validate_url)r,r-r.r/rrrIr rR__annotations__rKr rFrTrrDrrYr(r(r(r)rGms  rGc@sPeZdZUdZeejddZeed<edddZ e ed<edd dZ e ed <d S) RivaCommonConfigMixinz%A collection of common Riva settings.z!The encoding on the audio stream.)defaultrHencodingi@z*The sample rate frequency of audio stream.sample_rate_hertzzen-USzaThe [BCP-47 language code](https://www.rfc-editor.org/rfc/bcp/bcp47.txt) for the target language. language_codeN) r,r-r.r/rr0r4r]rZr^rEr_rRr(r(r(r)r[s r[c@szeZdZUdZejed<ejed<ddddZdddd Z ddd d Z e dd d Z ddddZ ddddZdS)_Eventz3A combined event that is threadsafe and async safe._event_aeventNr!cCst|_t|_dS)zInitialize the event.N) threadingEventraasynciorbrAr(r(r)__init__s z_Event.__init__cCs|j|jdSzSet the event.N)rasetrbrfr(r(r)ris z _Event.setcCs|j|jdSrh)raclearrbrfr(r(r)rjs z _Event.clearcCs |jS)zIndicate if the event is set.)rais_setrfr(r(r)rksz _Event.is_setcCs|jdS)zWait for the event to be set.N)rawaitrfr(r(r)rlsz _Event.waitcs|jIdHdS)z#Async wait for the event to be set.N)rbrlrfr(r(r) async_waitsz_Event.async_wait)r,r-r.r/rcrdrZrergrirjboolrkrlrmr(r(r(r)r`s   r`)output_directory sample_rater"cCsr|rntjddd|d}|j}Wdn1s20Yt|d}|d|d||||fSd S) zECreate a new wave file and return the wave write object and filename.bxz.wavF)modesuffixdeletedirNwbr8rN)NN)tempfileNamedTemporaryFilenamewaveopenZ setnchannelsZ setsampwidthZ setframerate)rorpf wav_file_namewav_filer(r(r) _mk_wave_files$    r TTSInputTyperVcCs.t|tr|St|tr&t|jSt|S)zAttempt to coerce the input value to a string. This is particularly useful for converting LangChain message to strings. )rXrZ to_stringrrRcontentrWr(r(r)_coerce_strings    r)inputsr"ccsd}|D]}t|}tD]0}||vr||d\}}|||Vd}qq||7}t|tkrtdt|tD]}|||dVqnd}q|r|VdS)z9Filter the input chunks are return strings ready for TTS.r8rN)r_SENTENCE_TERMINATORSrSlen_MAX_TEXT_LENGTHrange)rbufferchunk terminatorZ last_sentenceidxr(r(r)_process_chunkss rc@sJeZdZUdZejed<ejed<ejed<e ed<e ed<e ed<e ej ed<d)e d d d d Z eed d fdddZeedddZeedddZeedddZeedddZeedddZd*ee e d dddZd+ee e d dddZd,e e d d d!d"Zd-e e d d d#d$Zed%d d&d'd(Zd S). AudioStreamz%A message containing streaming audio. _put_lock_queueoutputhangup user_talking user_quiet_workerrN)maxsizer"cCsDt|_tj|d|_t|_t|_t|_ t|_ d|_ dS)zInitialize the queue.)rN) rcLockrqueueQueuerrr`rrrr)rArr(r(r)rg s  zAudioStream.__init__r!ccsLz|jdt}Wntjy*YqYn0|tkr6qH|V|jqdS)zReturn an error.TN)rget_QUEUE_GET_TIMEOUTrEmptyHANGUP task_donerAZnext_valr(r(r)__iter__s zAudioStream.__iter__cCs\z"td|jjdtIdH}Wntjy:YqYn0|tkrFqX|V|j qdS)z4Iterate through all items in the queue until HANGUP.NT) reget_event_looprun_in_executorrrrrrrrrr(r(r) __aiter__&s  zAudioStream.__aiter__cCs |jS)z(Indicate if the audio stream has hungup.)rrkrfr(r(r)hungup9szAudioStream.hungupcCs |jS)z-Indicate in the input stream buffer is empty.)remptyrfr(r(r)r>szAudioStream.emptycCs4|jo |j}|jduo*|j o*|j}|o2|S)z;Indicate if the audio stream has hungup and been processed.N)rrris_aliver)rAZ input_doneZ output_doner(r(r)completeCs   zAudioStream.completecCs|jr|jSdS)z&Indicate if the ASR stream is running.F)rrrfr(r(r)runningNs zAudioStream.running)itemtimeoutr"cCsZ|j@|jrtd|tur(|j|jj||dWdn1sL0YdS)zPut a new item into the queue.z?The audio stream has already been hungup. Cannot put more data.rN)rr RuntimeErrorrrrirput)rArrr(r(r)rUs zAudioStream.putcs*t}t|d|j||IdHdS)z$Async put a new item into the queue.N)rerwait_forrr)rArrloopr(r(r)aput`szAudioStream.aput)rr"cCs|t|dS)zSend the hangup signal.N)rrrArr(r(r)closeeszAudioStream.closecs|t|IdHdS)zAsync send the hangup signal.N)rrrr(r(r)acloseiszAudioStream.aclosezrasr.StreamingRecognizeResponse) responsesr"cs^jrtdtjdddddfdd }tj|d _d j_jdS) zIDrain the responses from the provided iterator and put them into a queue.z,An ASR instance has already been registered.rNrrNr!csD]x}|jsq |jD]d}|js*q|jrdjjtt |jdj }j |qj sjjqq dS)zConsume the ASR Generator.rN)rlresultsZ alternativesis_finalrrjrrirrR transcriptrrrk)responseresultrZ has_startedrrAr(r)workerts     z$AudioStream.register..worker)targetT) rrrcBarrierThreadrdaemonstartrl)rArrr(rr)registerms zAudioStream.register)r)N)N)N)N) r,r-r.r/rcrrZrrr`r rrErgrbytesrrStreamInputTyperrFrnrrrrrrrrr rr(r(r(r)rs0       rc@seZdZUdZdZeed<dZeed<edddZ e ed <ed d dZ e ed <ed d dZ e ed<ed deeeefeeefdddZeddddZddddZdeeeeedddZdS)RivaASRzNA runnable that performs Automatic Speech Recognition (ASR) using NVIDIA Riva.Znvidia_riva_asrryzA Runnable for converting audio bytes to a string.This is useful for feeding an audio stream into a chain andpreprocessing that audio to create an LLM prompt.rHr8z7The number of audio channels in the input audio stream.rJaudio_channel_countTz\Controls whether or not Riva should attempt to filter profanity out of the transcribed text.profanity_filterz]Controls whether Riva should attempt to correct senetence puncuation in the transcribed text.enable_automatic_punctuationrUvaluesr"cCs t}|Sz4Validate the Python environment and input arguments.r*r=r_r(r(r)_validate_environmentszRivaASR._validate_environmentz&riva.client.StreamingRecognitionConfigr!c Cs4t}|jd|j|j|j|jd|j|j|jddS)z)Create and return the riva config object.Tr8)r]r^rZmax_alternativesrrr_)Zinterim_resultsconfig) r*ZStreamingRecognitionConfigZRecognitionConfigr]r^rrrr_r@r(r(r)rszRivaASR.configzriva.client.ASRServicec CsHt}z||jWStyB}ztd|WYd}~n d}~00dS);Connect to the riva service and return the a client object.z5Error raised while connecting to the Riva ASR server.N)r*Z ASRServicerT Exception ValueErrorrArBr'r(r(r) _get_serviceszRivaASR._get_serviceNinputrkwargsr"cKs|js(|}|j||jd}||g}|js|jj|jjd}Wdn1s^0Y|r,|j sz||j g7}Wnt j yYqlYn0|j qltdt|d|Sq,dS)z3Transcribe the audio bytes into a string with Riva.)Z audio_chunksZstreaming_configg?NzRiva ASR returning: %s r)rrZstreaming_response_generatorrrrr not_emptyrlr get_nowaitrrr_LOGGERdebugreprjoinstrip)rArrrservicerZ full_responsereadyr(r(r)invokes*  ,   zRivaASR.invoke)N)r,r-r.r/ryrRrZrHrrrErrnrrrDrrrrFrr ASRInputTyper r ASROutputTyperr(r(r(r)rs8   $ rc@seZdZUdZdZeed<dZeed<edddZ eed <ed d dZ e eed <e d de eeefeeefdddZed e eedddZddddZd ee eeedddZd!eee ee eeedddZd"eee ee eeed fdddZd S)#RivaTTSz?A runnable that performs Text-to-Speech (TTS) with NVIDIA Riva.Znvidia_riva_ttsryz_A tool for converting text to speech.This is useful for converting LLM output into audio bytes.rHzEnglish-US.Female-1zThe voice model in Riva to use for speech. Pre-trained models are documented in [the Riva documentation](https://docs.nvidia.com/deeplearning/riva/user-guide/docs/tts/tts-overview.html).rJ voice_nameNzThe directory where all audio files should be saved. A null value indicates that wave files should not be saved. This is useful for debugging purposes.roTrrcCs t}|Srrrr(r(r)r szRivaTTS._validate_environment)vr"cCs,|r(t|}|jdddt|S|S)NT)parentsexist_ok)pathlibPathmkdirrRabsolute)r=rdirpathr(r(r)_output_directory_validator's   z#RivaTTS._output_directory_validatorz"riva.client.SpeechSynthesisServicer!c CsHt}z||jWStyB}ztd|WYd}~n d}~00dS)rz5Error raised while connecting to the Riva TTS server.N)r*ZSpeechSynthesisServicerTrrrr(r(r)r0szRivaTTS._get_servicercKsd|t|gS)zDPerform TTS by taking a string and outputting the entire audio file.)r transformiter)rArrrr(r(r)r:szRivaTTS.invokec ks|}t|j|j\}}t|D]X}td||j||j|j |j j |jd}|D]$} t t | j} |rr|| | VqTq"|r|td|dS)zHPerform TTS by taking a stream of characters and streaming output bytes.zRiva TTS chunk: %s)textrr_r]Zsample_rate_hzzRiva TTS wrote file: %sN)rrror^rrrZsynthesize_onlinerr_r]rCrraudioZwriteframesrawr) rArrrrr}r~rrresprr(r(r)rCs*      zRivaTTS.transformc  stttddfdd }ttdfdd ddfdd ddfd d }|}|}ztd IdH}Wntj j yYqYn0 |t urq|Vq|IdH|IdHdS) zGIntercept async transforms and route them to the synchronous transform.Nr!cs,2z3dHW}|q6tdS)z#Produce input into the input queue.N) put_nowait_TRANSFORM_ENDr)r input_queuer(r) _producerusz%RivaTTS.atransform.._producerc3s@zjdd}Wntjy(YqYn0|tkr4q<|VqdS)zIterate over the input_queue.rrN)rrrrr)rr(r)_input_iterator{s z+RivaTTS.atransform.._input_iteratorcs*D]}|q tdS)z!Consume the input with transform.N)rrrr)r out_queuerAr(r) _consumers z%RivaTTS.atransform.._consumercsdIdHdS)z"Coroutine that wraps the consumer.N)rr()rrr(r)_consumer_corosz*RivaTTS.atransform.._consumer_coror) reget_running_looprrr r create_taskrr exceptions TimeoutErrorrr) rArrrrrZproducerZconsumerrWr()rrrrrrrAr) atransformjs&     zRivaTTS.atransform)N)N)N)r,r-r.r/ryrRrZrHrrror rrDrrrrrrrr TTSOutputTyperr rrrrr(r(r(r)rsP     $   * r)Nr/reloggingrrrwrcrzenumrtypingrrrrrrr r r r r rZlangchain_core.messagesrrZlangchain_core.prompt_valuesrZlangchain_core.runnablesrrZpydanticrrrrrrr#r&Zriva.client.proto.riva_asr_pb2r%protoZ riva_asr_pb2Zrasr getLoggerr,rrrrr*r+rrrRr0rGr[r`floatZ Wave_writerrrrrZStreamOutputTyperrrrrrrZ NVIDIARivaASRZ NVIDIARivaTTSZNVIDIARivaStreamr(r(r(r)sn 8    '!#     h  $