atom-audio-engine 0.1.0__py3-none-any.whl

llm/base.py

@@ -0,0 +1,108 @@
+"""Abstract base class for LLM providers."""
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterator, Optional
+
+from core.types import ResponseChunk, ConversationContext
+
+
+class BaseLLM(ABC):
+    """
+    Abstract base class for Large Language Model providers.
+
+    All LLM implementations must inherit from this class and implement
+    the required methods for both batch and streaming text generation.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "gpt-4o",
+        temperature: float = 0.7,
+        max_tokens: int = 1024,
+        system_prompt: Optional[str] = None,
+        **kwargs
+    ):
+        """
+        Initialize the LLM provider.
+
+        Args:
+            api_key: API key for the provider
+            model: Model identifier to use
+            temperature: Sampling temperature (0.0-2.0)
+            max_tokens: Maximum tokens in response
+            system_prompt: Default system prompt
+            **kwargs: Additional provider-specific configuration
+        """
+        self.api_key = api_key
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.system_prompt = system_prompt
+        self.config = kwargs
+
+    @abstractmethod
+    async def generate(
+        self, prompt: str, context: Optional[ConversationContext] = None
+    ) -> str:
+        """
+        Generate a complete response to a prompt.
+
+        Args:
+            prompt: User's input text
+            context: Optional conversation history
+
+        Returns:
+            Complete response text
+        """
+        pass
+
+    @abstractmethod
+    async def generate_stream(
+        self, prompt: str, context: Optional[ConversationContext] = None
+    ) -> AsyncIterator[ResponseChunk]:
+        """
+        Generate a streaming response to a prompt.
+
+        Args:
+            prompt: User's input text
+            context: Optional conversation history
+
+        Yields:
+            ResponseChunk objects with partial text
+        """
+        pass
+
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.disconnect()
+
+    async def connect(self):
+        """
+        Initialize the LLM client.
+        Override in subclasses if needed.
+        """
+        pass
+
+    async def disconnect(self):
+        """
+        Clean up the LLM client.
+        Override in subclasses if needed.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the name of this LLM provider."""
+        pass
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Whether this provider supports streaming responses."""
+        return True

llm/groq.py

