andrewyng--aisuite
Rohit P
424 行
17 KiB
Python
424 行
17 KiB
Python
import os
|
|
import json
|
|
import numpy as np
|
|
import queue
|
|
import threading
|
|
import time
|
|
from typing import Union, BinaryIO, AsyncGenerator
|
|
|
|
from aisuite.provider import Provider, ASRError, Audio
|
|
from aisuite.framework.message import (
|
|
TranscriptionResult,
|
|
Segment,
|
|
Word,
|
|
Alternative,
|
|
Channel,
|
|
StreamingTranscriptionChunk,
|
|
)
|
|
|
|
|
|
class DeepgramProvider(Provider):
|
|
"""Deepgram ASR provider."""
|
|
|
|
def __init__(self, **config):
|
|
"""Initialize the Deepgram provider with the given configuration."""
|
|
super().__init__()
|
|
|
|
# Ensure API key is provided either in config or via environment variable
|
|
self.api_key = config.get("api_key") or os.getenv("DEEPGRAM_API_KEY")
|
|
if not self.api_key:
|
|
raise ValueError(
|
|
"Deepgram API key is missing. Please provide it in the config or set the DEEPGRAM_API_KEY environment variable."
|
|
)
|
|
|
|
# Initialize Deepgram client (v5.0.0+)
|
|
try:
|
|
from deepgram import DeepgramClient
|
|
|
|
self.client = DeepgramClient(api_key=self.api_key)
|
|
except ImportError:
|
|
raise ImportError(
|
|
"Deepgram SDK is required. Install it with: pip install deepgram-sdk"
|
|
)
|
|
|
|
# Initialize audio functionality
|
|
self.audio = DeepgramAudio(self.client)
|
|
|
|
def chat_completions_create(self, model, messages):
|
|
"""Deepgram does not support chat completions."""
|
|
raise NotImplementedError(
|
|
"Deepgram provider only supports audio transcription, not chat completions."
|
|
)
|
|
|
|
|
|
# Audio Classes
|
|
class DeepgramAudio(Audio):
|
|
"""Deepgram Audio functionality container."""
|
|
|
|
def __init__(self, client):
|
|
super().__init__()
|
|
self.transcriptions = self.Transcriptions(client)
|
|
|
|
class Transcriptions(Audio.Transcription):
|
|
"""Deepgram Audio Transcriptions functionality."""
|
|
|
|
def __init__(self, client):
|
|
self.client = client
|
|
|
|
def create(
|
|
self,
|
|
model: str,
|
|
file: Union[str, BinaryIO],
|
|
**kwargs,
|
|
) -> TranscriptionResult:
|
|
"""
|
|
Create audio transcription using Deepgram SDK v5.
|
|
|
|
All parameters are already validated and mapped by the Client layer.
|
|
This is a simple pass-through to the Deepgram API.
|
|
"""
|
|
try:
|
|
# Add model to params and set defaults
|
|
kwargs["model"] = model
|
|
kwargs.setdefault("smart_format", True)
|
|
kwargs.setdefault("punctuate", True)
|
|
kwargs.setdefault("language", "en")
|
|
|
|
# Get audio bytes
|
|
audio_bytes = self._prepare_audio_payload(file)
|
|
|
|
# Use v5 API: client.listen.v1.media.transcribe_file()
|
|
# All parameters passed as kwargs, no PrerecordedOptions needed
|
|
response = self.client.listen.v1.media.transcribe_file(
|
|
request=audio_bytes, **kwargs
|
|
)
|
|
|
|
# Convert Pydantic model to dict (v5 uses Pydantic v2)
|
|
if hasattr(response, "model_dump"):
|
|
response_dict = response.model_dump()
|
|
elif hasattr(response, "to_dict"):
|
|
response_dict = response.to_dict()
|
|
elif hasattr(response, "dict"):
|
|
response_dict = response.dict()
|
|
else:
|
|
response_dict = response
|
|
|
|
return self._parse_deepgram_response(response_dict)
|
|
|
|
except Exception as e:
|
|
raise ASRError(f"Deepgram transcription error: {e}") from e
|
|
|
|
async def create_stream_output(
|
|
self,
|
|
model: str,
|
|
file: Union[str, BinaryIO],
|
|
chunk_size_minutes: float = 3.0,
|
|
**kwargs,
|
|
) -> AsyncGenerator[StreamingTranscriptionChunk, None]:
|
|
"""
|
|
Create streaming audio transcription using Deepgram SDK v5 with chunked processing.
|
|
|
|
All parameters are already validated and mapped by the Client layer.
|
|
This implementation handles audio chunking and streaming.
|
|
"""
|
|
try:
|
|
# Load and prepare audio
|
|
audio_data, sample_rate = await self._load_and_prepare_audio(file)
|
|
|
|
# Calculate chunking strategy
|
|
duration_seconds = len(audio_data) / sample_rate
|
|
chunk_duration_seconds = chunk_size_minutes * 60
|
|
|
|
if duration_seconds <= chunk_duration_seconds:
|
|
chunks = [audio_data]
|
|
else:
|
|
chunk_size_samples = int(chunk_duration_seconds * sample_rate)
|
|
chunks = []
|
|
num_chunks = int(np.ceil(duration_seconds / chunk_duration_seconds))
|
|
for i in range(num_chunks):
|
|
start_sample = i * chunk_size_samples
|
|
end_sample = min(
|
|
start_sample + chunk_size_samples, len(audio_data)
|
|
)
|
|
chunks.append(audio_data[start_sample:end_sample])
|
|
|
|
# Setup API parameters for v5
|
|
kwargs["model"] = model
|
|
kwargs.setdefault("smart_format", "true")
|
|
kwargs.setdefault("punctuate", "true")
|
|
kwargs.setdefault("language", "en")
|
|
kwargs["interim_results"] = (
|
|
"true" # Enable interim results for streaming
|
|
)
|
|
|
|
# Remove parameters not supported by streaming
|
|
kwargs.pop("utterances", None)
|
|
|
|
# Add critical audio format parameters (as strings for v5)
|
|
kwargs["encoding"] = "linear16" # PCM16 format
|
|
kwargs["sample_rate"] = "16000" # Match our target sample rate
|
|
kwargs["channels"] = "1" # Mono audio
|
|
|
|
# Use thread-safe queue for cross-thread communication
|
|
transcript_queue = queue.Queue()
|
|
connection_closed = threading.Event()
|
|
|
|
def on_message(*args, **message_kwargs):
|
|
"""Handle transcript events"""
|
|
# Extract result from args or kwargs
|
|
result = None
|
|
if len(args) >= 2:
|
|
result = args[1]
|
|
elif "result" in message_kwargs:
|
|
result = message_kwargs["result"]
|
|
else:
|
|
return
|
|
|
|
if hasattr(result, "channel") and result.channel.alternatives:
|
|
alt = result.channel.alternatives[0]
|
|
if alt.transcript:
|
|
chunk = StreamingTranscriptionChunk(
|
|
text=alt.transcript,
|
|
is_final=getattr(result, "is_final", False),
|
|
confidence=getattr(alt, "confidence", None),
|
|
)
|
|
transcript_queue.put(chunk)
|
|
|
|
def on_error(*args, **error_kwargs):
|
|
"""Handle error events"""
|
|
error = None
|
|
if len(args) >= 2:
|
|
error = args[1]
|
|
elif "error" in error_kwargs:
|
|
error = error_kwargs["error"]
|
|
|
|
if error:
|
|
transcript_queue.put(
|
|
ASRError(f"Deepgram streaming error: {error}")
|
|
)
|
|
|
|
def on_close(*args, **close_kwargs):
|
|
"""Handle connection close events"""
|
|
connection_closed.set()
|
|
|
|
# Use v5 streaming API with context manager
|
|
from deepgram.core.events import EventType
|
|
|
|
async with self.client.listen.v1.connect(**kwargs) as connection:
|
|
# Register event handlers
|
|
connection.on(EventType.Transcript, on_message)
|
|
connection.on(EventType.Error, on_error)
|
|
connection.on(EventType.Close, on_close)
|
|
|
|
# Send all chunks through connection
|
|
for audio_chunk in chunks:
|
|
self._send_audio_chunk(connection, audio_chunk)
|
|
|
|
# Send CloseStream message to signal end
|
|
close_stream_message = json.dumps({"type": "CloseStream"})
|
|
connection.send(close_stream_message)
|
|
|
|
# Yield results until connection closes
|
|
while not connection_closed.is_set():
|
|
try:
|
|
chunk = transcript_queue.get(timeout=0.1)
|
|
if isinstance(chunk, Exception):
|
|
raise chunk
|
|
yield chunk
|
|
except queue.Empty:
|
|
continue
|
|
|
|
# Get any remaining results
|
|
while not transcript_queue.empty():
|
|
try:
|
|
chunk = transcript_queue.get_nowait()
|
|
if isinstance(chunk, Exception):
|
|
raise chunk
|
|
yield chunk
|
|
except queue.Empty:
|
|
break
|
|
|
|
except Exception as e:
|
|
raise ASRError(f"Deepgram streaming transcription error: {e}")
|
|
|
|
def _prepare_audio_payload(self, file: Union[str, BinaryIO]) -> bytes:
|
|
"""Prepare audio payload for Deepgram API v5.
|
|
|
|
Returns raw bytes instead of dict payload (v5 API change).
|
|
"""
|
|
if isinstance(file, str):
|
|
with open(file, "rb") as audio_file:
|
|
buffer_data = audio_file.read()
|
|
else:
|
|
if hasattr(file, "read"):
|
|
buffer_data = file.read()
|
|
else:
|
|
raise ValueError(
|
|
"File must be a file path string or file-like object"
|
|
)
|
|
return buffer_data
|
|
|
|
async def _load_and_prepare_audio(
|
|
self, file: Union[str, BinaryIO]
|
|
) -> tuple[np.ndarray, int]:
|
|
"""Load and prepare audio file for streaming.
|
|
|
|
Conversions performed only when necessary:
|
|
- Stereo to mono: Required for multi-channel audio
|
|
- Sample rate conversion: Required when input != 16kHz
|
|
- Other formats: Error out as unsupported
|
|
"""
|
|
try:
|
|
try:
|
|
import soundfile as sf
|
|
except ImportError:
|
|
raise ASRError(
|
|
"soundfile is required for audio processing. Install with: pip install soundfile"
|
|
)
|
|
|
|
if isinstance(file, str):
|
|
audio_data, original_sample_rate = sf.read(file)
|
|
else:
|
|
audio_data, original_sample_rate = sf.read(file)
|
|
|
|
audio_data = np.asarray(audio_data, dtype=np.float32)
|
|
|
|
# Convert to mono if stereo
|
|
if len(audio_data.shape) > 1:
|
|
if audio_data.shape[1] == 2:
|
|
audio_data = np.mean(audio_data, axis=1)
|
|
else:
|
|
raise ASRError(
|
|
f"Unsupported audio format: {audio_data.shape[1]} channels. Only mono and stereo are supported."
|
|
)
|
|
|
|
# Resample to 16kHz if needed
|
|
target_sample_rate = 16000
|
|
if original_sample_rate != target_sample_rate:
|
|
try:
|
|
from scipy import signal
|
|
|
|
num_samples = int(
|
|
len(audio_data) * target_sample_rate / original_sample_rate
|
|
)
|
|
audio_data = signal.resample(audio_data, num_samples)
|
|
except ImportError:
|
|
raise ASRError(
|
|
f"Audio resampling required but scipy not available. "
|
|
f"Input is {original_sample_rate}Hz, need {target_sample_rate}Hz. "
|
|
f"Install scipy or provide audio at {target_sample_rate}Hz."
|
|
)
|
|
|
|
return np.asarray(audio_data, dtype=np.float32), target_sample_rate
|
|
|
|
except Exception as e:
|
|
if isinstance(e, ASRError):
|
|
raise
|
|
raise ASRError(f"Error loading audio file: {e}")
|
|
|
|
def _send_audio_chunk(self, connection, audio_chunk: np.ndarray) -> None:
|
|
"""Send audio chunk data through the connection."""
|
|
streaming_chunk_size = 8000 # Match reference BLOCKSIZE (~0.5s @16kHz mono)
|
|
send_delay = 0.01
|
|
|
|
for i in range(0, len(audio_chunk), streaming_chunk_size):
|
|
piece = audio_chunk[i : i + streaming_chunk_size]
|
|
|
|
if len(piece) < streaming_chunk_size:
|
|
piece = np.pad(
|
|
piece, (0, streaming_chunk_size - len(piece)), mode="constant"
|
|
)
|
|
|
|
pcm16 = (piece * 32767).astype(np.int16).tobytes()
|
|
connection.send(pcm16)
|
|
time.sleep(send_delay) # Use synchronous sleep like reference
|
|
|
|
def _parse_deepgram_response(self, response_dict: dict) -> TranscriptionResult:
|
|
"""Convert Deepgram API response to unified TranscriptionResult."""
|
|
try:
|
|
results = response_dict.get("results", {})
|
|
channels = results.get("channels", [])
|
|
|
|
if not channels or not channels[0].get("alternatives"):
|
|
return TranscriptionResult(
|
|
text="", language=None, confidence=None, task="transcribe"
|
|
)
|
|
|
|
best_alternative = channels[0]["alternatives"][0]
|
|
text = best_alternative.get("transcript", "")
|
|
confidence = best_alternative.get("confidence", None)
|
|
|
|
words = [
|
|
Word(
|
|
word=word_data.get("word", ""),
|
|
start=word_data.get("start", None),
|
|
end=word_data.get("end", None),
|
|
confidence=word_data.get("confidence", None),
|
|
)
|
|
for word_data in best_alternative.get("words", [])
|
|
]
|
|
|
|
segments = []
|
|
paragraphs = results.get("paragraphs", {}).get("paragraphs", [])
|
|
for para in paragraphs:
|
|
for sentence in para.get("sentences", []):
|
|
segments.append(
|
|
Segment(
|
|
id=len(segments),
|
|
seek=0,
|
|
start=sentence.get("start", None),
|
|
end=sentence.get("end", None),
|
|
text=sentence.get("text", ""),
|
|
tokens=[],
|
|
temperature=0.0,
|
|
avg_logprob=0.0,
|
|
compression_ratio=0.0,
|
|
no_speech_prob=0.0,
|
|
)
|
|
)
|
|
|
|
alternatives_list = [
|
|
Alternative(
|
|
transcript=alt.get("transcript", ""),
|
|
confidence=alt.get("confidence", None),
|
|
)
|
|
for alt in channels[0]["alternatives"][1:]
|
|
]
|
|
|
|
channels_list = [
|
|
Channel(
|
|
alternatives=[
|
|
Alternative(
|
|
transcript=alt.get("transcript", ""),
|
|
confidence=alt.get("confidence", None),
|
|
)
|
|
for alt in channel.get("alternatives", [])
|
|
]
|
|
)
|
|
for channel in channels
|
|
]
|
|
|
|
metadata = response_dict.get("metadata", {})
|
|
|
|
return TranscriptionResult(
|
|
text=text,
|
|
language=results.get("language", None),
|
|
confidence=confidence,
|
|
task="transcribe",
|
|
duration=metadata.get("duration", None) if metadata else None,
|
|
segments=segments or None,
|
|
words=words or None,
|
|
channels=channels_list or None,
|
|
alternatives=alternatives_list or None,
|
|
utterances=results.get("utterances", []),
|
|
paragraphs=results.get("paragraphs", None),
|
|
topics=results.get("topics", []),
|
|
intents=results.get("intents", []),
|
|
sentiment=results.get("sentiment", None),
|
|
summary=results.get("summary", None),
|
|
metadata=metadata,
|
|
)
|
|
|
|
except (KeyError, TypeError, IndexError) as e:
|
|
raise ASRError(f"Error parsing Deepgram response: {e}")
|