import streamlit as st
from openai import OpenAI
import os

# Get API Key from environment variable
api_key = os.getenv("NVIDIA_API_KEY")

# Check if the API key is found
if api_key is None:
    st.error("NVIDIA_API_KEY environment variable not found.")
else:
    # Initialize the OpenAI client
    client = OpenAI(
        base_url="https://integrate.api.nvidia.com/v1",
        api_key=api_key
    )
def get_response(user_prompt):
    try:
        # Create and send the chat completion request
        completion = client.chat.completions.create(
            model="nvidia/llama-3.1-nemotron-70b-instruct",
            messages=[{"role": "user", "content": user_prompt}],
            temperature=0.5,  # Adjust temperature for creativity
            top_p=1,
            max_tokens=1024,
            stream=True  # Enable streaming response
        )

        # Display the streaming response
        response_container = st.empty()
        full_response = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                full_response += chunk.choices[0].delta.content                
                response_container.write(full_response)
    except Exception as e:
        st.error(f"An error occurred: {e}")

def main():
    # Streamlit app
    st.title("NVIDIA Nemotron Chatbot")

    # User Prompt Input
    user_prompt = st.text_area("Enter your prompt:")

    # When the user clicks "Submit"
    if st.button("Submit"):
        if user_prompt:
            # Call the OpenAI API
            get_response(user_prompt)            
        else:
            st.warning("Please enter a prompt.")

if __name__ == "__main__":
    main()