ff-aitoolkit 0.2.0__py3-none-any.whl

aitoolkit/retry.py

@@ -0,0 +1,51 @@
+"""Minimal async retry helper: exponential backoff + jitter for transient faults.
+
+Intentionally dependency-free (no tenacity) and small. The caller decides what
+is transient by passing the exception types in ``retry_on`` — nothing else is
+retried, so non-retriable errors (e.g. a 4xx / validation) surface immediately.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from typing import Awaitable, Callable, Tuple, Type, TypeVar
+
+from loguru import logger
+
+T = TypeVar("T")
+
+
+async def retry_async(
+    fn: Callable[[], Awaitable[T]],
+    *,
+    attempts: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 10.0,
+    retry_on: Tuple[Type[BaseException], ...] = (Exception,),
+    label: str = "request",
+) -> T:
+    """Call ``fn`` (an async, no-arg callable), retrying transient failures.
+
+    Backoff is exponential (``base_delay * 2**n``, capped at ``max_delay``) with
+    jitter (50–100% of the computed delay) to avoid synchronized retry storms.
+    Only exceptions in ``retry_on`` are retried; anything else propagates at once.
+    The final failure is re-raised after ``attempts`` tries.
+    """
+    last_exc: BaseException | None = None
+    for attempt in range(1, attempts + 1):
+        try:
+            return await fn()
+        except retry_on as exc:
+            last_exc = exc
+            if attempt >= attempts:
+                break
+            delay = min(max_delay, base_delay * (2 ** (attempt - 1)))
+            delay *= 0.5 + random.random() / 2  # jitter: 50–100% of the delay
+            logger.warning(
+                f"{label}: attempt {attempt}/{attempts} failed "
+                f"({type(exc).__name__}: {exc}); retrying in {delay:.1f}s"
+            )
+            await asyncio.sleep(delay)
+    assert last_exc is not None
+    raise last_exc

aitoolkit/stt/__init__.py

@@ -0,0 +1,5 @@
+"""Speech-to-text capability — OpenAI-compatible faster-whisper server."""
+
+from aitoolkit.stt.client import STTClient, get_stt_client
+
+__all__ = ["STTClient", "get_stt_client"]

aitoolkit/stt/client.py

@@ -0,0 +1,147 @@
+"""Speech-to-text client backed by an OpenAI-compatible transcription endpoint.
+
+Targets the self-hosted faster-whisper service which exposes
+``/v1/audio/transcriptions``. Replaces any local/in-process Whisper so no model
+weights are loaded inside the application.
+"""
+
+from __future__ import annotations
+
+import io
+from functools import lru_cache
+from pathlib import Path
+from typing import BinaryIO, Optional, Union
+
+from loguru import logger
+from openai import AsyncOpenAI, OpenAI
+
+from aitoolkit.config import AIToolkitSettings, get_settings
+from aitoolkit.exceptions import STTError
+from aitoolkit.types import TranscriptionResult
+
+AudioInput = Union[str, Path, bytes, BinaryIO]
+
+
+class STTClient:
+    """Transcribe audio via an OpenAI-compatible ``audio.transcriptions`` API."""
+
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        language: Optional[str] = None,
+        timeout: Optional[float] = None,
+        settings: Optional[AIToolkitSettings] = None,
+    ) -> None:
+        settings = settings or get_settings()
+        self.model = model or settings.stt_model
+        self.default_language = language or settings.stt_language
+        self._base_url = base_url or settings.stt_base_url
+        self._api_key = api_key or settings.stt_api_key
+        self._timeout = timeout if timeout is not None else settings.stt_timeout
+
+        self._aclient = AsyncOpenAI(
+            base_url=self._base_url, api_key=self._api_key, timeout=self._timeout
+        )
+        self._sclient: Optional[OpenAI] = None
+        logger.info(
+            f"STTClient ready (model={self.model}, base_url={self._base_url})"
+        )
+
+    @property
+    def sync_client(self) -> OpenAI:
+        if self._sclient is None:
+            self._sclient = OpenAI(
+                base_url=self._base_url, api_key=self._api_key, timeout=self._timeout
+            )
+        return self._sclient
+
+    @staticmethod
+    def _to_file(audio: AudioInput):
+        """Normalize various audio inputs into something the SDK accepts."""
+        if isinstance(audio, (str, Path)):
+            path = Path(audio)
+            return (path.name, path.read_bytes())
+        if isinstance(audio, bytes):
+            return ("audio.wav", audio)
+        if isinstance(audio, io.IOBase) or hasattr(audio, "read"):
+            data = audio.read()
+            name = getattr(audio, "name", "audio.wav")
+            return (Path(str(name)).name, data)
+        raise STTError(f"unsupported audio input type: {type(audio)!r}")
+
+    async def transcribe(
+        self,
+        audio: AudioInput,
+        *,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        **kwargs,
+    ) -> TranscriptionResult:
+        """Transcribe audio and return text plus optional metadata.
+
+        ``response_format`` defaults to ``"json"`` (text only). Pass
+        ``"verbose_json"`` to also populate ``language`` and ``duration`` on the
+        returned :class:`TranscriptionResult`.
+        """
+        file_arg = self._to_file(audio)
+        try:
+            resp = await self._aclient.audio.transcriptions.create(
+                file=file_arg,
+                model=self.model,
+                language=language or self.default_language,
+                prompt=prompt,
+                response_format=response_format,
+                **kwargs,
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise STTError(f"transcription failed: {exc}") from exc
+        return self._to_result(resp, language or self.default_language)
+
+    def transcribe_sync(
+        self,
+        audio: AudioInput,
+        *,
+        language: Optional[str] = None,
+        prompt: Optional[str] = None,
+        response_format: str = "json",
+        **kwargs,
+    ) -> TranscriptionResult:
+        """Synchronous counterpart of :meth:`transcribe`."""
+        file_arg = self._to_file(audio)
+        try:
+            resp = self.sync_client.audio.transcriptions.create(
+                file=file_arg,
+                model=self.model,
+                language=language or self.default_language,
+                prompt=prompt,
+                response_format=response_format,
+                **kwargs,
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise STTError(f"transcription failed: {exc}") from exc
+        return self._to_result(resp, language or self.default_language)
+
+    @staticmethod
+    def _to_result(resp, language: Optional[str]) -> TranscriptionResult:
+        text = getattr(resp, "text", None)
+        if text is None and isinstance(resp, dict):
+            text = resp.get("text", "")
+        return TranscriptionResult(
+            text=text or "",
+            language=getattr(resp, "language", None) or language,
+            duration=getattr(resp, "duration", None),
+        )
+
+    async def aclose(self) -> None:
+        await self._aclient.close()
+        if self._sclient is not None:
+            self._sclient.close()
+
+
+@lru_cache(maxsize=1)
+def get_stt_client() -> STTClient:
+    """Return the process-wide STT client singleton."""
+    return STTClient()

aitoolkit/tts/__init__.py

@@ -0,0 +1,10 @@
+"""Text-to-speech capability — custom self-hosted ``/api/tts`` server.
+
+Works with any TTS service exposing the ``/api/tts`` + ``/api/voices`` contract
+(``POST /api/tts`` returns raw audio bytes; ``GET /api/voices`` lists voices).
+"""
+
+from aitoolkit.tts.audio import concat_wav
+from aitoolkit.tts.client import TTSClient, get_tts_client
+
+__all__ = ["TTSClient", "get_tts_client", "concat_wav"]

aitoolkit/tts/audio.py

@@ -0,0 +1,68 @@
+"""WAV audio helpers — stitch multiple WAV clips into one.
+
+Pure standard library (``wave``); no third-party audio dependencies, keeping the
+core package light. All input clips must share the same format (channels, sample
+width, frame rate) — produce them with a single TTS engine / voice family.
+"""
+
+from __future__ import annotations
+
+import io
+import wave
+from typing import List, Optional, Sequence
+
+from aitoolkit.exceptions import TTSError
+
+
+def concat_wav(segments: Sequence[bytes], *, gap_ms: int = 0) -> bytes:
+    """Concatenate WAV byte clips into a single WAV.
+
+    Args:
+        segments: WAV-encoded audio clips. Empty clips are skipped. All
+            non-empty clips must share channels, sample width and frame rate.
+        gap_ms: silence inserted between consecutive clips, in milliseconds.
+
+    Returns:
+        A single WAV-encoded byte string.
+
+    Raises:
+        TTSError: if there are no usable segments or their formats differ.
+    """
+    clips = [clip for clip in segments if clip]
+    if not clips:
+        raise TTSError("concat_wav: no audio segments to concatenate")
+
+    params: Optional[wave._wave_params] = None
+    base_format: Optional[tuple] = None
+    frames: List[bytes] = []
+
+    for index, clip in enumerate(clips):
+        try:
+            with wave.open(io.BytesIO(clip), "rb") as reader:
+                clip_params = reader.getparams()
+                clip_frames = reader.readframes(clip_params.nframes)
+        except (wave.Error, EOFError) as exc:
+            raise TTSError(f"concat_wav: segment {index} is not valid WAV: {exc}") from exc
+
+        clip_format = (clip_params.nchannels, clip_params.sampwidth, clip_params.framerate)
+        if params is None:
+            params, base_format = clip_params, clip_format
+        elif clip_format != base_format:
+            raise TTSError(
+                f"concat_wav: segment {index} format {clip_format} != {base_format}; "
+                "all clips must share channels/width/rate (synthesize with one engine)"
+            )
+
+        if frames and gap_ms > 0:
+            silent_samples = int(params.framerate * gap_ms / 1000)
+            frames.append(b"\x00" * silent_samples * params.nchannels * params.sampwidth)
+        frames.append(clip_frames)
+
+    assert params is not None  # guaranteed: clips is non-empty
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as writer:
+        writer.setnchannels(params.nchannels)
+        writer.setsampwidth(params.sampwidth)
+        writer.setframerate(params.framerate)
+        writer.writeframes(b"".join(frames))
+    return buffer.getvalue()

aitoolkit/tts/client.py

@@ -0,0 +1,219 @@
+"""Text-to-speech client for the custom self-hosted TTS services.
+
+Unlike LLM/embeddings/STT, the TTS servers are **not** OpenAI-compatible: they
+expose ``POST /api/tts`` (returning raw audio bytes) and ``GET /api/voices``.
+A successful synthesis returns binary audio; errors return JSON ``{"detail": ...}``.
+The request requires either a ``voice_id`` or an ``instruct`` prompt.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from pathlib import Path
+from typing import List, Optional, Sequence, Union
+
+import httpx
+from loguru import logger
+
+from aitoolkit.config import AIToolkitSettings, get_settings
+from aitoolkit.exceptions import TTSError
+from aitoolkit.retry import retry_async
+from aitoolkit.tts.audio import concat_wav
+from aitoolkit.types import DialogueTurn
+
+# HTTP statuses worth retrying (transient). Other 4xx are caller errors.
+_RETRIABLE_STATUS = {408, 425, 429, 500, 502, 503, 504}
+
+
+class _RetriableTTS(Exception):
+    """Internal marker for a transient TTS failure (timeout / 5xx / 429)."""
+
+
+class TTSClient:
+    """Synthesize speech via a custom ``/api/tts`` endpoint."""
+
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        default_voice: Optional[str] = None,
+        timeout: Optional[float] = None,
+        tts_path: str = "/api/tts",
+        voices_path: str = "/api/voices",
+        settings: Optional[AIToolkitSettings] = None,
+    ) -> None:
+        settings = settings or get_settings()
+        self._base_url = (base_url or settings.tts_base_url).rstrip("/")
+        self.default_voice = default_voice or settings.tts_default_voice
+        self._timeout = timeout if timeout is not None else settings.tts_timeout
+        self._tts_path = tts_path
+        self._voices_path = voices_path
+        logger.info(f"TTSClient ready (base_url={self._base_url})")
+
+    async def synthesize(
+        self,
+        text: str,
+        *,
+        voice: Optional[str] = None,
+        language: str = "en",
+        instruct: Optional[str] = None,
+        speed: Optional[float] = None,
+        num_step: Optional[int] = None,
+        ref_text: Optional[str] = None,
+    ) -> bytes:
+        """Synthesize ``text`` and return raw audio bytes.
+
+        Either ``voice`` (resolved against ``default_voice``) or ``instruct``
+        must be provided, matching the server contract.
+        """
+        voice_id = voice or self.default_voice
+        if not voice_id and not instruct:
+            raise TTSError("either 'voice' (voice_id) or 'instruct' must be provided")
+
+        payload: dict = {"text": text, "language": language}
+        if voice_id:
+            payload["voice_id"] = voice_id
+        if instruct is not None:
+            payload["instruct"] = instruct
+        if speed is not None:
+            payload["speed"] = speed
+        if num_step is not None:
+            payload["num_step"] = num_step
+        if ref_text is not None:
+            payload["ref_text"] = ref_text
+
+        url = f"{self._base_url}{self._tts_path}"
+
+        # Fast connect (fail quickly + retry), longer read for the actual synthesis.
+        timeout = httpx.Timeout(self._timeout, connect=5.0)
+
+        async def _attempt() -> bytes:
+            try:
+                async with httpx.AsyncClient(timeout=timeout) as client:
+                    resp = await client.post(url, json=payload)
+            except httpx.TransportError as exc:
+                # Transient transport faults — covers all timeouts (TimeoutException)
+                # and network/connection errors (ConnectError, ReadError, …).
+                raise _RetriableTTS(f"{type(exc).__name__}: {exc}") from exc
+            except httpx.HTTPError as exc:  # other httpx errors — don't retry.
+                raise TTSError(
+                    f"TTS request failed: {type(exc).__name__}: {exc}"
+                ) from exc
+
+            if resp.status_code != 200:
+                detail = self._error_detail(resp)
+                if resp.status_code in _RETRIABLE_STATUS:
+                    raise _RetriableTTS(detail)
+                raise TTSError(detail)  # 4xx caller error — don't retry.
+
+            if not resp.content:
+                raise TTSError("TTS returned an empty audio response")
+            return resp.content
+
+        try:
+            return await retry_async(
+                _attempt, retry_on=(_RetriableTTS,), label="TTS synthesize"
+            )
+        except _RetriableTTS as exc:
+            # Exhausted retries on a transient fault — surface the real reason.
+            raise TTSError(f"TTS request failed after retries: {exc}") from exc
+
+    async def synthesize_to_file(
+        self, text: str, path: Union[str, Path], **kwargs
+    ) -> Path:
+        """Synthesize and write the audio to ``path``; returns the path."""
+        audio = await self.synthesize(text, **kwargs)
+        out = Path(path)
+        out.write_bytes(audio)
+        return out
+
+    async def synthesize_dialogue(
+        self,
+        turns: Sequence[DialogueTurn],
+        *,
+        language: str = "en",
+        gap_ms: int = 300,
+        **kwargs,
+    ) -> bytes:
+        """Synthesize a multi-speaker dialogue into a single WAV.
+
+        Each turn is synthesized with its own ``voice_id`` and the clips are
+        concatenated with a short silent gap between turns. All voices must live
+        on the same TTS engine so the clips share an audio format (see
+        :func:`aitoolkit.tts.audio.concat_wav`).
+
+        Args:
+            turns: ordered ``{"voice_id", "text"}`` turns. Empty texts are skipped.
+            language: language code passed to each synthesis.
+            gap_ms: silence inserted between turns, in milliseconds.
+            **kwargs: forwarded to :meth:`synthesize` (e.g. ``speed``).
+
+        Returns:
+            A single WAV-encoded byte string.
+
+        Raises:
+            TTSError: if no turns have text, or synthesis/concatenation fails.
+        """
+        spoken = [turn for turn in turns if turn.get("text", "").strip()]
+        if not spoken:
+            raise TTSError("synthesize_dialogue: no non-empty turns provided")
+
+        # Each turn already retries transient faults (see ``synthesize``). If a
+        # turn STILL fails, skip it rather than fail the whole dialogue — a
+        # podcast missing one line is far better than no podcast at all. We only
+        # fail if every turn failed.
+        clips: List[bytes] = []
+        failed = 0
+        for index, turn in enumerate(spoken):
+            try:
+                clips.append(
+                    await self.synthesize(
+                        turn["text"],
+                        voice=turn["voice_id"],
+                        language=language,
+                        **kwargs,
+                    )
+                )
+            except TTSError as exc:
+                failed += 1
+                logger.warning(
+                    f"synthesize_dialogue: skipping turn {index + 1}/{len(spoken)} "
+                    f"after retries failed ({exc})"
+                )
+
+        if not clips:
+            raise TTSError(
+                f"synthesize_dialogue: all {len(spoken)} turns failed to synthesize"
+            )
+        if failed:
+            logger.info(
+                f"synthesize_dialogue: produced {len(clips)}/{len(spoken)} turns "
+                f"({failed} skipped after retries)"
+            )
+        return concat_wav(clips, gap_ms=gap_ms)
+
+    async def list_voices(self) -> List[dict]:
+        """Return the available voices from the server."""
+        url = f"{self._base_url}{self._voices_path}"
+        try:
+            async with httpx.AsyncClient(timeout=self._timeout) as client:
+                resp = await client.get(url)
+                resp.raise_for_status()
+                data = resp.json()
+        except httpx.HTTPError as exc:
+            raise TTSError(f"failed to list voices: {exc}") from exc
+        return data.get("voices", data) if isinstance(data, dict) else data
+
+    @staticmethod
+    def _error_detail(resp: httpx.Response) -> str:
+        try:
+            body = resp.json()
+            detail = body.get("detail", body) if isinstance(body, dict) else body
+        except Exception:  # noqa: BLE001 - body may be non-JSON
+            detail = resp.text[:200]
+        return f"TTS failed (HTTP {resp.status_code}): {detail}"
+
+
+@lru_cache(maxsize=1)
+def get_tts_client() -> TTSClient:
+    """Return the process-wide TTS client singleton."""
+    return TTSClient()

