agent-cli 0.70.5__py3-none-any.whl

agent_cli/server/streaming.py

@@ -0,0 +1,67 @@
+"""Core streaming types for subprocess-based audio generation."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+    from multiprocessing import Queue
+
+
+@dataclass
+class StreamChunk:
+    """A chunk of streamed data from a subprocess worker."""
+
+    chunk_type: Literal["data", "error", "done"]
+    payload: bytes | str | None = None
+    metadata: dict[str, Any] | None = None
+
+
+class AsyncQueueReader:
+    """Async iterator over multiprocessing.Queue with timeout handling."""
+
+    def __init__(self, queue: Queue, *, timeout: float = 30.0) -> None:
+        """Initialize reader with queue and timeout."""
+        self._queue = queue
+        self._timeout = timeout
+
+    def __aiter__(self) -> AsyncIterator[StreamChunk]:
+        """Return async iterator."""
+        return self
+
+    async def __anext__(self) -> StreamChunk:
+        """Get the next chunk from the queue."""
+        loop = asyncio.get_running_loop()
+        try:
+            chunk = await asyncio.wait_for(
+                loop.run_in_executor(None, self._queue.get),
+                timeout=self._timeout,
+            )
+        except TimeoutError as e:
+            msg = f"Queue read timeout after {self._timeout}s"
+            raise TimeoutError(msg) from e
+        return chunk
+
+
+class QueueWriter:
+    """Helper for subprocess to send chunks/errors/done sentinel."""
+
+    def __init__(self, queue: Queue) -> None:
+        """Initialize writer with queue."""
+        self._queue = queue
+
+    def send_data(self, data: bytes, metadata: dict[str, Any] | None = None) -> None:
+        """Send a data chunk."""
+        self._queue.put(StreamChunk("data", data, metadata))
+
+    def send_error(self, error: str | Exception) -> None:
+        """Send an error chunk."""
+        error_msg = str(error) if isinstance(error, Exception) else error
+        self._queue.put(StreamChunk("error", error_msg))
+
+    def send_done(self, metadata: dict[str, Any] | None = None) -> None:
+        """Send the done sentinel."""
+        self._queue.put(StreamChunk("done", metadata=metadata))

agent_cli/server/tts/__init__.py

@@ -0,0 +1,3 @@
+"""TTS server module with TTL-based model management."""
+
+from __future__ import annotations

agent_cli/server/tts/api.py

@@ -0,0 +1,335 @@
+"""FastAPI application for TTS server with OpenAI-compatible API."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Annotated, Literal
+
+from fastapi import FastAPI, HTTPException, Query
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+from agent_cli import constants
+from agent_cli.core.audio_format import check_ffmpeg_available, convert_to_mp3
+from agent_cli.server.common import configure_app, create_lifespan
+from agent_cli.server.tts.backends.base import InvalidTextError
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from agent_cli.server.tts.model_registry import TTSModelRegistry
+
+logger = logging.getLogger(__name__)
+
+
+def _format_audio_response(
+    audio: bytes,
+    response_format: str,
+    sample_rate: int,
+    sample_width: int,
+    channels: int,
+) -> StreamingResponse:
+    """Format audio data as a streaming response."""
+    if response_format == "wav":
+        return StreamingResponse(iter([audio]), media_type="audio/wav")
+
+    if response_format == "pcm":
+        pcm_data = (
+            audio[constants.WAV_HEADER_SIZE :] if len(audio) > constants.WAV_HEADER_SIZE else audio
+        )
+        return StreamingResponse(
+            iter([pcm_data]),
+            media_type="audio/pcm",
+            headers={
+                "X-Sample-Rate": str(sample_rate),
+                "X-Sample-Width": str(sample_width),
+                "X-Channels": str(channels),
+            },
+        )
+
+    if response_format == "mp3":
+        if not check_ffmpeg_available():
+            raise HTTPException(
+                status_code=422,
+                detail="MP3 format requires ffmpeg to be installed",
+            )
+        try:
+            mp3_data = convert_to_mp3(audio, input_format="wav")
+        except RuntimeError as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+        return StreamingResponse(iter([mp3_data]), media_type="audio/mpeg")
+
+    # Unreachable due to early validation
+    msg = f"Unsupported response_format: {response_format}"
+    raise HTTPException(status_code=422, detail=msg)  # pragma: no cover
+
+
+# --- Pydantic Models ---
+
+
+class ModelStatusResponse(BaseModel):
+    """Status of a single model."""
+
+    name: str
+    loaded: bool
+    device: str | None
+    ttl_seconds: int
+    ttl_remaining: float | None
+    active_requests: int
+    # Stats
+    load_count: int
+    unload_count: int
+    total_requests: int
+    total_characters: int
+    total_audio_seconds: float
+    total_synthesis_seconds: float
+    last_load_time: float | None
+    last_request_time: float | None
+    load_duration_seconds: float | None
+
+
+class HealthResponse(BaseModel):
+    """Health check response."""
+
+    status: str
+    models: list[ModelStatusResponse]
+
+
+class UnloadResponse(BaseModel):
+    """Response from model unload request."""
+
+    status: str
+    model: str
+    was_loaded: bool
+
+
+class SpeechRequest(BaseModel):
+    """Request body for JSON speech synthesis endpoint."""
+
+    input: str
+    model: str = "tts-1"
+    voice: str = "alloy"
+    response_format: Literal["mp3", "wav", "pcm"] = "mp3"
+    speed: float = 1.0
+    stream_format: Literal["audio"] | None = None
+
+
+class VoiceInfo(BaseModel):
+    """Information about an available voice."""
+
+    voice_id: str
+    name: str
+    description: str
+    preview_url: str | None = None
+    labels: dict[str, str] | None = None
+
+
+class VoicesResponse(BaseModel):
+    """Response containing available voices."""
+
+    voices: list[VoiceInfo]
+
+
+# --- App Factory ---
+
+
+def create_app(
+    registry: TTSModelRegistry,
+    *,
+    enable_wyoming: bool = True,
+    wyoming_uri: str = "tcp://0.0.0.0:10200",
+) -> FastAPI:
+    """Create the FastAPI application.
+
+    Args:
+        registry: The model registry to use.
+        enable_wyoming: Whether to start Wyoming server.
+        wyoming_uri: URI for Wyoming server.
+
+    Returns:
+        Configured FastAPI application.
+
+    """
+    lifespan = create_lifespan(
+        registry,
+        wyoming_handler_module="agent_cli.server.tts.wyoming_handler",
+        enable_wyoming=enable_wyoming,
+        wyoming_uri=wyoming_uri,
+    )
+
+    app = FastAPI(
+        title="TTS Server",
+        description="OpenAI-compatible TTS server with TTL-based model unloading",
+        version="1.0.0",
+        lifespan=lifespan,
+    )
+
+    configure_app(app)
+
+    # --- Health & Status Endpoints ---
+
+    @app.get("/health", response_model=HealthResponse)
+    async def health_check() -> HealthResponse:
+        """Health check endpoint."""
+        models = [
+            ModelStatusResponse(
+                name=s.name,
+                loaded=s.loaded,
+                device=s.device,
+                ttl_seconds=s.ttl_seconds,
+                ttl_remaining=s.ttl_remaining,
+                active_requests=s.active_requests,
+                load_count=s.load_count,
+                unload_count=s.unload_count,
+                total_requests=s.total_requests,
+                total_characters=int(s.extra.get("total_characters", 0.0)),
+                total_audio_seconds=s.total_audio_seconds,
+                total_synthesis_seconds=s.extra.get("total_synthesis_seconds", 0.0),
+                last_load_time=s.last_load_time,
+                last_request_time=s.last_request_time,
+                load_duration_seconds=s.load_duration_seconds,
+            )
+            for s in registry.list_status()
+        ]
+        return HealthResponse(status="healthy", models=models)
+
+    @app.post("/v1/model/unload", response_model=UnloadResponse)
+    async def unload_model(
+        model: Annotated[str | None, Query(description="Model to unload")] = None,
+    ) -> UnloadResponse:
+        """Manually unload a model from memory."""
+        try:
+            manager = registry.get_manager(model)
+            was_loaded = await manager.unload()
+            return UnloadResponse(
+                status="success",
+                model=manager.config.model_name,
+                was_loaded=was_loaded,
+            )
+        except ValueError as e:
+            raise HTTPException(status_code=404, detail=str(e)) from e
+
+    @app.get("/v1/voices", response_model=VoicesResponse)
+    async def list_voices() -> VoicesResponse:
+        """List available voices (models).
+
+        For Piper TTS, each model IS a voice. This endpoint returns
+        the list of registered models as available voices.
+        """
+        voices = [
+            VoiceInfo(
+                voice_id=s.name,
+                name=s.name,
+                description=f"Piper TTS voice: {s.name}",
+                labels={"language": s.name.split("_")[0] if "_" in s.name else "en"},
+            )
+            for s in registry.list_status()
+        ]
+        return VoicesResponse(voices=voices)
+
+    # --- OpenAI-Compatible TTS Endpoint ---
+
+    async def _synthesize(
+        input_text: str,
+        model: str,
+        voice: str,
+        response_format: str,
+        speed: float,
+        stream_format: str | None,
+    ) -> StreamingResponse:
+        """Core synthesis logic shared by JSON and form endpoints."""
+        # Resolve model name - "tts-1" and "tts-1-hd" are OpenAI's model names
+        model_name = None if model in ("tts-1", "tts-1-hd") else model
+
+        try:
+            manager = registry.get_manager(model_name)
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e)) from e
+
+        if not input_text.strip():
+            raise HTTPException(status_code=400, detail="Input text cannot be empty")
+
+        # Clamp speed to valid range
+        speed = max(0.25, min(4.0, speed))
+
+        # Handle streaming mode (OpenAI uses stream_format=audio with response_format=pcm)
+        if stream_format is not None:
+            if stream_format != "audio":
+                raise HTTPException(
+                    status_code=422,
+                    detail="Only 'audio' stream_format is supported",
+                )
+            if response_format != "pcm":
+                raise HTTPException(
+                    status_code=422,
+                    detail="Streaming requires response_format=pcm",
+                )
+            if not manager.supports_streaming:
+                raise HTTPException(
+                    status_code=422,
+                    detail="This model does not support streaming synthesis",
+                )
+
+            async def generate_audio() -> AsyncIterator[bytes]:
+                async for chunk in manager.synthesize_stream(
+                    input_text,
+                    voice=voice,
+                    speed=speed,
+                ):
+                    yield chunk
+
+            return StreamingResponse(
+                generate_audio(),
+                media_type="audio/pcm",
+                headers={
+                    "X-Sample-Rate": str(constants.KOKORO_DEFAULT_SAMPLE_RATE),
+                    "X-Sample-Width": "2",
+                    "X-Channels": "1",
+                },
+            )
+
+        # Non-streaming mode: validate format and synthesize complete audio
+        valid_formats = ("wav", "pcm", "mp3")
+        if response_format not in valid_formats:
+            raise HTTPException(
+                status_code=422,
+                detail=f"Unsupported response_format: {response_format}. Supported: {', '.join(valid_formats)}",
+            )
+
+        try:
+            result = await manager.synthesize(
+                input_text,
+                voice=voice,
+                speed=speed,
+            )
+        except InvalidTextError as e:
+            raise HTTPException(status_code=400, detail=str(e)) from e
+        except Exception as e:
+            logger.exception("Synthesis failed")
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+        return _format_audio_response(
+            result.audio,
+            response_format,
+            result.sample_rate,
+            result.sample_width,
+            result.channels,
+        )
+
+    @app.post("/v1/audio/speech")
+    async def synthesize_speech(request: SpeechRequest) -> StreamingResponse:
+        """OpenAI-compatible text-to-speech endpoint.
+
+        Accepts JSON body with input, model, voice, response_format, speed,
+        and optional stream_format parameters.
+        """
+        return await _synthesize(
+            input_text=request.input,
+            model=request.model,
+            voice=request.voice,
+            response_format=request.response_format,
+            speed=request.speed,
+            stream_format=request.stream_format,
+        )
+
+    return app

