Refactor audio stream workers and enhance configuration management

- Introduced a base class `BaseStreamWorker` to standardize audio stream worker implementations, improving code reuse and maintainability.
- Updated `DeepgramStreamWorker` and `ParakeetStreamWorker` to inherit from `BaseStreamWorker`, streamlining their initialization and configuration validation processes.
- Enhanced error handling and logging for Redis connections and consumer initialization in the worker classes.
- Refactored memory processing logic to improve conversation text extraction and speaker identification.
- Added unit tests for audio stream workers to ensure reliability and correct functionality across various scenarios.

Author: 0xrushi

backends/advanced/src/advanced_omi_backend/controllers/conversation_controller.py

@@ -447,9 +447,6 @@ async def reprocess_memory(conversation_id: str, transcript_version_id: str, use
         from advanced_omi_backend.models.job import JobPriority
 
         job = enqueue_memory_processing(
-            client_id=conversation_model.client_id,
-            user_id=str(user.user_id),
-            user_email=user.email,
             conversation_id=conversation_id,
             priority=JobPriority.NORMAL
         )

backends/advanced/src/advanced_omi_backend/routers/modules/annotation_routes.py

@@ -73,9 +73,6 @@ async def create_annotation(
 
             # Trigger memory reprocessing
             enqueue_memory_processing(
-                client_id=conversation.client_id,
-                user_id=str(current_user.id),
-                user_email=current_user.email,
                 conversation_id=conversation.conversation_id,
                 priority=JobPriority.NORMAL
             )

backends/advanced/src/advanced_omi_backend/workers/audio_stream_deepgram_worker.py

@@ -5,76 +5,38 @@
 Starts a consumer that reads from audio:stream:deepgram and transcribes audio.
 """
 
-import asyncio
-import logging
 import os
-import signal
-import sys
-
-import redis.asyncio as redis
 
 from advanced_omi_backend.services.transcription.deepgram import DeepgramStreamConsumer
+from advanced_omi_backend.workers.base_audio_worker import BaseStreamWorker
 
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
-)
-
-logger = logging.getLogger(__name__)
-
-
-async def main():
-    """Main worker entry point."""
-    logger.info("🚀 Starting Deepgram audio stream worker")
-
-    # Check that config.yml has Deepgram configured
-    # The registry provider will load configuration from config.yml
-    api_key = os.getenv("DEEPGRAM_API_KEY")
-    if not api_key:
-        logger.warning("DEEPGRAM_API_KEY environment variable not set")
-        logger.warning("Ensure config.yml has a default 'stt' model configured for Deepgram")
-        logger.warning("Audio transcription will use alternative providers if configured in config.yml")
-
-    redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
-
-    # Create Redis client
-    redis_client = await redis.from_url(
-        redis_url,
-        encoding="utf-8",
-        decode_responses=False
-    )
-    logger.info("Connected to Redis")
-
-    # Create consumer with balanced buffer size
-    # 20 chunks = ~5 seconds of audio
-    # Balance between transcription accuracy and latency
-    # Consumer uses registry-driven provider from config.yml
-    consumer = DeepgramStreamConsumer(
-        redis_client=redis_client,
-        buffer_chunks=20  # 5 seconds - good context without excessive delay
-    )
-
-    # Setup signal handlers for graceful shutdown
-    def signal_handler(signum, frame):
-        logger.info(f"Received signal {signum}, shutting down...")
-        asyncio.create_task(consumer.stop())
 
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
+class DeepgramStreamWorker(BaseStreamWorker):
+    """Deepgram audio stream worker implementation."""
 
-    try:
-        logger.info("✅ Deepgram worker ready")
+    def __init__(self):
+        super().__init__(service_name="Deepgram audio stream worker")
 
-        # This blocks until consumer is stopped
-        await consumer.start_consuming()
+    def validate_config(self):
+        """Check that config.yml has Deepgram configured."""
+        # The registry provider will load configuration from config.yml
+        api_key = os.getenv("DEEPGRAM_API_KEY")
+        if not api_key:
+            self.logger.warning("DEEPGRAM_API_KEY environment variable not set")
+            self.logger.warning("Ensure config.yml has a default 'stt' model configured for Deepgram")
+            self.logger.warning("Audio transcription will use alternative providers if configured in config.yml")
 
-    except Exception as e:
-        logger.error(f"Worker error: {e}", exc_info=True)
-        sys.exit(1)
-    finally:
-        await redis_client.aclose()
-        logger.info("👋 Deepgram worker stopped")
+    def get_consumer(self, redis_client):
+        """Create Deepgram consumer with balanced buffer size."""
+        # Create consumer with balanced buffer size
+        # 20 chunks = ~5 seconds of audio
+        # Balance between transcription accuracy and latency
+        # Consumer uses registry-driven provider from config.yml
+        return DeepgramStreamConsumer(
+            redis_client=redis_client,
+            buffer_chunks=20  # 5 seconds - good context without excessive delay
+        )
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    DeepgramStreamWorker.start()

backends/advanced/src/advanced_omi_backend/workers/audio_stream_parakeet_worker.py

@@ -5,91 +5,39 @@
 Starts a consumer that reads from audio:stream:* and transcribes audio using Parakeet.
 """
 
-import asyncio
-import logging
 import os
-import signal
-import sys
-
-import redis.asyncio as redis
 
 from advanced_omi_backend.services.transcription.parakeet_stream_consumer import ParakeetStreamConsumer
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
-)
-
-logger = logging.getLogger(__name__)
-
-
-async def main():
-    """Main worker entry point."""
-    logger.info("🚀 Starting Parakeet audio stream worker")
-
-    # Check that config.yml has Parakeet configured
-    # The registry provider will load configuration from config.yml
-    service_url = os.getenv("PARAKEET_ASR_URL")
-    if not service_url:
-        logger.warning("PARAKEET_ASR_URL environment variable not set")
-        logger.warning("Ensure config.yml has a default 'stt' model configured for Parakeet")
-        logger.warning("Audio transcription will use alternative providers if configured in config.yml")
-
-    redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
-
-    # Create Redis client
-    redis_client = await redis.from_url(
-        redis_url,
-        encoding="utf-8",
-        decode_responses=False
-    )
-    logger.info("Connected to Redis")
-
-    # Create consumer with balanced buffer size
-    # 20 chunks = ~5 seconds of audio
-    # Balance between transcription accuracy and latency
-    # Consumer uses registry-driven provider from config.yml
-    consumer = ParakeetStreamConsumer(
-        redis_client=redis_client,
-        buffer_chunks=20  # 5 seconds - good context without excessive delay
-    )
-
-    # Setup signal handlers for graceful shutdown
-    shutdown_event = asyncio.Event()
-
-    def signal_handler(signum, _frame):
-        logger.info(f"Received signal {signum}, shutting down...")
-        shutdown_event.set()
-
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
-
-    try:
-        logger.info("✅ Parakeet worker ready")
-
-        # This blocks until consumer is stopped or shutdown signaled
-        consume_task = asyncio.create_task(consumer.start_consuming())
-        shutdown_task = asyncio.create_task(shutdown_event.wait())
-
-        done, pending = await asyncio.wait(
-            [consume_task, shutdown_task],
-            return_when=asyncio.FIRST_COMPLETED
+from advanced_omi_backend.workers.base_audio_worker import BaseStreamWorker
+
+
+class ParakeetStreamWorker(BaseStreamWorker):
+    """Parakeet audio stream worker implementation."""
+
+    def __init__(self):
+        super().__init__(service_name="Parakeet audio stream worker")
+
+    def validate_config(self):
+        """Check that config.yml has Parakeet configured."""
+        # The registry provider will load configuration from config.yml
+        service_url = os.getenv("PARAKEET_ASR_URL")
+        if not service_url:
+            self.logger.warning("PARAKEET_ASR_URL environment variable not set")
+            self.logger.warning("Ensure config.yml has a default 'stt' model configured for Parakeet")
+            self.logger.warning("Audio transcription will use alternative providers if configured in config.yml")
+
+    def get_consumer(self, redis_client):
+        """Create Parakeet consumer with balanced buffer size."""
+        # Create consumer with balanced buffer size
+        # 20 chunks = ~5 seconds of audio
+        # Balance between transcription accuracy and latency
+        # Consumer uses registry-driven provider from config.yml
+        return ParakeetStreamConsumer(
+            redis_client=redis_client,
+            buffer_chunks=20  # 5 seconds - good context without excessive delay
         )
 
-        # Cancel pending tasks
-        for task in pending:
-            task.cancel()
-
-        await consumer.stop()
-
-    except Exception as e:
-        logger.error(f"Worker error: {e}", exc_info=True)
-        sys.exit(1)
-    finally:
-        await redis_client.aclose()
-        logger.info("👋 Parakeet worker stopped")
-
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    ParakeetStreamWorker.start()
 

backends/advanced/src/advanced_omi_backend/workers/base_audio_worker.py

@@ -0,0 +1,156 @@
+"""
+Base audio stream worker.
+
+Provides a template for stream workers with consistent Redis connection,
+signal handling, and error management.
+"""
+
+import asyncio
+import logging
+import os
+import signal
+import sys
+from abc import ABC, abstractmethod
+
+import redis.asyncio as redis
+
+# Configure basic logging if not already configured
+if not logging.getLogger().handlers:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
+    )
+
+class BaseStreamWorker(ABC):
+    """
+    Base class for audio stream workers using the Template Method pattern.
+    
+    Subclasses must implement:
+    - validate_config(): Check environment/config requirements
+    - get_consumer(redis_client): Return the specific consumer instance
+    """
+
+    def __init__(self, service_name: str):
+        self.service_name = service_name
+        self.logger = logging.getLogger(self.__class__.__name__)
+        self.redis_client = None
+        self.consumer = None
+
+    @abstractmethod
+    def validate_config(self):
+        """
+        Check required environment variables or configuration.
+        Should log warnings/errors if configuration is missing.
+        """
+        pass
+
+    @abstractmethod
+    def get_consumer(self, redis_client):
+        """
+        Create and return the consumer instance.
+        
+        Args:
+            redis_client: Initialized Redis client
+            
+        Returns:
+            An instance complying with the BaseAudioStreamConsumer interface
+        """
+        pass
+
+    async def run(self):
+        """Main execution loop."""
+        self.logger.info(f"🚀 Starting {self.service_name}")
+        
+        self.validate_config()
+
+        redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0")
+        
+        try:
+            self.redis_client = await redis.from_url(
+                redis_url,
+                encoding="utf-8",
+                decode_responses=False
+            )
+            self.logger.info("Connected to Redis")
+        except Exception as e:
+            self.logger.error(f"Failed to connect to Redis: {e}")
+            sys.exit(1)
+
+        try:
+            self.consumer = self.get_consumer(self.redis_client)
+        except Exception as e:
+            self.logger.error(f"Failed to initialize consumer: {e}")
+            await self.redis_client.aclose()
+            sys.exit(1)
+
+        # Setup graceful shutdown
+        loop = asyncio.get_running_loop()
+        stop_event = asyncio.Event()
+
+        def signal_handler():
+            self.logger.info("Received stop signal, shutting down...")
+            stop_event.set()
+
+        # Register signal handlers
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            try:
+                loop.add_signal_handler(sig, signal_handler)
+            except NotImplementedError:
+                # Fallback for environments where add_signal_handler is not supported
+                # (e.g. some Windows environments or custom loops)
+                self.logger.warning(f"Could not add signal handler for {sig}")
+
+        try:
+            self.logger.info(f"✅ {self.service_name} ready")
+            
+            # Run consumer as a task
+            consumer_task = asyncio.create_task(self.consumer.start_consuming())
+            stop_wait_task = asyncio.create_task(stop_event.wait())
+
+            # Wait for either the consumer to finish (error/done) or stop signal
+            done, pending = await asyncio.wait(
+                [consumer_task, stop_wait_task],
+                return_when=asyncio.FIRST_COMPLETED
+            )
+
+            # Check if consumer failed
+            if consumer_task in done:
+                try:
+                    await consumer_task
+                except asyncio.CancelledError:
+                    pass
+                except Exception as e:
+                    self.logger.error(f"Consumer task failed: {e}", exc_info=True)
+                    # We continue to cleanup
+
+            # Trigger stop on consumer
+            self.logger.info("Stopping consumer...")
+            await self.consumer.stop()
+            
+            # Ensure consumer task finishes if it was running
+            if consumer_task in pending:
+                try:
+                    await consumer_task
+                except asyncio.CancelledError:
+                    pass
+                except Exception as e:
+                    # Ignore expected errors during shutdown if any
+                    self.logger.debug(f"Consumer shutdown exception: {e}")
+
+        except Exception as e:
+            self.logger.error(f"Worker runtime error: {e}", exc_info=True)
+            sys.exit(1)
+        finally:
+            if self.redis_client:
+                await self.redis_client.aclose()
+            self.logger.info(f"👋 {self.service_name} stopped")
+
+    @classmethod
+    def start(cls):
+        """Entry point for script execution."""
+        instance = cls()
+        try:
+            asyncio.run(instance.run())
+        except KeyboardInterrupt:
+            # Handle keyboard interrupt outside the loop if it propagates
+            pass