aitoolkit/types.py

@@ -0,0 +1,66 @@
+"""Plain, provider-agnostic data types used across aitoolkit.
+
+These are deliberately simple so the public surface does not leak any third-party
+SDK types to callers.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Literal, Optional, TypedDict
+
+from pydantic import BaseModel, Field
+
+Role = Literal["system", "user", "assistant", "tool"]
+
+
+class ChatMessage(TypedDict):
+    """A single chat message in OpenAI message format."""
+
+    role: Role
+    content: str
+
+
+class DialogueTurn(TypedDict):
+    """One turn of a multi-speaker dialogue to synthesize: a voice and its line."""
+
+    voice_id: str
+    text: str
+
+
+class TranscriptionResult(BaseModel):
+    """Result of a speech-to-text transcription."""
+
+    text: str
+    language: Optional[str] = None
+    duration: Optional[float] = None
+
+
+class RetrievedChunk(BaseModel):
+    """A single retrieved context chunk with score and metadata."""
+
+    text: str
+    score: float = 0.0
+    file_id: Optional[str] = None
+    metadata: Dict[str, object] = Field(default_factory=dict)
+
+
+def as_messages(
+    prompt: Optional[str] = None,
+    *,
+    system: Optional[str] = None,
+    messages: Optional[List[ChatMessage]] = None,
+) -> List[ChatMessage]:
+    """Normalize the various ways callers pass prompts into a message list.
+
+    Accepts either an explicit ``messages`` list, or a ``prompt`` (plus optional
+    ``system``) convenience form.
+    """
+    if messages is not None:
+        return list(messages)
+
+    out: List[ChatMessage] = []
+    if system:
+        out.append({"role": "system", "content": system})
+    if prompt:
+        out.append({"role": "user", "content": prompt})
+    return out