Update app.py

app.py

@@ -1,9 +1,7 @@
 import requests
-import os
 import json
 import streamlit as st
 from datetime import datetime
-import time
 
 # Page configuration
 st.set_page_config(
@@ -12,27 +10,18 @@ st.set_page_config(
     initial_sidebar_state="collapsed"
 )
 
-# White background
+# Minimal CSS for styling
 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;
@@ -47,12 +36,10 @@ HISTORY_FILE = "chat_history.json"
 def load_chat_history():
     """Load chat history from file"""
     try:
-        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 []
+        with open(HISTORY_FILE, 'r', encoding='utf-8') as f:
+            return json.load(f)
+    except (FileNotFoundError, json.JSONDecodeError):
+        return []
 
 def save_chat_history(messages):
     """Save chat history to file"""
@@ -65,13 +52,12 @@ def save_chat_history(messages):
 def clear_chat_history():
     """Clear chat history file"""
     try:
-        if os.path.exists(HISTORY_FILE):
-            os.remove(HISTORY_FILE)
+        open(HISTORY_FILE, 'w').close()
         st.session_state.messages = []
     except Exception as e:
         st.error(f"Error clearing chat history: {e}")
 
-# Initialize session state with saved history
+# Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = load_chat_history()
 
@@ -80,6 +66,7 @@ OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
 
 @st.cache_data(ttl=300)
 def check_api_status():
+    """Check OpenRouter API status"""
     if not OPENROUTER_API_KEY:
         return "No API Key"
     try:
@@ -87,22 +74,21 @@ 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:
+    except requests.RequestException:
         return "Error"
 
 def get_ai_response(messages, model="openai/gpt-3.5-turbo"):
+    """Get streaming AI response from OpenRouter"""
     if not OPENROUTER_API_KEY:
-        return "No API key found. Please add OPENROUTER_API_KEY to environment variables."
+        yield "No API key found. Please add OPENROUTER_API_KEY to environment variables."
+        return
     
     url = "https://openrouter.ai/api/v1/chat/completions"
     headers = {
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "HTTP-Referer": "http://localhost:8501",  # Optional: Your site URL
-        "X-Title": "Streamlit AI Assistant"  # Optional: Your app name
+        "Authorization": f"Bearer {OPENROUTER_API_KEY}"
     }
     
-    # 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)
     
@@ -119,49 +105,31 @@ 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 = ""
-            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}"
-            
+            error_detail = response.json().get('error', {}).get('message', f"HTTP {response.status_code}")
             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:
-                # 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."
+            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)}."
 
 # Header
 st.title("AI Assistant")
@@ -182,7 +150,7 @@ with st.sidebar:
     
     st.divider()
     
-    # All models including new ones
+    # Model list
     models = [
         ("GPT-3.5 Turbo", "openai/gpt-3.5-turbo"),
         ("LLaMA 3.1 8B", "meta-llama/llama-3.1-8b-instruct"),
@@ -200,26 +168,20 @@ 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):
@@ -227,51 +189,18 @@ with st.sidebar:
             st.success("History saved!")
     
     with col2:
-        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()
+        if st.button("Clear History", use_container_width=True):
+            clear_chat_history()
+            st.success("History cleared!")
 
-# Show welcome message when no messages
+# Show welcome message
 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:
@@ -285,44 +214,29 @@ 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 = ""
-        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)
+        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)
 
-# Show currently using model
-st.caption(f"Currently using: **{model_names[selected_index]}**")
+# Show current model
+st.caption(f"Currently using: {model_names[selected_index]}")