53 lines
2.4 KiB
Python
53 lines
2.4 KiB
Python
"""Add native SMART mode and vocabulary only to opted-in dictation calls."""
|
||
import json
|
||
from pathlib import Path
|
||
|
||
VOCABULARY = Path('/etc/litellm/personal-vocabulary.json')
|
||
|
||
|
||
def vocabulary(path=VOCABULARY):
|
||
words = json.loads(path.read_text())
|
||
if not isinstance(words, list) or not 1 <= len(words) <= 100:
|
||
raise ValueError('Personal vocabulary must contain 1–100 terms')
|
||
if any(not isinstance(word, str) or not word.strip() or len(word) > 100 for word in words):
|
||
raise ValueError('Invalid personal vocabulary term')
|
||
return list(dict.fromkeys(words))
|
||
|
||
|
||
def install():
|
||
from litellm.llms.gemini.audio_transcription.transformation import GeminiAudioTranscriptionConfig
|
||
original = GeminiAudioTranscriptionConfig.transform_audio_transcription_request
|
||
if getattr(original, '_personal_vocabulary_patch', False):
|
||
return
|
||
|
||
def transform(self, model, audio_file, optional_params, litellm_params):
|
||
enabled = optional_params.get('personal_vocabulary') is True
|
||
if not enabled:
|
||
return original(self, model, audio_file, optional_params, litellm_params)
|
||
if model.removeprefix('gemini/') != 'gemini-3.5-transcribe':
|
||
raise ValueError('Personal vocabulary is restricted to Gemini 3.5 Transcribe')
|
||
if optional_params.get('timestamp_granularities') or optional_params.get('response_format') not in (None, 'json', 'text'):
|
||
raise ValueError('Personal dictation does not support timestamps or diarization')
|
||
clean = {key: value for key, value in optional_params.items() if key != 'personal_vocabulary'}
|
||
result = original(self, model, audio_file, clean, litellm_params)
|
||
config = result.data.setdefault('generation_config', {}).setdefault('transcription_config', {})
|
||
# Fail closed if a future LiteLLM version enables conflicting features by default.
|
||
if config.get('mode'):
|
||
raise ValueError('Unexpected Gemini transcription mode')
|
||
config['mode'] = 'smart'
|
||
config['custom_vocabulary'] = vocabulary()
|
||
return result
|
||
|
||
transform._personal_vocabulary_patch = True
|
||
GeminiAudioTranscriptionConfig.transform_audio_transcription_request = transform
|
||
|
||
|
||
if __name__ == '__main__':
|
||
import sys
|
||
from litellm import run_server
|
||
from bailian_transcription import install as install_bailian_asr
|
||
install_bailian_asr()
|
||
install()
|
||
sys.argv[0] = 'litellm'
|
||
sys.exit(run_server())
|