agent-cli 0.70.5__py3-none-any.whl

agent_cli/services/llm.py

@@ -0,0 +1,228 @@
+"""Client for interacting with LLMs."""
+
+from __future__ import annotations
+
+import sys
+import time
+from typing import TYPE_CHECKING
+
+from rich.live import Live
+
+from agent_cli.core.utils import console, live_timer, print_error_message, print_output_panel
+
+if TYPE_CHECKING:
+    import logging
+
+    from pydantic_ai import Agent
+    from pydantic_ai.models.gemini import GeminiModel
+    from pydantic_ai.models.openai import OpenAIModel
+    from pydantic_ai.tools import Tool
+
+    from agent_cli import config
+
+
+def _openai_llm_model(openai_cfg: config.OpenAILLM) -> OpenAIModel:
+    from pydantic_ai.models.openai import OpenAIModel  # noqa: PLC0415
+    from pydantic_ai.providers.openai import OpenAIProvider  # noqa: PLC0415
+
+    # For custom base URLs (like llama-server), API key might not be required
+    if openai_cfg.openai_base_url:
+        # Custom endpoint - API key is optional
+        provider = OpenAIProvider(
+            api_key=openai_cfg.openai_api_key or "dummy",
+            base_url=openai_cfg.openai_base_url,
+        )
+    else:
+        # Standard OpenAI - API key is required
+        if not openai_cfg.openai_api_key:
+            msg = "OpenAI API key is not set."
+            raise ValueError(msg)
+        provider = OpenAIProvider(api_key=openai_cfg.openai_api_key)
+
+    model_name = openai_cfg.llm_openai_model
+    return OpenAIModel(model_name=model_name, provider=provider)
+
+
+def _ollama_llm_model(ollama_cfg: config.Ollama) -> OpenAIModel:
+    from pydantic_ai.models.openai import OpenAIModel  # noqa: PLC0415
+    from pydantic_ai.providers.openai import OpenAIProvider  # noqa: PLC0415
+
+    provider = OpenAIProvider(base_url=f"{ollama_cfg.llm_ollama_host}/v1")
+    model_name = ollama_cfg.llm_ollama_model
+    return OpenAIModel(model_name=model_name, provider=provider)
+
+
+def _gemini_llm_model(gemini_cfg: config.GeminiLLM) -> GeminiModel:
+    from pydantic_ai.models.gemini import GeminiModel  # noqa: PLC0415
+    from pydantic_ai.providers.google_gla import GoogleGLAProvider  # noqa: PLC0415
+
+    if not gemini_cfg.gemini_api_key:
+        msg = "Gemini API key is not set."
+        raise ValueError(msg)
+    provider = GoogleGLAProvider(api_key=gemini_cfg.gemini_api_key)
+    model_name = gemini_cfg.llm_gemini_model
+    return GeminiModel(model_name=model_name, provider=provider)
+
+
+def create_llm_agent(
+    provider_cfg: config.ProviderSelection,
+    ollama_cfg: config.Ollama,
+    openai_cfg: config.OpenAILLM,
+    gemini_cfg: config.GeminiLLM,
+    *,
+    system_prompt: str | None = None,
+    instructions: str | None = None,
+    tools: list[Tool] | None = None,
+) -> Agent:
+    """Construct and return a PydanticAI agent."""
+    from pydantic_ai import Agent  # noqa: PLC0415
+
+    if provider_cfg.llm_provider == "openai":
+        llm_model = _openai_llm_model(openai_cfg)
+    elif provider_cfg.llm_provider == "ollama":
+        llm_model = _ollama_llm_model(ollama_cfg)
+    elif provider_cfg.llm_provider == "gemini":
+        llm_model = _gemini_llm_model(gemini_cfg)
+
+    return Agent(
+        model=llm_model,
+        system_prompt=system_prompt or (),
+        instructions=instructions,
+        tools=tools or [],
+    )
+
+
+# --- LLM (Editing) Logic ---
+
+INPUT_TEMPLATE = """
+{context_block}<original-text>
+{original_text}
+</original-text>
+
+<instruction>
+{instruction}
+</instruction>
+"""
+
+
+async def get_llm_response(
+    *,
+    system_prompt: str,
+    agent_instructions: str,
+    user_input: str,
+    provider_cfg: config.ProviderSelection,
+    ollama_cfg: config.Ollama,
+    openai_cfg: config.OpenAILLM,
+    gemini_cfg: config.GeminiLLM,
+    logger: logging.Logger,
+    live: Live | None = None,
+    tools: list[Tool] | None = None,
+    quiet: bool = False,
+    clipboard: bool = False,
+    show_output: bool = False,
+    exit_on_error: bool = False,
+) -> str | None:
+    """Get a response from the LLM with optional clipboard and output handling."""
+    agent = create_llm_agent(
+        provider_cfg=provider_cfg,
+        ollama_cfg=ollama_cfg,
+        openai_cfg=openai_cfg,
+        gemini_cfg=gemini_cfg,
+        system_prompt=system_prompt,
+        instructions=agent_instructions,
+        tools=tools,
+    )
+
+    start_time = time.monotonic()
+
+    try:
+        if provider_cfg.llm_provider == "ollama":
+            model_name = ollama_cfg.llm_ollama_model
+        elif provider_cfg.llm_provider == "openai":
+            model_name = openai_cfg.llm_openai_model
+        elif provider_cfg.llm_provider == "gemini":
+            model_name = gemini_cfg.llm_gemini_model
+
+        async with live_timer(
+            live or Live(console=console),
+            f"🤖 Applying instruction with {model_name}",
+            style="bold yellow",
+            quiet=quiet,
+        ):
+            result = await agent.run(user_input)
+
+        elapsed = time.monotonic() - start_time
+        result_text = result.output
+
+        if clipboard:
+            import pyperclip  # noqa: PLC0415
+
+            pyperclip.copy(result_text)
+            logger.info("Copied result to clipboard.")
+
+        if show_output and not quiet:
+            print_output_panel(
+                result_text,
+                title="✨ Result (Copied to Clipboard)" if clipboard else "✨ Result",
+                subtitle=f"[dim]took {elapsed:.2f}s[/dim]",
+            )
+        elif quiet and clipboard:
+            print(result_text)
+
+        return result_text
+
+    except Exception as e:
+        logger.exception("An error occurred during LLM processing.")
+        if provider_cfg.llm_provider == "openai":
+            msg = "Please check your OpenAI API key."
+        elif provider_cfg.llm_provider == "gemini":
+            msg = "Please check your Gemini API key."
+        elif provider_cfg.llm_provider == "ollama":
+            msg = f"Please check your Ollama server at [cyan]{ollama_cfg.llm_ollama_host}[/cyan]"
+        print_error_message(f"An unexpected LLM error occurred: {e}", msg)
+        if exit_on_error:
+            sys.exit(1)
+        return None
+
+
+async def process_and_update_clipboard(
+    system_prompt: str,
+    agent_instructions: str,
+    *,
+    provider_cfg: config.ProviderSelection,
+    ollama_cfg: config.Ollama,
+    openai_cfg: config.OpenAILLM,
+    gemini_cfg: config.GeminiLLM,
+    logger: logging.Logger,
+    original_text: str,
+    instruction: str,
+    clipboard: bool,
+    quiet: bool,
+    live: Live | None,
+    context: str | None = None,
+) -> str | None:
+    """Processes the text with the LLM, updates the clipboard, and displays the result."""
+    context_block = ""
+    if context:
+        context_block = f"<context>\n{context}\n</context>\n\n"
+    user_input = INPUT_TEMPLATE.format(
+        context_block=context_block,
+        original_text=original_text,
+        instruction=instruction,
+    )
+
+    return await get_llm_response(
+        system_prompt=system_prompt,
+        agent_instructions=agent_instructions,
+        user_input=user_input,
+        provider_cfg=provider_cfg,
+        ollama_cfg=ollama_cfg,
+        openai_cfg=openai_cfg,
+        gemini_cfg=gemini_cfg,
+        logger=logger,
+        quiet=quiet,
+        clipboard=clipboard,
+        live=live,
+        show_output=True,
+        exit_on_error=False,  # Don't exit the server on LLM errors
+    )

