iukea/open-notebooklm

1"""2utils.py3 4Functions:5- generate_script: Get the dialogue from the LLM.6- call_llm: Call the LLM with the given prompt and dialogue format.7- parse_url: Parse the given URL and return the text content.8- generate_podcast_audio: Generate audio for podcast using TTS or advanced audio models.9"""10 11# Standard library imports12import time13from typing import Any, Union14 15# Third-party imports16import requests17from bark import SAMPLE_RATE, generate_audio, preload_models18from gradio_client import Client19from openai import OpenAI20from pydantic import ValidationError21from scipy.io.wavfile import write as write_wav22 23# Local imports24from constants import (25    FIREWORKS_API_KEY,26    FIREWORKS_BASE_URL,27    FIREWORKS_MODEL_ID,28    FIREWORKS_MAX_TOKENS,29    FIREWORKS_TEMPERATURE,30    FIREWORKS_JSON_RETRY_ATTEMPTS,31    MELO_API_NAME,32    MELO_TTS_SPACES_ID,33    MELO_RETRY_ATTEMPTS,34    MELO_RETRY_DELAY,35    JINA_READER_URL,36    JINA_RETRY_ATTEMPTS,37    JINA_RETRY_DELAY,38)39from schema import ShortDialogue, MediumDialogue40 41# Initialize clients42fw_client = OpenAI(base_url=FIREWORKS_BASE_URL, api_key=FIREWORKS_API_KEY)43hf_client = Client(MELO_TTS_SPACES_ID)44 45# Download and load all models for Bark46preload_models()47 48 49def generate_script(50    system_prompt: str,51    input_text: str,52    output_model: Union[ShortDialogue, MediumDialogue],53) -> Union[ShortDialogue, MediumDialogue]:54    """Get the dialogue from the LLM."""55 56    # Call the LLM57    response = call_llm(system_prompt, input_text, output_model)58    response_json = response.choices[0].message.content59 60    # Validate the response61    for attempt in range(FIREWORKS_JSON_RETRY_ATTEMPTS):62        try:63            first_draft_dialogue = output_model.model_validate_json(response_json)64            break65        except ValidationError as e:66            if attempt == FIREWORKS_JSON_RETRY_ATTEMPTS - 1:  # Last attempt67                raise ValueError(68                    f"Failed to parse dialogue JSON after {FIREWORKS_JSON_RETRY_ATTEMPTS} attempts: {e}"69                ) from e70            error_message = (71                f"Failed to parse dialogue JSON (attempt {attempt + 1}): {e}"72            )73            # Re-call the LLM with the error message74            system_prompt_with_error = f"{system_prompt}\n\nPlease return a VALID JSON object. This was the earlier error: {error_message}"75            response = call_llm(system_prompt_with_error, input_text, output_model)76            response_json = response.choices[0].message.content77            first_draft_dialogue = output_model.model_validate_json(response_json)78 79    # Call the LLM a second time to improve the dialogue80    system_prompt_with_dialogue = f"{system_prompt}\n\nHere is the first draft of the dialogue you provided:\n\n{first_draft_dialogue}."81 82    # Validate the response83    for attempt in range(FIREWORKS_JSON_RETRY_ATTEMPTS):84        try:85            response = call_llm(86                system_prompt_with_dialogue,87                "Please improve the dialogue. Make it more natural and engaging.",88                output_model,89            )90            final_dialogue = output_model.model_validate_json(91                response.choices[0].message.content92            )93            break94        except ValidationError as e:95            if attempt == FIREWORKS_JSON_RETRY_ATTEMPTS - 1:  # Last attempt96                raise ValueError(97                    f"Failed to improve dialogue after {FIREWORKS_JSON_RETRY_ATTEMPTS} attempts: {e}"98                ) from e99            error_message = f"Failed to improve dialogue (attempt {attempt + 1}): {e}"100            system_prompt_with_dialogue += f"\n\nPlease return a VALID JSON object. This was the earlier error: {error_message}"101    return final_dialogue102 103 104def call_llm(system_prompt: str, text: str, dialogue_format: Any) -> Any:105    """Call the LLM with the given prompt and dialogue format."""106    response = fw_client.chat.completions.create(107        messages=[108            {"role": "system", "content": system_prompt},109            {"role": "user", "content": text},110        ],111        model=FIREWORKS_MODEL_ID,112        max_tokens=FIREWORKS_MAX_TOKENS,113        temperature=FIREWORKS_TEMPERATURE,114        response_format={115            "type": "json_object",116            "schema": dialogue_format.model_json_schema(),117        },118    )119    return response120 121 122def parse_url(url: str) -> str:123    """Parse the given URL and return the text content."""124    for attempt in range(JINA_RETRY_ATTEMPTS):125        try:126            full_url = f"{JINA_READER_URL}{url}"127            response = requests.get(full_url, timeout=60)128            response.raise_for_status()  # Raise an exception for bad status codes129            break130        except requests.RequestException as e:131            if attempt == JINA_RETRY_ATTEMPTS - 1:  # Last attempt132                raise ValueError(133                    f"Failed to fetch URL after {JINA_RETRY_ATTEMPTS} attempts: {e}"134                ) from e135            time.sleep(JINA_RETRY_DELAY)  # Wait for X second before retrying136    return response.text137 138 139def generate_podcast_audio(140    text: str, speaker: str, language: str, use_advanced_audio: bool, random_voice_number: int141) -> str:142    """Generate audio for podcast using TTS or advanced audio models."""143    if use_advanced_audio:144        return _use_suno_model(text, speaker, language, random_voice_number)145    else:146        return _use_melotts_api(text, speaker, language)147 148 149def _use_suno_model(text: str, speaker: str, language: str, random_voice_number: int) -> str:150    """Generate advanced audio using Bark."""151    host_voice_num = str(random_voice_number)152    guest_voice_num = str(random_voice_number + 1)153    audio_array = generate_audio(154        text,155        history_prompt=f"v2/{language}_speaker_{host_voice_num if speaker == 'Host (Jane)' else guest_voice_num}",156    )157    file_path = f"audio_{language}_{speaker}.mp3"158    write_wav(file_path, SAMPLE_RATE, audio_array)159    return file_path160 161 162def _use_melotts_api(text: str, speaker: str, language: str) -> str:163    """Generate audio using TTS model."""164    accent, speed = _get_melo_tts_params(speaker, language)165 166    for attempt in range(MELO_RETRY_ATTEMPTS):167        try:168            return hf_client.predict(169                text=text,170                language=language,171                speaker=accent,172                speed=speed,173                api_name=MELO_API_NAME,174            )175        except Exception as e:176            if attempt == MELO_RETRY_ATTEMPTS - 1:  # Last attempt177                raise  # Re-raise the last exception if all attempts fail178            time.sleep(MELO_RETRY_DELAY)  # Wait for X second before retrying179 180 181def _get_melo_tts_params(speaker: str, language: str) -> tuple[str, float]:182    """Get TTS parameters based on speaker and language."""183    if speaker == "Guest":184        accent = "EN-US" if language == "EN" else language185        speed = 0.9186    else:  # host187        accent = "EN-Default" if language == "EN" else language188        speed = (189            1.1 if language != "EN" else 1190        )  # if the language is not English, try speeding up so it'll sound different from the host191        # for non-English, there is only one voice192    return accent, speed193