Krisp VIVA SDK Filter and Turn support. (#3261)
* Krisp VIVA SDK Filter and Turn support. * Reverted the krisp_filter.py as it's already deprectaed. * enabled test with krisp_audio mock. * More review comment fixes. reverted the state logic in viva filter to be similar to the existing impl on main branch. Fixed tests, ruff, etc. * More review comments for Turn detection. removed integration tests. * Moved the SDK init/deinit into start/stop
This commit is contained in:
Garegin Harutyunyan
committed by
GitHub
GitHub
parent
72a44c2fcd
commit
16819a5caa
@@ -97,7 +97,8 @@ INWORLD_API_KEY=...
|
||||
KRISP_MODEL_PATH=...
|
||||
|
||||
# Krisp Viva
|
||||
KRISP_VIVA_MODEL_PATH=...
|
||||
KRISP_VIVA_FILTER_MODEL_PATH=...
|
||||
KRISP_VIVA_TURN_MODEL_PATH=...
|
||||
|
||||
# LiveKit
|
||||
LIVEKIT_API_KEY=...
|
||||
|
||||
148
scripts/krisp/audio_file_utils.py
Normal file
148
scripts/krisp/audio_file_utils.py
Normal file
@@ -0,0 +1,148 @@
|
||||
"""Utility functions for reading and writing audio files in integration tests.
|
||||
|
||||
This module provides consistent audio file I/O operations for test scripts,
|
||||
handling format detection and conversion to int16 PCM format.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
|
||||
def read_audio_file(input_path: str, verbose: bool = False) -> Tuple[np.ndarray, int]:
|
||||
"""Read an audio file and convert to int16 mono format.
|
||||
|
||||
This function:
|
||||
- Detects the audio format from the file header
|
||||
- Reads PCM_16 files directly without conversion
|
||||
- Converts float formats by scaling to int16 range
|
||||
- Converts stereo to mono
|
||||
- Validates the audio data range
|
||||
|
||||
Args:
|
||||
input_path: Path to the input audio file
|
||||
verbose: If True, print detailed format information
|
||||
|
||||
Returns:
|
||||
Tuple of (audio_data, sample_rate) where audio_data is int16 mono
|
||||
|
||||
Raises:
|
||||
SystemExit: If the audio format is not supported
|
||||
"""
|
||||
if verbose:
|
||||
print(f"Loading audio from: {input_path}")
|
||||
|
||||
# Get audio file info to determine the format
|
||||
info = sf.info(input_path)
|
||||
if verbose:
|
||||
print(
|
||||
f"Audio file format: {info.subtype}, {info.channels} channel(s), {info.samplerate} Hz"
|
||||
)
|
||||
|
||||
# Read audio data based on the source format
|
||||
if info.subtype in ["PCM_16", "PCM_S16"]:
|
||||
# File is already int16, read directly to avoid unnecessary conversion
|
||||
audio_data, sample_rate = sf.read(input_path, dtype="int16")
|
||||
if verbose:
|
||||
print("Read as int16 (native format)")
|
||||
elif info.subtype in ["FLOAT", "DOUBLE"]:
|
||||
# File is float format, read as float32 and scale to int16
|
||||
audio_data, sample_rate = sf.read(input_path, dtype="float32")
|
||||
# Convert float32 (-1.0 to 1.0) to int16 (-32768 to 32767)
|
||||
audio_data = (audio_data * 32767).astype(np.int16)
|
||||
if verbose:
|
||||
print("Read as float32 and scaled to int16")
|
||||
else:
|
||||
print(f"Error: Unsupported audio format: {info.subtype}")
|
||||
print(f"Supported formats: PCM_16, PCM_S16, FLOAT, DOUBLE")
|
||||
sys.exit(1)
|
||||
|
||||
# Convert stereo to mono if needed
|
||||
if len(audio_data.shape) > 1:
|
||||
if verbose:
|
||||
print(f"Converting from {audio_data.shape[1]} channels to mono")
|
||||
if audio_data.dtype == np.int16:
|
||||
# For int16, convert to int32 for averaging to avoid overflow
|
||||
audio_data = audio_data.astype(np.int32).mean(axis=1).astype(np.int16)
|
||||
else:
|
||||
audio_data = audio_data.mean(axis=1).astype(np.int16)
|
||||
|
||||
# Verify the audio has proper range
|
||||
audio_max = abs(audio_data.max())
|
||||
audio_min = abs(audio_data.min())
|
||||
audio_range = max(audio_max, audio_min)
|
||||
|
||||
if audio_range < 100:
|
||||
print(
|
||||
f"⚠️ WARNING: Audio values are very small (max: {audio_data.max()}, min: {audio_data.min()})"
|
||||
)
|
||||
print(f" Expected int16 range: -32768 to 32767")
|
||||
print(f" This may indicate a format conversion issue.")
|
||||
elif verbose:
|
||||
print(f"Audio range: {audio_data.min()} to {audio_data.max()} ✓")
|
||||
|
||||
if verbose:
|
||||
print(
|
||||
f"Audio info: {len(audio_data)} samples, {sample_rate} Hz, {len(audio_data) / sample_rate:.2f} seconds"
|
||||
)
|
||||
|
||||
return audio_data, sample_rate
|
||||
|
||||
|
||||
def write_audio_file(
|
||||
output_path: str, audio_data: np.ndarray, sample_rate: int, verbose: bool = False
|
||||
) -> None:
|
||||
"""Write audio data to a file.
|
||||
|
||||
Args:
|
||||
output_path: Path to the output audio file
|
||||
audio_data: Audio data as numpy array (int16)
|
||||
sample_rate: Sample rate in Hz
|
||||
verbose: If True, print status information
|
||||
|
||||
Raises:
|
||||
ValueError: If output file extension is not supported
|
||||
"""
|
||||
# Validate output file extension
|
||||
valid_extensions = [".wav", ".flac", ".ogg"]
|
||||
output_ext = output_path[output_path.rfind(".") :].lower() if "." in output_path else ""
|
||||
|
||||
if output_ext not in valid_extensions:
|
||||
raise ValueError(
|
||||
f"Invalid output file extension: '{output_ext}'. "
|
||||
f"Supported formats: {', '.join(valid_extensions)}"
|
||||
)
|
||||
|
||||
if verbose:
|
||||
print(f"Saving audio to: {output_path}")
|
||||
print(f" - Format: {output_ext[1:].upper()}")
|
||||
print(f" - Samples: {len(audio_data)}")
|
||||
print(f" - Sample rate: {sample_rate} Hz")
|
||||
|
||||
# Write the audio file
|
||||
sf.write(output_path, audio_data, sample_rate)
|
||||
|
||||
if verbose:
|
||||
print(f"✓ Audio saved successfully")
|
||||
|
||||
|
||||
def calculate_audio_stats(audio_data: np.ndarray) -> dict:
|
||||
"""Calculate statistics for audio data.
|
||||
|
||||
Args:
|
||||
audio_data: Audio data as numpy array
|
||||
|
||||
Returns:
|
||||
Dictionary with audio statistics
|
||||
"""
|
||||
rms = np.sqrt(np.mean(audio_data.astype(np.float32) ** 2))
|
||||
return {
|
||||
"min": int(audio_data.min()),
|
||||
"max": int(audio_data.max()),
|
||||
"mean": float(audio_data.mean()),
|
||||
"std": float(audio_data.std()),
|
||||
"rms": float(rms),
|
||||
"samples": len(audio_data),
|
||||
}
|
||||
262
scripts/krisp/test_krisp_viva_filter_audiofile.py
Normal file
262
scripts/krisp/test_krisp_viva_filter_audiofile.py
Normal file
@@ -0,0 +1,262 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Standalone script to test Krisp VIVA filter with real audio files.
|
||||
|
||||
This script processes audio files through Krisp VIVA filter (noise reduction) and saves the output,
|
||||
allowing you to compare the original and filtered audio.
|
||||
|
||||
Usage:
|
||||
python test_krisp_viva_filter_audiofile.py input.wav output.wav
|
||||
python test_krisp_viva_filter_audiofile.py input.wav output.wav --level 80
|
||||
|
||||
Requirements:
|
||||
pip install soundfile numpy pipecat-ai[krisp]
|
||||
Set KRISP_VIVA_FILTER_MODEL_PATH environment variable to point to your .kef model file
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from audio_file_utils import calculate_audio_stats, read_audio_file, write_audio_file
|
||||
except ImportError as e:
|
||||
print(f"Error: Missing required dependencies: {e}")
|
||||
print("Install with: pip install soundfile numpy")
|
||||
sys.exit(1)
|
||||
|
||||
# Add src directory to Python path for development environment
|
||||
script_dir = Path(__file__).parent
|
||||
project_root = script_dir.parent.parent
|
||||
src_dir = project_root / "src"
|
||||
if src_dir.exists() and str(src_dir) not in sys.path:
|
||||
sys.path.insert(0, str(src_dir))
|
||||
|
||||
# Import Krisp VIVA filter
|
||||
try:
|
||||
from pipecat.audio.filters.krisp_viva_filter import KrispVivaFilter
|
||||
from pipecat.audio.krisp_instance import KRISP_SAMPLE_RATES
|
||||
except ImportError as e:
|
||||
print(f"Error: Could not import Krisp VIVA filter: {e}")
|
||||
print("Make sure pipecat-ai is installed: pip install pipecat-ai[krisp]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def validate_model_path():
|
||||
"""Validate that the Krisp VIVA model path is set and exists."""
|
||||
env_var = "KRISP_VIVA_FILTER_MODEL_PATH"
|
||||
|
||||
model_path = os.getenv(env_var)
|
||||
if not model_path:
|
||||
print(f"Error: {env_var} environment variable not set")
|
||||
print(f"Set it with: export {env_var}=/path/to/model.kef")
|
||||
print(f"Or in PowerShell: $env:{env_var}='C:\\path\\to\\model.kef'")
|
||||
sys.exit(1)
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
print(f"Error: Model file not found: {model_path}")
|
||||
sys.exit(1)
|
||||
|
||||
return model_path
|
||||
|
||||
|
||||
async def process_audio_file(
|
||||
input_path: str,
|
||||
output_path: str,
|
||||
noise_suppression_level: int = 100,
|
||||
frame_duration_ms: int = 10,
|
||||
verbose: bool = False,
|
||||
) -> None:
|
||||
"""Process an audio file through Krisp VIVA filter.
|
||||
|
||||
Args:
|
||||
input_path: Path to input audio file
|
||||
output_path: Path to save filtered audio
|
||||
noise_suppression_level: Noise suppression level (0-100)
|
||||
frame_duration_ms: Frame duration in milliseconds (for chunking input)
|
||||
verbose: Show detailed processing information
|
||||
"""
|
||||
# Read and convert audio file
|
||||
audio_data, sample_rate = read_audio_file(input_path, verbose=True)
|
||||
|
||||
# Validate model path
|
||||
model_path = validate_model_path()
|
||||
|
||||
# Check if sample rate is supported
|
||||
supported_rates = list(KRISP_SAMPLE_RATES.keys())
|
||||
if sample_rate not in supported_rates:
|
||||
print(f"Warning: Sample rate {sample_rate} not in supported rates {supported_rates}")
|
||||
print("Resampling may be required. Continuing anyway...")
|
||||
|
||||
print(f"\nInitializing VIVA filter:")
|
||||
print(f" - Model path: {model_path}")
|
||||
print(f" - Noise suppression level: {noise_suppression_level}")
|
||||
print(f" - Frame duration: {frame_duration_ms}ms (processing chunk size)")
|
||||
print(f" - Sample rate: {sample_rate}Hz")
|
||||
|
||||
# Create filter instance and measure preload time
|
||||
print("\nInitializing filter (preloading model)...")
|
||||
preload_start_time = time.time()
|
||||
filter_obj = KrispVivaFilter(
|
||||
model_path=model_path,
|
||||
noise_suppression_level=noise_suppression_level,
|
||||
)
|
||||
preload_duration = time.time() - preload_start_time
|
||||
print(f"Model preloaded in {preload_duration * 1000:.2f}ms")
|
||||
|
||||
try:
|
||||
# Measure filter start time
|
||||
print("\nStarting filter...")
|
||||
start_time = time.time()
|
||||
await filter_obj.start(sample_rate)
|
||||
start_duration = time.time() - start_time
|
||||
print(f"Filter started in {start_duration * 1000:.2f}ms")
|
||||
|
||||
print("\nProcessing audio...")
|
||||
filtered_samples = []
|
||||
total_frames = 0
|
||||
|
||||
# Use chunk size matching filter frame duration
|
||||
chunk_size = int(sample_rate * frame_duration_ms / 1000)
|
||||
print(f" - Chunk size: {chunk_size} samples ({frame_duration_ms}ms)")
|
||||
|
||||
if verbose:
|
||||
print(f" - Processing {len(audio_data)} samples in chunks of {chunk_size}")
|
||||
|
||||
for i in range(0, len(audio_data), chunk_size):
|
||||
chunk = audio_data[i : i + chunk_size]
|
||||
|
||||
if len(chunk) == 0:
|
||||
break
|
||||
|
||||
# Skip incomplete chunks
|
||||
if len(chunk) < chunk_size:
|
||||
if verbose:
|
||||
print(f"\n Skipping incomplete final chunk: {len(chunk)} samples")
|
||||
break
|
||||
|
||||
# Filter the chunk
|
||||
filtered_chunk_bytes = await filter_obj.filter(chunk.tobytes())
|
||||
|
||||
# Collect filtered samples
|
||||
if filtered_chunk_bytes:
|
||||
filtered_chunk = np.frombuffer(filtered_chunk_bytes, dtype=np.int16)
|
||||
filtered_samples.append(filtered_chunk)
|
||||
total_frames += 1
|
||||
|
||||
if verbose and total_frames <= 3:
|
||||
print(
|
||||
f" Frame {total_frames}: {len(chunk)} -> {len(filtered_chunk)} samples"
|
||||
)
|
||||
|
||||
# Progress indicator
|
||||
if i % (chunk_size * 50) == 0:
|
||||
progress = (i / len(audio_data)) * 100
|
||||
print(f" Progress: {progress:.1f}%", end="\r")
|
||||
|
||||
print(f" Progress: 100.0% - Processed {total_frames} frames")
|
||||
|
||||
# Concatenate all filtered samples
|
||||
if filtered_samples:
|
||||
filtered_audio = np.concatenate(filtered_samples)
|
||||
print(f"\nFiltered audio: {len(filtered_audio)} samples")
|
||||
|
||||
# Save the filtered audio
|
||||
write_audio_file(output_path, filtered_audio, sample_rate, verbose=True)
|
||||
|
||||
# Calculate statistics
|
||||
original_stats = calculate_audio_stats(audio_data)
|
||||
filtered_stats = calculate_audio_stats(filtered_audio)
|
||||
|
||||
print("\nAudio Statistics:")
|
||||
print(f" Original RMS: {original_stats['rms']:.2f}")
|
||||
print(f" Filtered RMS: {filtered_stats['rms']:.2f}")
|
||||
print(f" RMS Ratio: {filtered_stats['rms'] / original_stats['rms']:.2f}")
|
||||
|
||||
if filtered_stats["rms"] < 0.01:
|
||||
print("\n ⚠️ WARNING: Filtered audio is very quiet or silent!")
|
||||
print(" This may indicate a processing issue.")
|
||||
|
||||
print("\n✅ Processing complete!")
|
||||
print(f" Original: {input_path}")
|
||||
print(f" Filtered: {output_path}")
|
||||
print("\nListen to both files to compare the results.")
|
||||
|
||||
else:
|
||||
print("Error: No filtered audio produced")
|
||||
sys.exit(1)
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await filter_obj.stop()
|
||||
print("Filter stopped.")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Test Krisp VIVA filter with real audio files",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python test_krisp_viva_audiofile.py noisy_input.wav clean_output.wav
|
||||
python test_krisp_viva_audiofile.py input.wav output.wav --level 80
|
||||
|
||||
Supported audio formats: WAV, FLAC, OGG, etc. (via soundfile)
|
||||
Supported sample rates: 8000, 16000, 24000, 32000, 44100, 48000 Hz
|
||||
|
||||
Note: Set KRISP_VIVA_FILTER_MODEL_PATH environment variable to point to your .kef model file
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument("input", help="Input audio file path")
|
||||
parser.add_argument("output", help="Output audio file path")
|
||||
parser.add_argument(
|
||||
"--level",
|
||||
type=int,
|
||||
default=100,
|
||||
help="Noise suppression level (0-100, default: 100)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frame-duration",
|
||||
type=int,
|
||||
default=10,
|
||||
choices=[10, 15, 20, 30, 32],
|
||||
help="Frame duration in milliseconds (default: 10)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
help="Show detailed processing information",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Validate input file exists
|
||||
if not os.path.exists(args.input):
|
||||
print(f"Error: Input file not found: {args.input}")
|
||||
sys.exit(1)
|
||||
|
||||
# Create output directory if needed
|
||||
output_dir = os.path.dirname(args.output)
|
||||
if output_dir and not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
# Process the audio
|
||||
asyncio.run(
|
||||
process_audio_file(
|
||||
args.input,
|
||||
args.output,
|
||||
noise_suppression_level=args.level,
|
||||
frame_duration_ms=args.frame_duration,
|
||||
verbose=args.verbose,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
333
scripts/krisp/test_krisp_viva_turn_audiofile.py
Normal file
333
scripts/krisp/test_krisp_viva_turn_audiofile.py
Normal file
@@ -0,0 +1,333 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Standalone script to test Krisp VIVA turn analyzer with real audio files.
|
||||
|
||||
This script processes audio files through Krisp VIVA turn analyzer and analyzes
|
||||
turn detection, allowing you to test turn detection on real audio data.
|
||||
|
||||
Usage:
|
||||
python test_krisp_viva_turn_audiofile.py input.wav
|
||||
python test_krisp_viva_turn_audiofile.py input.wav --threshold 0.7
|
||||
python test_krisp_viva_turn_audiofile.py input.wav --frame-duration 20
|
||||
|
||||
Requirements:
|
||||
pip install soundfile numpy pipecat-ai[krisp]
|
||||
Set KRISP_VIVA_TURN_MODEL_PATH environment variable to point to your .kef model file
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from audio_file_utils import read_audio_file
|
||||
except ImportError as e:
|
||||
print(f"Error: Missing required dependencies: {e}")
|
||||
print("Install with: pip install soundfile numpy")
|
||||
sys.exit(1)
|
||||
|
||||
# Add src directory to Python path for development environment
|
||||
script_dir = Path(__file__).parent
|
||||
project_root = script_dir.parent.parent
|
||||
src_dir = project_root / "src"
|
||||
if src_dir.exists() and str(src_dir) not in sys.path:
|
||||
sys.path.insert(0, str(src_dir))
|
||||
|
||||
# Import Krisp VIVA turn analyzer
|
||||
try:
|
||||
from pipecat.audio.krisp_instance import KRISP_SAMPLE_RATES
|
||||
from pipecat.audio.turn.krisp_viva_turn import KrispTurnParams, KrispVivaTurn
|
||||
except ImportError as e:
|
||||
print(f"Error: Could not import Krisp VIVA turn analyzer: {e}")
|
||||
print("Make sure pipecat-ai is installed: pip install pipecat-ai[krisp]")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def validate_model_path():
|
||||
"""Validate that the Krisp VIVA turn model path is set and exists."""
|
||||
env_var = "KRISP_VIVA_TURN_MODEL_PATH"
|
||||
|
||||
model_path = os.getenv(env_var)
|
||||
if not model_path:
|
||||
print(f"Error: {env_var} environment variable not set")
|
||||
print(f"Set it with: export {env_var}=/path/to/model.kef")
|
||||
print(f"Or in PowerShell: $env:{env_var}='C:\\path\\to\\model.kef'")
|
||||
sys.exit(1)
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
print(f"Error: Model file not found: {model_path}")
|
||||
sys.exit(1)
|
||||
|
||||
return model_path
|
||||
|
||||
|
||||
async def analyze_audio_file(
|
||||
input_path: str,
|
||||
threshold: float = 0.5,
|
||||
frame_duration_ms: int = 20,
|
||||
chunk_duration_ms: int = 20,
|
||||
verbose: bool = False,
|
||||
output_file: str = None,
|
||||
) -> None:
|
||||
"""Analyze an audio file for turn detection using Krisp VIVA turn analyzer.
|
||||
|
||||
Args:
|
||||
input_path: Path to input audio file
|
||||
threshold: Probability threshold for turn completion (0.0 to 1.0)
|
||||
frame_duration_ms: Frame duration in milliseconds for turn detection model
|
||||
chunk_duration_ms: Processing chunk size in milliseconds
|
||||
verbose: Show detailed processing information
|
||||
output_file: Optional path to save turn probabilities (one per line)
|
||||
"""
|
||||
# Read and convert audio file
|
||||
audio_data, sample_rate = read_audio_file(input_path, verbose=True)
|
||||
|
||||
# Validate model path
|
||||
model_path = validate_model_path()
|
||||
|
||||
# Check if sample rate is supported
|
||||
supported_rates = list(KRISP_SAMPLE_RATES.keys())
|
||||
if sample_rate not in supported_rates:
|
||||
print(f"Warning: Sample rate {sample_rate} not in supported rates {supported_rates}")
|
||||
print("Resampling may be required. Continuing anyway...")
|
||||
|
||||
print(f"\nInitializing VIVA turn analyzer:")
|
||||
print(f" - Model path: {model_path}")
|
||||
print(f" - Threshold: {threshold}")
|
||||
print(f" - Frame duration: {frame_duration_ms}ms")
|
||||
print(f" - Sample rate: {sample_rate}Hz")
|
||||
print(f" - Processing chunk size: {chunk_duration_ms}ms")
|
||||
|
||||
# Create turn analyzer instance
|
||||
print("\nInitializing turn analyzer...")
|
||||
init_start_time = time.time()
|
||||
params = KrispTurnParams(threshold=threshold, frame_duration_ms=frame_duration_ms)
|
||||
turn_analyzer = KrispVivaTurn(model_path=model_path, params=params)
|
||||
init_duration = time.time() - init_start_time
|
||||
print(f"Turn analyzer initialized in {init_duration * 1000:.2f}ms")
|
||||
|
||||
try:
|
||||
# Set sample rate
|
||||
print("\nSetting sample rate...")
|
||||
set_rate_start_time = time.time()
|
||||
turn_analyzer.set_sample_rate(sample_rate)
|
||||
set_rate_duration = time.time() - set_rate_start_time
|
||||
print(f"Sample rate set to {turn_analyzer.sample_rate}Hz")
|
||||
print(f"set_sample_rate latency: {set_rate_duration * 1000:.2f}ms")
|
||||
|
||||
print("\nProcessing audio for turn detection...")
|
||||
|
||||
# Calculate exact frame size based on frame duration
|
||||
# The Krisp Tt processor requires exact frame sizes matching the configured frame duration
|
||||
frame_size_samples = int(sample_rate * frame_duration_ms / 1000)
|
||||
print(f" Frame size: {frame_size_samples} samples ({frame_duration_ms}ms)")
|
||||
|
||||
turn_events = []
|
||||
speech_segments = []
|
||||
current_speech_start = None
|
||||
all_probabilities = [] # Store all probabilities for output file
|
||||
|
||||
# Simple energy-based VAD (for demonstration)
|
||||
energy_threshold = np.std(audio_data) * 0.1
|
||||
|
||||
# Buffer for incomplete frames - we need to send exact frame sizes
|
||||
audio_buffer = np.array([], dtype=np.int16)
|
||||
frames_processed = 0
|
||||
|
||||
# Process audio in chunks, buffering to ensure exact frame sizes
|
||||
read_chunk_size = max(frame_size_samples, int(sample_rate * chunk_duration_ms / 1000))
|
||||
|
||||
for i in range(0, len(audio_data), read_chunk_size):
|
||||
chunk = audio_data[i : i + read_chunk_size]
|
||||
|
||||
if len(chunk) == 0:
|
||||
break
|
||||
|
||||
# Add chunk to buffer
|
||||
audio_buffer = np.concatenate([audio_buffer, chunk])
|
||||
|
||||
# Process complete frames from buffer
|
||||
while len(audio_buffer) >= frame_size_samples:
|
||||
# Extract exactly one frame
|
||||
frame_samples = audio_buffer[:frame_size_samples].copy()
|
||||
audio_buffer = audio_buffer[frame_size_samples:]
|
||||
|
||||
# Calculate timestamp for this frame
|
||||
timestamp = frames_processed * frame_duration_ms / 1000.0
|
||||
frames_processed += 1
|
||||
|
||||
# Simple VAD: check if frame has significant energy
|
||||
frame_energy = np.sqrt(np.mean(frame_samples.astype(np.float32) ** 2))
|
||||
is_speech = frame_energy > energy_threshold
|
||||
|
||||
# Process frame through turn analyzer
|
||||
frame_bytes = frame_samples.tobytes()
|
||||
end_of_turn_state = turn_analyzer.append_audio(frame_bytes, is_speech)
|
||||
|
||||
# Collect all probabilities from this call
|
||||
# The TT model processes frames and returns probabilities per 100ms
|
||||
# append_audio may process multiple frames, so collect all of them
|
||||
all_probabilities.extend(turn_analyzer.frame_probabilities)
|
||||
|
||||
# Track speech segments
|
||||
if is_speech:
|
||||
if current_speech_start is None:
|
||||
current_speech_start = timestamp
|
||||
else:
|
||||
if current_speech_start is not None:
|
||||
speech_segments.append((current_speech_start, timestamp))
|
||||
current_speech_start = None
|
||||
|
||||
# Track turn completion events
|
||||
if end_of_turn_state.value == 1: # EndOfTurnState.COMPLETE
|
||||
turn_events.append(
|
||||
{
|
||||
"timestamp": timestamp,
|
||||
"speech_triggered": turn_analyzer.speech_triggered,
|
||||
}
|
||||
)
|
||||
if verbose:
|
||||
print(f" Turn completed at {timestamp:.2f}s")
|
||||
|
||||
# Progress indicator
|
||||
if i % (read_chunk_size * 50) == 0:
|
||||
progress = (i / len(audio_data)) * 100
|
||||
print(f" Progress: {progress:.1f}%", end="\r")
|
||||
|
||||
# Process any remaining incomplete frame (if buffer has data)
|
||||
if len(audio_buffer) > 0:
|
||||
if verbose:
|
||||
print(
|
||||
f"\n Warning: {len(audio_buffer)} samples remaining (incomplete frame, will be discarded)"
|
||||
)
|
||||
|
||||
print(f" Progress: 100.0%")
|
||||
|
||||
# Final speech segment if still speaking
|
||||
if current_speech_start is not None:
|
||||
speech_segments.append((current_speech_start, len(audio_data) / sample_rate))
|
||||
|
||||
# Print results
|
||||
print("\n" + "=" * 60)
|
||||
print("Turn Detection Results:")
|
||||
print("=" * 60)
|
||||
|
||||
print(f"\nSpeech Segments Detected: {len(speech_segments)}")
|
||||
for i, (start, end) in enumerate(speech_segments, 1):
|
||||
duration = end - start
|
||||
print(f" Segment {i}: {start:.2f}s - {end:.2f}s (duration: {duration:.2f}s)")
|
||||
|
||||
print(f"\nTurn Completion Events: {len(turn_events)}")
|
||||
for i, event in enumerate(turn_events, 1):
|
||||
print(f" Turn {i} completed at {event['timestamp']:.2f}s")
|
||||
|
||||
print(f"\nFinal State:")
|
||||
print(f" Speech triggered: {turn_analyzer.speech_triggered}")
|
||||
print(f" Sample rate: {turn_analyzer.sample_rate}Hz")
|
||||
print(f" Total probabilities collected: {len(all_probabilities)}")
|
||||
|
||||
if len(turn_events) == 0:
|
||||
print("\n ⚠️ No turn completion events detected.")
|
||||
print(" This could mean:")
|
||||
print(" - The audio doesn't contain clear turn boundaries")
|
||||
print(" - The threshold is too high")
|
||||
print(" - The model needs different parameters")
|
||||
|
||||
# Save probabilities to output file if specified
|
||||
if output_file:
|
||||
with open(output_file, "w") as f:
|
||||
for prob in all_probabilities:
|
||||
f.write(f"{prob}\n")
|
||||
print(f"\n📄 Turn probabilities saved to: {output_file}")
|
||||
print(f" Total frames: {len(all_probabilities)}")
|
||||
|
||||
print("\n✅ Analysis complete!")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
turn_analyzer.clear()
|
||||
print("Turn analyzer cleared.")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Test Krisp VIVA turn analyzer with real audio files",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
python test_krisp_viva_turn_audiofile.py conversation.wav
|
||||
python test_krisp_viva_turn_audiofile.py input.wav --threshold 0.7
|
||||
python test_krisp_viva_turn_audiofile.py input.wav --frame-duration 20
|
||||
|
||||
Supported audio formats: WAV, FLAC, OGG, etc. (via soundfile)
|
||||
Supported sample rates: 8000, 16000, 24000, 32000, 44100, 48000 Hz
|
||||
|
||||
Note: Set KRISP_VIVA_TURN_MODEL_PATH environment variable to point to your .kef model file
|
||||
""",
|
||||
)
|
||||
|
||||
parser.add_argument("input", help="Input audio file path")
|
||||
parser.add_argument(
|
||||
"--threshold",
|
||||
type=float,
|
||||
default=0.5,
|
||||
help="Probability threshold for turn completion (0.0 to 1.0, default: 0.5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frame-duration",
|
||||
type=int,
|
||||
default=20,
|
||||
choices=[10, 15, 20, 30, 32],
|
||||
help="Frame duration in milliseconds for turn detection model (default: 20)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chunk-duration",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Processing chunk size in milliseconds (default: 20)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-v",
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
help="Show detailed processing information",
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Output file path to save turn probabilities (.tt format, one probability per line)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Validate input file exists
|
||||
if not os.path.exists(args.input):
|
||||
print(f"Error: Input file not found: {args.input}")
|
||||
sys.exit(1)
|
||||
|
||||
# Validate threshold
|
||||
if not 0.0 <= args.threshold <= 1.0:
|
||||
print(f"Error: Threshold must be between 0.0 and 1.0, got {args.threshold}")
|
||||
sys.exit(1)
|
||||
|
||||
# Process the audio
|
||||
asyncio.run(
|
||||
analyze_audio_file(
|
||||
args.input,
|
||||
threshold=args.threshold,
|
||||
frame_duration_ms=args.frame_duration,
|
||||
chunk_duration_ms=args.chunk_duration,
|
||||
verbose=args.verbose,
|
||||
output_file=args.output,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -61,7 +61,6 @@ class KrispFilter(BaseAudioFilter):
|
||||
Provides real-time noise reduction for audio streams using Krisp's
|
||||
proprietary noise suppression algorithms. Requires a Krisp model file
|
||||
for operation.
|
||||
|
||||
.. deprecated:: 0.0.94
|
||||
The KrispFilter is deprecated and will be removed in a future version.
|
||||
Use KrispVivaFilter instead.
|
||||
|
||||
@@ -9,111 +9,121 @@
|
||||
This module provides an audio filter implementation using Krisp VIVA SDK.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from loguru import logger
|
||||
|
||||
from pipecat.audio.filters.base_audio_filter import BaseAudioFilter
|
||||
from pipecat.audio.krisp_instance import (
|
||||
KrispVivaSDKManager,
|
||||
int_to_krisp_frame_duration,
|
||||
int_to_krisp_sample_rate,
|
||||
)
|
||||
from pipecat.frames.frames import FilterControlFrame, FilterEnableFrame
|
||||
|
||||
try:
|
||||
import krisp_audio
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
logger.error("In order to use the Krisp filter, you need to install krisp_audio.")
|
||||
logger.error("In order to use KrispVivaFilter, you need to install krisp_audio.")
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
def _log_callback(log_message, log_level):
|
||||
logger.info(f"[{log_level}] {log_message}")
|
||||
|
||||
|
||||
class KrispVivaFilter(BaseAudioFilter):
|
||||
"""Audio filter using the Krisp VIVA SDK.
|
||||
|
||||
Provides real-time noise reduction for audio streams using Krisp's
|
||||
proprietary noise suppression algorithms. This filter requires a
|
||||
valid Krisp model file to operate.
|
||||
|
||||
Supported sample rates:
|
||||
- 8000 Hz
|
||||
- 16000 Hz
|
||||
- 24000 Hz
|
||||
- 32000 Hz
|
||||
- 44100 Hz
|
||||
- 48000 Hz
|
||||
"""
|
||||
|
||||
# Initialize Krisp Audio SDK globally
|
||||
krisp_audio.globalInit("", _log_callback, krisp_audio.LogLevel.Off)
|
||||
SDK_VERSION = krisp_audio.getVersion()
|
||||
logger.debug(
|
||||
f"Krisp Audio Python SDK Version: {SDK_VERSION.major}."
|
||||
f"{SDK_VERSION.minor}.{SDK_VERSION.patch}"
|
||||
)
|
||||
|
||||
SAMPLE_RATES = {
|
||||
8000: krisp_audio.SamplingRate.Sr8000Hz,
|
||||
16000: krisp_audio.SamplingRate.Sr16000Hz,
|
||||
24000: krisp_audio.SamplingRate.Sr24000Hz,
|
||||
32000: krisp_audio.SamplingRate.Sr32000Hz,
|
||||
44100: krisp_audio.SamplingRate.Sr44100Hz,
|
||||
48000: krisp_audio.SamplingRate.Sr48000Hz,
|
||||
}
|
||||
|
||||
FRAME_SIZE_MS = 10 # Krisp requires audio frames of 10ms duration for processing.
|
||||
|
||||
def __init__(self, model_path: str = None, noise_suppression_level: int = 100) -> None:
|
||||
def __init__(
|
||||
self, model_path: str = None, frame_duration: int = 10, noise_suppression_level: int = 100
|
||||
) -> None:
|
||||
"""Initialize the Krisp noise reduction filter.
|
||||
|
||||
Args:
|
||||
model_path: Path to the Krisp model file (.kef extension).
|
||||
If None, uses KRISP_VIVA_MODEL_PATH environment variable.
|
||||
If None, uses KRISP_VIVA_FILTER_MODEL_PATH environment variable.
|
||||
frame_duration: Frame duration in milliseconds.
|
||||
noise_suppression_level: Noise suppression level.
|
||||
|
||||
Raises:
|
||||
ValueError: If model_path is not provided and KRISP_VIVA_MODEL_PATH is not set.
|
||||
ValueError: If model_path is not provided and KRISP_VIVA_FILTER_MODEL_PATH is not set.
|
||||
Exception: If model file doesn't have .kef extension.
|
||||
FileNotFoundError: If model file doesn't exist.
|
||||
RuntimeError: If Krisp SDK initialization fails.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Set model path, checking environment if not specified
|
||||
self._model_path = model_path or os.getenv("KRISP_VIVA_MODEL_PATH")
|
||||
if not self._model_path:
|
||||
logger.error("Model path is not provided and KRISP_VIVA_MODEL_PATH is not set.")
|
||||
raise ValueError("Model path for KrispAudioProcessor must be provided.")
|
||||
try:
|
||||
# Set model path, checking environment if not specified
|
||||
if model_path:
|
||||
self._model_path = model_path
|
||||
else:
|
||||
# Check new environment variable first
|
||||
self._model_path = os.getenv("KRISP_VIVA_FILTER_MODEL_PATH")
|
||||
# Fall back to old environment variable for backward compatibility
|
||||
if not self._model_path:
|
||||
self._model_path = os.getenv("KRISP_VIVA_MODEL_PATH")
|
||||
if self._model_path:
|
||||
logger.warning(
|
||||
"KRISP_VIVA_MODEL_PATH is deprecated. "
|
||||
"Please use KRISP_VIVA_FILTER_MODEL_PATH instead."
|
||||
)
|
||||
if not self._model_path:
|
||||
logger.error(
|
||||
"Model path is not provided and KRISP_VIVA_FILTER_MODEL_PATH is not set."
|
||||
)
|
||||
raise ValueError("Model path for KrispAudioProcessor must be provided.")
|
||||
|
||||
if not self._model_path.endswith(".kef"):
|
||||
raise Exception("Model is expected with .kef extension")
|
||||
if not self._model_path.endswith(".kef"):
|
||||
raise Exception("Model is expected with .kef extension")
|
||||
|
||||
if not os.path.isfile(self._model_path):
|
||||
raise FileNotFoundError(f"Model file not found: {self._model_path}")
|
||||
if not os.path.isfile(self._model_path):
|
||||
raise FileNotFoundError(f"Model file not found: {self._model_path}")
|
||||
|
||||
self._filtering = True
|
||||
self._session = None
|
||||
self._samples_per_frame = None
|
||||
self._noise_suppression_level = noise_suppression_level
|
||||
self._session = None
|
||||
self._samples_per_frame = None
|
||||
self._noise_suppression_level = noise_suppression_level
|
||||
self._frame_duration_ms = frame_duration
|
||||
self._audio_buffer = bytearray()
|
||||
self._filtering = True
|
||||
|
||||
# Audio buffer to accumulate samples for complete frames
|
||||
self._audio_buffer = bytearray()
|
||||
except Exception:
|
||||
# If initialization fails, release the SDK reference
|
||||
KrispVivaSDKManager.release()
|
||||
raise
|
||||
|
||||
def _int_to_sample_rate(self, sample_rate):
|
||||
"""Convert integer sample rate to krisp_audio SamplingRate enum.
|
||||
def _create_session(self, sample_rate: int, frame_duration: int):
|
||||
"""Create a Krisp session with a specific sample rate.
|
||||
|
||||
Args:
|
||||
sample_rate: Sample rate as integer
|
||||
|
||||
Returns:
|
||||
krisp_audio.SamplingRate enum value
|
||||
sample_rate: Sample rate for the session
|
||||
frame_duration: Frame duration in milliseconds
|
||||
|
||||
Raises:
|
||||
ValueError: If sample rate is not supported
|
||||
Exception: If session creation fails
|
||||
"""
|
||||
if sample_rate not in self.SAMPLE_RATES:
|
||||
raise ValueError("Unsupported sample rate")
|
||||
return self.SAMPLE_RATES[sample_rate]
|
||||
try:
|
||||
model_info = krisp_audio.ModelInfo()
|
||||
model_info.path = self._model_path
|
||||
|
||||
nc_cfg = krisp_audio.NcSessionConfig()
|
||||
nc_cfg.inputSampleRate = int_to_krisp_sample_rate(sample_rate)
|
||||
nc_cfg.inputFrameDuration = int_to_krisp_frame_duration(frame_duration)
|
||||
nc_cfg.outputSampleRate = nc_cfg.inputSampleRate
|
||||
nc_cfg.modelInfo = model_info
|
||||
|
||||
self._samples_per_frame = int((sample_rate * frame_duration) / 1000)
|
||||
self._current_sample_rate = sample_rate
|
||||
session = krisp_audio.NcInt16.create(nc_cfg)
|
||||
return session
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create Krisp session: {e}", exc_info=True)
|
||||
raise RuntimeError(f"Failed to create Krisp processing session: {e}") from e
|
||||
|
||||
async def start(self, sample_rate: int):
|
||||
"""Initialize the Krisp processor with the transport's sample rate.
|
||||
@@ -121,21 +131,24 @@ class KrispVivaFilter(BaseAudioFilter):
|
||||
Args:
|
||||
sample_rate: The sample rate of the input transport in Hz.
|
||||
"""
|
||||
model_info = krisp_audio.ModelInfo()
|
||||
model_info.path = self._model_path
|
||||
|
||||
nc_cfg = krisp_audio.NcSessionConfig()
|
||||
nc_cfg.inputSampleRate = self._int_to_sample_rate(sample_rate)
|
||||
nc_cfg.inputFrameDuration = krisp_audio.FrameDuration.Fd10ms
|
||||
nc_cfg.outputSampleRate = nc_cfg.inputSampleRate
|
||||
nc_cfg.modelInfo = model_info
|
||||
|
||||
self._samples_per_frame = int((sample_rate * self.FRAME_SIZE_MS) / 1000)
|
||||
self._session = krisp_audio.NcInt16.create(nc_cfg)
|
||||
try:
|
||||
# Acquire SDK reference (will initialize on first call)
|
||||
KrispVivaSDKManager.acquire()
|
||||
self._session = self._create_session(sample_rate, self._frame_duration_ms)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start Krisp session: {e}", exc_info=True)
|
||||
self._session = None
|
||||
raise RuntimeError(f"Failed to create Krisp processing session: {e}") from e
|
||||
|
||||
async def stop(self):
|
||||
"""Clean up the Krisp processor when stopping."""
|
||||
self._session = None
|
||||
try:
|
||||
self._session = None
|
||||
self._audio_buffer.clear()
|
||||
KrispVivaSDKManager.release()
|
||||
except Exception as e:
|
||||
logger.error(f"Error in stop: {e}", exc_info=True)
|
||||
raise RuntimeError(f"Failed to stop Krisp processor: {e}") from e
|
||||
|
||||
async def process_frame(self, frame: FilterControlFrame):
|
||||
"""Process control frames to enable/disable filtering.
|
||||
@@ -158,36 +171,41 @@ class KrispVivaFilter(BaseAudioFilter):
|
||||
if not self._filtering:
|
||||
return audio
|
||||
|
||||
# Add incoming audio to our buffer
|
||||
self._audio_buffer.extend(audio)
|
||||
try:
|
||||
# Add incoming audio to our buffer
|
||||
self._audio_buffer.extend(audio)
|
||||
|
||||
# Calculate how many complete frames we can process
|
||||
total_samples = len(self._audio_buffer) // 2 # 2 bytes per int16 sample
|
||||
num_complete_frames = total_samples // self._samples_per_frame
|
||||
# Calculate how many complete frames we can process
|
||||
total_samples = len(self._audio_buffer) // 2 # 2 bytes per int16 sample
|
||||
num_complete_frames = total_samples // self._samples_per_frame
|
||||
|
||||
if num_complete_frames == 0:
|
||||
# Not enough samples for a complete frame yet, return empty
|
||||
return b""
|
||||
if num_complete_frames == 0:
|
||||
# Not enough samples for a complete frame yet, return empty
|
||||
return b""
|
||||
|
||||
# Calculate how many bytes we need for complete frames
|
||||
complete_samples_count = num_complete_frames * self._samples_per_frame
|
||||
bytes_to_process = complete_samples_count * 2 # 2 bytes per sample
|
||||
# Calculate how many bytes we need for complete frames
|
||||
complete_samples_count = num_complete_frames * self._samples_per_frame
|
||||
bytes_to_process = complete_samples_count * 2 # 2 bytes per sample
|
||||
|
||||
# Extract the bytes we can process
|
||||
audio_to_process = bytes(self._audio_buffer[:bytes_to_process])
|
||||
# Extract the bytes we can process
|
||||
audio_to_process = bytes(self._audio_buffer[:bytes_to_process])
|
||||
|
||||
# Remove processed bytes from buffer, keep the remainder
|
||||
self._audio_buffer = self._audio_buffer[bytes_to_process:]
|
||||
# Remove processed bytes from buffer, keep the remainder
|
||||
self._audio_buffer = self._audio_buffer[bytes_to_process:]
|
||||
|
||||
# Process the complete frames
|
||||
samples = np.frombuffer(audio_to_process, dtype=np.int16)
|
||||
frames = samples.reshape(-1, self._samples_per_frame)
|
||||
processed_samples = np.empty_like(samples)
|
||||
# Process the complete frames
|
||||
samples = np.frombuffer(audio_to_process, dtype=np.int16)
|
||||
frames = samples.reshape(-1, self._samples_per_frame)
|
||||
processed_samples = np.empty_like(samples)
|
||||
|
||||
for i, frame in enumerate(frames):
|
||||
cleaned_frame = self._session.process(frame, self._noise_suppression_level)
|
||||
processed_samples[i * self._samples_per_frame : (i + 1) * self._samples_per_frame] = (
|
||||
cleaned_frame
|
||||
)
|
||||
for i, frame in enumerate(frames):
|
||||
cleaned_frame = self._session.process(frame, self._noise_suppression_level)
|
||||
processed_samples[
|
||||
i * self._samples_per_frame : (i + 1) * self._samples_per_frame
|
||||
] = cleaned_frame
|
||||
|
||||
return processed_samples.tobytes()
|
||||
return processed_samples.tobytes()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during Krisp filtering: {e}", exc_info=True)
|
||||
return audio
|
||||
|
||||
183
src/pipecat/audio/krisp_instance.py
Normal file
183
src/pipecat/audio/krisp_instance.py
Normal file
@@ -0,0 +1,183 @@
|
||||
#
|
||||
# Copyright (c) 2024–2025, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""Krisp Instance manager for pipecat audio."""
|
||||
|
||||
import atexit
|
||||
from threading import Lock
|
||||
|
||||
from loguru import logger
|
||||
|
||||
try:
|
||||
import krisp_audio
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
logger.error("In order to use the Krisp instance, you need to install krisp_audio.")
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
# Mapping of sample rates (Hz) to Krisp SDK SamplingRate enums
|
||||
KRISP_SAMPLE_RATES = {
|
||||
8000: krisp_audio.SamplingRate.Sr8000Hz,
|
||||
16000: krisp_audio.SamplingRate.Sr16000Hz,
|
||||
24000: krisp_audio.SamplingRate.Sr24000Hz,
|
||||
32000: krisp_audio.SamplingRate.Sr32000Hz,
|
||||
44100: krisp_audio.SamplingRate.Sr44100Hz,
|
||||
48000: krisp_audio.SamplingRate.Sr48000Hz,
|
||||
}
|
||||
|
||||
KRISP_FRAME_DURATIONS = {
|
||||
10: krisp_audio.FrameDuration.Fd10ms,
|
||||
15: krisp_audio.FrameDuration.Fd15ms,
|
||||
20: krisp_audio.FrameDuration.Fd20ms,
|
||||
30: krisp_audio.FrameDuration.Fd30ms,
|
||||
32: krisp_audio.FrameDuration.Fd32ms,
|
||||
}
|
||||
|
||||
|
||||
def int_to_krisp_sample_rate(sample_rate: int):
|
||||
"""Convert integer sample rate to Krisp SDK enum value.
|
||||
|
||||
Args:
|
||||
sample_rate: Sample rate in Hz (e.g., 16000, 24000, 48000).
|
||||
|
||||
Returns:
|
||||
Corresponding Krisp SDK SampleRate enum value.
|
||||
|
||||
Raises:
|
||||
ValueError: If the sample rate is not supported by Krisp SDK.
|
||||
"""
|
||||
if sample_rate not in KRISP_SAMPLE_RATES:
|
||||
supported_rates = ", ".join(str(rate) for rate in sorted(KRISP_SAMPLE_RATES.keys()))
|
||||
raise ValueError(
|
||||
f"Unsupported sample rate: {sample_rate} Hz. Supported rates: {supported_rates} Hz"
|
||||
)
|
||||
return KRISP_SAMPLE_RATES[sample_rate]
|
||||
|
||||
|
||||
def int_to_krisp_frame_duration(frame_duration_ms: int):
|
||||
"""Convert integer frame duration to Krisp SDK enum value.
|
||||
|
||||
Args:
|
||||
frame_duration_ms: Frame duration in milliseconds (e.g., 10, 20, 30).
|
||||
|
||||
Returns:
|
||||
Corresponding Krisp SDK FrameDuration enum value.
|
||||
|
||||
Raises:
|
||||
ValueError: If the frame duration is not supported by Krisp SDK.
|
||||
"""
|
||||
if frame_duration_ms not in KRISP_FRAME_DURATIONS:
|
||||
supported_durations = ", ".join(
|
||||
str(duration) for duration in sorted(KRISP_FRAME_DURATIONS.keys())
|
||||
)
|
||||
raise ValueError(
|
||||
f"Unsupported frame duration: {frame_duration_ms} ms. "
|
||||
f"Supported durations: {supported_durations} ms"
|
||||
)
|
||||
return KRISP_FRAME_DURATIONS[frame_duration_ms]
|
||||
|
||||
|
||||
class KrispVivaSDKManager:
|
||||
"""Singleton manager for Krisp VIVA SDK with reference counting."""
|
||||
|
||||
_initialized = False
|
||||
_lock = Lock()
|
||||
_reference_count = 0
|
||||
|
||||
@staticmethod
|
||||
def _log_callback(log_message, log_level):
|
||||
"""Thread-safe callback for Krisp SDK logging."""
|
||||
logger.info(f"[{log_level}] {log_message}")
|
||||
|
||||
@classmethod
|
||||
def acquire(cls):
|
||||
"""Acquire a reference to the SDK (initializes if needed).
|
||||
|
||||
Call this when creating a filter instance.
|
||||
|
||||
Raises:
|
||||
Exception: If SDK initialization fails (propagated from krisp_audio)
|
||||
"""
|
||||
with cls._lock:
|
||||
# Initialize SDK on first acquire
|
||||
if cls._reference_count == 0:
|
||||
try:
|
||||
krisp_audio.globalInit("", cls._log_callback, krisp_audio.LogLevel.Off)
|
||||
|
||||
cls._initialized = True
|
||||
|
||||
SDK_VERSION = krisp_audio.getVersion()
|
||||
logger.debug(
|
||||
f"Krisp Audio Python SDK initialized - Version: "
|
||||
f"{SDK_VERSION.major}.{SDK_VERSION.minor}.{SDK_VERSION.patch}"
|
||||
)
|
||||
|
||||
# Register cleanup on program exit (failsafe)
|
||||
atexit.register(cls._force_cleanup)
|
||||
|
||||
except Exception as e:
|
||||
cls._initialized = False
|
||||
logger.error(f"Krisp SDK initialization failed: {e}")
|
||||
raise
|
||||
|
||||
cls._reference_count += 1
|
||||
logger.debug(f"Krisp SDK reference count: {cls._reference_count}")
|
||||
|
||||
@classmethod
|
||||
def release(cls):
|
||||
"""Release a reference to the SDK (destroys if last reference).
|
||||
|
||||
Call this when destroying a filter instance.
|
||||
"""
|
||||
with cls._lock:
|
||||
if cls._reference_count > 0:
|
||||
cls._reference_count -= 1
|
||||
logger.debug(f"Krisp SDK reference count: {cls._reference_count}")
|
||||
|
||||
# Destroy SDK when last reference is released
|
||||
if cls._reference_count == 0 and cls._initialized:
|
||||
try:
|
||||
krisp_audio.globalDestroy()
|
||||
cls._initialized = False
|
||||
logger.debug("Krisp Audio SDK destroyed (all references released)")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during Krisp SDK cleanup: {e}")
|
||||
cls._initialized = False
|
||||
|
||||
@classmethod
|
||||
def get_reference_count(cls) -> int:
|
||||
"""Get the current reference count.
|
||||
|
||||
Returns:
|
||||
Current number of active references to the SDK.
|
||||
"""
|
||||
with cls._lock:
|
||||
return cls._reference_count
|
||||
|
||||
@classmethod
|
||||
def is_initialized(cls) -> bool:
|
||||
"""Check if the SDK is currently initialized.
|
||||
|
||||
Returns:
|
||||
True if SDK is initialized, False otherwise.
|
||||
"""
|
||||
with cls._lock:
|
||||
return cls._initialized
|
||||
|
||||
@classmethod
|
||||
def _force_cleanup(cls):
|
||||
"""Force cleanup on program exit (failsafe)."""
|
||||
with cls._lock:
|
||||
if cls._initialized:
|
||||
try:
|
||||
logger.warning(
|
||||
f"Force cleaning up Krisp SDK at exit (ref count: {cls._reference_count})"
|
||||
)
|
||||
krisp_audio.globalDestroy()
|
||||
cls._initialized = False
|
||||
except Exception as e:
|
||||
logger.error(f"Error during forced Krisp SDK cleanup: {e}")
|
||||
353
src/pipecat/audio/turn/krisp_viva_turn.py
Normal file
353
src/pipecat/audio/turn/krisp_viva_turn.py
Normal file
@@ -0,0 +1,353 @@
|
||||
#
|
||||
# Copyright (c) 2024–2025, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""Krisp turn analyzer for end-of-turn detection using Krisp VIVA SDK.
|
||||
|
||||
This module provides a turn analyzer implementation using Krisp's turn detection
|
||||
(Tt) API to determine when a user has finished speaking in a conversation.
|
||||
|
||||
Note: This analyzer uses a different model than KrispVivaFilter. The model path
|
||||
can be specified via the KRISP_VIVA_TURN_MODEL_PATH environment variable or
|
||||
passed directly to the constructor.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
from loguru import logger
|
||||
|
||||
from pipecat.audio.krisp_instance import (
|
||||
KrispVivaSDKManager,
|
||||
int_to_krisp_frame_duration,
|
||||
int_to_krisp_sample_rate,
|
||||
)
|
||||
from pipecat.audio.turn.base_turn_analyzer import BaseTurnAnalyzer, BaseTurnParams, EndOfTurnState
|
||||
from pipecat.metrics.metrics import MetricsData
|
||||
|
||||
try:
|
||||
import krisp_audio
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
logger.error("In order to use KrispVivaTurn, you need to install krisp_audio.")
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
class KrispTurnParams(BaseTurnParams):
|
||||
"""Configuration parameters for Krisp turn analysis.
|
||||
|
||||
Parameters:
|
||||
threshold: Probability threshold for turn completion (0.0 to 1.0).
|
||||
Higher values require more confidence before marking turn as complete.
|
||||
frame_duration_ms: Frame duration in milliseconds for turn detection.
|
||||
Supported values: 10, 15, 20, 30, 32.
|
||||
"""
|
||||
|
||||
threshold: float = 0.5
|
||||
frame_duration_ms: int = 20
|
||||
|
||||
|
||||
class KrispVivaTurn(BaseTurnAnalyzer):
|
||||
"""Turn analyzer using Krisp VIVA SDK for end-of-turn detection.
|
||||
|
||||
Uses Krisp's turn detection (Tt) API to determine when a user has finished
|
||||
speaking. This analyzer requires a valid Krisp model file to operate.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model_path: Optional[str] = None,
|
||||
sample_rate: Optional[int] = None,
|
||||
params: Optional[KrispTurnParams] = None,
|
||||
) -> None:
|
||||
"""Initialize the Krisp turn analyzer.
|
||||
|
||||
Args:
|
||||
model_path: Path to the Krisp turn detection model file (.kef extension).
|
||||
If None, uses KRISP_VIVA_TURN_MODEL_PATH environment variable.
|
||||
sample_rate: Optional initial sample rate for audio processing.
|
||||
If provided, this will be used as the fixed sample rate.
|
||||
params: Configuration parameters for turn analysis behavior.
|
||||
|
||||
Raises:
|
||||
ValueError: If model_path is not provided and KRISP_VIVA_TURN_MODEL_PATH is not set.
|
||||
Exception: If model file doesn't have .kef extension.
|
||||
FileNotFoundError: If model file doesn't exist.
|
||||
RuntimeError: If Krisp SDK initialization fails.
|
||||
"""
|
||||
super().__init__(sample_rate=sample_rate)
|
||||
|
||||
# Acquire SDK reference (will initialize on first call)
|
||||
try:
|
||||
KrispVivaSDKManager.acquire()
|
||||
self._sdk_acquired = True
|
||||
except Exception as e:
|
||||
self._sdk_acquired = False
|
||||
raise RuntimeError(f"Failed to initialize Krisp SDK: {e}")
|
||||
|
||||
try:
|
||||
# Set model path, checking environment if not specified
|
||||
self._model_path = model_path or os.getenv("KRISP_VIVA_TURN_MODEL_PATH")
|
||||
if not self._model_path:
|
||||
logger.error(
|
||||
"Model path is not provided and KRISP_VIVA_TURN_MODEL_PATH is not set."
|
||||
)
|
||||
raise ValueError("Model path for KrispVivaTurn must be provided.")
|
||||
|
||||
if not self._model_path.endswith(".kef"):
|
||||
raise Exception("Model is expected with .kef extension")
|
||||
|
||||
if not os.path.isfile(self._model_path):
|
||||
raise FileNotFoundError(f"Model file not found: {self._model_path}")
|
||||
|
||||
self._params = params or KrispTurnParams()
|
||||
self._tt_session = None
|
||||
self._preload_tt_session = None
|
||||
self._samples_per_frame = None
|
||||
self._audio_buffer = bytearray()
|
||||
|
||||
# State tracking
|
||||
self._speech_triggered = False
|
||||
self._last_probability = None
|
||||
self._frame_probabilities = []
|
||||
self._last_state = EndOfTurnState.INCOMPLETE
|
||||
|
||||
# Create session with provided sample rate or default to 16000 Hz
|
||||
# This preloads the model to improve latency when set_sample_rate is called later
|
||||
preload_sample_rate = sample_rate if sample_rate else 16000
|
||||
try:
|
||||
self._preload_tt_session = self._create_tt_session(preload_sample_rate)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create turn detection session: {e}", exc_info=True)
|
||||
self._preload_tt_session = None
|
||||
raise RuntimeError(f"Failed to create turn detection session: {e}") from e
|
||||
|
||||
except Exception:
|
||||
# If initialization fails, release the SDK reference
|
||||
if self._sdk_acquired:
|
||||
KrispVivaSDKManager.release()
|
||||
self._sdk_acquired = False
|
||||
raise
|
||||
|
||||
def __del__(self):
|
||||
"""Release SDK reference when analyzer is destroyed."""
|
||||
if self._sdk_acquired:
|
||||
try:
|
||||
# Clean up session first
|
||||
if hasattr(self, "_tt_session") and self._tt_session is not None:
|
||||
self._tt_session = None
|
||||
if hasattr(self, "_preload_tt_session") and self._preload_tt_session is not None:
|
||||
self._preload_tt_session = None
|
||||
|
||||
KrispVivaSDKManager.release()
|
||||
self._sdk_acquired = False
|
||||
except Exception as e:
|
||||
logger.error(f"Error in __del__: {e}", exc_info=True)
|
||||
|
||||
def _create_tt_session(self, sample_rate: int):
|
||||
"""Create a turn detection session with the specified sample rate.
|
||||
|
||||
Args:
|
||||
sample_rate: Sample rate for the session
|
||||
|
||||
Returns:
|
||||
krisp_audio.TtFloat instance
|
||||
|
||||
Raises:
|
||||
ValueError: If sample rate or frame duration is not supported
|
||||
RuntimeError: If session creation fails
|
||||
"""
|
||||
try:
|
||||
model_info = krisp_audio.ModelInfo()
|
||||
model_info.path = self._model_path
|
||||
|
||||
tt_cfg = krisp_audio.TtSessionConfig()
|
||||
tt_cfg.inputSampleRate = int_to_krisp_sample_rate(sample_rate)
|
||||
tt_cfg.inputFrameDuration = int_to_krisp_frame_duration(self._params.frame_duration_ms)
|
||||
tt_cfg.modelInfo = model_info
|
||||
|
||||
# Calculate samples per frame for this sample rate
|
||||
self._samples_per_frame = int((sample_rate * self._params.frame_duration_ms) / 1000)
|
||||
|
||||
tt_instance = krisp_audio.TtFloat.create(tt_cfg)
|
||||
return tt_instance
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create Krisp turn detection session: {e}", exc_info=True)
|
||||
raise RuntimeError(f"Failed to create Krisp turn detection session: {e}") from e
|
||||
|
||||
def set_sample_rate(self, sample_rate: int):
|
||||
"""Set the sample rate and create/update the turn detection session.
|
||||
|
||||
Args:
|
||||
sample_rate: The sample rate to set.
|
||||
"""
|
||||
if self._sample_rate == sample_rate:
|
||||
return
|
||||
|
||||
super().set_sample_rate(sample_rate)
|
||||
# Create session when sample rate is set
|
||||
try:
|
||||
self._tt_session = self._create_tt_session(self._sample_rate)
|
||||
# Clear buffer when sample rate changes
|
||||
self._audio_buffer.clear()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create turn detection session: {e}", exc_info=True)
|
||||
self._tt_session = None
|
||||
|
||||
@property
|
||||
def frame_probabilities(self) -> list:
|
||||
"""Get all probabilities from the last append_audio call.
|
||||
|
||||
Returns:
|
||||
List of probability values for each frame processed in the last append_audio call.
|
||||
"""
|
||||
return self._frame_probabilities
|
||||
|
||||
@property
|
||||
def last_probability(self) -> Optional[float]:
|
||||
"""Get the last turn probability value computed.
|
||||
|
||||
Returns:
|
||||
Last probability value, or None if no frames have been processed yet.
|
||||
"""
|
||||
return self._last_probability
|
||||
|
||||
@property
|
||||
def speech_triggered(self) -> bool:
|
||||
"""Check if speech has been detected and triggered analysis.
|
||||
|
||||
Returns:
|
||||
True if speech has been detected and turn analysis is active.
|
||||
"""
|
||||
return self._speech_triggered
|
||||
|
||||
@property
|
||||
def params(self) -> KrispTurnParams:
|
||||
"""Get the current turn analyzer parameters.
|
||||
|
||||
Returns:
|
||||
Current turn analyzer configuration parameters.
|
||||
"""
|
||||
return self._params
|
||||
|
||||
def append_audio(self, buffer: bytes, is_speech: bool) -> EndOfTurnState:
|
||||
"""Append audio data for turn analysis.
|
||||
|
||||
Args:
|
||||
buffer: Raw audio data bytes to append for analysis.
|
||||
is_speech: Whether the audio buffer contains detected speech.
|
||||
|
||||
Returns:
|
||||
Current end-of-turn state after processing the audio.
|
||||
"""
|
||||
if self._tt_session is None:
|
||||
logger.warning("Turn detection session not initialized, returning INCOMPLETE")
|
||||
self._last_state = EndOfTurnState.INCOMPLETE
|
||||
return EndOfTurnState.INCOMPLETE
|
||||
|
||||
if self._samples_per_frame is None:
|
||||
logger.warning("Samples per frame not initialized, returning INCOMPLETE")
|
||||
self._last_state = EndOfTurnState.INCOMPLETE
|
||||
return EndOfTurnState.INCOMPLETE
|
||||
|
||||
try:
|
||||
# Add incoming audio to our buffer
|
||||
self._audio_buffer.extend(buffer)
|
||||
|
||||
# Clear frame probabilities from previous call
|
||||
self._frame_probabilities = []
|
||||
|
||||
total_samples = len(self._audio_buffer) // 2 # 2 bytes per int16 sample
|
||||
num_complete_frames = total_samples // self._samples_per_frame
|
||||
|
||||
if num_complete_frames == 0:
|
||||
# Not enough samples for a complete frame yet, return current state
|
||||
self._last_state = EndOfTurnState.INCOMPLETE
|
||||
return EndOfTurnState.INCOMPLETE
|
||||
|
||||
complete_samples_count = num_complete_frames * self._samples_per_frame
|
||||
bytes_to_process = complete_samples_count * 2 # 2 bytes per sample
|
||||
|
||||
audio_to_process = bytes(self._audio_buffer[:bytes_to_process])
|
||||
|
||||
self._audio_buffer = self._audio_buffer[bytes_to_process:]
|
||||
|
||||
audio_int16 = np.frombuffer(audio_to_process, dtype=np.int16)
|
||||
audio_float32 = audio_int16.astype(np.float32) / 32768.0
|
||||
|
||||
frames = audio_float32.reshape(-1, self._samples_per_frame)
|
||||
|
||||
state = EndOfTurnState.INCOMPLETE
|
||||
|
||||
# Process each complete frame
|
||||
for frame in frames:
|
||||
if is_speech:
|
||||
# Track speech start time
|
||||
if not self._speech_triggered:
|
||||
logger.trace("Speech detected, turn analysis started")
|
||||
self._speech_triggered = True
|
||||
# Note: We don't immediately mark as complete on silence detection.
|
||||
# Instead, we wait for the model's probability check below to confirm
|
||||
# end-of-turn based on the threshold.
|
||||
|
||||
prob = self._tt_session.process(frame.tolist())
|
||||
|
||||
# Negative values indicate the model is not ready yet (working with 100ms data)
|
||||
# Skip processing until we get positive probabilities
|
||||
if prob < 0:
|
||||
continue
|
||||
|
||||
# Store the probability for external access
|
||||
self._last_probability = prob
|
||||
self._frame_probabilities.append(prob)
|
||||
|
||||
# Check if turn is complete based on probability threshold
|
||||
# Only mark as complete if we've detected speech and the model
|
||||
# confirms with sufficient confidence
|
||||
if self._speech_triggered and prob >= self._params.threshold:
|
||||
state = EndOfTurnState.COMPLETE
|
||||
self._clear(state)
|
||||
break
|
||||
|
||||
# Store the last state for analyze_end_of_turn()
|
||||
self._last_state = state
|
||||
return state
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error during Krisp turn detection: {e}", exc_info=True)
|
||||
error_state = EndOfTurnState.INCOMPLETE
|
||||
self._last_state = error_state
|
||||
return error_state
|
||||
|
||||
async def analyze_end_of_turn(self) -> Tuple[EndOfTurnState, Optional[MetricsData]]:
|
||||
"""Analyze the current audio state to determine if turn has ended.
|
||||
|
||||
Returns:
|
||||
Tuple containing the end-of-turn state and optional metrics data.
|
||||
Returns the last state determined by append_audio().
|
||||
"""
|
||||
# For real-time processing, the state is determined in append_audio
|
||||
# Return the last state that was computed
|
||||
return self._last_state, None
|
||||
|
||||
def clear(self):
|
||||
"""Reset the turn analyzer to its initial state."""
|
||||
self._clear(EndOfTurnState.COMPLETE)
|
||||
|
||||
def _clear(self, turn_state: EndOfTurnState):
|
||||
"""Clear internal state based on turn completion status.
|
||||
|
||||
Args:
|
||||
turn_state: The end-of-turn state to use for clearing.
|
||||
"""
|
||||
# If the state is still incomplete, keep the _speech_triggered as True
|
||||
self._speech_triggered = turn_state == EndOfTurnState.INCOMPLETE
|
||||
# Clear audio buffer on turn completion
|
||||
if turn_state == EndOfTurnState.COMPLETE:
|
||||
self._audio_buffer.clear()
|
||||
# Reset last state when clearing
|
||||
self._last_state = EndOfTurnState.INCOMPLETE
|
||||
196
tests/test_krisp_sdk_manager.py
Normal file
196
tests/test_krisp_sdk_manager.py
Normal file
@@ -0,0 +1,196 @@
|
||||
#
|
||||
# Copyright (c) 2024-2025 Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""Unit tests for Krisp SDK Manager (singleton with reference counting)."""
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
# Mock package version check before importing pipecat
|
||||
# This allows tests to run in development mode without installed package
|
||||
_version_patcher = patch("importlib.metadata.version", return_value="0.0.0-dev")
|
||||
_version_patcher.start()
|
||||
|
||||
# Mock krisp_audio module BEFORE any pipecat imports
|
||||
# This allows tests to run without krisp_audio installed
|
||||
mock_krisp_audio = MagicMock()
|
||||
mock_krisp_audio.SamplingRate.Sr8000Hz = 8000
|
||||
mock_krisp_audio.SamplingRate.Sr16000Hz = 16000
|
||||
mock_krisp_audio.SamplingRate.Sr24000Hz = 24000
|
||||
mock_krisp_audio.SamplingRate.Sr32000Hz = 32000
|
||||
mock_krisp_audio.SamplingRate.Sr44100Hz = 44100
|
||||
mock_krisp_audio.SamplingRate.Sr48000Hz = 48000
|
||||
mock_krisp_audio.FrameDuration.Fd10ms = "10ms"
|
||||
mock_krisp_audio.FrameDuration.Fd15ms = "15ms"
|
||||
mock_krisp_audio.FrameDuration.Fd20ms = "20ms"
|
||||
mock_krisp_audio.FrameDuration.Fd30ms = "30ms"
|
||||
mock_krisp_audio.FrameDuration.Fd32ms = "32ms"
|
||||
mock_krisp_audio.LogLevel.Off = 0
|
||||
|
||||
# Mock getVersion to return a version object
|
||||
mock_version = MagicMock()
|
||||
mock_version.major = 1
|
||||
mock_version.minor = 0
|
||||
mock_version.patch = 0
|
||||
mock_krisp_audio.getVersion.return_value = mock_version
|
||||
|
||||
# Install the mock in sys.modules before importing
|
||||
sys.modules["krisp_audio"] = mock_krisp_audio
|
||||
|
||||
# Mock pipecat_ai_krisp package
|
||||
mock_pipecat_krisp = MagicMock()
|
||||
sys.modules["pipecat_ai_krisp"] = mock_pipecat_krisp
|
||||
sys.modules["pipecat_ai_krisp.audio"] = MagicMock()
|
||||
sys.modules["pipecat_ai_krisp.audio.krisp_processor"] = MagicMock()
|
||||
|
||||
# Now we can safely import
|
||||
from pipecat.audio.krisp_instance import (
|
||||
KRISP_SAMPLE_RATES,
|
||||
KrispVivaSDKManager,
|
||||
int_to_krisp_sample_rate,
|
||||
)
|
||||
|
||||
|
||||
class TestKrispVivaSDKManager:
|
||||
"""Tests for KrispVivaSDKManager singleton."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Reset mocks and SDK state before each test."""
|
||||
mock_krisp_audio.reset_mock()
|
||||
mock_krisp_audio.getVersion.return_value = mock_version
|
||||
|
||||
# Reset the SDK manager state for clean tests
|
||||
# We access internal state to ensure tests are isolated
|
||||
with KrispVivaSDKManager._lock:
|
||||
# Release any leftover references from previous tests
|
||||
while KrispVivaSDKManager._reference_count > 0:
|
||||
KrispVivaSDKManager._reference_count -= 1
|
||||
KrispVivaSDKManager._initialized = False
|
||||
|
||||
def test_reference_counting(self):
|
||||
"""Test that SDK manager properly tracks references."""
|
||||
# Initial state
|
||||
initial_count = KrispVivaSDKManager.get_reference_count()
|
||||
assert initial_count == 0
|
||||
|
||||
# Acquire first reference
|
||||
KrispVivaSDKManager.acquire()
|
||||
assert KrispVivaSDKManager.get_reference_count() == initial_count + 1
|
||||
assert KrispVivaSDKManager.is_initialized()
|
||||
|
||||
# Verify globalInit was called
|
||||
mock_krisp_audio.globalInit.assert_called_once()
|
||||
|
||||
# Acquire second reference
|
||||
KrispVivaSDKManager.acquire()
|
||||
assert KrispVivaSDKManager.get_reference_count() == initial_count + 2
|
||||
assert KrispVivaSDKManager.is_initialized()
|
||||
|
||||
# globalInit should NOT be called again
|
||||
assert mock_krisp_audio.globalInit.call_count == 1
|
||||
|
||||
# Release first reference
|
||||
KrispVivaSDKManager.release()
|
||||
assert KrispVivaSDKManager.get_reference_count() == initial_count + 1
|
||||
assert KrispVivaSDKManager.is_initialized()
|
||||
|
||||
# globalDestroy should NOT be called yet
|
||||
mock_krisp_audio.globalDestroy.assert_not_called()
|
||||
|
||||
# Release second reference
|
||||
KrispVivaSDKManager.release()
|
||||
assert KrispVivaSDKManager.get_reference_count() == initial_count
|
||||
|
||||
# globalDestroy should be called now
|
||||
mock_krisp_audio.globalDestroy.assert_called_once()
|
||||
|
||||
def test_multiple_acquire_release_cycles(self):
|
||||
"""Test multiple acquire/release cycles."""
|
||||
initial_count = KrispVivaSDKManager.get_reference_count()
|
||||
|
||||
for i in range(3):
|
||||
KrispVivaSDKManager.acquire()
|
||||
assert KrispVivaSDKManager.get_reference_count() > initial_count
|
||||
assert KrispVivaSDKManager.is_initialized()
|
||||
KrispVivaSDKManager.release()
|
||||
assert KrispVivaSDKManager.get_reference_count() == initial_count
|
||||
|
||||
# Verify globalInit/globalDestroy were called for each cycle
|
||||
assert mock_krisp_audio.globalInit.call_count == 3
|
||||
assert mock_krisp_audio.globalDestroy.call_count == 3
|
||||
|
||||
def test_sdk_initialization_failure(self):
|
||||
"""Test that SDK initialization failures are handled properly."""
|
||||
mock_krisp_audio.globalInit.side_effect = Exception("SDK init failed")
|
||||
|
||||
with pytest.raises(Exception, match="SDK init failed"):
|
||||
KrispVivaSDKManager.acquire()
|
||||
|
||||
# Verify SDK is not initialized after failure
|
||||
assert not KrispVivaSDKManager.is_initialized()
|
||||
assert KrispVivaSDKManager.get_reference_count() == 0
|
||||
|
||||
# Reset the side effect for other tests
|
||||
mock_krisp_audio.globalInit.side_effect = None
|
||||
|
||||
def test_release_without_acquire(self):
|
||||
"""Test that release without acquire is safe."""
|
||||
initial_count = KrispVivaSDKManager.get_reference_count()
|
||||
|
||||
# Release without acquire should be safe (no-op)
|
||||
KrispVivaSDKManager.release()
|
||||
|
||||
assert KrispVivaSDKManager.get_reference_count() == initial_count
|
||||
mock_krisp_audio.globalDestroy.assert_not_called()
|
||||
|
||||
def test_is_initialized_state(self):
|
||||
"""Test is_initialized state transitions."""
|
||||
# Initially not initialized
|
||||
assert not KrispVivaSDKManager.is_initialized()
|
||||
|
||||
# After acquire, should be initialized
|
||||
KrispVivaSDKManager.acquire()
|
||||
assert KrispVivaSDKManager.is_initialized()
|
||||
|
||||
# After release, should not be initialized
|
||||
KrispVivaSDKManager.release()
|
||||
assert not KrispVivaSDKManager.is_initialized()
|
||||
|
||||
|
||||
class TestSampleRateConversion:
|
||||
"""Tests for sample rate conversion utilities."""
|
||||
|
||||
def test_supported_sample_rates(self):
|
||||
"""Test conversion of all supported sample rates."""
|
||||
for rate_hz, krisp_enum in KRISP_SAMPLE_RATES.items():
|
||||
result = int_to_krisp_sample_rate(rate_hz)
|
||||
assert result == krisp_enum
|
||||
|
||||
def test_unsupported_sample_rate(self):
|
||||
"""Test that unsupported rates raise ValueError."""
|
||||
with pytest.raises(ValueError, match="Unsupported sample rate"):
|
||||
int_to_krisp_sample_rate(22050) # Not supported
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported sample rate"):
|
||||
int_to_krisp_sample_rate(96000) # Not supported
|
||||
|
||||
def test_sample_rate_error_message(self):
|
||||
"""Test that error message includes helpful information."""
|
||||
try:
|
||||
int_to_krisp_sample_rate(11025)
|
||||
except ValueError as e:
|
||||
assert "11025" in str(e)
|
||||
assert "Supported rates" in str(e)
|
||||
# Should list at least some supported rates
|
||||
assert "16000" in str(e)
|
||||
|
||||
def test_all_krisp_sample_rates_defined(self):
|
||||
"""Test that all expected sample rates are in KRISP_SAMPLE_RATES."""
|
||||
expected_rates = [8000, 16000, 24000, 32000, 44100, 48000]
|
||||
for rate in expected_rates:
|
||||
assert rate in KRISP_SAMPLE_RATES
|
||||
777
tests/test_krisp_viva_filter.py
Normal file
777
tests/test_krisp_viva_filter.py
Normal file
@@ -0,0 +1,777 @@
|
||||
#
|
||||
# Copyright (c) 2024-2025 Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Mock package version check before importing pipecat
|
||||
# This allows tests to run in development mode without installed package
|
||||
_version_patcher = patch("importlib.metadata.version", return_value="0.0.0-dev")
|
||||
_version_patcher.start()
|
||||
|
||||
# Mock krisp_audio module BEFORE any pipecat imports
|
||||
# This allows tests to run without krisp_audio installed
|
||||
mock_krisp_audio = MagicMock()
|
||||
mock_krisp_audio.SamplingRate.Sr8000Hz = 8000
|
||||
mock_krisp_audio.SamplingRate.Sr16000Hz = 16000
|
||||
mock_krisp_audio.SamplingRate.Sr24000Hz = 24000
|
||||
mock_krisp_audio.SamplingRate.Sr32000Hz = 32000
|
||||
mock_krisp_audio.SamplingRate.Sr44100Hz = 44100
|
||||
mock_krisp_audio.SamplingRate.Sr48000Hz = 48000
|
||||
mock_krisp_audio.FrameDuration.Fd10ms = "10ms"
|
||||
mock_krisp_audio.FrameDuration.Fd15ms = "15ms"
|
||||
mock_krisp_audio.FrameDuration.Fd20ms = "20ms"
|
||||
mock_krisp_audio.FrameDuration.Fd30ms = "30ms"
|
||||
mock_krisp_audio.FrameDuration.Fd32ms = "32ms"
|
||||
|
||||
# Install the mock in sys.modules before importing
|
||||
sys.modules["krisp_audio"] = mock_krisp_audio
|
||||
|
||||
# Mock pipecat_ai_krisp package
|
||||
mock_pipecat_krisp = MagicMock()
|
||||
sys.modules["pipecat_ai_krisp"] = mock_pipecat_krisp
|
||||
sys.modules["pipecat_ai_krisp.audio"] = MagicMock()
|
||||
sys.modules["pipecat_ai_krisp.audio.krisp_processor"] = MagicMock()
|
||||
|
||||
# Now we can safely import
|
||||
from pipecat.audio.filters.krisp_viva_filter import KrispVivaFilter
|
||||
from pipecat.frames.frames import FilterEnableFrame
|
||||
|
||||
|
||||
class TestKrispVivaFilter(unittest.IsolatedAsyncioTestCase):
|
||||
"""Test suite for KrispVivaFilter audio filter."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures before each test method."""
|
||||
# Create a temporary .kef model file for testing
|
||||
self.temp_model_file = tempfile.NamedTemporaryFile(suffix=".kef", delete=False)
|
||||
self.temp_model_file.write(b"dummy model data")
|
||||
self.temp_model_file.close()
|
||||
self.model_path = self.temp_model_file.name
|
||||
|
||||
# Use the global mock_krisp_audio that was set up before imports
|
||||
self.mock_krisp_audio = mock_krisp_audio
|
||||
|
||||
# Reset all mocks to clear call counts from previous tests
|
||||
self.mock_krisp_audio.reset_mock()
|
||||
self.mock_krisp_audio.ModelInfo.reset_mock()
|
||||
self.mock_krisp_audio.NcSessionConfig.reset_mock()
|
||||
self.mock_krisp_audio.NcInt16.reset_mock()
|
||||
|
||||
# Mock ModelInfo
|
||||
self.mock_model_info = MagicMock()
|
||||
self.mock_krisp_audio.ModelInfo.return_value = self.mock_model_info
|
||||
|
||||
# Mock NcSessionConfig
|
||||
self.mock_nc_cfg = MagicMock()
|
||||
self.mock_krisp_audio.NcSessionConfig.return_value = self.mock_nc_cfg
|
||||
|
||||
# Mock session
|
||||
self.mock_session = MagicMock()
|
||||
self.mock_session.process = MagicMock(side_effect=lambda x, level: x)
|
||||
self.mock_krisp_audio.NcInt16.create.return_value = self.mock_session
|
||||
|
||||
# Patch krisp_audio in the module
|
||||
self.sample_rates_patch = patch(
|
||||
"pipecat.audio.filters.krisp_viva_filter.krisp_audio", self.mock_krisp_audio
|
||||
)
|
||||
self.sample_rates_patch.start()
|
||||
|
||||
# Patch KrispVivaSDKManager
|
||||
self.sdk_manager_patcher = patch(
|
||||
"pipecat.audio.filters.krisp_viva_filter.KrispVivaSDKManager"
|
||||
)
|
||||
self.mock_sdk_manager = self.sdk_manager_patcher.start()
|
||||
self.mock_sdk_manager.acquire = MagicMock()
|
||||
self.mock_sdk_manager.release = MagicMock()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up test fixtures after each test method."""
|
||||
# Stop all patchers
|
||||
self.sample_rates_patch.stop()
|
||||
self.sdk_manager_patcher.stop()
|
||||
|
||||
# Remove temporary model file
|
||||
if os.path.exists(self.model_path):
|
||||
os.unlink(self.model_path)
|
||||
|
||||
async def test_initialization_with_model_path(self):
|
||||
"""Test filter initialization with explicit model path."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# Verify SDK was NOT acquired during initialization (happens in start())
|
||||
self.mock_sdk_manager.acquire.assert_not_called()
|
||||
|
||||
# Verify filter attributes
|
||||
self.assertEqual(filter_instance._model_path, self.model_path)
|
||||
self.assertTrue(filter_instance._filtering) # Filtering starts enabled
|
||||
self.assertEqual(filter_instance._noise_suppression_level, 100)
|
||||
self.assertIsNotNone(filter_instance._audio_buffer)
|
||||
|
||||
async def test_initialization_with_env_variable(self):
|
||||
"""Test filter initialization using KRISP_VIVA_FILTER_MODEL_PATH environment variable."""
|
||||
with patch.dict(os.environ, {"KRISP_VIVA_FILTER_MODEL_PATH": self.model_path}):
|
||||
filter_instance = KrispVivaFilter()
|
||||
|
||||
# Verify SDK was NOT acquired during initialization (happens in start())
|
||||
self.mock_sdk_manager.acquire.assert_not_called()
|
||||
self.assertEqual(filter_instance._model_path, self.model_path)
|
||||
|
||||
async def test_initialization_without_model_path(self):
|
||||
"""Test filter initialization fails without model path."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with self.assertRaises(ValueError) as context:
|
||||
KrispVivaFilter()
|
||||
|
||||
self.assertIn("Model path", str(context.exception))
|
||||
# SDK acquire not called during initialization (happens in start())
|
||||
# But release() is called in exception handler even though acquire() wasn't called
|
||||
self.mock_sdk_manager.acquire.assert_not_called()
|
||||
self.mock_sdk_manager.release.assert_called_once()
|
||||
|
||||
async def test_initialization_with_invalid_extension(self):
|
||||
"""Test filter initialization fails with non-.kef file."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as tmp:
|
||||
tmp.write(b"dummy")
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
with self.assertRaises(Exception) as context:
|
||||
KrispVivaFilter(model_path=tmp_path)
|
||||
|
||||
self.assertIn(".kef extension", str(context.exception))
|
||||
# SDK acquire not called during initialization (happens in start())
|
||||
# But release() is called in exception handler even though acquire() wasn't called
|
||||
self.mock_sdk_manager.acquire.assert_not_called()
|
||||
self.mock_sdk_manager.release.assert_called_once()
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
async def test_initialization_with_nonexistent_file(self):
|
||||
"""Test filter initialization fails with non-existent model file."""
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
KrispVivaFilter(model_path="/nonexistent/path/model.kef")
|
||||
|
||||
# SDK acquire not called during initialization (happens in start())
|
||||
# But release() is called in exception handler even though acquire() wasn't called
|
||||
self.mock_sdk_manager.acquire.assert_not_called()
|
||||
self.mock_sdk_manager.release.assert_called_once()
|
||||
|
||||
async def test_initialization_with_custom_noise_level(self):
|
||||
"""Test filter initialization with custom noise suppression level."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path, noise_suppression_level=50)
|
||||
|
||||
self.assertEqual(filter_instance._noise_suppression_level, 50)
|
||||
|
||||
async def test_initialization_with_default_noise_level(self):
|
||||
"""Test filter initialization with default noise suppression level."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
self.assertEqual(filter_instance._noise_suppression_level, 100)
|
||||
|
||||
async def test_start_with_supported_sample_rate(self):
|
||||
"""Test starting filter with a supported sample rate."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Verify SDK was acquired during start()
|
||||
self.mock_sdk_manager.acquire.assert_called_once()
|
||||
|
||||
# Verify session was created
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 16000)
|
||||
self.assertEqual(filter_instance._samples_per_frame, 160) # 16000 * 10ms / 1000
|
||||
|
||||
# Verify NcSessionConfig was created and configured
|
||||
# Note: Called once in start() (no preload session anymore)
|
||||
self.assertEqual(self.mock_krisp_audio.NcSessionConfig.call_count, 1)
|
||||
# Verify frame duration was set (hardcoded to 10ms in filter)
|
||||
self.assertEqual(self.mock_nc_cfg.inputFrameDuration, "10ms")
|
||||
# inputSampleRate and outputSampleRate are now set to the enum value
|
||||
from pipecat.audio.krisp_instance import int_to_krisp_sample_rate
|
||||
|
||||
expected_sample_rate = int_to_krisp_sample_rate(16000)
|
||||
self.assertEqual(self.mock_nc_cfg.inputSampleRate, expected_sample_rate)
|
||||
self.assertEqual(self.mock_nc_cfg.outputSampleRate, expected_sample_rate)
|
||||
|
||||
async def test_start_with_unsupported_sample_rate(self):
|
||||
"""Test starting filter with an unsupported sample rate raises RuntimeError."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
with self.assertRaises(RuntimeError) as context:
|
||||
await filter_instance.start(12000) # Unsupported sample rate
|
||||
|
||||
self.assertIn("Unsupported sample rate", str(context.exception))
|
||||
|
||||
async def test_start_multiple_sample_rates(self):
|
||||
"""Test starting filter with multiple different sample rates."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
for sample_rate in [8000, 16000, 24000, 32000, 44100, 48000]:
|
||||
# Reset mock config for each iteration to verify frame duration is always set
|
||||
mock_nc_cfg = MagicMock()
|
||||
self.mock_krisp_audio.NcSessionConfig.return_value = mock_nc_cfg
|
||||
|
||||
await filter_instance.start(sample_rate)
|
||||
self.assertEqual(filter_instance._current_sample_rate, sample_rate)
|
||||
expected_samples = int((sample_rate * 10) / 1000)
|
||||
self.assertEqual(filter_instance._samples_per_frame, expected_samples)
|
||||
|
||||
# Verify frame duration is always set to 10ms (hardcoded in filter)
|
||||
self.assertEqual(mock_nc_cfg.inputFrameDuration, "10ms")
|
||||
|
||||
async def test_stop(self):
|
||||
"""Test stopping the filter."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
await filter_instance.stop()
|
||||
|
||||
# Verify session was cleared
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
async def test_process_frame_enable(self):
|
||||
"""Test processing FilterEnableFrame to enable filtering."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
# Disable filtering first
|
||||
filter_instance._filtering = False
|
||||
|
||||
enable_frame = FilterEnableFrame(enable=True)
|
||||
await filter_instance.process_frame(enable_frame)
|
||||
|
||||
self.assertTrue(filter_instance._filtering)
|
||||
|
||||
async def test_process_frame_disable(self):
|
||||
"""Test processing FilterEnableFrame to disable filtering."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
# After start, filtering should be enabled
|
||||
self.assertTrue(filter_instance._filtering)
|
||||
|
||||
disable_frame = FilterEnableFrame(enable=False)
|
||||
await filter_instance.process_frame(disable_frame)
|
||||
|
||||
self.assertFalse(filter_instance._filtering)
|
||||
|
||||
async def test_filter_when_disabled(self):
|
||||
"""Test that filter returns audio unchanged when filtering is disabled."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
# Disable filtering
|
||||
filter_instance._filtering = False
|
||||
|
||||
input_audio = b"\x00\x01\x02\x03\x04\x05"
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
self.assertEqual(output_audio, input_audio)
|
||||
|
||||
async def test_filter_with_complete_frame(self):
|
||||
"""Test filtering audio with exactly one complete frame."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create audio data for exactly one 10ms frame (160 samples = 320 bytes)
|
||||
samples = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
# Verify audio was processed
|
||||
self.assertIsInstance(output_audio, bytes)
|
||||
self.assertEqual(len(output_audio), len(input_audio))
|
||||
|
||||
# Verify session.process was called
|
||||
self.mock_session.process.assert_called()
|
||||
|
||||
async def test_filter_with_multiple_frames(self):
|
||||
"""Test filtering audio with multiple complete frames."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create audio data for 3 complete 10ms frames (480 samples = 960 bytes)
|
||||
samples = np.random.randint(-32768, 32767, size=480, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
# Verify audio was processed
|
||||
self.assertIsInstance(output_audio, bytes)
|
||||
self.assertEqual(len(output_audio), len(input_audio))
|
||||
|
||||
# Verify session.process was called 3 times
|
||||
self.assertEqual(self.mock_session.process.call_count, 3)
|
||||
|
||||
async def test_filter_with_incomplete_frame(self):
|
||||
"""Test filtering audio with incomplete frame data."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create audio data for less than one frame (100 samples = 200 bytes)
|
||||
samples = np.random.randint(-32768, 32767, size=100, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
# Should return empty bytes since no complete frame
|
||||
self.assertEqual(output_audio, b"")
|
||||
|
||||
# Verify session.process was NOT called
|
||||
self.mock_session.process.assert_not_called()
|
||||
|
||||
async def test_filter_with_buffering(self):
|
||||
"""Test that filter properly buffers incomplete frames."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# First call: Send 100 samples (incomplete frame)
|
||||
samples1 = np.random.randint(-32768, 32767, size=100, dtype=np.int16)
|
||||
input_audio1 = samples1.tobytes()
|
||||
output_audio1 = await filter_instance.filter(input_audio1)
|
||||
|
||||
# Should buffer and return empty
|
||||
self.assertEqual(output_audio1, b"")
|
||||
self.assertEqual(len(filter_instance._audio_buffer), 200)
|
||||
|
||||
# Second call: Send 60 more samples (now we have 160 total = 1 complete frame)
|
||||
samples2 = np.random.randint(-32768, 32767, size=60, dtype=np.int16)
|
||||
input_audio2 = samples2.tobytes()
|
||||
output_audio2 = await filter_instance.filter(input_audio2)
|
||||
|
||||
# Should process one frame and return 320 bytes
|
||||
self.assertEqual(len(output_audio2), 320)
|
||||
self.assertEqual(len(filter_instance._audio_buffer), 0)
|
||||
self.mock_session.process.assert_called_once()
|
||||
|
||||
async def test_filter_with_partial_buffering(self):
|
||||
"""Test that filter keeps remainder in buffer after processing."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Send 250 samples (1 complete frame + 90 samples remainder)
|
||||
samples = np.random.randint(-32768, 32767, size=250, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
# Should process one frame (320 bytes)
|
||||
self.assertEqual(len(output_audio), 320)
|
||||
|
||||
# Should keep remainder (90 samples = 180 bytes) in buffer
|
||||
self.assertEqual(len(filter_instance._audio_buffer), 180)
|
||||
|
||||
self.mock_session.process.assert_called_once()
|
||||
|
||||
async def test_filter_error_handling(self):
|
||||
"""Test that filter handles processing errors gracefully."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Make session.process raise an exception
|
||||
self.mock_session.process.side_effect = Exception("Processing error")
|
||||
|
||||
# Create audio data for one complete frame
|
||||
samples = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
# Should return original audio on error
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
self.assertEqual(output_audio, input_audio)
|
||||
|
||||
async def test_filter_different_sample_rates(self):
|
||||
"""Test filtering with different sample rates."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
test_cases = [
|
||||
(8000, 80), # 8kHz: 80 samples per 10ms frame
|
||||
(16000, 160), # 16kHz: 160 samples per 10ms frame
|
||||
(48000, 480), # 48kHz: 480 samples per 10ms frame
|
||||
]
|
||||
|
||||
for sample_rate, expected_samples in test_cases:
|
||||
await filter_instance.start(sample_rate)
|
||||
|
||||
# Create audio data for exactly one frame
|
||||
samples = np.random.randint(-32768, 32767, size=expected_samples, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
# Verify correct processing
|
||||
self.assertEqual(len(output_audio), len(input_audio))
|
||||
|
||||
async def test_stop_releases_sdk(self):
|
||||
"""Test that stop() properly releases SDK reference."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Stop the filter
|
||||
await filter_instance.stop()
|
||||
|
||||
# Verify SDK was released
|
||||
self.mock_sdk_manager.release.assert_called_once()
|
||||
|
||||
async def test_int_to_sample_rate_conversion(self):
|
||||
"""Test sample rate conversion using the shared utility function."""
|
||||
from pipecat.audio.krisp_instance import KRISP_SAMPLE_RATES, int_to_krisp_sample_rate
|
||||
|
||||
# Test valid sample rates - verify they return the correct enum values
|
||||
for rate in [8000, 16000, 24000, 32000, 44100, 48000]:
|
||||
result = int_to_krisp_sample_rate(rate)
|
||||
# Check that result is from the KRISP_SAMPLE_RATES dict
|
||||
self.assertEqual(result, KRISP_SAMPLE_RATES[rate])
|
||||
|
||||
# Test invalid sample rate
|
||||
with self.assertRaises(ValueError) as context:
|
||||
int_to_krisp_sample_rate(12000)
|
||||
|
||||
self.assertIn("Unsupported sample rate", str(context.exception))
|
||||
|
||||
async def test_noise_suppression_level_applied(self):
|
||||
"""Test that noise suppression level is passed to processing."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path, noise_suppression_level=75)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create audio data for one frame
|
||||
samples = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
await filter_instance.filter(input_audio)
|
||||
|
||||
# Verify noise suppression level was passed to process()
|
||||
call_args = self.mock_session.process.call_args
|
||||
self.assertEqual(call_args[0][1], 75) # Second argument should be the level
|
||||
|
||||
async def test_start_acquires_sdk(self):
|
||||
"""Test that start() acquires SDK reference and creates session."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# Verify no session exists before start
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
# Start the filter
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Verify SDK was acquired
|
||||
self.mock_sdk_manager.acquire.assert_called_once()
|
||||
|
||||
# Verify session was created
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
|
||||
# Verify NcSessionConfig was created and frame duration was set
|
||||
self.mock_krisp_audio.NcSessionConfig.assert_called_once()
|
||||
# Verify frame duration was set to 10ms (hardcoded in filter)
|
||||
self.assertEqual(self.mock_nc_cfg.inputFrameDuration, "10ms")
|
||||
|
||||
async def test_filter_preserves_audio_data_integrity(self):
|
||||
"""Test that filter processing preserves data integrity."""
|
||||
# Make mock session return the same data
|
||||
self.mock_session.process.side_effect = lambda x, level: x.copy()
|
||||
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create deterministic audio data
|
||||
samples = np.arange(160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
output_audio = await filter_instance.filter(input_audio)
|
||||
|
||||
# Verify output matches input (since mock returns same data)
|
||||
output_samples = np.frombuffer(output_audio, dtype=np.int16)
|
||||
np.testing.assert_array_equal(output_samples, samples)
|
||||
|
||||
# ==================== Concurrency & Thread Safety Tests ====================
|
||||
|
||||
async def test_concurrent_filter_calls(self):
|
||||
"""Test that concurrent filter calls are handled safely."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create audio data for one frame
|
||||
samples = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
# Create multiple concurrent filter calls
|
||||
async def filter_audio():
|
||||
return await filter_instance.filter(input_audio)
|
||||
|
||||
# Run 10 concurrent filter operations
|
||||
tasks = [filter_audio() for _ in range(10)]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all calls completed successfully
|
||||
self.assertEqual(len(results), 10)
|
||||
for result in results:
|
||||
self.assertIsInstance(result, bytes)
|
||||
self.assertEqual(len(result), len(input_audio))
|
||||
|
||||
# Verify session.process was called for each frame
|
||||
self.assertEqual(self.mock_session.process.call_count, 10)
|
||||
|
||||
async def test_concurrent_enable_disable(self):
|
||||
"""Test rapid enable/disable toggling during filtering."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Create audio data
|
||||
samples = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
# Concurrently toggle enable/disable while filtering
|
||||
async def toggle_and_filter(toggle_enable):
|
||||
enable_frame = FilterEnableFrame(enable=toggle_enable)
|
||||
await filter_instance.process_frame(enable_frame)
|
||||
return await filter_instance.filter(input_audio)
|
||||
|
||||
# Run concurrent enable/disable operations
|
||||
tasks = [
|
||||
toggle_and_filter(True),
|
||||
toggle_and_filter(False),
|
||||
toggle_and_filter(True),
|
||||
toggle_and_filter(False),
|
||||
]
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
# Verify all operations completed
|
||||
self.assertEqual(len(results), 4)
|
||||
|
||||
# Verify final state is consistent (last operation was disable)
|
||||
self.assertFalse(filter_instance._filtering)
|
||||
|
||||
async def test_concurrent_start_stop(self):
|
||||
"""Test concurrent start/stop operations."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
async def start_filter():
|
||||
await filter_instance.start(16000)
|
||||
|
||||
async def stop_filter():
|
||||
await filter_instance.stop()
|
||||
|
||||
# Run start and stop concurrently
|
||||
await asyncio.gather(start_filter(), stop_filter())
|
||||
|
||||
# Verify final state (stop should clear session)
|
||||
# Note: This tests that operations don't crash, final state may vary
|
||||
# depending on which completes first
|
||||
|
||||
async def test_concurrent_filter_with_state_changes(self):
|
||||
"""Test filtering while state changes occur concurrently."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
samples = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
|
||||
async def filter_operation():
|
||||
return await filter_instance.filter(input_audio)
|
||||
|
||||
async def toggle_filtering():
|
||||
# Toggle based on current filtering state
|
||||
is_filtering = filter_instance._filtering
|
||||
enable_frame = FilterEnableFrame(enable=not is_filtering)
|
||||
await filter_instance.process_frame(enable_frame)
|
||||
|
||||
# Run filtering and toggling concurrently
|
||||
filter_tasks = [filter_operation() for _ in range(5)]
|
||||
toggle_tasks = [toggle_filtering() for _ in range(3)]
|
||||
|
||||
results = await asyncio.gather(*filter_tasks + toggle_tasks)
|
||||
|
||||
# Verify all operations completed without errors
|
||||
self.assertEqual(len(results), 8)
|
||||
|
||||
# ==================== State Transition Tests ====================
|
||||
|
||||
async def test_multiple_start_stop_cycles(self):
|
||||
"""Test multiple start/stop cycles."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# First cycle
|
||||
await filter_instance.start(16000)
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 16000)
|
||||
|
||||
await filter_instance.stop()
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
# Second cycle
|
||||
await filter_instance.start(24000)
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 24000)
|
||||
|
||||
await filter_instance.stop()
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
# Third cycle
|
||||
await filter_instance.start(48000)
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 48000)
|
||||
|
||||
await filter_instance.stop()
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
# Verify session was created multiple times
|
||||
self.assertGreaterEqual(self.mock_krisp_audio.NcInt16.create.call_count, 3)
|
||||
|
||||
async def test_sample_rate_change_during_operation(self):
|
||||
"""Test changing sample rate between start/stop cycles."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# Start with 16kHz
|
||||
await filter_instance.start(16000)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 16000)
|
||||
self.assertEqual(filter_instance._samples_per_frame, 160)
|
||||
|
||||
# Process some audio
|
||||
samples_16k = np.random.randint(-32768, 32767, size=160, dtype=np.int16)
|
||||
output_16k = await filter_instance.filter(samples_16k.tobytes())
|
||||
self.assertEqual(len(output_16k), 320) # 160 samples * 2 bytes
|
||||
|
||||
# Stop and change to 48kHz
|
||||
await filter_instance.stop()
|
||||
await filter_instance.start(48000)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 48000)
|
||||
self.assertEqual(filter_instance._samples_per_frame, 480)
|
||||
|
||||
# Process audio at new sample rate
|
||||
samples_48k = np.random.randint(-32768, 32767, size=480, dtype=np.int16)
|
||||
output_48k = await filter_instance.filter(samples_48k.tobytes())
|
||||
self.assertEqual(len(output_48k), 960) # 480 samples * 2 bytes
|
||||
|
||||
await filter_instance.stop()
|
||||
|
||||
async def test_start_after_stop_with_different_sample_rate(self):
|
||||
"""Test starting with different sample rate after stop."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# Start with 8kHz
|
||||
await filter_instance.start(8000)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 8000)
|
||||
await filter_instance.stop()
|
||||
|
||||
# Start with 32kHz
|
||||
await filter_instance.start(32000)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 32000)
|
||||
await filter_instance.stop()
|
||||
|
||||
# Start with 44.1kHz
|
||||
await filter_instance.start(44100)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 44100)
|
||||
await filter_instance.stop()
|
||||
|
||||
async def test_filter_state_persistence_across_start_stop(self):
|
||||
"""Test that filtering state persists across start/stop cycles."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# Filter starts with filtering enabled
|
||||
self.assertTrue(filter_instance._filtering)
|
||||
|
||||
# Start the filter
|
||||
await filter_instance.start(16000)
|
||||
self.assertTrue(filter_instance._filtering)
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
|
||||
# Disable filtering
|
||||
disable_frame = FilterEnableFrame(enable=False)
|
||||
await filter_instance.process_frame(disable_frame)
|
||||
self.assertFalse(filter_instance._filtering)
|
||||
|
||||
# Stop the filter (cleanup)
|
||||
await filter_instance.stop()
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
# Enable filtering again
|
||||
enable_frame = FilterEnableFrame(enable=True)
|
||||
await filter_instance.process_frame(enable_frame)
|
||||
self.assertTrue(filter_instance._filtering)
|
||||
|
||||
# Start the filter again
|
||||
await filter_instance.start(16000)
|
||||
self.assertTrue(filter_instance._filtering)
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
|
||||
async def test_noise_suppression_level_persistence(self):
|
||||
"""Test that noise suppression level persists across start/stop."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path, noise_suppression_level=75)
|
||||
|
||||
self.assertEqual(filter_instance._noise_suppression_level, 75)
|
||||
|
||||
# Start and stop
|
||||
await filter_instance.start(16000)
|
||||
await filter_instance.stop()
|
||||
|
||||
# Verify noise suppression level persisted
|
||||
self.assertEqual(filter_instance._noise_suppression_level, 75)
|
||||
|
||||
async def test_buffer_cleared_on_stop(self):
|
||||
"""Test that audio buffer is cleared when stopping."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
await filter_instance.start(16000)
|
||||
|
||||
# Add incomplete frame to buffer
|
||||
samples = np.random.randint(-32768, 32767, size=100, dtype=np.int16)
|
||||
input_audio = samples.tobytes()
|
||||
await filter_instance.filter(input_audio)
|
||||
|
||||
# Verify buffer has data
|
||||
self.assertGreater(len(filter_instance._audio_buffer), 0)
|
||||
|
||||
# Stop should clear buffer (or at least not cause issues)
|
||||
await filter_instance.stop()
|
||||
|
||||
# Buffer state after stop - verify no errors on next start
|
||||
await filter_instance.start(16000)
|
||||
# Should be able to filter after restart
|
||||
output = await filter_instance.filter(input_audio)
|
||||
self.assertIsInstance(output, bytes)
|
||||
|
||||
async def test_multiple_starts_without_stop(self):
|
||||
"""Test behavior when start is called multiple times without stop."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# First start
|
||||
await filter_instance.start(16000)
|
||||
session1 = filter_instance._session
|
||||
self.assertIsNotNone(session1)
|
||||
|
||||
# Second start without stop (should replace session)
|
||||
await filter_instance.start(24000)
|
||||
session2 = filter_instance._session
|
||||
self.assertIsNotNone(session2)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 24000)
|
||||
|
||||
# Third start
|
||||
await filter_instance.start(48000)
|
||||
session3 = filter_instance._session
|
||||
self.assertIsNotNone(session3)
|
||||
self.assertEqual(filter_instance._current_sample_rate, 48000)
|
||||
|
||||
await filter_instance.stop()
|
||||
|
||||
async def test_stop_without_start(self):
|
||||
"""Test that stop can be called safely without start."""
|
||||
filter_instance = KrispVivaFilter(model_path=self.model_path)
|
||||
|
||||
# Stop without starting should not raise an error
|
||||
await filter_instance.stop()
|
||||
|
||||
# Verify session is None
|
||||
self.assertIsNone(filter_instance._session)
|
||||
|
||||
# Should be able to start after stop without start
|
||||
await filter_instance.start(16000)
|
||||
self.assertIsNotNone(filter_instance._session)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user