Update app.py

app.py

@@ -1,139 +1,110 @@
 import streamlit as st
 import cohere
+import os
 
-st.set_page_config(page_title="Cohere Chat Interface", layout="wide")
+st.set_page_config(page_title="Cohere Chat", layout="wide")
 
-MODEL_PFPS = {
-    "command-a-03-2025": "/media/pfp/cohere-pfp.png",
-    "command-r7b-12-2024": "/media/pfp/cohere-pfp.png",
-    "command-r-plus-04-2024": "/media/pfp/cohere-pfp.png",
-    "command-r-plus": "/media/pfp/cohere-pfp.png",
-    "command-r-08-2024": "/media/pfp/cohere-pfp.png",
-    "command-r-03-2024": "/media/pfp/cohere-pfp.png",
-    "command-r": "/media/pfp/cohere-pfp.png",
-    "command": "/media/pfp/cohere-pfp.png",
-    "command-nightly": "/media/pfp/cohere-pfp.png",
-    "command-light": "/media/pfp/cohere-pfp.png",
-    "command-light-nightly": "/media/pfp/cohere-pfp.png"
-}
+AI_PFP = "media/pfps/cohere-pfp.png"
+USER_PFP = "media/pfps/user-pfp.jpg"
 
-USER_PFP = "https://example.com/user-default.png"
+if not os.path.exists(AI_PFP) or not os.path.exists(USER_PFP):
+    st.error("Missing profile pictures in media/pfps directory")
+    st.stop()
 
-MODEL_INFO = {
+model_info = {
     "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_window": "256K tokens",
-        "output_tokens": "8K tokens"
+        "description": "Command A is our most performant model to date, excelling at tool use, agents, retrieval augmented generation (RAG), and multilingual use cases.",
+        "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_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Small, fast update excelling at RAG, tool use, and complex reasoning tasks.",
+        "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_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Instruction-following model for complex RAG workflows and multi-step tool use.",
+        "context": "128K",
+        "output": "4K"
     },
     "command-r-plus": {
-        "description": "command-r-plus is an alias for command-r-plus-04-2024, so if you use command-r-plus in the API, that's the model you're pointing to.",
-        "context_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Alias for command-r-plus-04-2024.",
+        "context": "128K",
+        "output": "4K"
     },
     "command-r-08-2024": {
-        "description": "command-r-08-2024 is an update of the Command R model, delivered in August 2024.",
-        "context_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Updated Command R model from August 2024.",
+        "context": "128K",
+        "output": "4K"
     },
     "command-r-03-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 can be used for complex workflows like code generation, retrieval augmented generation (RAG), tool use, and agents.",
-        "context_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Instruction-following model for code generation, RAG, and agents.",
+        "context": "128K",
+        "output": "4K"
     },
     "command-r": {
-        "description": "command-r is an alias for command-r-03-2024, so if you use command-r in the API, that's the model you're pointing to.",
-        "context_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Alias for command-r-03-2024.",
+        "context": "128K",
+        "output": "4K"
     },
     "command": {
-        "description": "An instruction-following conversational model that performs language tasks with high quality, more reliably and with a longer context than our base generative models.",
-        "context_window": "4K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Conversational model with long context capabilities.",
+        "context": "4K",
+        "output": "4K"
     },
     "command-nightly": {
-        "description": "Nightly version of command - experimental and unstable. Not recommended for production use.",
-        "context_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Experimental nightly build (not for production).",
+        "context": "128K",
+        "output": "4K"
     },
     "command-light": {
-        "description": "Smaller, faster version of command with similar capabilities.",
-        "context_window": "4K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Faster lightweight version of command.",
+        "context": "4K",
+        "output": "4K"
     },
     "command-light-nightly": {
-        "description": "Nightly version of command-light - experimental and unstable. Not for production use.",
-        "context_window": "128K tokens",
-        "output_tokens": "4K tokens"
+        "description": "Experimental nightly build of command-light.",
+        "context": "128K",
+        "output": "4K"
     }
 }
 
 with st.sidebar:
-    st.title("Configuration")
+    st.title("Settings")
     api_key = st.text_input("Cohere API Key", type="password")
-    
-    selected_model = st.selectbox(
-        "Select Model",
-        options=list(MODEL_INFO.keys()),
-        format_func=lambda x: x.upper()
-    )
-    
+    selected_model = st.selectbox("Model", options=list(model_info.keys()))
     st.divider()
-    st.subheader("Model Details")
-    st.image(MODEL_PFPS[selected_model], width=80)
-    st.markdown(f"**{selected_model}**")
-    st.markdown(MODEL_INFO[selected_model]["description"])
-    st.markdown(f"**Context Window:** {MODEL_INFO[selected_model]['context_window']}")
-    st.markdown(f"**Max Output:** {MODEL_INFO[selected_model]['output_tokens']}")
+    st.image(AI_PFP, width=60)
+    st.subheader(selected_model)
+    st.markdown(model_info[selected_model]["description"])
+    st.caption(f"Context: {model_info[selected_model]['context']}")
+    st.caption(f"Output: {model_info[selected_model]['output']}")
 
-st.title(f"Chat with {selected_model.upper()}")
-st.image(MODEL_PFPS[selected_model], width=50)
+st.title(f"Chat - {selected_model}")
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-for message in st.session_state.messages:
-    avatar = USER_PFP if message["role"] == "user" else MODEL_PFPS[selected_model]
-    with st.chat_message(message["role"], avatar=avatar):
-        st.markdown(message["content"])
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"], avatar=USER_PFP if msg["role"] == "user" else AI_PFP):
+        st.markdown(msg["content"])
 
-if prompt := st.chat_input("Type your message..."):
+if prompt := st.chat_input("Message..."):
     if not api_key:
-        st.error("API key required - enter in sidebar")
+        st.error("API key required")
         st.stop()
     
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user", avatar=USER_PFP):
         st.markdown(prompt)
-    
+
     try:
         co = cohere.ClientV2(api_key)
-        
-        with st.chat_message("assistant", avatar=MODEL_PFPS[selected_model]):
-            response = co.chat(
-                model=selected_model,
-                messages=st.session_state.messages
-            )
-            
-            if hasattr(response, 'text'):
-                full_response = response.text
-            else:
-                full_response = "Error: Unexpected API response format"
-            
-            st.markdown(full_response)
-        
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
+        with st.chat_message("assistant", avatar=AI_PFP):
+            response = co.chat(model=selected_model, messages=st.session_state.messages)
+            reply = response.text if hasattr(response, 'text') else "Error processing response"
+            st.markdown(reply)
+        st.session_state.messages.append({"role": "assistant", "content": reply})
     
-    except cohere.CohereError as e:
-        st.error(f"Cohere API Error: {str(e)}")
     except Exception as e:
-        st.error(f"General Error: {str(e)}")
+        st.error(f"Error: {str(e)}")