o 5i,@sdZddlZddlZddlmZddlmZmZmZm Z ddl m Z ddl m Z ddlmZdd lmZmZmZdd lmZmZmZmZmZmZdd lmZdd lmZe Gd ddZGdddZ    d#deede!de!de!dee"deef ddZ#e$dkrddl%Z%e&e%j'dkre%j'dZ(e&e%j'dkre%j'dndZ)edde)dZe Z*e*+e(eZ,e-d ddl.Z.e-e.j/e,dd!d"dSdSdS)$z Gemini transcription module using Google AI Studio endpoints. Handles audio transcription with structured output and multiple model support. N)Path)ListOptionalDictAny) dataclass)genai)types)GEMINI_API_KEY GEMINI_MODELSget_model_name)TranscriptionOutputTranscriptionResult SpeakerMetaget_transcription_promptget_user_promptTRANSCRIPTION_JSON_SCHEMA) AudioChunk)romanize_for_alignmentc@sReZdZUdZdZeed<dZeeed<dZ e ed<dZ eed <d Z e ed <d S) TranscriptionConfigz&Configuration for a transcription job.gemini-3-flash-previewmodellowthinking_levelg temperatureTelugulanguage< timeout_secN)__name__ __module__ __qualname____doc__rstr__annotations__rrrfloatrrintr(r(?/home/ubuntu/maya3_transcribe/src/backend/gemini_transcriber.pyrs    rc @seZdZdZddeefddZdedefdd Zdedefd d Z d e de j fd dZ ded e deeeffddZded e defddZ  ddeed e deedeedeef ddZdS)GeminiTranscriberzO Handles audio transcription using Gemini models via Google AI Studio. Napi_keycCs,|pt|_|js tdtj|jd|_dS)z~ Initialize the transcriber. Args: api_key: Gemini API key (defaults to env var) zGEMINI_API_KEY not set)r+N)r r+ ValueErrorrClientclient)selfr+r(r(r)__init__.s zGeminiTranscriber.__init__ file_pathreturncCs,t|j}ddddddd}||dS)z(Determine MIME type from file extension.z audio/flacz audio/wavz audio/mpegz audio/oggz audio/mp4z audio/aac)z.flacz.wavz.mp3z.oggz.m4az.aac)rsuffixlowerget)r/r1ext mime_typesr(r(r)_get_mime_type;s z GeminiTranscriber._get_mime_typecCs6t|d }|WdS1swYdS)zLoad audio file as bytes.rbN)openread)r/r1fr(r(r)_load_audio_bytesHs $z#GeminiTranscriber._load_audio_bytesconfigcCszt|j}t|i}d}|dr#|jr#|dr#tj|jd}tj |j dt tj j t|jdgd}|r;||_|S)z-Build the generation config for the API call.Nsupports_thinkingzgemini-3)rzapplication/jsontext)rresponse_mime_typeresponse_json_schemasystem_instruction)r rr r5r startswithr ThinkingConfigupperGenerateContentConfigrrPart from_textrrthinking_config)r/r> model_name model_inforK gen_configr(r(r) _build_configMs$    zGeminiTranscriber._build_config audio_pathc CsTt|j}||}||}tjdtjj||dtjjt dgdg}| |}t }z[|j j j|||d} t |} | jrddl} | | j} | | d<|| d <|j| d <| d d } | rkd | vrkt| | d <d| vr| dd }tdd || d<| WSd| |dWSty}zt |} t|| |dWYd}~Sd}~ww)z Transcribe a single audio file. Args: audio_path: Path to the audio file config: Transcription configuration Returns: Dictionary with transcription results user) mime_typedatar@)roleparts)rcontentsr>rN_processing_time_sec_model_thinking_level transcription romanized code_switchtaggedz \[[\w_]+\]zEmpty response from model)errorrWrX)r rr=r8r ContentrI from_bytesrJrrOtimer.modelsgenerate_contentrAjsonloadsrr5rresubstrip Exceptionr$)r/rPr>rL audio_bytesrRrVrN start_timeresponseprocessing_timereresultnativer^er(r(r)transcribe_audiohs^            z"GeminiTranscriber.transcribe_audiochunkc Cs||j|}|d}|r'td|dd|dd|dd|dd}n1|d}t|tr8tdi|nd}t|dd|d d|d d|d d||d dd }t|j|j |j |j |j ||d|j |d|dd S)a  Transcribe an AudioChunk and return structured result. Args: chunk: AudioChunk to transcribe config: Transcription configuration Returns: TranscriptionResult with full metadata r_z[ERROR: ])rZr]r\r^speakerNrZr[r]r\r^detected_language)rZr]r\r^rurvrXrYrW) segment_id chunk_index total_chunks duration_secrrZ model_usedrprocessing_time_secr()rrr1r5r isinstancedictrroriginal_segmentrxryrzrr)r/rsr> raw_result error_msgrZ speaker_datarur(r(r)transcribe_chunks<             z"GeminiTranscriber.transcribe_chunkchunks max_chunksprogress_callbackc Cs|r|d|}g}t|}td|d|jdt|D]?\}}|r,||d|ntd|dd|d|jd|jdd|jd |||} || ||dkr]t d qtd t|d |S) a Transcribe multiple audio chunks. Args: chunks: List of AudioChunks to transcribe config: Transcription configuration max_chunks: Maximum chunks to process (for testing) progress_callback: Optional callback(current, total) for progress updates Returns: List of TranscriptionResults Nz[Transcriber] Processing z chunks with z...r /z: z (chunk )g?z[Transcriber] Completed z transcriptions) lenprintr enumeraterrxryrappendrbsleep) r/rr>rrresultstotalirsror(r(r)transcribe_batchs(     z"GeminiTranscriber.transcribe_batch)N)NN)r r!r"r#rr$r0r8bytesr=rr rHrOrrrrrrrrr'callablerr(r(r(r)r*)sB    P 9r*rrhighsegmentsrrr max_segmentsr2cCs$t|||d}t}|j|||dS)a Convenience function to transcribe a list of audio segments. Args: segments: List of AudioChunk objects language: Primary language of the audio model: Gemini model to use thinking_level: Thinking level for Gemini 3 models max_segments: Max segments to process (for testing) Returns: List of TranscriptionResult objects rrr)r)rr*r)rrrrrr> transcriberr(r(r)transcribe_segmentssr__main__rrz Transcription Result:F)indent ensure_ascii)rrrN)0r#rgrbpathlibrtypingrrrr dataclassesrgoogler google.genair r>r r r transcription_schemarrrrrraudio_processorrsrc.romanizationrrr*r$r'rr sysrargvrPrrrrrorredumpsr(r(r(r)sd        t