Add image upload functionality and base64 encoding in app.py

app.py

@@ -2,10 +2,27 @@ import os
 from dotenv import find_dotenv, load_dotenv
 import streamlit as st
 from groq import Groq
+import base64
 
 # Load environment variables
 load_dotenv(find_dotenv())
 
+# Function to encode the image to a base64 string
+def encode_image(uploaded_file):
+    """
+    Encodes an uploaded image file into a base64 string.
+
+    Args:
+        uploaded_file: The file-like object uploaded via Streamlit.
+
+    Returns:
+        str: The base64 encoded string of the image.
+    """
+    return base64.b64encode(uploaded_file.read()).decode('utf-8')
+
+# Initialize the Groq client using the API key from the environment variables
+client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
 # Set up Streamlit page configuration
 st.set_page_config(
     page_icon="📃",
@@ -16,9 +33,6 @@ st.set_page_config(
 # App Title
 st.title("Groq Chat with LLaMA3x")
 
-# Initialize the Groq client using the API key from the environment variables
-client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
-
 # Cache the model fetching function to improve performance
 @st.cache_data
 def fetch_available_models():
@@ -77,6 +91,9 @@ with st.sidebar:
     # Define a function to clear messages when the model changes
     def reset_chat_on_model_change():
         st.session_state.messages = []
+        st.session_state.image_used = False
+        uploaded_file = None
+        base64_image = None
 
     # Model selection dropdown
     if models:
@@ -110,6 +127,35 @@ with st.sidebar:
     # Button to clear the chat
     if st.button("Clear Chat"):
         st.session_state.messages = []
+        st.session_state.image_used = False
+
+    # Initialize session state for tracking uploaded image usage
+    if "image_used" not in st.session_state:
+        st.session_state.image_used = False  # Flag to track image usage
+
+    # Check if the selected model supports vision
+    base64_image = None
+    uploaded_file = None
+    if model_option and "vision" in model_option.lower():
+        st.markdown(
+            "### Upload an Image"
+            "\n\n*One per conversation*"
+            )
+
+        # File uploader for images (only if image hasn't been used yet)
+        if not st.session_state.image_used:
+            uploaded_file = st.file_uploader(
+                "Upload an image for the model to process:",
+                type=["png", "jpg", "jpeg"],
+                help="Upload an image if the model supports vision tasks.",
+                accept_multiple_files=False
+            )
+            if uploaded_file:
+                base64_image = encode_image(uploaded_file)
+                st.image(uploaded_file, caption="Uploaded Image")
+    else:
+        base64_image = None
+
 
     st.markdown("### Usage Summary")
     usage_box = st.empty()
@@ -132,16 +178,56 @@ st.markdown("### Chat Interface")
 for message in st.session_state.messages:
     avatar = "🔋" if message["role"] == "assistant" else "🧑‍💻"
     with st.chat_message(message["role"], avatar=avatar):
-        st.markdown(message["content"])
+        # Check if the content is a list (text and image combined)
+        if isinstance(message["content"], list):
+            for item in message["content"]:
+                if item["type"] == "text":
+                    st.markdown(item["text"])
+                elif item["type"] == "image_url":
+                    # Handle base64-encoded image URLs
+                    if item["image_url"]["url"].startswith("data:image"):
+                        st.image(item["image_url"]["url"], caption="Uploaded Image")
+                        st.session_state.image_used = True
+                    else:
+                        st.warning("Invalid image format or unsupported URL.")
+        else:
+            # For plain text content
+            st.markdown(message["content"])
 
-# Capture user input
-user_input = st.chat_input("Enter your message here...")
 
-if user_input:
+# Capture user input
+if user_input:=st.chat_input("Enter your message here..."):
     # Append the user input to the session state
-    st.session_state.messages.append({"role": "user", "content": user_input})
+    # including the image if uploaded
+    if base64_image and not st.session_state.image_used:
+        # Append the user message with the image to session state
+        st.session_state.messages.append(
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": user_input},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{base64_image}",
+                            },
+                        },
+                    ],
+                }
+            )
+        st.session_state.image_used = True
+    else:
+        st.session_state.messages.append({"role": "user", "content": user_input})
+
+    # Display the uploaded image and user query in the chat
     with st.chat_message("user", avatar="🧑‍💻"):
+        # Display the user input
         st.markdown(user_input)
+        
+        # Display the uploaded image only if it's included in the current message
+        if base64_image and st.session_state.image_used:
+            st.image(uploaded_file, caption="Uploaded Image")
+            base64_image = None
 
     # Generate a response using the selected model
     try: