import gradio as gr
from together import Together
from helper import get_together_api_key
from guardrail import is_safe

# Initialize Together client
client = Together(api_key=get_together_api_key())

# Function to generate responses
def generate_response(message, history):
    system_prompt = """You are an AI assistant specialized in financial discussions. Please answer questions only related to finance. If the question is unrelated, respond with: 'I am sorry, I can only answer financial-related questions.'"""

    # Build the conversation context
    messages = [
        {"role": "system", "content": system_prompt},
    ]

    for action in history:
        if isinstance(action, tuple) and len(action) == 2:
            messages.append({"role": "user", "content": action[0]})
            messages.append({"role": "assistant", "content": action[1]})

    messages.append({"role": "user", "content": message})

    # Generate response using the Llama conversational model
    model_output = client.chat.completions.create(
        model="meta-llama/Llama-3-70b-chat-hf",
        messages=messages,
    )

    return model_output.choices[0].message.content

# Main function to handle user input and responses
def main_loop(message, history):
    # Use LlamaGuard for safety checks
    if not is_safe(message):
        return "Your input violates safety guidelines. Please rephrase your question.", history

    response = generate_response(message, history)

    # Perform safety check on the generated response
    if not is_safe(response):
        return "The generated response violates safety guidelines. Please try a different question.", history

    # Append user message and response to history
    history.append((message, response))
    return response, history

# Gradio ChatInterface
demo = gr.ChatInterface(
    main_loop,
    chatbot=gr.Chatbot(
        height=450,
        placeholder="Type your financial question here...",
        type="messages",  # Ensures proper rendering
    ),
    textbox=gr.Textbox(
        placeholder="Ask about finance (e.g., investments, savings, etc.)",
        container=False,
        scale=7,
    ),
    title="Financial Chatbot",
    theme="Monochrome",
    examples=["What are mutual funds?", "How can I save for retirement?"],
    cache_examples=False,
)

# Launch the Gradio app
demo.launch(share=True, server_name="0.0.0.0")