app/pipeline.py

import os
import uuid
from app.config import settings
from app.modules.asr_whisper import ArabicASRWhisper
from app.modules.normalize_eg import EgyptianNormalizer
from app.modules.llm_jais import JaisChatbot
from app.modules.postprocess import ResponsePostProcessor
from app.modules.tts_egtts import EGTTS

class FullEgyptianVoiceChatbot:
    def __init__(self):
        # 1) ASR
        self.asr = ArabicASRWhisper(settings.WHISPER_MODEL_SIZE)

        # 2) Normalize
        self.norm = EgyptianNormalizer()

        # 3) LLM
        self.llm = JaisChatbot(settings.JAIS_MODEL_ID)

        # 4) Post-process
        self.post = ResponsePostProcessor()

        # 5) TTS
        self.tts = EGTTS(
            config_path=settings.XTTS_CONFIG_PATH,
            vocab_path=settings.XTTS_VOCAB_PATH,
            model_dir=settings.XTTS_MODEL_DIR,
            speaker_wav=settings.XTTS_SPEAKER_WAV,
            use_deepspeed=True,
        )

        os.makedirs("outputs", exist_ok=True)

    def chat_text(self, user_text: str) -> dict:
        # Normalize user text
        user_norm = self.norm.normalize(user_text)

        # LLM response
        raw_answer = self.llm.generate(
            user_norm,
            max_new_tokens=settings.JAIS_MAX_NEW_TOKENS,
            temperature=settings.JAIS_TEMPERATURE,
            top_p=settings.JAIS_TOP_P,
        )

        # Clean for TTS
        answer_clean = self.post.clean_for_tts(raw_answer)

        # Synthesize
        out_id = str(uuid.uuid4())
        out_wav = f"outputs/{out_id}.wav"
        self.tts.synthesize_to_file(answer_clean, out_wav)

        return {
            "user_text": user_text,
            "user_normalized": user_norm,
            "assistant_text_raw": raw_answer,
            "assistant_text_tts": answer_clean,
            "audio_path": out_wav,
            "audio_id": out_id,
        }

    def chat_voice(self, audio_path: str) -> dict:
        user_text = self.asr.transcribe(audio_path)
        if not user_text:
            return {
                "error": "ASR returned empty text",
                "audio_path": audio_path,
            }
        return self.chat_text(user_text)