import streamlit as st
import os
from openai import OpenAI

# Initialize the OpenAI client
client = OpenAI(
    base_url="https://api.studio.nebius.ai/v1/",
    api_key=os.environ.get("NEBIUS_API_KEY")  # Make sure to set this in Hugging Face Secrets
)

# Streamlit app title
st.title("AI Title Generator")

# Text Area Input Bar
user_input = st.text_area(
    label="Enter a description for generating titles:",
    placeholder="e.g., Man who went to jail for no reason"
)

# Generate Button
if st.button("Generate Titles"):
    if user_input.strip():
        try:
            # API call to OpenAI
            completion = client.chat.completions.create(
                model="nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
                messages=[
                    {"role": "system", "content": "Your task is to generate 3 very short titles based on the user input."},
                    {"role": "user", "content": user_input}
                ],
                temperature=0.6,
                maxTokens=512,
                topP=0.9,
                topK=50
            )

            # Extracting the generated titles
            response_content = completion.choices[0].message["content"]

            # Output Text Area
            st.text_area(
                label="Generated Titles:",
                value=response_content,
                height=200,
                disabled=True
            )

        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.warning("Please provide input before clicking Generate.")