vibe-bot/vibe_bot/services/speech_service.py

"""Speech service: parse the voice flag, dispatch bot-vs-plain, run TTS."""

from __future__ import annotations

import asyncio
import logging
import re
from collections.abc import Callable
from io import BytesIO
from typing import TYPE_CHECKING, Any

from vibe_bot import llm_client
from vibe_bot.config import (
    CHAT_MODEL,
    MAX_COMPLETION_TOKENS,
    TTS_SPEED,
    TTS_VOICE,
    VOICES_LIST,
)
from vibe_bot.prompts import build_system_prompt, get_user_info
from vibe_bot.tts import DEFAULT_LANG

if TYPE_CHECKING:
    from discord.ext.commands import Bot, Context

    from vibe_bot.database import ChatDatabase, CustomBotManager
    from vibe_bot.tts import TTSEngine

logger = logging.getLogger(__name__)

# Input size bound: reject oversized text before it reaches the TTS engine.
MAX_SPEAK_LENGTH = 5000

# Injected factory that turns an in-memory audio buffer into a discord.File.
FileFactory = Callable[[BytesIO, str], Any]

# Precomputed voice -> language lookup (replaces a per-call VOICES_LIST scan).
VOICE_LANGUAGES: dict[str, str] = {
    voice: str(category["language"])
    for category in VOICES_LIST.values()
    for voice in category["voices"]
}

# Trailing-anchored voice flag: only a `--voice <name>` at the very end of the
# message is treated as a flag, so the flag mid-text is preserved as speech.
_VOICE_FLAG_RE = re.compile(r"^(?P<text>.*)\s+--voice\s+(?P<voice>\S+)$")


def parse_voice_flag(message: str) -> tuple[str, str | None]:
    """Split a trailing `--voice <name>` flag off a speak message.

    Returns:
        (text, voice) where voice is None when no trailing flag is present and
        text is then returned unchanged.
    """
    match = _VOICE_FLAG_RE.match(message)
    if not match:
        return message, None
    return match["text"].rstrip(), match["voice"]