@@ -0,0 +1,210 @@
+"""Groq API implementation for LLM (Language Model)."""
+
+import logging
+from typing import AsyncIterator, Optional
+
+from groq import Groq
+
+from core.types import ResponseChunk, ConversationContext
+from .base import BaseLLM
+
+logger = logging.getLogger(__name__)
+
+
+class GroqLLM(BaseLLM):
+    """
+    Groq API client for language model text generation.
+
+    Supports both batch and streaming text generation with excellent
+    latency for conversational AI. Uses Groq's optimized inference engine.
+
+    Example:
+        llm = GroqLLM(
+            api_key="gsk_...",
+            model="llama-3.1-8b-instant"
+        )
+
+        # Batch generation
+        response = await llm.generate("Hello", context=conversation)
+
+        # Streaming generation
+        async for chunk in llm.generate_stream("Hello", context=conversation):
+            print(chunk.text, end="", flush=True)
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: str = "llama-3.1-8b-instant",
+        temperature: float = 0.7,
+        max_tokens: int = 1024,
+        system_prompt: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Initialize Groq LLM provider.
+
+        Args:
+            api_key: Groq API key
+            model: Model to use (default "llama-3.1-8b-instant", alternatives: "llama-3.3-70b-versatile", "mixtral-8x7b-32768")
+            temperature: Sampling temperature (0.0-2.0)
+            max_tokens: Maximum tokens in response
+            system_prompt: Default system prompt
+            **kwargs: Additional config (stored in self.config)
+        """
+        super().__init__(
+            api_key=api_key,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            system_prompt=system_prompt,
+            **kwargs,
+        )
+        self.client = None
+
+    @property
+    def name(self) -> str:
+        """Return provider name."""
+        return "groq"
+
+    async def connect(self):
+        """Initialize Groq client."""
+        try:
+            self.client = Groq(api_key=self.api_key)
+            logger.debug("Groq client initialized")
+        except Exception as e:
+            logger.error(f"Failed to initialize Groq client: {e}")
+            raise
+
+    async def disconnect(self):
+        """Close Groq client connection."""
+        if self.client:
+            try:
+                # Groq client cleanup (if supported)
+                pass
+            except Exception as e:
+                logger.error(f"Error disconnecting Groq: {e}")
+
+    async def generate(
+        self, prompt: str, context: Optional[ConversationContext] = None
+    ) -> str:
+        """
+        Generate a complete response to a prompt.
+
+        Args:
+            prompt: User's input text
+            context: Optional conversation history
+
+        Returns:
+            Complete response text
+        """
+        if not self.client:
+            await self.connect()
+
+        try:
+            # Build message list from context
+            messages = []
+
+            # Add system prompt
+            system = self.system_prompt or context.system_prompt if context else None
+            if system:
+                messages.append({"role": "system", "content": system})
+
+            # Add conversation history
+            if context:
+                for msg in context.get_messages_for_llm():
+                    if msg["role"] != "system":  # Avoid duplicate system prompt
+                        messages.append(msg)
+
+            # Add current prompt
+            messages.append({"role": "user", "content": prompt})
+
+            logger.debug(f"Generating response with {len(messages)} messages")
+
+            # Call Groq API
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                stream=False,
+            )
+
+            # Extract text
+            if response.choices and response.choices[0].message:
+                text = response.choices[0].message.content
+                logger.debug(f"Generated response: {text[:100]}...")
+                return text
+
+            return ""
+
+        except Exception as e:
+            logger.error(f"Groq generation error: {e}")
+            raise
+
+    async def generate_stream(
+        self, prompt: str, context: Optional[ConversationContext] = None
+    ) -> AsyncIterator[ResponseChunk]:
+        """
+        Generate a streaming response to a prompt.
+
+        Yields text chunks as they are generated for real-time display.
+
+        Args:
+            prompt: User's input text
+            context: Optional conversation history
+
+        Yields:
+            ResponseChunk objects with partial and final text
+        """
+        if not self.client:
+            await self.connect()
+
+        try:
+            # Build message list from context
+            messages = []
+
+            # Add system prompt
+            system = self.system_prompt or context.system_prompt if context else None
+            if system:
+                messages.append({"role": "system", "content": system})
+
+            # Add conversation history
+            if context:
+                for msg in context.get_messages_for_llm():
+                    if msg["role"] != "system":  # Avoid duplicate system prompt
+                        messages.append(msg)
+
+            # Add current prompt
+            messages.append({"role": "user", "content": prompt})
+
+            logger.debug(f"Streaming response with {len(messages)} messages")
+
+            # Call Groq API with streaming
+            with self.client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                stream=True,
+            ) as response:
+
+                full_text = ""
+                for chunk in response:
+                    if chunk.choices and chunk.choices[0].delta.content:
+                        delta = chunk.choices[0].delta.content
+                        full_text += delta
+
+                        # Check if this is the last chunk
+                        is_final = (
+                            chunk.choices[0].finish_reason is not None
+                            and chunk.choices[0].finish_reason != "length"
+                        )
+
+                        yield ResponseChunk(text=delta, is_final=is_final)
+
+                logger.debug(f"Streaming complete. Total: {full_text[:100]}...")
+
+        except Exception as e:
+            logger.error(f"Groq streaming error: {e}")
+            raise

pipelines/__init__.py

@@ -0,0 +1 @@
+"""Pipeline implementations for audio-engine."""

streaming/__init__.py

@@ -0,0 +1,5 @@
+"""Streaming and WebSocket server components."""
+
+from streaming.websocket_server import WebSocketServer
+
+__all__ = ["WebSocketServer"]

streaming/websocket_server.py