agent_cli/server/tts/backends/__init__.py

@@ -0,0 +1,82 @@
+"""TTS backend factory with platform auto-detection."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from agent_cli.server.tts.backends.base import TTSBackend
+
+from agent_cli.server.tts.backends.base import (
+    BackendConfig,
+    SynthesisResult,
+    has_gpu,
+)
+
+logger = logging.getLogger(__name__)
+
+BackendType = Literal["piper", "kokoro", "auto"]
+
+
+def detect_backend() -> Literal["piper", "kokoro"]:
+    """Detect the best backend for the current platform.
+
+    Returns:
+        "kokoro" if GPU is available and kokoro is installed, otherwise "piper".
+
+    """
+    if has_gpu():
+        try:
+            import kokoro  # noqa: F401, PLC0415
+
+            return "kokoro"
+        except ImportError:
+            pass
+    return "piper"
+
+
+def create_backend(
+    config: BackendConfig,
+    backend_type: BackendType = "auto",
+) -> TTSBackend:
+    """Create a TTS backend instance.
+
+    Args:
+        config: Backend configuration.
+        backend_type: Backend to use, or "auto" for platform detection.
+
+    Returns:
+        Configured TTSBackend instance.
+
+    Raises:
+        ImportError: If the required backend package is not installed.
+        ValueError: If an unknown backend type is specified.
+
+    """
+    if backend_type == "auto":
+        backend_type = detect_backend()
+
+    logger.debug("Creating %s backend for model %s", backend_type, config.model_name)
+
+    if backend_type == "piper":
+        from agent_cli.server.tts.backends.piper import PiperBackend  # noqa: PLC0415
+
+        return PiperBackend(config)
+
+    if backend_type == "kokoro":
+        from agent_cli.server.tts.backends.kokoro import KokoroBackend  # noqa: PLC0415
+
+        return KokoroBackend(config)
+
+    msg = f"Unknown backend type: {backend_type}"
+    raise ValueError(msg)
+
+
+__all__ = [
+    "BackendConfig",
+    "BackendType",
+    "SynthesisResult",
+    "create_backend",
+    "detect_backend",
+]

agent_cli/server/tts/backends/base.py

@@ -0,0 +1,139 @@
+"""Base types and protocol for TTS backends."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+
+def get_torch_device() -> str:
+    """Detect the best available PyTorch device."""
+    try:
+        import torch  # noqa: PLC0415
+
+        if torch.cuda.is_available():
+            return "cuda"
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return "mps"
+    except ImportError:
+        pass
+    return "cpu"
+
+
+def has_gpu() -> bool:
+    """Check if a GPU (CUDA or MPS) is available."""
+    return get_torch_device() in ("cuda", "mps")
+
+
+def get_backend_cache_dir(backend_name: str) -> Path:
+    """Get default cache directory for a TTS backend."""
+    cache_dir = Path.home() / ".cache" / backend_name
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+@dataclass
+class SynthesisResult:
+    """Result of a synthesis operation."""
+
+    audio: bytes
+    sample_rate: int
+    sample_width: int
+    channels: int
+    duration: float
+
+
+@dataclass
+class BackendConfig:
+    """Configuration for a TTS backend."""
+
+    model_name: str
+    device: str = "auto"
+    cache_dir: Path | None = None
+
+
+class InvalidTextError(ValueError):
+    """Raised when the input text is invalid or unsupported."""
+
+
+@runtime_checkable
+class TTSBackend(Protocol):
+    """Protocol for TTS synthesis backends.
+
+    Backends handle model loading, unloading, and synthesis.
+    The ModelManager handles TTL, stats, and lifecycle.
+    """
+
+    @property
+    def is_loaded(self) -> bool:
+        """Check if the model is currently loaded."""
+        ...
+
+    @property
+    def device(self) -> str | None:
+        """Get the device the model is loaded on, or None if not loaded."""
+        ...
+
+    async def load(self) -> float:
+        """Load the model into memory.
+
+        Returns:
+            Load duration in seconds.
+
+        """
+        ...
+
+    async def unload(self) -> None:
+        """Unload the model and free memory."""
+        ...
+
+    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, uses model default if not specified).
+            speed: Speech speed multiplier (0.25 to 4.0).
+
+        Returns:
+            SynthesisResult with audio data and metadata.
+
+        """
+        ...
+
+    @property
+    def supports_streaming(self) -> bool:
+        """Check if backend supports streaming synthesis."""
+        return False
+
+    def synthesize_stream(
+        self,
+        text: str,
+        *,
+        voice: str | None = None,
+        speed: float = 1.0,
+    ) -> AsyncIterator[bytes]:
+        """Stream synthesized audio chunks as they are generated.
+
+        Implementations should be async generators (async def with yield).
+
+        Args:
+            text: Text to synthesize.
+            voice: Voice to use (optional).
+            speed: Speech speed multiplier (0.25 to 4.0).
+
+        Yields:
+            Raw PCM audio chunks (int16, mono).
+
+        """
+        ...