fixed the handling of the streamed response

app.py

@@ -3,6 +3,7 @@ import base64
 import requests
 import streamlit as st
 import json
+import tempfile
 
 if "stream" not in st.session_state:
     st.session_state.stream = True
@@ -16,40 +17,36 @@ def encode_image(image_path):
     with open(image_path, "rb") as image_file:
         return base64.b64encode(image_file.read()).decode('utf-8')
 
-def extract_content(chunk):
-    try:
-        decoded_chunk = chunk.decode('utf-8')
-        json_data = decoded_chunk.split('data: ')[1]
-        parsed_data = json.loads(json_data)
-        content = parsed_data['choices'][0]['delta']['content']
-        return content
-    except json.JSONDecodeError as e:
-        #ignore the error
-        return ""
-    
-  
 def main():
-    st.title("Multimodal Image Analysis with " + MODEL_ID)
+    st.title(f"Multimodal Image Analysis with {MODEL_ID}")
 
-    text = """Prof. Louie F. Cervantes, M. Eng. (Information Engineering)
+    # Display about section
+    about_text = """Prof. Louie F. Cervantes, M. Eng. (Information Engineering)
     CCS 229 - Intelligent Systems
     Department of Computer Science
     College of Information and Communications Technology
     West Visayas State University
     """
     with st.expander("About"):
-        st.text(text)
+        st.text(about_text)
 
     st.write("Upload an image and select the image analysis task.")
 
     # File upload for image
     uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+    temp_file_path = None
+
     if uploaded_image is not None:
-        # Encode the uploaded image to base64
-        base64_image = base64.b64encode(uploaded_image.getvalue()).decode('utf-8')
+        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+            temp_file.write(uploaded_image.getvalue())
+            temp_file_path = temp_file.name
+
+        # Encode image as Base64
+        with open(temp_file_path, "rb") as f:
+            base64_image = base64.b64encode(f.read()).decode()
 
         # Display the uploaded image
-        st.image(uploaded_image, caption="Uploaded Image",  use_container_width=True)
+        st.image(uploaded_image, caption="Uploaded Image", use_container_width=True)
 
     # List of image analysis tasks
     analysis_tasks = [
@@ -66,83 +63,93 @@ def main():
     ]
 
     # Task selection dropdown
-    selected_task = st.selectbox("Select an image analysis task:", analysis_tasks)
-    
-    
+    selected_task = st.selectbox("Select an image analysis task:", [""] + analysis_tasks)
+
+    # Checkbox for streaming
+    stream = st.checkbox("Begin streaming the AI response as soon as it is available.", value=st.session_state.stream)
 
     if st.button("Generate Response"):
-        st.session_state.stream = st.checkbox("Begin streaming the AI response as soon as it is available.", value=True)    
-        stream = st.session_state.stream
+        if not api_key:
+            st.error("API key not found. Please set the NVIDIA_VISION_API_KEY environment variable.")
+            return
 
-        if uploaded_image is None or selected_task == "":
-            st.error("Please upload an image and select a task.")
+        if uploaded_image is None:
+            st.error("Please upload an image.")
             return
 
-        else:
-            headers = {
-                "Authorization": f"Bearer {api_key}",
-                "Accept": "text/event-stream" if stream else "application/json"
-            }
-
-            # Prepare the multimodal prompt
-            payload = {
-                "model": MODEL_ID,
-                "messages": [
-                    {
-                        "role": "user",
-                        "content": f'{selected_task} <img src="data:image/png;base64,{base64_image}" />'
-                    }
-                ],
-                "max_tokens": 512,
-                "temperature": 1.00,
-                "top_p": 1.00,
-                "stream": stream  
-            }
+        if not selected_task:
+            st.error("Please select an image analysis task.")
+            return
 
+        # Headers for the API call
+        headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Accept": "text/event-stream" if stream else "application/json"
+        }
+
+        # Prepare the multimodal prompt
+        payload = {
+            "model": MODEL_ID,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": f'{selected_task} <img src="data:image/png;base64,{base64_image}" />'
+                }
+            ],
+            "max_tokens": 512,
+            "temperature": 1.0,
+            "top_p": 1.0,
+            "stream": stream
+        }
+
+        try:
             with st.spinner("Processing..."):
                 response = requests.post(
                     invoke_url,
                     headers=headers,
                     json=payload,
-                    stream=stream  # Important for streaming
+                    stream=stream
                 )
+                response.raise_for_status()  # Raise exception for HTTP errors
 
                 if stream:
+                    # Handle streaming response
                     response_container = st.empty()
                     content = ""
-                    # Efficiently handle streaming response
-                    for chunk in response.iter_lines(): 
-                        if len(chunk) > 0:                             
-                            # Decode the bytes object into a string
-                            chunk_str = chunk.decode('utf-8')
-                            # Remove the "data: " prefix
-                            if chunk_str.startswith("data: "):
-                                chunk_str = chunk_str[6:]
-                            if chunk_str.strip() == "[DONE]":
-                                break
-                            # Check if the string is not empty
-                            if chunk_str.strip() != "":
+                    for chunk in response.iter_lines(decode_unicode=True):
+                        if chunk:
+                            if "[DONE]" in chunk:
+                                # Handle the end chunk
+                                st.write("Response generation complete.")
+                                break                            
+                            # Check if the chunk is a JSON string
+                            elif chunk.startswith("data:"):
+                                chunk = chunk[5:].strip()  # Remove the "data:" prefix
                                 try:
-                                    # Attempt to parse the string as JSON
-                                    chunk_dict = json.loads(chunk_str)
-                                    # Now you can access the 'choices' key
-                                    content += chunk_dict['choices'][0]['delta']['content']
-                                    response_container.markdown(content)
+                                    if len(chunk) > 0:
+                                        chunk_dict = json.loads(chunk)
+                                        if "choices" in chunk_dict and chunk_dict["choices"]:
+                                            delta_content = chunk_dict["choices"][0]["delta"]["content"]
+                                            content += delta_content
+                                            response_container.write(content)
+
                                 except json.JSONDecodeError as e:
-                                    # Handle the error if the string is not valid JSON
-                                    print(f"Error parsing JSON: {e}")
-                                    print(f"Invalid JSON string: {chunk_str}")
+                                    st.error(f"Error parsing JSON: {e}")
 
                 else:
-                    try:
-                        content = response.json()
-                        content_string = content.get('choices', [{}])[0].get('message', {}).get('content', '')
-                        st.write(f"AI Response: {content_string}")
-
-                        st.success("Response generated!")
+                    # Handle non-streaming response
+                    content = response.json()
+                    content_string = content.get("choices", [{}])[0].get("message", {}).get("content", "")
+                    st.write(f"AI Response: {content_string}")
+                    st.success("Response generated!")
+        except requests.exceptions.RequestException as e:
+            st.error(f"An error occurred while making the API call: {e}")
+        except Exception as e:
+            st.error(f"An unexpected error occurred: {e}")
+        finally:
+            # Clean up temporary file
+            if temp_file_path and os.path.exists(temp_file_path):
+                os.remove(temp_file_path)
 
-                    except Exception as e:
-                        st.error(f"An error occurred: {e}")
- 
 if __name__ == "__main__":
-    main()
+    main()