from dotenv import load_dotenv
import gradio as gr
from transformers import pipeline
import os

load_dotenv()
# Load the Hugging Face model and tokenizer for text generation
hf_token = os.getenv('HF_TOKEN')  # Hugging Face Token for authentication
model_name = "meta-llama/Llama-3-70b-chat-hf"  # Hugging Face model
chat_pipeline = pipeline("text-generation", model=model_name, use_auth_token=hf_token)

# Function to handle the chatbot's response to user queries
# You can only answer finance-related queries.
# - Do not answer non-finance questions.
def run_action(message, history):
    system_prompt = """You are a financial assistant. 
    - Answer in 50 words.
    - Ensure responses adhere to the safety policy."""

    messages = [{"role": "system", "content": system_prompt}]

    # Convert history into the appropriate format
    for entry in history:
        if entry["role"] == "user":
            messages.append({"role": "user", "content": entry["content"]})
        elif entry["role"] == "assistant":
            messages.append({"role": "assistant", "content": entry["content"]})

    # Add the user's current action
    messages.append({"role": "user", "content": message})

    # Generate the model output using Hugging Face's pipeline
    response = chat_pipeline(messages)

    return response[0]['generated_text']

# Main loop for the chatbot to handle user input
def main_loop(message, history):
    """
    Main loop for the chatbot to handle user input.
    """
    # Validate the user's input for safety
    if not is_safe(message):
        return "Your input violates our safety policy. Please try again with a finance-related query."
    
    # Generate and validate the response
    return run_action(message, history)

# Gradio Chat Interface
demo = gr.ChatInterface(
    main_loop,
    chatbot=gr.Chatbot(
        height=450,
        placeholder="Ask a finance-related question. Type 'exit' to quit.",
        type="messages",  # Proper rendering of chat format
    ),
    textbox=gr.Textbox(
        placeholder="What do you want to ask about finance?",
        container=False,
        scale=7,
    ),
    title="Finance Chatbot",
    theme="Monochrome",
    examples=["What is compound interest?", "How to save for retirement?", "What are tax-saving options?"],
    cache_examples=False,
)

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