Update src/streamlit_app.py

src/streamlit_app.py

@@ -11,16 +11,10 @@ endpoint_data = json.load(open(f"{working_dir}/model_info.json"))
 def clear_chat():
     st.session_state.messages = []
 
-def get_api_key():
-    # Retrieve API key from environment variable or prompt user
-    return os.getenv("OPENAI_API_KEY") or st.text_input("Enter your API Key", type="password")    
-
-st.title("AIaaS on Intel® Gaudi® Demo")
-
-# Extract endpoint and model names from JSON data
-endpoint = endpoint_data['endpoint']
-model_names = endpoint_data['models']
+st.title("Inference as a Service Playground")
 
+# Extract the keys (model names) from the JSON data
+model_names = list(endpoint_data.keys())
 
 with st.sidebar:
     modelname = st.selectbox("Select a LLM model (Running on Intel® Gaudi®) ", model_names)
@@ -28,7 +22,7 @@ with st.sidebar:
     st.button("Start New Chat", on_click=clear_chat)
 
     # Add a text input for the API key
-    api_key = get_api_key()
+    api_key = st.text_input("Enter your API Key", type="password")
     if api_key:
         st.session_state.api_key = api_key    
 
@@ -37,12 +31,15 @@ if "api_key" not in st.session_state or not st.session_state.api_key:
     st.error("Please enter your API Key in the sidebar.")
 else:
     try:
+        endpoint = endpoint_data[modelname]
+        
         api_key = st.session_state.api_key
         base_url = endpoint
         client = OpenAI(api_key=api_key, base_url=base_url)
 
-        print(f"Selected Model --> {modelname}")
-        st.write(f"**Model Info:** `{modelname}`")
+        # Extract the model name
+        models = client.models.list()
+        modelname = models.data[0].id
 
         if "messages" not in st.session_state:
             st.session_state.messages = []
@@ -65,10 +62,8 @@ else:
                             for m in st.session_state.messages
                         ],
                         max_tokens=1024,
-						temperature=0,
                         stream=True,
                     )
-                    
                     response = st.write_stream(stream)
                 except Exception as e:
                     st.error(f"An error occurred while generating the response: {e}")