Update app.py

app.py

@@ -55,7 +55,6 @@ def interact_with_huggingface_model(messages, model):
 # Function to interact with the Together API model
 def interact_with_together_api(messages):
     all_messages = []
-    entire_assistant_response = ""  # Initialize the variable
 
     if not messages:  # If history is empty
         all_messages.append({"role": "user", "content": ""})
@@ -76,7 +75,6 @@ def interact_with_together_api(messages):
         "repetition_penalty": 1,
         "n": 1,
         "messages": all_messages,
-        "stream_tokens": True,
     }
 
     TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
@@ -86,34 +84,14 @@ def interact_with_together_api(messages):
         "Authorization": f"Bearer {TOGETHER_API_KEY}",
     }
 
-    response = requests.post(url, json=payload, headers=headers, stream=True)
+    response = requests.post(url, json=payload, headers=headers)
     response.raise_for_status()  # Ensure HTTP request was successful
 
-    for line in response.iter_lines():
-        if line:
-            decoded_line = line.decode('utf-8')
-
-            # Check for the completion signal
-            if decoded_line == "data: [DONE]":
-                if entire_assistant_response:
-                    yield entire_assistant_response  # Yield the entire response at the end
-                break
-
-            try:
-                # Decode and strip any SSE format specific prefix ("data: ")
-                if decoded_line.startswith("data: "):
-                    decoded_line = decoded_line.replace("data: ", "")
-                    chunk_data = json.loads(decoded_line)
-                    content = chunk_data['choices'][0]['delta']['content']
-                    entire_assistant_response += content  # Aggregate content
-                    yield content  # Yield each chunk separately
-
-            except json.JSONDecodeError:
-                print(f"Invalid JSON received: {decoded_line}")
-                continue
-            except KeyError as e:
-                print(f"KeyError encountered: {e}")
-                continue
+    # Extract response from JSON
+    response_data = response.json()
+    assistant_response = response_data["choices"][0]["message"]
+
+    return assistant_response
 
 # Create sidebar with model selection dropdown and temperature slider
 selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
@@ -147,15 +125,15 @@ if prompt := st.chat_input(f"Hi, I'm {selected_model}, ask me a question"):
 
     # Interact with selected model
     if selected_model == "Nous-Hermes-2-Yi-34B":
-        stream = interact_with_together_api(st.session_state.messages)
+        assistant_response = interact_with_together_api(st.session_state.messages)
     else:
-        stream = interact_with_huggingface_model(st.session_state.messages, model_links[selected_model])
+        assistant_response = interact_with_huggingface_model(st.session_state.messages, model_links[selected_model])
 
     # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        response = ""
-        if stream:
-            for chunk in stream:
-                response = chunk
-                st.markdown(response)
-    st.session_state.messages.append(("assistant", response))
+    with st.empty():
+        st.markdown("AI is typing...")
+        st.empty()
+        st.markdown(assistant_response)
+
+    # Add assistant response to chat history
+    st.session_state.messages.append(("assistant", assistant_response))