class SpeechService:
    """Speaks text (plain or via a custom bot) using the Kokoro TTS engine."""

    def __init__(
        self,
        db: ChatDatabase,
        manager: CustomBotManager,
        tts: TTSEngine | None,
        make_file: FileFactory,
    ) -> None:
        self._db = db
        self._manager = manager
        self._tts = tts
        self._make_file = make_file

    async def speak(self, ctx: Context[Bot], *, message: str) -> None:
        """Have the bot speak the given text, or have a custom bot respond+speaks."""
        if self._tts is None:
            await ctx.send(
                "TTS engine not initialized. "
                "Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.",
            )
            return

        text, voice = parse_voice_flag(message)

        if not text or not text.strip():
            await ctx.send("Please provide text to speak.")
            return

        if len(text) > MAX_SPEAK_LENGTH:
            logger.warning(
                "Speak text too long from user %s: length=%d",
                ctx.author.id,
                len(text),
            )
            await ctx.send(
                f"Text too long to speak. Max {MAX_SPEAK_LENGTH} characters."
            )
            return

        # Validate the voice if one was requested.
        if voice is not None and voice not in VOICE_LANGUAGES:
            await ctx.send(
                f"Unknown voice '{voice}'. Use `!voices` to see available voices."
            )
            return

        custom_bots = await asyncio.to_thread(self._manager.list_custom_bots)
        bot_names = [b[0] for b in custom_bots]

        first_word = text.split(maxsplit=1)[0] if text.split() else ""
        if first_word in bot_names:
            await self._speak_with_bot(ctx, first_word, text, voice)
        else:
            await self._speak_plain(ctx, text, voice)

    async def _speak_with_bot(
        self,
        ctx: Context[Bot],
        bot_name: str,
        message: str,
        voice: str | None,
    ) -> None:
        """Have a custom bot respond to the message and speak the response."""
        text_to_speak = message[len(bot_name) :].lstrip()
        if not text_to_speak:
            await ctx.send("Please provide text for the bot to respond to.")
            return

        await ctx.send(f"**{bot_name}** is thinking...")

        bot_info = await asyncio.to_thread(self._manager.get_custom_bot, bot_name)
        if not bot_info:
            await ctx.send(f"Custom bot '{bot_name}' not found.")
            return

        _, system_prompt, _, _ = bot_info
        system_prompt_edit = build_system_prompt(
            system_prompt, get_user_info(ctx.author)
        )

        engine = self._tts
        if engine is None:
            await ctx.send(
                "TTS engine not initialized. "
                "Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.",
            )
            return

        # Determine language for the chosen voice.
        chosen_voice = voice or TTS_VOICE
        lang = VOICE_LANGUAGES.get(chosen_voice, DEFAULT_LANG)

        try:
            context = await asyncio.to_thread(
                self._db.get_conversation_context,
                user_id=str(ctx.author.id),
                current_message=text_to_speak,
                max_context=5,
            )

            prompts: list[dict[str, str]] = [{"role": "user", "content": text_to_speak}]
            if context:
                prompts = context + prompts

            # Tools come from the shared registry (schema + dispatch).
            registry = llm_client.get_tool_registry()
            speak_tools = registry.to_openai_tools()

            def speak_tool_executor(tool_name: str, tool_args: dict[str, str]) -> str:
                """Dispatch a tool call through the registry."""
                return registry.execute(tool_name, tool_args, channel=ctx.channel)

            async def speak_tool_call_notifier(
                tool_name: str, tool_args: dict[str, str]
            ) -> None:
                """Send a notification message when a tool is called."""
                if tool_name == "get_channel_members":
                    await ctx.send(
                        f"**{bot_name}** is looking at the channel members..."
                    )

            bot_response = await llm_client.chat_completion_with_tools(
                system_prompt=system_prompt_edit,
                prompts=prompts,
                tools=speak_tools,
                tool_executor=speak_tool_executor,
                tool_call_notifier=speak_tool_call_notifier,
                model=CHAT_MODEL,
                max_tokens=MAX_COMPLETION_TOKENS,
            )

            if not bot_response:
                await ctx.send(f"**{bot_name}** failed to generate a response.")
                return

            await asyncio.to_thread(
                self._db.add_message,
                message_id=f"{ctx.message.id}",
                user_id=str(ctx.author.id),
                username=ctx.author.name,
                content=f"User: {text_to_speak}",
                bot_name=bot_name,
                channel_id=str(ctx.channel.id),
                guild_id=str(ctx.guild.id) if ctx.guild else None,
            )

            if ctx.bot.user is not None:
                await asyncio.to_thread(
                    self._db.add_message,
                    message_id=f"{ctx.message.id}_response",
                    user_id=str(ctx.bot.user.id),
                    username=ctx.bot.user.name,
                    content=bot_response,
                    bot_name=bot_name,
                    channel_id=str(ctx.channel.id),
                    guild_id=str(ctx.guild.id) if ctx.guild else None,
                    role="assistant",
                    embed=False,
                )

            await ctx.send(f"**{bot_name}**: {bot_response}")
            await ctx.send(f"Generating speech for **{bot_name}**...")
            result = await asyncio.to_thread(
                engine.generate_audio,
                bot_response,
                voice=chosen_voice,
                speed=TTS_SPEED,
                lang=lang,
            )
            if result.partial:
                await ctx.send("Some audio chunks failed; audio may be incomplete.")

            audio_file = self._make_file(result.audio, "speech.mp3")
            await ctx.send(file=audio_file)
        except Exception:
            logger.exception(
                "Error in speak command with bot %r",
                bot_name,
            )
            await ctx.send("Error generating speech.")

    async def _speak_plain(
        self,
        ctx: Context[Bot],
        message: str,
        voice: str | None,
    ) -> None:
        """Speak plain text (no custom bot involved)."""
        engine = self._tts
        if engine is None:
            await ctx.send(
                "TTS engine not initialized. "
                "Make sure kokoro-v1.0.onnx and voices-v1.0.bin are present.",
            )
            return

        chosen_voice = voice or TTS_VOICE
        lang = VOICE_LANGUAGES.get(chosen_voice, DEFAULT_LANG)

        try:
            await ctx.send("Generating speech...")
            result = await asyncio.to_thread(
                engine.generate_audio,
                message,
                voice=chosen_voice,
                speed=TTS_SPEED,
                lang=lang,
            )
            if result.partial:
                await ctx.send("Some audio chunks failed; audio may be incomplete.")

            audio_file = self._make_file(result.audio, "speech.mp3")
            await ctx.send(file=audio_file)
        except Exception:
            logger.exception("Error in speak command")
            await ctx.send("Error generating speech.")