Update app.py

app.py

@@ -1,8 +1,9 @@
 import requests
-import json
 import os
+import json
 import streamlit as st
 from datetime import datetime
+import time
 
 # Page configuration
 st.set_page_config(
@@ -11,18 +12,27 @@ st.set_page_config(
     initial_sidebar_state="collapsed"
 )
 
-# Minimal CSS for styling
+# White background
 st.markdown("""
 <style>
     .stApp {
         background: white;
+    }
+    
+    .main .block-container {
         max-width: 800px;
-        margin: 0 auto;
     }
+    
+    #MainMenu {visibility: hidden;}
+    footer {visibility: hidden;}
+    header {visibility: hidden;}
+    .stDeployButton {display: none;}
+    
     .model-id {
         color: #28a745;
         font-family: monospace;
     }
+    
     .model-attribution {
         color: #28a745;
         font-size: 0.8em;
@@ -37,10 +47,12 @@ HISTORY_FILE = "chat_history.json"
 def load_chat_history():
     """Load chat history from file"""
     try:
-        with open(HISTORY_FILE, 'r', encoding='utf-8') as f:
-            return json.load(f)
-    except (FileNotFoundError, json.JSONDecodeError):
-        return []
+        if os.path.exists(HISTORY_FILE):
+            with open(HISTORY_FILE, 'r', encoding='utf-8') as f:
+                return json.load(f)
+    except Exception as e:
+        st.error(f"Error loading chat history: {e}")
+    return []
 
 def save_chat_history(messages):
     """Save chat history to file"""
@@ -53,24 +65,21 @@ def save_chat_history(messages):
 def clear_chat_history():
     """Clear chat history file"""
     try:
-        open(HISTORY_FILE, 'w').close()
+        if os.path.exists(HISTORY_FILE):
+            os.remove(HISTORY_FILE)
         st.session_state.messages = []
     except Exception as e:
         st.error(f"Error clearing chat history: {e}")
 
-# Initialize session state
+# Initialize session state with saved history
 if "messages" not in st.session_state:
     st.session_state.messages = load_chat_history()
 
 # Get API key
 OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
-if not OPENROUTER_API_KEY:
-    st.error("API key not found. Please set OPENROUTER_API_KEY.")
-    st.stop()
 
 @st.cache_data(ttl=300)
 def check_api_status():
-    """Check OpenRouter API status"""
     if not OPENROUTER_API_KEY:
         return "No API Key"
     try:
@@ -78,21 +87,22 @@ def check_api_status():
         headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}"}
         response = requests.get(url, headers=headers, timeout=10)
         return "Connected" if response.status_code == 200 else "Error"
-    except requests.RequestException:
+    except:
         return "Error"
 
 def get_ai_response(messages, model="openai/gpt-3.5-turbo"):
-    """Get streaming AI response from OpenRouter"""
     if not OPENROUTER_API_KEY:
-        yield "No API key found. Please add OPENROUTER_API_KEY to environment variables."
-        return
+        return "No API key found. Please add OPENROUTER_API_KEY to environment variables."
     
     url = "https://openrouter.ai/api/v1/chat/completions"
     headers = {
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}"
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+        "HTTP-Referer": "http://localhost:8501",  # Optional: Your site URL
+        "X-Title": "Streamlit AI Assistant"  # Optional: Your app name
     }
     
+    # Create system message and user messages
     api_messages = [{"role": "system", "content": "You are a helpful AI assistant. Provide clear and helpful responses."}]
     api_messages.extend(messages)
     
@@ -109,31 +119,49 @@ def get_ai_response(messages, model="openai/gpt-3.5-turbo"):
     
     try:
         response = requests.post(url, headers=headers, json=data, stream=True, timeout=60)
+        
+        # Better error handling
         if response.status_code != 200:
-            error_detail = response.json().get('error', {}).get('message', f"HTTP {response.status_code}")
+            error_detail = ""
+            try:
+                error_data = response.json()
+                error_detail = error_data.get('error', {}).get('message', f"HTTP {response.status_code}")
+            except:
+                error_detail = f"HTTP {response.status_code}: {response.reason}"
+            
             yield f"API Error: {error_detail}. Please try a different model or check your API key."
             return
         
         full_response = ""
+        buffer = ""
+        
+        # Using your working streaming logic
         for line in response.iter_lines():
-            if line and line.startswith(b"data: "):
-                data_str = line[len(b"data: "):].decode("utf-8")
-                if data_str.strip() == "[DONE]":
-                    break
-                try:
-                    data = json.loads(data_str)
-                    delta = data["choices"][0]["delta"].get("content", "")
-                    if delta:
-                        full_response += delta
-                        yield full_response
-                except json.JSONDecodeError:
-                    continue
-    except requests.Timeout:
-        yield "Request timed out. Please try again."
-    except requests.ConnectionError:
-        yield "Connection error. Please check your internet."
-    except requests.RequestException as e:
-        yield f"Request error: {str(e)}."
+            if line:
+                # The server sends lines starting with "data: ..."
+                if line.startswith(b"data: "):
+                    data_str = line[len(b"data: "):].decode("utf-8")
+                    if data_str.strip() == "[DONE]":
+                        break
+                    try:
+                        data = json.loads(data_str)
+                        delta = data["choices"][0]["delta"].get("content", "")
+                        if delta:
+                            full_response += delta
+                            yield full_response
+                    except json.JSONDecodeError:
+                        continue
+                    except (KeyError, IndexError):
+                        continue
+                            
+    except requests.exceptions.Timeout:
+        yield "Request timed out. Please try again with a shorter message or different model."
+    except requests.exceptions.ConnectionError:
+        yield "Connection error. Please check your internet connection and try again."
+    except requests.exceptions.RequestException as e:
+        yield f"Request error: {str(e)}. Please try again."
+    except Exception as e:
+        yield f"Unexpected error: {str(e)}. Please try again or contact support."
 
 # Header
 st.title("AI Assistant")
@@ -154,7 +182,7 @@ with st.sidebar:
     
     st.divider()
     
-    # Model list
+    # All models including new ones
     models = [
         ("GPT-3.5 Turbo", "openai/gpt-3.5-turbo"),
         ("LLaMA 3.1 8B", "meta-llama/llama-3.1-8b-instruct"),
@@ -172,20 +200,26 @@ with st.sidebar:
     model_ids = [model_id for _, model_id in models]
     
     selected_index = st.selectbox("Model", range(len(model_names)), 
-                                format_func=lambda x: model_names[x], index=0)
+                                format_func=lambda x: model_names[x], 
+                                index=0)
     selected_model = model_ids[selected_index]
     
+    # Show selected model ID in green
     st.markdown(f"**Model ID:** <span class='model-id'>{selected_model}</span>", unsafe_allow_html=True)
     
     st.divider()
     
     # Chat History Controls
     st.header("Chat History")
+    
+    # Show number of messages
     if st.session_state.messages:
         st.info(f"Messages stored: {len(st.session_state.messages)}")
     
+    # Auto-save toggle
     auto_save = st.checkbox("Auto-save messages", value=True)
     
+    # Manual save/load buttons
     col1, col2 = st.columns(2)
     with col1:
         if st.button("Save History", use_container_width=True):
@@ -193,18 +227,51 @@ with st.sidebar:
             st.success("History saved!")
     
     with col2:
-        if st.button("Clear History", use_container_width=True):
-            clear_chat_history()
-            st.success("History cleared!")
+        if st.button("Load History", use_container_width=True):
+            st.session_state.messages = load_chat_history()
+            st.success("History loaded!")
+            st.rerun()
+    
+    st.divider()
+    
+    # View History
+    if st.button("View History File", use_container_width=True):
+        if os.path.exists(HISTORY_FILE):
+            with open(HISTORY_FILE, 'r', encoding='utf-8') as f:
+                history_content = f.read()
+            st.text_area("Chat History (JSON)", history_content, height=200)
+        else:
+            st.warning("No history file found")
+    
+    # Download History
+    if os.path.exists(HISTORY_FILE):
+        with open(HISTORY_FILE, 'rb') as f:
+            st.download_button(
+                label="Download History",
+                data=f.read(),
+                file_name=f"chat_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
+                mime="application/json",
+                use_container_width=True
+            )
+    
+    st.divider()
+    
+    # Clear controls
+    if st.button("Clear Chat", use_container_width=True, type="secondary"):
+        clear_chat_history()
+        st.success("Chat cleared!")
+        st.rerun()
 
-# Show welcome message
+# Show welcome message when no messages
 if not st.session_state.messages:
     st.info("How can I help you today?")
 
 # Display chat messages
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
+        # Check if this is an assistant message with attribution
         if message["role"] == "assistant" and "Response created by:" in message["content"]:
+            # Split content and attribution
             parts = message["content"].split("\n\n---\n*Response created by:")
             main_content = parts[0]
             if len(parts) > 1:
@@ -218,29 +285,44 @@ for message in st.session_state.messages:
 
 # Chat input
 if prompt := st.chat_input("Ask anything..."):
+    # Add user message
     user_message = {"role": "user", "content": prompt}
     st.session_state.messages.append(user_message)
     
+    # Auto-save if enabled
     if auto_save:
         save_chat_history(st.session_state.messages)
     
+    # Display user message
     with st.chat_message("user"):
         st.markdown(prompt)
     
+    # Get AI response
     with st.chat_message("assistant"):
         placeholder = st.empty()
-        full_response = ""
-        for response in get_ai_response(st.session_state.messages, selected_model):
-            full_response = response
-            placeholder.markdown(full_response + "▌")
-        placeholder.markdown(full_response)
         
-        full_response_with_attribution = full_response + f"\n\n---\n*Response created by: **{model_names[selected_index]}***"
-        assistant_message = {"role": "assistant", "content": full_response_with_attribution}
-        st.session_state.messages.append(assistant_message)
-        
-        if auto_save:
-            save_chat_history(st.session_state.messages)
+        full_response = ""
+        try:
+            for response in get_ai_response(st.session_state.messages, selected_model):
+                full_response = response
+                placeholder.markdown(full_response + "▌")
+            
+            # Remove cursor and show final response
+            placeholder.markdown(full_response)
+            
+        except Exception as e:
+            error_msg = f"An error occurred: {str(e)}"
+            placeholder.markdown(error_msg)
+            full_response = error_msg
+    
+    # Add AI response to messages with attribution
+    full_response_with_attribution = full_response + f"\n\n---\n*Response created by: **{model_names[selected_index]}***"
+    assistant_message = {"role": "assistant", "content": full_response_with_attribution}
+    st.session_state.messages.append(assistant_message)
+    
+    # Auto-save if enabled
+    if auto_save:
+        save_chat_history(st.session_state.messages)
 
-# Show current model
-st.caption(f"Currently using: {model_names[selected_index]}")
+# Show currently using model
+st.caption(f"Currently using: **{model_names[selected_index]}**")