app.py

import gradio as gr
import requests
import numpy as np
import soundfile as sf
import io
import tempfile
import os
from groq import Groq
from crdt import warn

# Language configuration
LANGUAGE_CONFIG = {
    "English": {
        "voice_id": "vits-eng-1",
        "whisper_lang": "en",
        "system_prompt": "You are a language tutor. Your goal is to help students practice language skills in English by engaging in conversations with them. As part of your responses, give them feedback on their English proficiency if needed. If they make a clear mistake (vocabulary or grammar, not spelling), indicate it gently and explain how to correct them. Please respond in English. Keep your answers very concise. Note that you are used in a school setting and should refuse to answer or produce any content that is violent, sexual, discriminatory, or inappropriate for children under 13."
    },
    "French": {
        "voice_id": "vits-fra-1",
        "whisper_lang": "fr",
        "system_prompt": "You are a language tutor. Your goal is to help students practice language skills in French by engaging in conversations with them. As part of your responses, give them feedback on their French proficiency if needed. If they make a clear mistake (vocabulary or grammar, not spelling), indicate it gently and explain how to correct them. Please respond in French. Keep your answers very concise. Note that you are used in a school setting and should refuse to answer or produce any content that is violent, sexual, discriminatory, or inappropriate for children under 13."
    },
    "Spanish": {
        "voice_id": "vits-spa-1",
        "whisper_lang": "es",
        "system_prompt": "You are a language tutor. Your goal is to help students practice language skills in Spanish by engaging in conversations with them. As part of your responses, give them feedback on their Spanish proficiency if needed. If they make a clear mistake (vocabulary or grammar, not spelling), indicate it gently and explain how to correct them. Please respond in Spanish. Keep your answers very concise. Note that you are used in a school setting and should refuse to answer or produce any content that is violent, sexual, discriminatory, or inappropriate for children under 13."
    }
}

def generate_audio(text: str, neets_api_key: str, language: str) -> tuple[int, np.ndarray]:
    """Generate audio from text using Neets API"""
    print(f"Generating audio for text in {language}:", text)
    
    try:
        # Make request with simplified params
        response = requests.post(
            url="https://api.neets.ai/v1/tts",
            headers={
                "Content-Type": "application/json",
                "X-API-Key": neets_api_key
            },
            json={
                "text": text,
                "voice_id": LANGUAGE_CONFIG[language]["voice_id"],
                "params": {
                    "model": "vits"
                }
            }
        )
        
        print(f"TTS Response status: {response.status_code}")
        
        if response.status_code != 200:
            print(f"TTS Response content: {response.content.decode()}")
            raise ValueError(f"TTS API returned status code {response.status_code}")

        # Save the audio to a temporary file
        with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as temp_file:
            temp_file.write(response.content)
            temp_path = temp_file.name

        # Read the audio file
        audio_data, sample_rate = sf.read(temp_path)
        
        # Convert to mono if stereo
        if len(audio_data.shape) > 1:
            audio_data = np.mean(audio_data, axis=1)
        
        # Convert to float32 if needed
        if audio_data.dtype != np.float32:
            audio_data = audio_data.astype(np.float32)
            
        # Clean up the temporary file
        os.unlink(temp_path)
        
        return sample_rate, audio_data
        
    except Exception as e:
        print(f"Audio generation error: {str(e)}")
        if hasattr(response, 'content'):
            print(f"Response content: {response.content[:100]}")  # Print first 100 bytes
        raise

def transcribe_audio(audio_path: str, groq_api_key: str, language: str) -> str:
    """Transcribe audio using Groq's Whisper API"""
    client = Groq(api_key=groq_api_key)
    
    try:
        with open(audio_path, "rb") as audio_file:
            transcription = client.audio.transcriptions.create(
                file=(audio_path, audio_file.read()),
                model="whisper-large-v3-turbo",
                response_format="json",
                language=LANGUAGE_CONFIG[language]["whisper_lang"],
                temperature=0.0
            )
            return transcription.text
    except Exception as e:
        print(f"Transcription error: {str(e)}")
        raise

def chat_with_groq(messages: list, groq_api_key: str, language: str) -> str:
    """Send chat request to Groq API"""
    client = Groq(api_key=groq_api_key)
    
    try:
        response = client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=[
                {"role": "system", "content": LANGUAGE_CONFIG[language]["system_prompt"]},
                *messages
            ],
            temperature=0.7
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"Groq chat error: {str(e)}")
        raise

