import streamlit as st
from transformers import pipeline
from ldclient import LDClient, Config, Context
import os

# Retrieve the LaunchDarkly SDK key from environment variables
ld_sdk_key = os.getenv("LAUNCHDARKLY_SDK_KEY")

# Initialize LaunchDarkly client with the correct configuration
ld_client = LDClient(Config(ld_sdk_key))

# Function to get the AI model configuration from LaunchDarkly
def get_model_config():
    flag_key = "swap-sentiment-models"  # Replace with your flag key
    # Create context using Context builder
    context = Context.builder("context-key-123abc").name("Sandy").build()
    flag_variation = ld_client.variation(flag_key, context, default={})
    
    model_id = flag_variation.get("modelID", "distilbert-base-uncased")
    return model_id


# Streamlit app
st.title("Sentiment Analysis Demo with AI Model Flags")

user_input = st.text_area("Enter text for sentiment analysis:")

if st.button("Analyze"):
    model_id = get_model_config()
    model = pipeline("sentiment-analysis", model=model_id)
    
    # Display model details
    st.write(f"Using model: {model_id}")
    
    # Perform sentiment analysis
    result = model(user_input)
    st.write(result)

# Closing the LD client
ld_client.close()