agent-cli 0.70.5__py3-none-any.whl

agent_cli/server/tts/model_manager.py

@@ -0,0 +1,201 @@
+"""TTS model manager with TTL-based unloading."""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from agent_cli import constants
+from agent_cli.server.model_manager import ModelConfig, ModelManager, ModelStats
+from agent_cli.server.tts.backends import (
+    BackendConfig,
+    BackendType,
+    SynthesisResult,
+    create_backend,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from agent_cli.server.tts.backends.base import TTSBackend
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TTSModelConfig(ModelConfig):
+    """Configuration for a TTS model."""
+
+    backend_type: BackendType = "auto"
+
+
+class TTSModelManager:
+    """Manages a TTS model with TTL-based unloading.
+
+    Thin wrapper around ModelManager that adds the synthesize() method.
+    """
+
+    def __init__(self, config: TTSModelConfig) -> None:
+        """Initialize the TTS model manager."""
+        self.config = config
+        backend = create_backend(
+            BackendConfig(
+                model_name=config.model_name,
+                device=config.device,
+                cache_dir=config.cache_dir,
+            ),
+            backend_type=config.backend_type,
+        )
+        self._manager = ModelManager(backend, config)
+
+    @property
+    def stats(self) -> ModelStats:
+        """Get the model statistics."""
+        return self._manager.stats
+
+    @property
+    def is_loaded(self) -> bool:
+        """Check if the model is currently loaded."""
+        return self._manager.is_loaded
+
+    @property
+    def device(self) -> str | None:
+        """Get the device the model is loaded on."""
+        return self._manager.device
+
+    @property
+    def active_requests(self) -> int:
+        """Get the number of active requests."""
+        return self._manager.active_requests
+
+    @property
+    def ttl_remaining(self) -> float | None:
+        """Get seconds remaining before model unloads."""
+        return self._manager.ttl_remaining
+
+    async def start(self) -> None:
+        """Start the TTL unload watcher."""
+        await self._manager.start()
+
+    async def stop(self) -> None:
+        """Stop the manager and unload the model."""
+        await self._manager.stop()
+
+    async def get_model(self) -> TTSBackend:
+        """Get the backend, loading it if necessary."""
+        return await self._manager.get_model()
+
+    async def unload(self) -> bool:
+        """Unload the model from memory."""
+        return await self._manager.unload()
+
+    def _update_stats(self, text: str, synthesis_duration: float) -> None:
+        """Update synthesis statistics."""
+        stats = self._manager.stats
+        stats.total_requests += 1
+        stats.total_processing_seconds += synthesis_duration
+        stats.extra["total_characters"] = stats.extra.get("total_characters", 0.0) + len(text)
+        stats.extra["total_synthesis_seconds"] = (
+            stats.extra.get("total_synthesis_seconds", 0.0) + synthesis_duration
+        )
+
+    async def synthesize(
+        self,
+        text: str,
+        *,
+        voice: str | None = None,
+        speed: float = 1.0,
+    ) -> SynthesisResult:
+        """Synthesize text to audio.
+
+        Args:
+            text: Text to synthesize.
+            voice: Voice to use (optional).
+            speed: Speech speed multiplier (0.25 to 4.0).
+
+        Returns:
+            SynthesisResult with audio data and metadata.
+
+        """
+        start_time = time.time()
+
+        async with self._manager.request():
+            backend: TTSBackend = self._manager.backend  # type: ignore[assignment]
+            result = await backend.synthesize(
+                text,
+                voice=voice,
+                speed=speed,
+            )
+
+        synthesis_duration = time.time() - start_time
+
+        self._update_stats(text, synthesis_duration)
+        self._manager.stats.total_audio_seconds += result.duration
+
+        logger.debug(
+            "Synthesized %d chars to %.1fs audio in %.2fs (model=%s)",
+            len(text),
+            result.duration,
+            synthesis_duration,
+            self.config.model_name,
+        )
+
+        return result
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Check if the backend supports streaming synthesis."""
+        backend: TTSBackend = self._manager.backend  # type: ignore[assignment]
+        return backend.supports_streaming
+
+    async def synthesize_stream(
+        self,
+        text: str,
+        *,
+        voice: str | None = None,
+        speed: float = 1.0,
+    ) -> AsyncIterator[bytes]:
+        """Stream synthesized audio chunks as they are generated."""
+        start_time = time.time()
+        chunk_count = 0
+        total_bytes = 0
+
+        async with self._manager.request():
+            backend: TTSBackend = self._manager.backend  # type: ignore[assignment]
+
+            if not backend.supports_streaming:
+                msg = "Backend does not support streaming"
+                raise RuntimeError(msg)
+
+            async for chunk in backend.synthesize_stream(
+                text,
+                voice=voice,
+                speed=speed,
+            ):
+                chunk_count += 1
+                total_bytes += len(chunk)
+                yield chunk
+
+        synthesis_duration = time.time() - start_time
+
+        # Calculate audio duration from PCM bytes (16-bit mono)
+        bytes_per_second = constants.KOKORO_DEFAULT_SAMPLE_RATE * 2  # 2 bytes per sample
+        audio_seconds = total_bytes / bytes_per_second
+
+        self._update_stats(text, synthesis_duration)
+        self._manager.stats.total_audio_seconds += audio_seconds
+        self._manager.stats.extra["streaming_requests"] = (
+            self._manager.stats.extra.get("streaming_requests", 0) + 1
+        )
+
+        logger.debug(
+            "Streamed %d chars to %.1fs audio in %d chunks (%d bytes) in %.2fs (model=%s)",
+            len(text),
+            audio_seconds,
+            chunk_count,
+            total_bytes,
+            synthesis_duration,
+            self.config.model_name,
+        )

agent_cli/server/tts/model_registry.py

@@ -0,0 +1,28 @@
+"""Registry for managing multiple TTS models."""
+
+from __future__ import annotations
+
+from agent_cli.server.model_registry import ModelRegistry
+from agent_cli.server.tts.model_manager import TTSModelConfig, TTSModelManager
+
+
+def create_tts_registry(
+    default_model: str | None = None,
+) -> ModelRegistry[TTSModelManager, TTSModelConfig]:
+    """Create a TTS model registry.
+
+    Args:
+        default_model: Name of the default model to use when not specified.
+
+    Returns:
+        Configured ModelRegistry for TTS models.
+
+    """
+    return ModelRegistry(
+        manager_factory=TTSModelManager,
+        default_model=default_model,
+    )
+
+
+# Alias for type hints
+TTSModelRegistry = ModelRegistry[TTSModelManager, TTSModelConfig]

agent_cli/server/tts/wyoming_handler.py

@@ -0,0 +1,249 @@
+"""Wyoming protocol handler for TTS server."""
+
+from __future__ import annotations
+
+import logging
+from functools import partial
+from typing import TYPE_CHECKING
+
+from wyoming.audio import AudioChunk, AudioStart, AudioStop
+from wyoming.info import Attribution, Describe, Info, TtsProgram, TtsVoice
+from wyoming.server import AsyncEventHandler, AsyncServer
+from wyoming.tts import Synthesize
+
+from agent_cli import constants
+
+if TYPE_CHECKING:
+    from wyoming.event import Event
+
+    from agent_cli.server.tts.model_manager import TTSModelManager
+    from agent_cli.server.tts.model_registry import TTSModelRegistry
+
+logger = logging.getLogger(__name__)
+
+
+class WyomingTTSHandler(AsyncEventHandler):
+    """Wyoming event handler for TTS.
+
+    Handles the Wyoming protocol for TTS (Text-to-Speech):
+    - Receives Synthesize event with text
+    - Synthesizes audio
+    - Returns AudioStart, AudioChunk(s), AudioStop
+    """
+
+    def __init__(
+        self,
+        registry: TTSModelRegistry,
+        *args: object,
+        **kwargs: object,
+    ) -> None:
+        """Initialize the handler.
+
+        Args:
+            registry: Model registry for getting TTS models.
+            *args: Passed to parent class.
+            **kwargs: Passed to parent class.
+
+        """
+        super().__init__(*args, **kwargs)
+        self._registry = registry
+
+    async def handle_event(self, event: Event) -> bool:
+        """Handle a Wyoming event.
+
+        Args:
+            event: The event to handle.
+
+        Returns:
+            True to continue processing events, False to stop.
+
+        """
+        if Synthesize.is_type(event.type):
+            return await self._handle_synthesize(event)
+
+        if Describe.is_type(event.type):
+            return await self._handle_describe()
+
+        return True
+
+    async def _handle_synthesize(self, event: Event) -> bool:
+        """Handle synthesize event - synthesize text to audio."""
+        synthesize = Synthesize.from_event(event)
+        text = synthesize.text
+
+        logger.debug("Synthesize: %s", text[:100] if text else "")
+
+        if not text:
+            logger.warning("Empty text received")
+            # Send empty audio response
+            await self.write_event(
+                AudioStart(
+                    rate=constants.PIPER_DEFAULT_SAMPLE_RATE,
+                    width=2,
+                    channels=1,
+                ).event(),
+            )
+            await self.write_event(AudioStop().event())
+            return False
+
+        try:
+            manager = self._registry.get_manager()
+
+            if manager.supports_streaming:
+                await self._synthesize_streaming(manager, text, synthesize.voice)
+            else:
+                await self._synthesize_complete(manager, text, synthesize.voice)
+
+        except Exception:
+            logger.exception("Wyoming synthesis failed")
+            # Send empty audio on error
+            await self.write_event(
+                AudioStart(
+                    rate=constants.PIPER_DEFAULT_SAMPLE_RATE,
+                    width=2,
+                    channels=1,
+                ).event(),
+            )
+            await self.write_event(AudioStop().event())
+
+        return False
+
+    async def _synthesize_streaming(
+        self,
+        manager: TTSModelManager,
+        text: str,
+        voice: str | None,
+    ) -> None:
+        """Stream audio chunks as they're generated."""
+        sample_rate = constants.KOKORO_DEFAULT_SAMPLE_RATE
+
+        # Send audio start
+        await self.write_event(
+            AudioStart(rate=sample_rate, width=2, channels=1).event(),
+        )
+
+        chunk_count = 0
+        total_bytes = 0
+        async for chunk in manager.synthesize_stream(text, voice=voice, speed=1.0):
+            await self.write_event(
+                AudioChunk(audio=chunk, rate=sample_rate, width=2, channels=1).event(),
+            )
+            chunk_count += 1
+            total_bytes += len(chunk)
+
+        await self.write_event(AudioStop().event())
+
+        # Calculate duration from PCM bytes (16-bit mono)
+        duration = total_bytes / (sample_rate * 2)
+        logger.info(
+            "Wyoming streaming synthesis: %d chars -> %.1fs audio in %d chunks",
+            len(text),
+            duration,
+            chunk_count,
+        )
+
+    async def _synthesize_complete(
+        self,
+        manager: TTSModelManager,
+        text: str,
+        voice: str | None,
+    ) -> None:
+        """Synthesize complete audio then send in chunks."""
+        result = await manager.synthesize(text, voice=voice, speed=1.0)
+
+        # Send audio start
+        await self.write_event(
+            AudioStart(
+                rate=result.sample_rate,
+                width=result.sample_width,
+                channels=result.channels,
+            ).event(),
+        )
+
+        # Send audio data - skip WAV header to get raw PCM
+        pcm_data = (
+            result.audio[constants.WAV_HEADER_SIZE :]
+            if len(result.audio) > constants.WAV_HEADER_SIZE
+            else result.audio
+        )
+
+        # Send in chunks
+        chunk_size = 4096
+        for i in range(0, len(pcm_data), chunk_size):
+            chunk = pcm_data[i : i + chunk_size]
+            await self.write_event(
+                AudioChunk(
+                    audio=chunk,
+                    rate=result.sample_rate,
+                    width=result.sample_width,
+                    channels=result.channels,
+                ).event(),
+            )
+
+        await self.write_event(AudioStop().event())
+
+        logger.info(
+            "Wyoming synthesis: %d chars -> %.1fs audio",
+            len(text),
+            result.duration,
+        )
+
+    async def _handle_describe(self) -> bool:
+        """Handle describe event - return server capabilities."""
+        logger.debug("Describe event")
+
+        # Get list of available models as voices
+        voices = [
+            TtsVoice(
+                name=status.name,
+                description=f"Piper TTS {status.name}",
+                attribution=Attribution(
+                    name="Piper",
+                    url="https://github.com/rhasspy/piper",
+                ),
+                installed=True,
+                # Extract language from model name (e.g., "en_US-lessac-medium" -> "en")
+                languages=[status.name.split("_")[0] if "_" in status.name else "en"],
+                version="1.0",
+            )
+            for status in self._registry.list_status()
+        ]
+
+        await self.write_event(
+            Info(
+                tts=[
+                    TtsProgram(
+                        name="agent-cli-tts",
+                        description="Agent CLI TTS Server with TTL-based model unloading",
+                        attribution=Attribution(
+                            name="agent-cli",
+                            url="https://github.com/basnijholt/agent-cli",
+                        ),
+                        installed=True,
+                        version="1.0",
+                        voices=voices,
+                    ),
+                ],
+            ).event(),
+        )
+        return True
+
+
+async def start_wyoming_server(
+    registry: TTSModelRegistry,
+    uri: str = "tcp://0.0.0.0:10200",
+) -> None:
+    """Start the Wyoming TTS server.
+
+    Args:
+        registry: Model registry for synthesis.
+        uri: URI to bind the server to (e.g., "tcp://0.0.0.0:10200").
+
+    """
+    server = AsyncServer.from_uri(uri)
+    logger.debug("Wyoming TTS server listening on %s", uri)
+
+    # Create handler factory with registry
+    handler_factory = partial(WyomingTTSHandler, registry)
+
+    await server.run(handler_factory)

agent_cli/server/whisper/__init__.py

@@ -0,0 +1,3 @@
+"""Whisper ASR server with TTL-based model unloading."""
+
+from __future__ import annotations