@@ -0,0 +1,341 @@
+"""WebSocket server for real-time audio streaming."""
+
+import asyncio
+import json
+import logging
+from typing import Optional, Callable, Any
+
+import websockets
+
+from core.pipeline import Pipeline
+from core.types import AudioChunk, AudioFormat
+from core.config import AudioEngineConfig
+
+logger = logging.getLogger(__name__)
+
+# Type alias for WebSocket connection
+WebSocketServerProtocol = Any
+
+
+class WebSocketServer:
+    """
+    WebSocket server for real-time audio-to-audio streaming.
+
+    Protocol:
+        Client sends:
+        - Binary messages: Raw audio chunks (PCM 16-bit, 16kHz mono)
+        - JSON messages: Control commands {"type": "end_of_speech"} or {"type": "reset"}
+
+        Server sends:
+        - Binary messages: Response audio chunks
+        - JSON messages: Events {"type": "transcript", "text": "..."} etc.
+
+    Example:
+        ```python
+        server = WebSocketServer(
+            pipeline=pipeline,
+            host="0.0.0.0",
+            port=8765
+        )
+        await server.start()
+        ```
+    """
+
+    def __init__(
+        self,
+        pipeline: Pipeline,
+        host: str = "0.0.0.0",
+        port: int = 8765,
+        input_sample_rate: int = 16000,
+        on_connect: Optional[Callable[[str], Any]] = None,
+        on_disconnect: Optional[Callable[[str], Any]] = None,
+    ):
+        """
+        Initialize the WebSocket server.
+
+        Args:
+            pipeline: Configured Pipeline instance
+            host: Host to bind to
+            port: Port to listen on
+            input_sample_rate: Expected sample rate of input audio
+            on_connect: Callback when client connects
+            on_disconnect: Callback when client disconnects
+        """
+        if websockets is None:
+            raise ImportError(
+                "websockets package required. Install with: pip install websockets"
+            )
+
+        self.pipeline = pipeline
+        self.host = host
+        self.port = port
+        self.input_sample_rate = input_sample_rate
+        self.on_connect = on_connect
+        self.on_disconnect = on_disconnect
+
+        self._server = None
+        self._clients: dict[str, WebSocketServerProtocol] = {}
+
+    async def start(self):
+        """Start the WebSocket server."""
+        await self.pipeline.connect()
+
+        self._server = await websockets.serve(
+            self._handle_client,
+            self.host,
+            self.port,
+        )
+
+        logger.info(f"WebSocket server started on ws://{self.host}:{self.port}")
+
+    async def stop(self):
+        """Stop the WebSocket server."""
+        if self._server:
+            self._server.close()
+            await self._server.wait_closed()
+            self._server = None
+
+        await self.pipeline.disconnect()
+        logger.info("WebSocket server stopped")
+
+    async def _handle_client(self, websocket: WebSocketServerProtocol):
+        """Handle a single client connection."""
+        client_id = str(id(websocket))
+        self._clients[client_id] = websocket
+
+        logger.info(f"Client connected: {client_id}")
+        if self.on_connect:
+            self.on_connect(client_id)
+
+        # Send welcome message
+        await websocket.send(
+            json.dumps(
+                {
+                    "type": "connected",
+                    "client_id": client_id,
+                    "providers": self.pipeline.providers,
+                }
+            )
+        )
+
+        try:
+            await self._process_client_stream(websocket, client_id)
+        except websockets.exceptions.ConnectionClosed:
+            logger.info(f"Client disconnected: {client_id}")
+        except Exception as e:
+            logger.error(f"Error handling client {client_id}: {e}")
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "error",
+                        "message": str(e),
+                    }
+                )
+            )
+        finally:
+            del self._clients[client_id]
+            if self.on_disconnect:
+                self.on_disconnect(client_id)
+
+    async def _process_client_stream(
+        self, websocket: WebSocketServerProtocol, client_id: str
+    ):
+        """Process streaming audio from a client."""
+        audio_queue: asyncio.Queue[AudioChunk] = asyncio.Queue()
+        end_of_speech = asyncio.Event()
+
+        async def audio_stream():
+            """Yield audio chunks from the queue."""
+            while True:
+                if end_of_speech.is_set() and audio_queue.empty():
+                    break
+                try:
+                    chunk = await asyncio.wait_for(audio_queue.get(), timeout=0.1)
+                    yield chunk
+                    if chunk.is_final:
+                        break
+                except asyncio.TimeoutError:
+                    if end_of_speech.is_set():
+                        break
+                    continue
+
+        async def receive_audio():
+            """Receive audio from WebSocket and queue it."""
+            async for message in websocket:
+                if isinstance(message, bytes):
+                    # Binary audio data
+                    chunk = AudioChunk(
+                        data=message,
+                        sample_rate=self.input_sample_rate,
+                        format=AudioFormat.PCM_16K,
+                    )
+                    await audio_queue.put(chunk)
+
+                elif isinstance(message, str):
+                    # JSON control message
+                    try:
+                        data = json.loads(message)
+                        msg_type = data.get("type")
+
+                        if msg_type == "end_of_speech":
+                            # Mark final chunk
+                            final_chunk = AudioChunk(
+                                data=b"",
+                                is_final=True,
+                            )
+                            await audio_queue.put(final_chunk)
+                            end_of_speech.set()
+                            break
+
+                        elif msg_type == "reset":
+                            self.pipeline.reset_context()
+                            await websocket.send(
+                                json.dumps(
+                                    {
+                                        "type": "context_reset",
+                                    }
+                                )
+                            )
+
+                    except json.JSONDecodeError:
+                        logger.warning(f"Invalid JSON from client: {message}")
+
+        async def send_response():
+            """Stream response audio back to client."""
+            # Set up callbacks to send events
+            original_on_transcript = self.pipeline.on_transcript
+            original_on_llm_response = self.pipeline.on_llm_response
+
+            async def send_transcript(text: str):
+                await websocket.send(
+                    json.dumps(
+                        {
+                            "type": "transcript",
+                            "text": text,
+                        }
+                    )
+                )
+                if original_on_transcript:
+                    original_on_transcript(text)
+
+            async def send_llm_response(text: str):
+                await websocket.send(
+                    json.dumps(
+                        {
+                            "type": "response_text",
+                            "text": text,
+                        }
+                    )
+                )
+                if original_on_llm_response:
+                    original_on_llm_response(text)
+
+            # Temporarily override callbacks
+            self.pipeline.on_transcript = lambda t: asyncio.create_task(
+                send_transcript(t)
+            )
+            self.pipeline.on_llm_response = lambda t: asyncio.create_task(
+                send_llm_response(t)
+            )
+
+            try:
+                # Wait for some audio to arrive
+                await asyncio.sleep(0.1)
+
+                # Stream response
+                await websocket.send(json.dumps({"type": "response_start"}))
+
+                async for audio_chunk in self.pipeline.stream(audio_stream()):
+                    await websocket.send(audio_chunk.data)
+
+                await websocket.send(json.dumps({"type": "response_end"}))
+
+            finally:
+                # Restore original callbacks
+                self.pipeline.on_transcript = original_on_transcript
+                self.pipeline.on_llm_response = original_on_llm_response
+
+        # Run receive and send concurrently
+        receive_task = asyncio.create_task(receive_audio())
+        send_task = asyncio.create_task(send_response())
+
+        try:
+            await asyncio.gather(receive_task, send_task)
+        except Exception as e:
+            receive_task.cancel()
+            send_task.cancel()
+            raise
+
+    async def broadcast(self, message: str):
+        """Broadcast a message to all connected clients."""
+        if self._clients:
+            await asyncio.gather(*[ws.send(message) for ws in self._clients.values()])
+
+    @property
+    def client_count(self) -> int:
+        """Return number of connected clients."""
+        return len(self._clients)
+
+
+async def run_server(
+    pipeline: Pipeline,
+    host: str = "0.0.0.0",
+    port: int = 8765,
+):
+    """
+    Convenience function to run the WebSocket server.
+
+    Args:
+        pipeline: Configured Pipeline instance
+        host: Host to bind to
+        port: Port to listen on
+    """
+    server = WebSocketServer(pipeline, host, port)
+    await server.start()
+
+    try:
+        await asyncio.Future()  # Run forever
+    finally:
+        await server.stop()
+
+
+async def run_server_from_config(
+    config: Optional["AudioEngineConfig"] = None,
+    host: Optional[str] = None,
+    port: Optional[int] = None,
+    system_prompt: Optional[str] = None,
+):
+    """
+    Create and run WebSocket server from AudioEngineConfig.
+
+    Approach:
+    1. Load config from environment (or use provided config)
+    2. Create Pipeline with providers from config
+    3. Initialize and run WebSocket server
+
+    Rationale: Single entry point to run full audio pipeline server.
+
+    Args:
+        config: AudioEngineConfig instance (loads from env if None)
+        host: Host to bind to (default: from config)
+        port: Port to listen on (default: from config)
+        system_prompt: Optional system prompt override
+    """
+    from core.config import AudioEngineConfig
+
+    if config is None:
+        config = AudioEngineConfig.from_env()
+
+    pipeline = config.create_pipeline(system_prompt=system_prompt)
+
+    host = host or config.streaming.host
+    port = port or config.streaming.port
+
+    logger.info(
+        f"Starting audio engine server with providers: "
+        f"ASR={config.asr.provider}, "
+        f"LLM={config.llm.provider}, "
+        f"TTS={config.tts.provider}"
+    )
+
+    await run_server(pipeline, host, port)

tts/__init__.py

@@ -0,0 +1,37 @@
+"""TTS (Text-to-Speech) providers."""
+
+from core.config import TTSConfig
+
+from .base import BaseTTS
+from .cartesia import CartesiaTTS
+
+__all__ = ["BaseTTS", "CartesiaTTS", "get_tts_from_config"]
+
+
+def get_tts_from_config(config: TTSConfig) -> BaseTTS:
+    """
+    Instantiate TTS provider from config.
+
+    Args:
+        config: TTSConfig object with provider name and settings
+
+    Returns:
+        Initialized BaseTTS provider instance
+
+    Raises:
+        ValueError: If provider name is not recognized
+    """
+    provider_name = config.provider.lower()
+
+    if provider_name == "cartesia":
+        return CartesiaTTS(
+            api_key=config.api_key,
+            voice_id=config.voice_id,  # None will use DEFAULT_VOICE_ID in CartesiaTTS
+            model=config.model or "sonic-3",
+            speed=config.speed,
+            **config.extra,
+        )
+    else:
+        raise ValueError(
+            f"Unknown TTS provider: {config.provider}. " f"Supported: cartesia"
+        )