update with appropriate keys + context builder

app.py

@@ -1,38 +1,38 @@
 import streamlit as st
 from transformers import pipeline
 import ldclient
-from ldclient.config import Config
+from ldclient import LDClient, Context
 
 # Initialize LaunchDarkly client
-ld_client = launchdarkly_api.LDClient("LAUNCHDARKLY_SDK_KEY")
+ldclient.set_config(Config("LAUNCHDARKLY_SDK_KEY"))
+client = ldclient.get()
 
-# Model descriptions
-model_descriptions = {
-    "bert-base-uncased": "BERT base model (uncased)",
-    "roberta-base": "RoBERTa base model",
-    "distilbert-base-uncased": "DistilBERT base model (uncased)",
-    "albert-base-v2": "ALBERT base model v2"
-}
+# 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
 
-# Create a function to get the active model from LaunchDarkly
-def get_active_model():
-    if ld_client.variation("use_bert", {"key": "user"}):
-        return pipeline("sentiment-analysis", model="bert-base-uncased"), "bert-base-uncased"
-    elif ld_client.variation("use_roberta", {"key": "user"}):
-        return pipeline("sentiment-analysis", model="roberta-base"), "roberta-base"
-    elif ld_client.variation("use_distilbert", {"key": "user"}):
-        return pipeline("sentiment-analysis", model="distilbert-base-uncased"), "distilbert-base-uncased"
-    elif ld_client.variation("use_albert", {"key": "user"}):
-        return pipeline("sentiment-analysis", model="albert-base-v2"), "albert-base-v2"
-    else:
-        return pipeline("sentiment-analysis", model="distilbert-base-uncased"), "distilbert-base-uncased"  # Default model
 
 # Streamlit app
-st.title("Sentiment Analysis Demo")
+st.title("Sentiment Analysis Demo with AI Model Flags")
+
 user_input = st.text_area("Enter text for sentiment analysis:")
 
 if st.button("Analyze"):
-    model, model_name = get_active_model()
+    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(f"Model used: {model_descriptions[model_name]}")
-    st.write(result)
+    st.write(result)
+
+# Closing the LD client
+ld_client.close()