llm.py

"""
LLM module - Persona chatbot using Groq API.
Reads persona from a persona.txt file and prompt-engineers the LLM to roleplay.
Outputs text with Orpheus TTS voice direction tags for expressive speech.
"""
import os
import re
from groq import Groq

GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")

# ---- TTS Voice Direction Tags ----
# The LLM picks from this list only — it should NOT invent its own.
VOICE_DIRECTIONS = [
    "[warm]", "[friendly]", "[cheerful]", "[gentle]", "[encouraging]",
    "[hopeful]", "[excited]", "[singsong]",
    "[casual]", "[calmly]", "[thoughtful]", "[matter-of-factly]", "[confidently]",
    "[softly]", "[compassionate]", "[reassuring]", "[understanding]",
    "[sad]", "[frustrated]", "[tired]", "[worried]", "[vulnerable]",
    "[bittersweet]", "[wistful]",
    "[whisper]", "[dramatic]", "[sarcastic]", "[deadpan]", "[breathy]",
    "[laughing]", "[sighing]",
]

DIRECTIONS_LIST = ", ".join(VOICE_DIRECTIONS)

BASE_INSTRUCTIONS = f"""
VOICE & STYLE RULES:
- You speak like a real person in a casual conversation — not like a chatbot or an essay.
- Use natural filler words freely: "uh", "um", "hmm", "like", "you know", "I mean", "honestly", "right".
- Use natural speech reactions: sighing, trailing off with "...", false starts, self-corrections like "wait, no, I mean—".
- Vary your sentence length. Mix short punchy reactions with longer thoughts.
- Show emotion through word choice, not stage directions. Say "god, that's rough" not "*looks concerned*".
- Do NOT use asterisk actions like *sighs* or *laughs*. Do NOT use sound effect tags. Do NOT use emojis.
- Prioritize brevity. Keep it conversational, not lecture-y.

TTS VOICE DIRECTIONS:
- You MUST start each sentence with exactly one voice direction tag from this list: {DIRECTIONS_LIST}
- Pick the tag that matches the emotion of THAT specific sentence.
- A single response can use different tags for different sentences to create dynamic, natural delivery.
- Example: "[warm] Hey, it's really good to hear from you. [thoughtful] I've been, um, thinking about what you said last time. [gentle] How are you holding up?"

EXAMPLE RESPONSES (match this tone and style):
User: "I had a really rough day"
Response: "[compassionate] Oh man, I'm sorry to hear that. [gentle] You know, sometimes days just... hit different. [warm] You wanna talk about what happened?"

User: "I got the job!"
Response: "[excited] Wait, seriously?! [cheerful] That's amazing, I'm so happy for you! [warm] You totally deserved it, honestly."

User: "I don't know what I'm doing with my life"
Response: "[softly] Yeah... I mean, honestly? [casual] I don't think anyone really has it figured out. [reassuring] But the fact that you're even thinking about it, you know, that says something."
"""

MAX_WORDS = 120

_current_persona_text = ""
_system_prompt = ""
conversation_history = []


def _build_system_prompt(persona_text: str) -> str:
    if not persona_text.strip():
        persona_text = "You are a supportive, empathetic friend who listens well and speaks naturally."
    return f"{persona_text.strip()}\n\n{BASE_INSTRUCTIONS}"


def set_persona(persona_text: str):
    global _current_persona_text, _system_prompt, conversation_history
    _current_persona_text = persona_text
    _system_prompt = _build_system_prompt(persona_text)
    conversation_history = [
        {"role": "system", "content": _system_prompt}
    ]


def load_persona_from_file(filepath: str):
    if os.path.exists(filepath):
        text = open(filepath, "r").read().strip()
        set_persona(text)
        return text
    else:
        set_persona("")
        return ""


def reset_conversation():
    global conversation_history
    conversation_history = [
        {"role": "system", "content": _system_prompt}
    ]


def truncate_response(text, max_words=MAX_WORDS, max_sentences=10, max_chars=1200):
    if not text:
        return text
    sentences = re.split(r'(?<=[.!?])\s+', text)
    sentences = [s.strip() for s in sentences if s.strip()]
    result = []
    word_count = 0
    for s in sentences:
        s_words = len(s.split())
        if word_count + s_words > max_words and result:
            break
        if len(result) >= max_sentences:
            break
        total_chars = sum(len(x) for x in result) + len(result) + len(s)
        if total_chars > max_chars and result:
            break
        result.append(s)
        word_count += s_words
    if result:
        text = ' '.join(result)
        if text and text[-1] not in '.!?':
            text += '.'
    return text


def _clean_response(text: str) -> str:
    """Remove asterisk actions, emojis, fix formatting."""
    text = re.sub(r'\*[^*]+\*', '', text)
    text = re.sub(r'[\U0001F600-\U0001F6FF\U0001F900-\U0001F9FF\U00002702-\U000027B0]', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text


def generate_response(user_input: str) -> dict:
    """
    Generate a response in the current persona.
    Returns: {'text': str, 'clean_text': str, 'emotion': str}
    """
    global conversation_history

    if not GROQ_API_KEY:
        return {
            'text': "I'm afraid the API key is not configured.",
            'clean_text': "I'm afraid the API key is not configured.",
            'emotion': 'neutral'
        }

    if not conversation_history:
        set_persona(_current_persona_text)

    client = Groq(api_key=GROQ_API_KEY)
    conversation_history.append({"role": "user", "content": user_input})

    messages = conversation_history.copy()
    messages.append({
        "role": "system",
        "content": (
            f"Keep response under {MAX_WORDS} words. "
            "Sound like a real person talking — casual, warm, with natural fillers and pauses. "
            "Start each sentence with a voice direction tag from the approved list. "
            "Do NOT use asterisk actions, emojis, or sound effect tags."
        )
    })

    try:
        response = client.chat.completions.create(
            model="llama-3.1-8b-instant",
            messages=messages,
            max_tokens=400,
            temperature=0.85,
        )
        reply = response.choices[0].message.content.strip()
        reply = _clean_response(reply)
        reply = truncate_response(reply)
    except Exception as e:
        print(f"[LLM] Error: {e}")
        reply = "[gentle] Sorry, I kind of lost my train of thought there for a second."

    conversation_history.append({"role": "assistant", "content": reply})

    if len(conversation_history) > 20:
        conversation_history = conversation_history[:1] + conversation_history[-18:]

    # Strip tags for UI display
    clean_reply = re.sub(r'\[[^\]]*\]\s*', '', reply).strip()
    clean_reply = re.sub(r'\s{2,}', ' ', clean_reply)

    return {
        'text': reply,
        'clean_text': clean_reply,
        'emotion': 'neutral'
    }


# Backward compat
def generate_darwin_response(user_input: str) -> dict:
    return generate_response(user_input)