agent_cli/services/tts.py

@@ -0,0 +1,450 @@
+"""Module for Text-to-Speech using Wyoming or OpenAI."""
+
+from __future__ import annotations
+
+import asyncio
+import importlib.util
+import io
+from functools import partial
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from rich.live import Live
+
+from agent_cli import config, constants
+from agent_cli.core.audio import open_audio_stream, setup_output_stream
+from agent_cli.core.audio_format import extract_pcm_from_wav
+from agent_cli.core.utils import (
+    InteractiveStopEvent,
+    live_timer,
+    manage_send_receive_tasks,
+    print_error_message,
+    print_with_style,
+)
+from agent_cli.services import pcm_to_wav, synthesize_speech_gemini, synthesize_speech_openai
+from agent_cli.services._wyoming_utils import wyoming_client_context
+
+if TYPE_CHECKING:
+    import logging
+    from collections.abc import Awaitable, Callable
+
+    from rich.live import Live
+    from wyoming.client import AsyncClient
+    from wyoming.tts import Synthesize
+
+
+has_audiostretchy = importlib.util.find_spec("audiostretchy") is not None
+
+
+def create_synthesizer(
+    provider_cfg: config.ProviderSelection,
+    audio_output_cfg: config.AudioOutput,
+    wyoming_tts_cfg: config.WyomingTTS,
+    openai_tts_cfg: config.OpenAITTS,
+    kokoro_tts_cfg: config.KokoroTTS,
+    gemini_tts_cfg: config.GeminiTTS | None = None,
+) -> Callable[..., Awaitable[bytes | None]]:
+    """Return the appropriate synthesizer based on the config."""
+    if not audio_output_cfg.enable_tts:
+        return _dummy_synthesizer
+    if provider_cfg.tts_provider == "openai":
+        return partial(
+            _synthesize_speech_openai,
+            openai_tts_cfg=openai_tts_cfg,
+        )
+    if provider_cfg.tts_provider == "kokoro":
+        return partial(
+            _synthesize_speech_kokoro,
+            kokoro_tts_cfg=kokoro_tts_cfg,
+        )
+    if provider_cfg.tts_provider == "gemini":
+        assert gemini_tts_cfg is not None, "Gemini TTS config required"
+        return partial(_synthesize_speech_gemini, gemini_tts_cfg=gemini_tts_cfg)
+    if provider_cfg.tts_provider == "wyoming":
+        return partial(_synthesize_speech_wyoming, wyoming_tts_cfg=wyoming_tts_cfg)
+    msg = f"Unknown TTS provider: {provider_cfg.tts_provider}"
+    raise NotImplementedError(msg)
+
+
+async def handle_tts_playback(
+    *,
+    text: str,
+    provider_cfg: config.ProviderSelection,
+    audio_output_cfg: config.AudioOutput,
+    wyoming_tts_cfg: config.WyomingTTS,
+    openai_tts_cfg: config.OpenAITTS,
+    kokoro_tts_cfg: config.KokoroTTS,
+    gemini_tts_cfg: config.GeminiTTS | None = None,
+    save_file: Path | None,
+    quiet: bool,
+    logger: logging.Logger,
+    play_audio: bool = True,
+    status_message: str = "🔊 Speaking...",
+    description: str = "Audio",
+    stop_event: InteractiveStopEvent | None = None,
+    live: Live,
+) -> bytes | None:
+    """Handle TTS synthesis, playback, and file saving."""
+    try:
+        if not quiet and status_message:
+            print_with_style(status_message, style="blue")
+
+        audio_data = await _speak_text(
+            text=text,
+            provider_cfg=provider_cfg,
+            audio_output_cfg=audio_output_cfg,
+            wyoming_tts_cfg=wyoming_tts_cfg,
+            openai_tts_cfg=openai_tts_cfg,
+            kokoro_tts_cfg=kokoro_tts_cfg,
+            gemini_tts_cfg=gemini_tts_cfg,
+            logger=logger,
+            quiet=quiet,
+            play_audio_flag=play_audio,
+            stop_event=stop_event,
+            live=live,
+        )
+
+        if save_file and audio_data:
+            await _save_audio_file(
+                audio_data,
+                save_file,
+                quiet,
+                logger,
+                description=description,
+            )
+
+        return audio_data
+
+    except (OSError, ConnectionError, TimeoutError) as e:
+        logger.warning("Failed TTS operation: %s", e)
+        if not quiet:
+            print_with_style(f"⚠️ TTS failed: {e}", style="yellow")
+        return None
+
+
+# --- Helper Functions ---
+
+
+def _create_synthesis_request(
+    text: str,
+    *,
+    voice_name: str | None = None,
+    language: str | None = None,
+    speaker: str | None = None,
+) -> Synthesize:
+    """Create a synthesis request with optional voice parameters."""
+    from wyoming.tts import Synthesize, SynthesizeVoice  # noqa: PLC0415
+
+    synthesize_event = Synthesize(text=text)
+
+    # Add voice parameters if specified
+    if voice_name or language or speaker:
+        synthesize_event.voice = SynthesizeVoice(
+            name=voice_name,
+            language=language,
+            speaker=speaker,
+        )
+
+    return synthesize_event
+
+
+async def _process_audio_events(
+    client: AsyncClient,
+    logger: logging.Logger,
+) -> tuple[bytes, int | None, int | None, int | None]:
+    """Process audio events from TTS server and return audio data with metadata."""
+    from wyoming.audio import AudioChunk, AudioStart, AudioStop  # noqa: PLC0415
+
+    audio_data = io.BytesIO()
+    sample_rate = None
+    sample_width = None
+    channels = None
+
+    while True:
+        event = await client.read_event()
+        if event is None:
+            logger.warning("Connection to TTS server lost.")
+            break
+
+        if AudioStart.is_type(event.type):
+            audio_start = AudioStart.from_event(event)
+            sample_rate = audio_start.rate
+            sample_width = audio_start.width
+            channels = audio_start.channels
+            logger.debug(
+                "Audio stream started: %dHz, %d channels, %d bytes/sample",
+                sample_rate,
+                channels,
+                sample_width,
+            )
+
+        elif AudioChunk.is_type(event.type):
+            chunk = AudioChunk.from_event(event)
+            audio_data.write(chunk.audio)
+            logger.debug("Received %d bytes of audio", len(chunk.audio))
+
+        elif AudioStop.is_type(event.type):
+            logger.debug("Audio stream completed")
+            break
+        else:
+            logger.debug("Ignoring event type: %s", event.type)
+
+    return audio_data.getvalue(), sample_rate, sample_width, channels
+
+
+async def _dummy_synthesizer(**_kwargs: object) -> bytes | None:
+    """A dummy synthesizer that does nothing."""
+    return None
+
+
+async def _synthesize_speech_openai(
+    *,
+    text: str,
+    openai_tts_cfg: config.OpenAITTS,
+    logger: logging.Logger,
+    **_kwargs: object,
+) -> bytes | None:
+    """Synthesize speech from text using OpenAI-compatible TTS server."""
+    return await synthesize_speech_openai(
+        text=text,
+        openai_tts_cfg=openai_tts_cfg,
+        logger=logger,
+    )
+
+
+async def _synthesize_speech_kokoro(
+    *,
+    text: str,
+    kokoro_tts_cfg: config.KokoroTTS,
+    logger: logging.Logger,
+    **_kwargs: object,
+) -> bytes | None:
+    """Synthesize speech from text using Kokoro TTS server via OpenAI client."""
+    openai_tts_cfg = config.OpenAITTS(
+        tts_openai_model=kokoro_tts_cfg.tts_kokoro_model,
+        tts_openai_voice=kokoro_tts_cfg.tts_kokoro_voice,
+        tts_openai_base_url=kokoro_tts_cfg.tts_kokoro_host,
+    )
+    try:
+        return await synthesize_speech_openai(
+            text=text,
+            openai_tts_cfg=openai_tts_cfg,
+            logger=logger,
+        )
+    except Exception:
+        logger.exception("Error during Kokoro speech synthesis")
+        return None
+
+
+async def _synthesize_speech_gemini(
+    *,
+    text: str,
+    gemini_tts_cfg: config.GeminiTTS,
+    logger: logging.Logger,
+    **_kwargs: object,
+) -> bytes | None:
+    """Synthesize speech from text using Gemini TTS."""
+    return await synthesize_speech_gemini(text=text, gemini_tts_cfg=gemini_tts_cfg, logger=logger)
+
+
+async def _synthesize_speech_wyoming(
+    *,
+    text: str,
+    wyoming_tts_cfg: config.WyomingTTS,
+    logger: logging.Logger,
+    quiet: bool = False,
+    live: Live,
+    **_kwargs: object,
+) -> bytes | None:
+    """Synthesize speech from text using Wyoming TTS server."""
+    try:
+        async with wyoming_client_context(
+            wyoming_tts_cfg.tts_wyoming_ip,
+            wyoming_tts_cfg.tts_wyoming_port,
+            "TTS",
+            logger,
+            quiet=quiet,
+        ) as client:
+            async with live_timer(live, "🔊 Synthesizing text", style="blue", quiet=quiet):
+                synthesize_event = _create_synthesis_request(
+                    text,
+                    voice_name=wyoming_tts_cfg.tts_wyoming_voice,
+                    language=wyoming_tts_cfg.tts_wyoming_language,
+                    speaker=wyoming_tts_cfg.tts_wyoming_speaker,
+                )
+                _send_task, recv_task = await manage_send_receive_tasks(
+                    client.write_event(synthesize_event.event()),
+                    _process_audio_events(client, logger),
+                )
+                audio_data, sample_rate, sample_width, channels = recv_task.result()
+            if sample_rate and sample_width and channels and audio_data:
+                wav_data = pcm_to_wav(
+                    audio_data,
+                    sample_rate=sample_rate,
+                    sample_width=sample_width,
+                    channels=channels,
+                )
+                logger.info("Speech synthesis completed: %d bytes", len(wav_data))
+                return wav_data
+            logger.warning("No audio data received from TTS server")
+            return None
+    except (ConnectionRefusedError, Exception):
+        return None
+
+
+def _apply_speed_adjustment(
+    audio_data: io.BytesIO,
+    speed: float,
+) -> tuple[io.BytesIO, bool]:
+    """Apply speed adjustment to audio data."""
+    if speed == 1.0 or not has_audiostretchy:
+        return audio_data, False
+    from audiostretchy.stretch import AudioStretch  # noqa: PLC0415
+
+    audio_data.seek(0)
+    input_copy = io.BytesIO(audio_data.read())
+    audio_stretch = AudioStretch()
+    audio_stretch.open(file=input_copy, format="wav")
+    audio_stretch.stretch(ratio=1 / speed)
+    out = io.BytesIO()
+    audio_stretch.save_wav(out, close=False)
+    out.seek(0)
+    return out, True
+
+
+async def _play_audio(
+    audio_data: bytes,
+    logger: logging.Logger,
+    *,
+    audio_output_cfg: config.AudioOutput,
+    quiet: bool = False,
+    stop_event: InteractiveStopEvent | None = None,
+    live: Live,
+) -> None:
+    """Play WAV audio data using SoundDevice."""
+    import numpy as np  # noqa: PLC0415
+
+    try:
+        wav_io = io.BytesIO(audio_data)
+        speed = audio_output_cfg.tts_speed
+        wav_io, speed_changed = _apply_speed_adjustment(wav_io, speed)
+        wav = extract_pcm_from_wav(wav_io.read())
+        sample_rate = wav.sample_rate if speed_changed else int(wav.sample_rate * speed)
+        base_msg = f"🔊 Playing audio at {speed}x speed" if speed != 1.0 else "🔊 Playing audio"
+        async with live_timer(live, base_msg, style="blue", quiet=quiet):
+            stream_config = setup_output_stream(
+                audio_output_cfg.output_device_index,
+                sample_rate=sample_rate,
+                sample_width=wav.sample_width,
+                channels=wav.num_channels,
+            )
+            dtype = stream_config.dtype
+
+            with open_audio_stream(stream_config) as stream:
+                chunk_size_frames = constants.AUDIO_CHUNK_SIZE
+                bytes_per_frame = wav.num_channels * wav.sample_width
+                chunk_bytes = chunk_size_frames * bytes_per_frame
+
+                for i in range(0, len(wav.pcm_data), chunk_bytes):
+                    if stop_event and stop_event.is_set():
+                        logger.info("Audio playback interrupted")
+                        if not quiet:
+                            print_with_style("⏹️ Audio playback interrupted", style="yellow")
+                        break
+                    chunk = wav.pcm_data[i : i + chunk_bytes]
+
+                    # Convert bytes to numpy array for sounddevice
+                    audio_array = np.frombuffer(chunk, dtype=dtype)
+                    if wav.num_channels > 1:
+                        audio_array = audio_array.reshape(-1, wav.num_channels)
+
+                    stream.write(audio_array)
+                    await asyncio.sleep(0)
+        if not (stop_event and stop_event.is_set()):
+            logger.info("Audio playback completed (speed: %.1fx)", speed)
+            if not quiet:
+                print_with_style("✅ Audio playback finished")
+    except Exception as e:
+        logger.exception("Error during audio playback")
+        if not quiet:
+            print_error_message(f"Playback error: {e}")
+
+
+async def _speak_text(
+    *,
+    text: str,
+    provider_cfg: config.ProviderSelection,
+    audio_output_cfg: config.AudioOutput,
+    wyoming_tts_cfg: config.WyomingTTS,
+    openai_tts_cfg: config.OpenAITTS,
+    kokoro_tts_cfg: config.KokoroTTS,
+    gemini_tts_cfg: config.GeminiTTS | None = None,
+    logger: logging.Logger,
+    quiet: bool = False,
+    play_audio_flag: bool = True,
+    stop_event: InteractiveStopEvent | None = None,
+    live: Live,
+) -> bytes | None:
+    """Synthesize and optionally play speech from text."""
+    synthesizer = create_synthesizer(
+        provider_cfg,
+        audio_output_cfg,
+        wyoming_tts_cfg,
+        openai_tts_cfg,
+        kokoro_tts_cfg,
+        gemini_tts_cfg,
+    )
+    audio_data = None
+    try:
+        async with live_timer(live, "🔊 Synthesizing text", style="blue", quiet=quiet):
+            audio_data = await synthesizer(
+                text=text,
+                wyoming_tts_cfg=wyoming_tts_cfg,
+                openai_tts_cfg=openai_tts_cfg,
+                kokoro_tts_cfg=kokoro_tts_cfg,
+                gemini_tts_cfg=gemini_tts_cfg,
+                logger=logger,
+                quiet=quiet,
+                live=live,
+            )
+    except Exception:
+        logger.exception("Error during speech synthesis")
+        return None
+
+    if audio_data and play_audio_flag:
+        await _play_audio(
+            audio_data,
+            logger,
+            audio_output_cfg=audio_output_cfg,
+            quiet=quiet,
+            stop_event=stop_event,
+            live=live,
+        )
+
+    return audio_data
+
+
+async def _save_audio_file(
+    audio_data: bytes,
+    save_file: Path,
+    quiet: bool,
+    logger: logging.Logger,
+    *,
+    description: str = "Audio",
+) -> None:
+    try:
+        save_path = Path(save_file)
+        await asyncio.to_thread(save_path.write_bytes, audio_data)
+        if not quiet:
+            print_with_style(f"💾 {description} saved to {save_file}")
+        logger.info("%s saved to %s", description, save_file)
+    except (OSError, PermissionError) as e:
+        logger.exception("Failed to save %s", description.lower())
+        if not quiet:
+            print_with_style(
+                f"❌ Failed to save {description.lower()}: {e}",
+                style="red",
+            )
+
+
+__all__ = ["handle_tts_playback"]