Update app.py

app.py

@@ -14,13 +14,13 @@ if not os.path.exists(AI_PFP) or not os.path.exists(USER_PFP):
     st.stop()
 
 model_info = {
-    "c4ai-aya-expanse-8b": {"description": "Aya Expanse is a highly performant 8B multilingual model, designed to rival monolingual performance through innovations in instruction tuning with data arbitrage, preference training, and model merging. Serves 23 languages.", "context": "4K", "output": "4K"},
-    "c4ai-aya-expanse-32b": {"description": "Aya Expanse is a highly performant 32B multilingual model, designed to rival monolingual performance through innovations in instruction tuning with data arbitrage, preference training, and model merging. Serves 23 languages.", "context": "128K", "output": "4K"},
-    "c4ai-aya-vision-8b": {"description": "Aya Vision is a state-of-the-art multimodal model excelling at a variety of critical benchmarks for language, text, and image capabilities. This 8 billion parameter variant is focused on low latency and best-in-class performance.", "context": "16K", "output": "4K"},
-    "c4ai-aya-vision-32b": {"description": "Aya Vision is a state-of-the-art multimodal model excelling at a variety of critical benchmarks for language, text, and image capabilities. Serves 23 languages. This 32 billion parameter variant is focused on state-of-art multilingual performance.", "context": "16k", "output": "4K"},
-    "command-a-03-2025": {"description": "Command A is our most performant model to date, excelling at tool use, agents, retrieval augmented generation (RAG), and multilingual use cases. Command A has a context length of 256K, only requires two GPUs to run, and has 150% higher throughput compared to Command R+ 08-2024.", "context": "256K", "output": "8K"},
-    "command-r7b-12-2024": {"description": "command-r7b-12-2024 is a small, fast update delivered in December 2024. It excels at RAG, tool use, agents, and similar tasks requiring complex reasoning and multiple steps.", "context": "128K", "output": "4K"},
-    "command-r-plus-04-2024": {"description": "Command R+ is an instruction-following conversational model that performs language tasks at a higher quality, more reliably, and with a longer context than previous models. It is best suited for complex RAG workflows and multi-step tool use.", "context": "128K", "output": "4K"},
+    "c4ai-aya-expanse-8b": {"description": "Aya Expanse is a highly performant 8B multilingual model...", "context": "4K", "output": "4K"},
+    "c4ai-aya-expanse-32b": {"description": "Aya Expanse is a highly performant 32B multilingual model...", "context": "128K", "output": "4K"},
+    "command-a-03-2025": {"description": "Command A is our most performant model to date...", "context": "256K", "output": "8K"},
+    "command-r7b-12-2024": {"description": "command-r7b-12-2024 is a small, fast update...", "context": "128K", "output": "4K"},
+    "command-r-plus-04-2024": {"description": "Command R+ is an instruction-following conversational model...", "context": "128K", "output": "4K"},
+    "c4ai-aya-vision-8b": {"description": "Aya Vision is an 8B vision-language model...", "context": "4K", "output": "4K"},
+    "c4ai-aya-vision-32b": {"description": "Aya Vision is a 32B vision-language model...", "context": "128K", "output": "4K"}
 }
 
 with st.sidebar:
@@ -29,10 +29,10 @@ with st.sidebar:
     st.title("Settings")
     api_key = st.text_input("Cohere API Key", type="password")
     selected_model = st.selectbox("Model", options=list(model_info.keys()))
-    if st.button("Clear Chat"):
+    def clear_chat():
         st.session_state.messages = []
         st.session_state.first_message_sent = False
-        st.experimental_rerun()
+    st.button("Clear Chat", on_click=clear_chat)
     st.divider()
     st.image(AI_PFP, width=60)
     st.subheader(selected_model)
@@ -45,53 +45,57 @@ 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 "uploaded_image" not in st.session_state:
+    st.session_state.uploaded_image = None
 
-main = st.container()
-with main:
-    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:
-        avatar = USER_PFP if msg["role"] == "user" else AI_PFP
-        with st.chat_message(msg["role"], avatar=avatar):
-            st.markdown(msg["content"])
+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
+    )
 
-    col1, col2 = st.columns([1, 4])
-    with col1:
-        if selected_model.startswith("c4ai-aya-vision"):
-            uploaded = st.file_uploader("Upload image", type=["png", "jpg", "jpeg"])
-        else:
-            uploaded = None
-    with col2:
-        prompt = st.chat_input("Message...")
+for msg in st.session_state.messages:
+    avatar = USER_PFP if msg["role"] == "user" else AI_PFP
+    with st.chat_message(msg["role"], avatar=avatar):
+        st.markdown(msg["content"])
 
-if prompt:
+col1, col2 = st.columns([1, 8])
+with col1:
+    if selected_model.startswith("c4ai-aya-vision"):
+        img = st.file_uploader(label="📷", key="uploader", type=["png","jpg","jpeg"], accept_multiple_files=False)
+        if img is not None:
+            st.session_state.uploaded_image = img
+            st.image(img, width=80)
+    else:
+        st.write("")
+with col2:
+    prompt = st.chat_input("Message...")
+
+if prompt or st.session_state.uploaded_image:
     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})
-    with st.chat_message("user", avatar=USER_PFP):
-        st.markdown(prompt)
-
+    user_items = []
+    if prompt:
+        st.session_state.first_message_sent = True
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user", avatar=USER_PFP):
+            st.markdown(prompt)
+        user_items.append({"type": "text", "text": prompt})
+    if st.session_state.uploaded_image:
+        raw = st.session_state.uploaded_image.read()
+        b64 = base64.b64encode(raw).decode("utf-8")
+        url = f"data:image/jpeg;base64,{b64}"
+        user_items.append({"type": "image_url", "image_url": {"url": url}})
+        with st.chat_message("user", avatar=USER_PFP):
+            st.image(raw, width=200)
+        st.session_state.uploaded_image = None
     try:
         co = cohere.ClientV2(api_key)
-        user_content = [{"type": "text", "text": prompt}]
-        if uploaded:
-            raw = uploaded.read()
-            b64 = base64.b64encode(raw).decode("utf-8")
-            data_url = f"data:image/jpeg;base64,{b64}"
-            user_content.append({"type": "image_url", "image_url": {"url": data_url}})
-        response = co.chat(
-            model=selected_model,
-            messages=[{"role": "user", "content": user_content}]
-        )
-        content_items = response.message.content
-        reply = "".join(getattr(item, 'text', '') for item in content_items)
-        st.session_state.messages.append({"role": "assistant", "content": reply})
+        response = co.chat(model=selected_model, messages=[{"role":"user","content":user_items}])
+        reply = "".join(getattr(item,'text','') for item in response.message.content)
         with st.chat_message("assistant", avatar=AI_PFP):
             st.markdown(reply)
+        st.session_state.messages.append({"role":"assistant","content":reply})
     except Exception as e:
-        st.error(f"Error: {str(e)}")
+        st.error(f"Error: {e}")