def process_voice_message(audio, history, neets_api_key: str, groq_api_key: str, 
                         english: bool, french: bool, spanish: bool):
    """Process recorded voice message"""
    if not all([neets_api_key, groq_api_key]):
        return "", [{"role": "error", "content": "Please provide all API keys."}], None
    
    # Check language selection
    selected_languages = []
    if english: selected_languages.append("English")
    if french: selected_languages.append("French")
    if spanish: selected_languages.append("Spanish")
    
    if not selected_languages:
        return "", [{"role": "error", "content": "Please select at least one language."}], None
    if len(selected_languages) > 1:
        return "", [{"role": "error", "content": "Please select only one language."}], None
    
    selected_language = selected_languages[0]
    print(f"Selected language: {selected_language}")
    
    try:
        # Save the recorded audio to a temporary file
        if isinstance(audio, tuple):
            sr, data = audio
            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
                sf.write(temp_file.name, data, sr)
                temp_path = temp_file.name
        else:
            temp_path = audio
        
        # Transcribe the audio using Groq
        transcribed_text = transcribe_audio(temp_path, groq_api_key, selected_language)
        print(f"Transcribed text: {transcribed_text}")
        
        # Clean up temporary file if we created one
        if isinstance(audio, tuple):
            os.unlink(temp_path)
        
        # Prepare messages
        messages = []
        for msg in (history or []):
            if isinstance(msg, dict):
                messages.append({
                    "role": msg["role"],
                    "content": msg["content"]
                })
        
        # Add transcribed message
        messages.append({
            "role": "user",
            "content": transcribed_text
        })
        
        # Get chat completion from Groq
        bot_response = chat_with_groq(messages, groq_api_key, selected_language)
        print(f"Bot response: {bot_response}")
        
        # Generate audio for the response
        try:
            sample_rate, audio_data = generate_audio(bot_response, neets_api_key, selected_language)
            audio_output = (sample_rate, audio_data)
        except Exception as audio_error:
            print(f"Audio generation error: {audio_error}")
            audio_output = None
        
        # Update history with new messages
        new_history = messages + [
            {"role": "assistant", "content": bot_response}
        ]
        
        return transcribed_text, new_history, audio_output
    
    except Exception as e:
        print(f"Processing error: {str(e)}")
        return "", [{"role": "error", "content": f"Error: {str(e)}"}], None

# Create Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        # Logo
        logo_url = "https://i.postimg.cc/cHmNBmg1/ED-COACH-1.jpg"  
        gr.Image(logo_url, label="", show_label=False, height=200)
    gr.Markdown("# BABELlama Multilingual Tutor")
    
    with gr.Row():
        neets_api_key = gr.Textbox(
            label="Neets API Key",
            type="password",
            placeholder="Enter your Neets API key here"
        )
        groq_api_key = gr.Textbox(
            label="Groq API Key",
            type="password",
            placeholder="Enter your Groq API key here"
        )
    
    with gr.Row():
        gr.Markdown("### Select Language")
        english_checkbox = gr.Checkbox(label="English", value=True)
        french_checkbox = gr.Checkbox(label="French")
        spanish_checkbox = gr.Checkbox(label="Spanish")
    
    chatbot = gr.Chatbot(
        type="messages",
        show_copy_button=True,
        layout="bubble"
    )
    
    with gr.Row():
        # Audio recorder for voice input
        audio_input = gr.Audio(
            label="Record Message",
            sources=["microphone"],
            type="numpy"
        )
        
        # Display transcribed text
        transcribed_msg = gr.Textbox(
            label="Transcribed Message",
            placeholder="Your message will appear here after recording...",
            interactive=False
        )
    
    # Audio output for bot response
    audio_output = gr.Audio(
        label="Bot Voice",
        autoplay=True,
        type="numpy"
    )
    
    # Add clear button
    clear = gr.Button("Clear Conversation")
    
    # Handle audio submission
    audio_input.stop_recording(
        process_voice_message,
        inputs=[
            audio_input, chatbot, neets_api_key, groq_api_key,
            english_checkbox, french_checkbox, spanish_checkbox
        ],
        outputs=[transcribed_msg, chatbot, audio_output]
    )
    
    # Make checkboxes mutually exclusive
    english_checkbox.change(
        lambda x: (False, False) if x else (False, False),
        inputs=english_checkbox,
        outputs=[french_checkbox, spanish_checkbox]
    )
    french_checkbox.change(
        lambda x: (False, False) if x else (False, False),
        inputs=french_checkbox,
        outputs=[english_checkbox, spanish_checkbox]
    )
    spanish_checkbox.change(
        lambda x: (False, False) if x else (False, False),
        inputs=spanish_checkbox,
        outputs=[english_checkbox, french_checkbox]
    )
    
    # Handle clear button
    clear.click(
        lambda: ("", [], None),
        outputs=[transcribed_msg, chatbot, audio_output]
    )
    gr.Markdown(warn)

    # "Powered by Groq" badge 
    with gr.Row():
        groq_url = "https://i.ibb.co/FxPsxKF/PBG-mark1-color.png"  
        gr.Image(groq_url, label="", show_label=False, height=50)

if __name__ == "__main__":
    demo.launch()