Update bot manager and noisegate functions to try and improve the logic

This commit is contained in:
Logan Cusano
2025-07-14 21:04:14 -04:00
parent b5191ef4d0
commit e7ff28da6e
2 changed files with 163 additions and 176 deletions

View File

@@ -1,22 +1,17 @@
import audioop
import math
import time
import pyaudio
import discord
import numpy
from internal.logger import create_logger
voice_connection = None
LOGGER = create_logger(__name__)
# noinspection PyUnresolvedReferences
class AudioStream:
def __init__(self, _channels: int = 2, _sample_rate: int = 48000, _frames_per_buffer: int = 1024,
def __init__(self, _channels: int = 2, _sample_rate: int = 48000, _frames_per_buffer: int = 960,
_input_device_index: int = None, _output_device_index: int = None, _input: bool = True,
_output: bool = True, _init_on_startup: bool = True):
# NOTE: frames_per_buffer changed to 960 to match Discord's 20ms frame size
self.paInstance_kwargs = {
'format': pyaudio.paInt16,
'channels': _channels,
@@ -26,14 +21,14 @@ class AudioStream:
'frames_per_buffer': _frames_per_buffer
}
if _input_device_index:
if _input_device_index is not None:
if _input:
self.paInstance_kwargs['input_device_index'] = _input_device_index
else:
LOGGER.warning("[AudioStream.__init__]:\tInput was not enabled."
" Reinitialize with '_input=True'")
if _output_device_index:
if _output_device_index is not None:
if _output:
self.paInstance_kwargs['output_device_index'] = _output_device_index
else:
@@ -41,28 +36,24 @@ class AudioStream:
" Reinitialize with '_output=True'")
if _init_on_startup:
# Init PyAudio instance
LOGGER.info("Creating PyAudio instance")
self.paInstance = pyaudio.PyAudio()
# Define and initialize stream object if we have been passed a device ID (pyaudio.open)
self.stream = None
if _output_device_index or _input_device_index:
if _output_device_index is not None or _input_device_index is not None:
if _init_on_startup:
LOGGER.info("Init stream")
self.init_stream()
def init_stream(self, _new_output_device_index: int = None, _new_input_device_index: int = None):
# Check what device was asked to be changed (or set)
if _new_input_device_index:
if _new_input_device_index is not None:
if self.paInstance_kwargs['input']:
self.paInstance_kwargs['input_device_index'] = _new_input_device_index
else:
LOGGER.warning("[AudioStream.init_stream]:\tInput was not enabled when initialized."
" Reinitialize with '_input=True'")
if _new_output_device_index:
if _new_output_device_index is not None:
if self.paInstance_kwargs['output']:
self.paInstance_kwargs['output_device_index'] = _new_output_device_index
else:
@@ -70,143 +61,115 @@ class AudioStream:
" Reinitialize with '_output=True'")
self.close_if_open()
# Open the stream
self.stream = self.paInstance.open(**self.paInstance_kwargs)
def close_if_open(self):
# Stop the stream if it is started
if self.stream:
if self.stream.is_active():
self.stream.stop_stream()
self.stream.close()
LOGGER.debug("[ReopenStream.close_if_open]:\t Stream was open; It was closed.")
if self.stream and self.stream.is_active():
self.stream.stop_stream()
self.stream.close()
LOGGER.debug("[ReopenStream.close_if_open]:\t Stream was open; It was closed.")
def list_devices(self, _display_input_devices: bool = True, _display_output_devices: bool = True):
LOGGER.info('Getting a list of the devices connected')
info = self.paInstance.get_host_api_info_by_index(0)
numdevices = info.get('deviceCount')
devices = {
'Input': {},
'Output': {}
}
devices = {'Input': {}, 'Output': {}}
for i in range(0, numdevices):
if (self.paInstance.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
input_device = self.paInstance.get_device_info_by_host_api_device_index(0, i).get('name')
device_info = self.paInstance.get_device_info_by_host_api_device_index(0, i)
if (device_info.get('maxInputChannels')) > 0:
input_device = device_info.get('name')
devices['Input'][i] = input_device
if _display_input_devices:
LOGGER.debug(f"Input Device id {i} - {input_device}")
if (self.paInstance.get_device_info_by_host_api_device_index(0, i).get('maxOutputChannels')) > 0:
output_device = self.paInstance.get_device_info_by_host_api_device_index(0, i).get('name')
if (device_info.get('maxOutputChannels')) > 0:
output_device = device_info.get('name')
devices['Output'][i] = output_device
if _display_output_devices:
LOGGER.debug(f"Output Device id {i} - {output_device}")
return devices
async def stop(self):
await voice_connection.disconnect()
self.close_if_open()
self.stream.close()
self.paInstance.terminate()
# noinspection PyUnresolvedReferences
class NoiseGate(AudioStream):
def __init__(self, _voice_connection, _noise_gate_threshold: int, **kwargs):
super(NoiseGate, self).__init__(_init_on_startup=True, **kwargs)
global voice_connection
voice_connection = _voice_connection
self.voice_connection = _voice_connection
self.THRESHOLD = _noise_gate_threshold
self.NGStream = NoiseGateStream(self)
self.Voice_Connection_Thread = None
def run(self) -> None:
global voice_connection
# Start the audio stream
LOGGER.debug("Starting stream")
self.stream.start_stream()
# Start the stream to discord
self.core()
def core(self, error=None):
if error:
LOGGER.warning(error)
LOGGER.warning(f"Audio stream stopped unexpectedly with error: {error}")
return # Avoid recursion on error
if voice_connection.is_connected() and not voice_connection.is_playing():
if self.voice_connection.is_connected() and not self.voice_connection.is_playing():
LOGGER.debug("Playing stream to discord")
voice_connection.play(self.NGStream, after=self.core)
# The 'after' callback can be prone to loops, simplified here
self.voice_connection.play(self.NGStream, after=lambda e: self.core(e))
async def close(self):
LOGGER.debug("Closing")
await voice_connection.disconnect()
if self.stream.is_active:
self.stream.stop_stream()
LOGGER.debug("Stopping stream")
LOGGER.debug("Closing NoiseGate resources...")
if self.voice_connection and self.voice_connection.is_connected():
self.voice_connection.stop() # Stop sending audio
self.close_if_open() # Close PyAudio stream
if self.paInstance:
self.paInstance.terminate()
LOGGER.debug("NoiseGate resources closed.")
# noinspection PyUnresolvedReferences
class NoiseGateStream(discord.AudioSource):
def __init__(self, _stream):
def __init__(self, noise_gate_instance: NoiseGate):
super(NoiseGateStream, self).__init__()
self.stream = _stream # The actual audio stream object
self.NG_fadeout = 240 / 20 # Fadeout value used to hold the noisegate after de-triggering
self.NG_fadeout_count = 0 # A count set when the noisegate is triggered and was de-triggered
self.process_set_count = 0 # Counts how many processes have been made
self.noise_gate = noise_gate_instance
self.NG_fadeout = 12 # Equivalent to 240ms of audio frames (240 / 20ms)
self.NG_fadeout_count = 0
self.process_set_count = 0
def read(self):
try:
while voice_connection.is_connected():
curr_buffer = bytearray(self.stream.stream.read(960))
buffer_rms = audioop.rms(curr_buffer, 2)
if buffer_rms > 0:
buffer_decibel = 20 * math.log10(buffer_rms)
# Check connection status via the parent NoiseGate instance
if not self.noise_gate.voice_connection.is_connected():
return b''
if self.process_set_count % 10 == 0:
if buffer_decibel >= self.stream.THRESHOLD:
LOGGER.debug(f"[Noisegate Open] {buffer_decibel} db")
else:
LOGGER.debug(f"[Noisegate Closed] {buffer_decibel} db")
# Read from the PyAudio stream, also via the parent instance
curr_buffer = self.noise_gate.stream.read(960, exception_on_overflow=False)
if buffer_decibel >= self.stream.THRESHOLD:
self.NG_fadeout_count = self.NG_fadeout
self.process_set_count += 1
if curr_buffer:
return bytes(curr_buffer)
if not curr_buffer:
return b''
else:
if self.NG_fadeout_count > 0:
self.NG_fadeout_count -= 1
LOGGER.debug(f"Frames in fadeout remaining: {self.NG_fadeout_count}")
self.process_set_count += 1
if curr_buffer:
return bytes(curr_buffer)
buffer_rms = audioop.rms(curr_buffer, 2)
if buffer_rms > 0:
buffer_decibel = 20 * math.log10(buffer_rms)
if self.process_set_count % 10 == 0:
log_msg = f"[{'Open' if buffer_decibel >= self.noise_gate.THRESHOLD else 'Closed'}]"
LOGGER.debug(f"[NoiseGate {log_msg}] {buffer_decibel:.2f} dB")
if buffer_decibel >= self.noise_gate.THRESHOLD:
self.NG_fadeout_count = self.NG_fadeout
self.process_set_count += 1
return bytes(curr_buffer)
elif self.NG_fadeout_count > 0:
self.NG_fadeout_count -= 1
self.process_set_count += 1
return bytes(curr_buffer)
return b'' # Return silence if below threshold and not in fadeout
except IOError as e:
LOGGER.error(f"PyAudio IOError in read(): {e}")
return b'' # Return silence to not break the Discord audio pump
except Exception as e:
LOGGER.warning(e)
pass
def audio_datalist_set_volume(self, datalist, volume):
""" Change value of list of audio chunks """
sound_level = (volume / 100.)
for i in range(len(datalist)):
chunk = numpy.fromstring(datalist[i], numpy.int16)
chunk = chunk * sound_level
datalist[i] = chunk.astype(numpy.int16)
if __name__ == '__main__':
input_index = int(input("Input:\t"))
output_index = int(input("Output:\t"))
ng = NoiseGate(_input_device_index=input_index, _output_device_index=output_index)
ng.list_devices()
ng.start()
LOGGER.error(f"Unhandled exception in NoiseGateStream.read: {e}", exc_info=True)
return b''

View File

@@ -18,7 +18,8 @@ class DiscordBotManager:
def __init__(self):
self.bot: Optional[commands.Bot] = None
self.bot_task: Optional[asyncio.Task] = None
self.voice_clients: Dict[int, VoiceClient] = {}
# This dictionary will hold both the client and its audio stream handler
self.voice_connections: Dict[int, Dict] = {}
self.token: Optional[str] = None
self.loop = asyncio.get_event_loop()
self.lock = asyncio.Lock()
@@ -38,33 +39,43 @@ class DiscordBotManager:
@self.bot.event
async def on_ready():
LOGGER.info(f'Logged in as {self.bot.user}')
# Set the event when on_ready is called
self._ready_event.set()
@self.bot.event
async def on_voice_state_update(member, before, after):
if member == self.bot.user and before.channel is None and after.channel is not None:
print(f"{member.name} joined voice channel {after.channel.name}")
if member != self.bot.user:
return
if before.channel is None and after.channel is not None:
LOGGER.info(f"{member.name} joined voice channel {after.channel.name}")
self._voice_ready_event.set()
elif before.channel is not None and after.channel is not None and before.channel != after.channel:
LOGGER.info(f"{member.name} was moved to voice channel {after.channel.name}")
if not self._voice_ready_event.is_set():
self._voice_ready_event.set()
elif before.channel is not None and after.channel is None:
LOGGER.warning(f"{member.name} left voice channel {before.channel.name}")
guild_id = before.channel.guild.id
if guild_id in self.voice_connections:
LOGGER.warning(f"Bot was disconnected from {guild_id} unexpectedly. Cleaning up...")
await self.leave_voice_channel(guild_id)
self._voice_ready_event.clear()
@self.bot.event
async def on_disconnect():
LOGGER.warning("Bot has been disconnected from Discord.")
# Load Opus for the current CPU
await self.load_opus()
# Create the task to run the bot in the background
self.bot_task = self.loop.create_task(self.bot.start(token))
# Wait for the on_ready event to be set by the bot task
LOGGER.info("Waiting for bot to become ready...")
try:
await asyncio.wait_for(self._ready_event.wait(), timeout=60.0)
LOGGER.info("Bot is ready, start_bot returning.")
return
LOGGER.info("Bot is ready.")
except asyncio.TimeoutError:
LOGGER.error("Timeout waiting for bot to become ready. Bot might have failed to start.")
LOGGER.error("Timeout waiting for bot to become ready.")
if self.bot_task and not self.bot_task.done():
self.bot_task.cancel()
raise RuntimeError("Bot failed to become ready within timeout.")
@@ -72,92 +83,105 @@ class DiscordBotManager:
async def stop_bot(self):
async with self.lock:
if self.bot:
# Disconnect from all voice channels cleanly
for guild_id in list(self.voice_connections.keys()):
await self.leave_voice_channel(guild_id)
await self.bot.close()
self.bot = None
if self.bot_task:
await self.bot_task
self.bot_task.cancel()
self.bot_task = None
self.voice_clients.clear()
self.voice_connections.clear()
self._ready_event.clear()
LOGGER.info("Bot has been stopped.")
async def join_voice_channel(self, guild_id: int, channel_id: int, ng_threshold: int = 50, device_id: int = 4):
if not self.bot:
raise RuntimeError("Bot is not running.")
if not self.bot: raise RuntimeError("Bot is not running.")
guild = self.bot.get_guild(guild_id)
if not guild:
raise ValueError("Guild not found.")
if not opus.is_loaded():
raise RuntimeError("Opus is not loaded.")
if not guild: raise ValueError("Guild not found.")
if not opus.is_loaded(): raise RuntimeError("Opus is not loaded.")
channel = guild.get_channel(channel_id)
if not isinstance(channel, VoiceChannel):
raise ValueError("Channel is not a voice channel.")
if guild_id in self.voice_clients:
raise RuntimeError("Already connected to this guild's voice channel.")
if not isinstance(channel, VoiceChannel): raise ValueError("Channel is not a voice channel.")
if guild_id in self.voice_connections: raise RuntimeError("Already connected to this guild's voice channel.")
try:
# 1. Connect to the channel first
self._voice_ready_event.clear()
voice_client = await channel.connect(timeout=60.0, reconnect=True)
LOGGER.debug("Voice Connected.")
streamHandler = NoiseGate(
LOGGER.debug("Voice client connecting...")
# 2. Wait for the on_voice_state_update event to confirm readiness
await asyncio.wait_for(self._voice_ready_event.wait(), timeout=15.0)
LOGGER.info("Bot voice connection is ready.")
# 3. NOW, create and start the audio stream handler
stream_handler = NoiseGate(
_input_device_index=device_id,
_voice_connection=voice_client,
_noise_gate_threshold=ng_threshold)
streamHandler.run()
LOGGER.debug("Stream is running.")
self.voice_clients[guild_id] = voice_client
LOGGER.info(f"Joined guild {guild_id} voice channel {channel_id} and stream is running.")
except Exception as e:
LOGGER.error(f"Failed to connect to voice channel: {e}")
stream_handler.run()
# 4. Store both client and stream handler for proper management
self.voice_connections[guild_id] = {
"client": voice_client,
"stream": stream_handler
}
LOGGER.info(f"Joined guild {guild_id} and audio stream is now running.")
LOGGER.info("Waiting for bot to join voice...")
try:
await asyncio.wait_for(self._voice_ready_event.wait(), timeout=60.0)
LOGGER.info("Bot joined voice, returning.")
return
except asyncio.TimeoutError:
LOGGER.error("Timeout waiting for bot to join voice.")
raise RuntimeError("Bot failed to join voice within timeout.")
LOGGER.error(f"Timeout waiting for bot to join voice channel {channel_id}.")
raise RuntimeError("Bot failed to confirm voice connection within timeout.")
except Exception as e:
LOGGER.error(f"Failed to connect to voice channel: {e}", exc_info=True)
raise
async def leave_voice_channel(self, guild_id: int):
if not self.bot:
raise RuntimeError("Bot is not running.")
if not self.bot: raise RuntimeError("Bot is not running.")
connection_info = self.voice_connections.get(guild_id)
if not connection_info: raise RuntimeError("Not connected to the specified guild's voice channel.")
voice_client = self.voice_clients.get(guild_id)
if not voice_client:
raise RuntimeError("Not connected to the specified guild's voice channel.")
# Cleanly stop the associated audio stream first
stream_handler = connection_info.get('stream')
if stream_handler:
LOGGER.info(f"Stopping audio stream for guild {guild_id}.")
await stream_handler.close()
await voice_client.disconnect()
del self.voice_clients[guild_id]
# Disconnect the voice client
voice_client = connection_info.get('client')
if voice_client and voice_client.is_connected():
await voice_client.disconnect()
del self.voice_connections[guild_id]
LOGGER.info(f"Left guild {guild_id} voice channel.")
async def load_opus(self):
""" Load the proper OPUS library for the device being used """
processor = platform.machine()
script_dir = os.path.dirname(os.path.abspath(__file__))
LOGGER.debug("Processor: ", processor)
if os.name == 'nt':
if processor == "AMD64":
opus.load_opus(os.path.join(script_dir, './opus/libopus_amd64.dll'))
LOGGER.info("Loaded OPUS library for AMD64")
return "AMD64"
else:
if processor == "aarch64":
opus.load_opus(os.path.join(script_dir, './opus/libopus_aarcch64.so'))
LOGGER.info("Loaded OPUS library for aarch64")
return "aarch64"
elif processor == "armv7l":
opus.load_opus(os.path.join(script_dir, './opus/libopus_armv7l.so'))
LOGGER.info("Loaded OPUS library for armv7l")
return "armv7l"
LOGGER.debug(f"Processor: {processor}, OS: {os.name}")
try:
if os.name == 'nt': # Windows
if processor == "AMD64":
opus.load_opus(os.path.join(script_dir, './opus/libopus_amd64.dll'))
LOGGER.info("Loaded OPUS library for AMD64")
else: # Linux / other
if processor == "aarch64":
opus.load_opus(os.path.join(script_dir, './opus/libopus_aarcch64.so'))
LOGGER.info("Loaded OPUS library for aarch64")
elif processor == "armv7l":
opus.load_opus(os.path.join(script_dir, './opus/libopus_armv7l.so'))
LOGGER.info("Loaded OPUS library for armv7l")
else:
# Fallback for other Linux archs like x86_64
opus.load_opus('libopus.so.0')
LOGGER.info(f"Loaded system OPUS library for {processor}")
except Exception as e:
LOGGER.error(f"Failed to load OPUS library: {e}")
raise RuntimeError("Could not load a valid Opus library. Voice functionality will fail.")
async def set_presence(self, system_name: str):
""" Set the presence (activity) of the bot """
if not self.bot:
LOGGER.warning("Bot is not running, cannot set presence.")
if not self.bot or not self.bot.is_ready():
LOGGER.warning("Bot is not ready, cannot set presence.")
return
try:
@@ -165,4 +189,4 @@ class DiscordBotManager:
await self.bot.change_presence(activity=activity)
LOGGER.info(f"Bot presence set to 'Listening to {system_name}'")
except Exception as pe:
LOGGER.error(f"Unable to set presence: '{pe}'")
LOGGER.error(f"Unable to set presence: '{pe}'")