mOrpheus/modules/lm_client.py

# modules/lm_client.py
import os
import time
import json
import wave
import requests
import re
from typing import Optional
from modules.logging import logger
from modules.audio import segment_text, combine_audio_files
from modules.snac_decoder import tokens_decoder_sync
from modules.config import load_config

def clean_text_for_tts(text: str) -> str:
    """
    Clean the text to be sent to the TTS engine by:
      • Removing newline characters and excessive whitespace.
      • Removing markdown symbols (e.g., asterisks).
      • Removing non-ASCII characters (e.g., emojis).
    """
    # Remove newline characters
    text = text.replace('\n', ' ')
    # Remove markdown formatting
    text = re.sub(r'\*+', '', text)
    # Remove non-ASCII characters (e.g., emojis)
    text = re.sub(r'[^\x00-\x7F]+', '', text)
    # Normalize whitespace
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

class LMStudioClient:
    def __init__(self, config):
        self.config = config
        lm_config = config["lm"]

        self.api_url = lm_config["api_url"]
        self.chat_endpoint = lm_config["chat"]["endpoint"]
        self.tts_endpoint = lm_config["tts"]["endpoint"]

        # Chat parameters
        chat_config = lm_config["chat"]
        self.chat_model = chat_config["model"]
        self.system_prompt = chat_config["system_prompt"]
        self.chat_max_tokens = chat_config["max_tokens"]
        self.chat_temperature = chat_config["temperature"]
        self.chat_top_p = chat_config["top_p"]
        self.chat_repetition_penalty = chat_config["repetition_penalty"]
        self.max_response_time = chat_config["max_response_time"]

        # TTS parameters
        tts_config = lm_config["tts"]
        self.tts_model = tts_config["model"]
        self.default_voice = tts_config["default_voice"]
        self.tts_max_tokens = tts_config["max_tokens"]
        self.tts_temperature = tts_config["temperature"]
        self.tts_top_p = tts_config["top_p"]
        self.tts_repetition_penalty = tts_config["repetition_penalty"]
        self.speed = tts_config["speed"]
        self.max_segment_duration = tts_config["max_segment_duration"]

        self.headers = {"Content-Type": "application/json"}
        self.retries = config["speech"]["max_retries"]
        self.tts_sample_rate = config["tts"]["sample_rate"]

        # Use a session for connection pooling
        self.session = requests.Session()

    def chat(self, user_input: str) -> str:
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_input}
        ]
        payload = {
            "model": self.chat_model,
            "messages": messages,
            "max_tokens": self.chat_max_tokens,
            "temperature": self.chat_temperature,
            "top_p": self.chat_top_p,
            "repeat_penalty": self.chat_repetition_penalty,
            "stream": False
        }
        url = self.api_url + self.chat_endpoint
        logger.debug("Chat request payload: %s", payload)
        for attempt in range(self.retries):
            try:
                start_time = time.time()
                response = self.session.post(
                    url,
                    headers=self.headers,
                    json=payload,
                    timeout=self.max_response_time
                )
                elapsed = time.time() - start_time
                logger.info("Chat response received in %.2f seconds", elapsed)
                logger.debug("LM Studio full response: %s", response.text)
                if response.status_code != 200:
                    logger.error("Chat API error: %s %s", response.status_code, response.text)
                    if attempt < self.retries - 1:
                        delay = 2 ** attempt
                        logger.warning("Chat API error, retrying in %d seconds...", delay)
                        time.sleep(delay)
                        continue
                    raise RuntimeError(f"Chat API error: {response.status_code} {response.text}")
                data = response.json()
                generated_text = data.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
                token_count = len(generated_text.split()) * 1.33
                logger.info("Generated %d tokens: %s", int(token_count), generated_text[:50])
                return generated_text
            except requests.exceptions.Timeout:
                delay = 2 ** attempt
                logger.warning("Chat API timeout (attempt %d), retrying in %d seconds", attempt + 1, delay)
                time.sleep(delay)
                if attempt == self.retries - 1:
                    return "I need more time to think about that. Could you ask again?"
            except Exception as e:
                delay = 2 ** attempt
                logger.error("Chat API failed (attempt %d): %s", attempt + 1, str(e))
                time.sleep(delay)
                if attempt == self.retries - 1:
                    return "I'm having trouble responding right now. Please try again later."

    def synthesize_speech(self, text: str, voice: Optional[str] = None, output_file: Optional[str] = None) -> str:
        voice = voice if voice else self.default_voice
        if not voice.isalpha() or len(voice) > 20:
            logger.warning("Invalid voice name '%s', using default", voice)
            voice = self.default_voice

        # Clean the text for TTS
        cleaned_text = clean_text_for_tts(text)
        prompt = f"<|audio|>{voice}: {cleaned_text}<|eot_id|>"
        payload = {
            "model": self.tts_model,
            "prompt": prompt,
            "max_tokens": self.tts_max_tokens,
            "temperature": self.tts_temperature,
            "top_p": self.tts_top_p,
            "repeat_penalty": self.tts_repetition_penalty,
            "speed": self.speed,
            "stream": True
        }
        url = self.api_url + self.tts_endpoint
        logger.debug("TTS request payload: %s", payload)
        for attempt in range(self.retries):
            try:
                response = self.session.post(
                    url,
                    headers=self.headers,
                    json=payload,
                    stream=True,
                    timeout=self.max_segment_duration + 5
                )
                if response.status_code != 200:
                    logger.error("TTS API error: %s %s", response.status_code, response.text)
                    if attempt < self.retries - 1:
                        delay = 2 ** attempt
                        logger.warning("TTS API error, retrying in %d seconds...", delay)
                        time.sleep(delay)
                        continue
                    raise RuntimeError(f"TTS API error: {response.status_code} {response.text}")

                def token_generator():
                    for line in response.iter_lines():
                        if line:
                            decoded_line = line.decode("utf-8")
                            if decoded_line.startswith("data: "):
                                data_str = decoded_line[6:]
                                if data_str.strip() == "[DONE]":
                                    break
                                try:
                                    data = json.loads(data_str)
                                    token_text = data.get("choices", [{}])[0].get("text", "")
                                    yield token_text
                                except json.JSONDecodeError as e:
                                    logger.error("JSON decode error: %s", e)

                audio_bytes = tokens_decoder_sync(token_generator())
                if not output_file:
                    timestamp = int(time.time())
                    output_file = f"outputs/{voice}_{timestamp}.wav"
                os.makedirs("outputs", exist_ok=True)
                with wave.open(output_file, "wb") as wf:
                    wf.setnchannels(1)
                    wf.setsampwidth(2)
                    wf.setframerate(self.tts_sample_rate)
                    wf.writeframes(audio_bytes)
                logger.info("Audio saved to %s", output_file)
                return output_file
            except Exception as e:
                delay = 2 ** attempt
                logger.error("TTS synthesis failed (attempt %d): %s", attempt + 1, str(e))
                time.sleep(delay)
                if attempt == self.retries - 1:
                    raise

    def synthesize_long_text(self, text: str, voice: Optional[str] = None) -> str:
        voice = voice if voice else self.default_voice
        segments = segment_text(text, max_words=self.config["segmentation"]["max_words"])
        logger.info("Text segmented into %d parts", len(segments))
        file_list = []
        for i, seg in enumerate(segments):
            seg_filename = f"outputs/{voice}_{int(time.time())}_{i}.wav"
            try:
                self.synthesize_speech(seg, voice=voice, output_file=seg_filename)
                file_list.append(seg_filename)
                time.sleep(0.2)  # Small delay between segments
            except Exception as e:
                logger.error("Failed to synthesize segment %d: %s", i, str(e))
                if file_list:
                    break
                raise
        combined_filename = f"outputs/{voice}_{int(time.time())}_combined.wav"
        combine_audio_files(file_list, combined_filename)
        for f in file_list:
            try:
                os.remove(f)
            except Exception as e:
                logger.warning("Could not remove temporary file %s: %s", f, str(e))
        return combined_filename