tools/generate_lyrics.py

import json
import os
import sys
from pathlib import Path
from typing import List

from dotenv import load_dotenv
from google import genai

from schemas.lyrics import LyricsSection, SongStructure

sys.path.append(str(Path(__file__).parent.parent.parent))

load_dotenv()


def generate_structured_lyrics(conversation: List[dict], genre: str, mood: str, theme: str) -> SongStructure:
    """

    Generate structured lyrics using Gemini API based on conversation history and song parameters.



    This function takes a conversation history and song preferences, then uses Google's

    Gemini AI to generate structured song lyrics following a specific format.



    Args:

        conversation: List of conversation messages with 'role' and 'content' keys

        genre: Musical genre for the song (e.g., "pop", "rock")

        mood: Emotional mood for the song (e.g., "romantic", "sad")

        theme: Subject matter or theme of the song (e.g., "love", "friendship")



    Returns:

        SongStructure: A structured representation of the generated song



    Raises:

        ValueError: If no API key is found or if the model returns an empty response

        Exception: For any other errors during API communication or response parsing

    """
    api_key = os.environ.get("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("Please set the GEMINI_API_KEY environment variable")

    genai_client = genai.Client(api_key=api_key)
    model = "gemini-2.0-flash"

    # Convert conversation history into a single prompt string
    conversation_text = "\n".join([
        f"{'User' if msg['role'] == 'user' else 'Assistant'}: {msg['content']}"
        for msg in conversation
    ])

    prompt = f"""Based on the following conversation:



{conversation_text}



Create song lyrics with these parameters:

- Genre: {genre}

- Mood: {mood}

- Theme: {theme}

"""

    system_instruction = """Generate a complete song with the following structure:

1. A title

2. At least one verse

3. A chorus

4. Optional bridge

5. Optional outro



The output must follow the exact JSON structure with these section types: VERSE, CHORUS, BRIDGE, OUTRO.

"""

    try:
        config = {
            'response_mime_type': 'application/json',
            'response_schema': SongStructure,
            'system_instruction': system_instruction
        }

        print(prompt)

        # Pass the prompt as a single string instead of conversation list
        response = genai_client.models.generate_content(
            contents=prompt,
            model=model,
            config=config
        )

        if not response.text:
            raise ValueError("No response generated from the model")

        # Parse the JSON string into a dictionary
        lyrics_data = json.loads(response.text)

        sections = []
        for section in lyrics_data["sections"]:
            sections.append(LyricsSection(
                section_type=section["section_type"],
                content=section["content"]
            ))

        song_structure = SongStructure(
            title=lyrics_data["title"],
            sections=sections,
        )

        return song_structure

    except Exception as e:
        print(f"Error generating lyrics: {str(e)}")
        raise


def format_lyrics(song_structure: SongStructure) -> str:
    """

    Convert the structured lyrics into a formatted string for display.



    This function takes a SongStructure object and creates a human-readable

    text representation with appropriate section headings and spacing.



    Args:

        song_structure: A SongStructure object containing the song title and sections



    Returns:

        str: Formatted lyrics text ready for display

    """
    formatted = f"TITLE: {song_structure.title}\n\n"

    for section in song_structure.sections:
        formatted += f"{section.section_type}:\n{section.content}\n\n"

    return formatted.strip()


def format_lyrics_for_yue(song_structure: SongStructure, genre: str, mood: str, theme: str) -> str:
    """

    Format lyrics in a specialized structure for the YUE music generation system.



    This function converts a SongStructure object into a specific format expected by

    the YUE music generation system, including genre and mood descriptors to guide

    the musical style generation.



    Args:

        song_structure: A SongStructure object containing the song title and sections

        genre: Musical genre that influences the instrumentation descriptors

        mood: Emotional mood that influences the style descriptors

        theme: Subject matter of the song (not directly used in formatting)



    Returns:

        str: Formatted text with appropriate section markers and descriptors for YUE

    """
    genre_descriptors = {
        "pop": "pop vocal clear melodic synthesizer",
        "rock": "rock electric-guitar drums powerful energetic",
        "jazz": "jazz piano smooth saxophone melodic",
        "hip-hop": "rap hip-hop beats vocal rhythmic",
        "electronic": "electronic synthesizer beats modern"
    }

    mood_descriptors = {
        "upbeat": "energetic bright positive",
        "sad": "melancholic emotional soft",
        "energetic": "dynamic powerful strong",
        "chill": "relaxed smooth gentle",
        "romantic": "soft emotional intimate"
    }

    # Create combined genre description
    base_genre = genre_descriptors.get(genre.lower(), "")
    mood_desc = mood_descriptors.get(mood.lower(), "")

    formatted = "Generate music from the given lyrics segment by segment.\n"

    formatted += f"[Genre] {base_genre} {mood_desc} clear vocal\n\n"

    formatted += f"[Title] {song_structure.title}\n\n"

    for section in song_structure.sections:
        section_type = section.section_type.value.lower()

        formatted += f"[{section_type}]\n"

        lines = section.content.strip().split('\n')
        formatted += '\n'.join(line.strip() for line in lines if line.strip())
        formatted += '\n\n'

    return formatted.strip()


# Example usage:
if __name__ == "__main__":
    """

    Test module functionality by generating sample lyrics and displaying them.



    This test code creates a mock conversation about writing a love song,

    generates structured lyrics using the Gemini API, and displays both

    the raw structure and formatted lyrics.

    """
    # Test conversation
    test_conversation = [
        {"role": "user", "content": "I want to write a love song"},
        {"role": "assistant", "content": "I'll help you create a love song. What style are you thinking of?"},
        {"role": "user", "content": "Something romantic and modern"},
        {"role": "assistant",
         "content": "Perfect! Here are the lyrics I've created:\n\nVERSE:\nSoft city lights paint the evening sky\nAs I think about you and I\nEvery moment we've shared feels right\nLike stars aligned in the night\n\nCHORUS:\nThis modern love, it's all we need\nBreaking rules and setting us free\nEvery text, every call, every memory\nMakes this love our reality"}
    ]

    try:
        song = generate_structured_lyrics(
            conversation=test_conversation,
            genre="pop",
            mood="romantic",
            theme="love"
        )

        print("Generated Song Structure:")
        print(song.json(indent=2))

        print("\nFormatted Lyrics:")
        print(format_lyrics(song))

    except Exception as e:
        print(f"Error in test: {str(e)}")