Update app.py

app.py

@@ -25,15 +25,9 @@ with st.sidebar:
     api_key = st.text_input("Cohere API Key", type="password")
     selected_model = st.selectbox("Model", options=list(model_info.keys()))
     if selected_model.startswith("c4ai-aya-vision"):
-        uploaded = st.file_uploader("Upload image", type=["png","jpg","jpeg"])
-        if uploaded:
-            data = uploaded.read()
-            session_image = base64.b64encode(data).decode('utf-8')
-            st.session_state.image_data = session_image
+        uploaded = st.file_uploader("Upload image", type=["png","jpg","jpeg"], key="uploaded")
     if st.button("Clear Chat"):
-        st.session_state.messages = []
-        st.session_state.first_message_sent = False
-        st.session_state.image_data = None
+        st.session_state.clear()
         st.rerun()
     st.divider()
     st.subheader(selected_model)
@@ -46,38 +40,51 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 if "first_message_sent" not in st.session_state:
     st.session_state.first_message_sent = False
-if "image_data" not in st.session_state:
-    st.session_state.image_data = None
 
 if not st.session_state.first_message_sent:
     st.markdown("<h1 style='text-align: center; color: #4a4a4a; margin-top: 100px;'>How can Cohere help you today?</h1>", unsafe_allow_html=True)
+
 for msg in st.session_state.messages:
-    with st.chat_message(msg["role"], avatar=USER_PFP if msg["role"]=="user" else AI_PFP):
+    with st.chat_message(msg["role"], avatar=USER_PFP if msg["role"] == "user" else AI_PFP):
         content = msg["content"]
         if isinstance(content, list):
             for item in content:
-                if item.get("type")=="text":
-                    st.markdown(item.get("text"))
-                if item.get("type")=="image_url":
-                    st.image(item.get("image_url").get("url"))
+                if getattr(item, 'type', None) == "text":
+                    st.markdown(item.text)
+                elif getattr(item, 'type', None) == "image_url":
+                    st.image(item.image_url.url)
         else:
             st.markdown(content)
+
 if prompt := st.chat_input("Message..."):
     if not api_key:
         st.error("API key required")
         st.stop()
+
     st.session_state.first_message_sent = True
-    st.session_state.messages.append({"role":"user","content":prompt})
+    st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user", avatar=USER_PFP):
         st.markdown(prompt)
+
     co = cohere.ClientV2(api_key)
-    content = [{"type":"text","text":prompt}]
-    if st.session_state.image_data and selected_model.startswith("c4ai-aya-vision"):
-        data_url = f"data:image/jpeg;base64,{st.session_state.image_data}"
-        content.append({"type":"image_url","image_url":{"url":data_url}})
-    response = co.chat(model=selected_model, messages=[*st.session_state.messages, {"role":"user","content":content}], temperature=0.3)
+    content = [{"type": "text", "text": prompt}]
+    if selected_model.startswith("c4ai-aya-vision") and st.session_state.get("uploaded"):
+        data = st.session_state.uploaded.read()
+        data_url = f"data:image/jpeg;base64,{base64.b64encode(data).decode('utf-8')}"
+        content.append({"type": "image_url", "image_url": {"url": data_url}})
+
+    response = co.chat(
+        model=selected_model,
+        messages=[*st.session_state.messages, {"role": "user", "content": content}],
+        temperature=0.3
+    )
     items = response.message.content
-    reply = "".join([getattr(i,'text','') for i in items])
+    reply = "".join(getattr(i, 'text', '') for i in items)
+
     with st.chat_message("assistant", avatar=AI_PFP):
         st.markdown(reply)
-    st.session_state.messages.append({"role":"assistant","content":items})
+
+    st.session_state.messages.append({"role": "assistant", "content": items})
+
+    if "uploaded" in st.session_state:
+        st.session_state.uploaded = None