Update app.py

app.py

@@ -1,7 +1,9 @@
 import requests
+import os
 import json
 import streamlit as st
 from datetime import datetime
+import time
 
 # Page configuration
 st.set_page_config(
@@ -10,18 +12,31 @@ 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;
-        font-size: 0.9em;
+    }
+    
+    .model-attribution {
+        color: #28a745;
+        font-size: 0.8em;
+        font-style: italic;
     }
 </style>
 """, unsafe_allow_html=True)
@@ -32,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"""
@@ -48,47 +65,45 @@ 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 = st.secrets.get("OPENROUTER_API_KEY", None)
+OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")
 
-@st.cache_data(ttl=3600)
-def fetch_models():
-    """Fetch available models from OpenRouter API"""
+@st.cache_data(ttl=300)
+def check_api_status():
     if not OPENROUTER_API_KEY:
-        return [("GPT-3.5 Turbo", "openai/gpt-3.5-turbo")]
+        return "No API Key"
     try:
         url = "https://openrouter.ai/api/v1/models"
         headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}"}
         response = requests.get(url, headers=headers, timeout=10)
-        if response.status_code == 200:
-            models = response.json().get("data", [])
-            return [(model["name"], model["id"]) for model in models if "name" in model and "id" in model]
-        return [("GPT-3.5 Turbo", "openai/gpt-3.5-turbo")]
-    except requests.RequestException:
-        return [("GPT-3.5 Turbo", "openai/gpt-3.5-turbo")]
+        return "Connected" if response.status_code == 200 else "Error"
+    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 secrets."
-        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
     }
     
-    api_messages = [{"role": "system", "content": "You are a helpful AI assistant."}]
+    # 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)
     
     data = {
@@ -96,36 +111,57 @@ def get_ai_response(messages, model="openai/gpt-3.5-turbo"):
         "messages": api_messages,
         "stream": True,
         "max_tokens": 2000,
-        "temperature": 0.7
+        "temperature": 0.7,
+        "top_p": 1,
+        "frequency_penalty": 0,
+        "presence_penalty": 0
     }
     
     try:
         response = requests.post(url, headers=headers, json=data, stream=True, timeout=60)
+        
+        # Better error handling
         if response.status_code != 200:
-            error_data = response.json().get('error', {})
-            yield f"API Error: {error_data.get('message', response.reason)}"
+            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. Try again."
-    except requests.ConnectionError:
-        yield "Connection error. 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")
@@ -136,32 +172,54 @@ with st.sidebar:
     st.header("Settings")
     
     # API Status
-    status = "Connected" if OPENROUTER_API_KEY and requests.get(
-        "https://openrouter.ai/api/v1/models",
-        headers={"Authorization": f"Bearer {OPENROUTER_API_KEY}"},
-        timeout=10
-    ).status_code == 200 else "No API Key"
+    status = check_api_status()
+    if status == "Connected":
+        st.success("🟢 API Connected")
+    elif status == "No API Key":
+        st.error("No API Key")
+    else:
+        st.warning("Connection Issue")
     
-    st.success("🟢 API Connected" if status == "Connected" else "🔴 No API Key")
+    st.divider()
+    
+    # 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"),
+        ("LLaMA 3.1 70B", "meta-llama/llama-3.1-70b-instruct"),
+        ("DeepSeek Chat v3", "deepseek/deepseek-chat-v3-0324:free"),
+        ("DeepSeek R1", "deepseek/deepseek-r1-0528:free"),
+        ("Qwen3 Coder", "qwen/qwen3-coder:free"),
+        ("Microsoft MAI DS R1", "microsoft/mai-ds-r1:free"),
+        ("Gemma 3 27B", "google/gemma-3-27b-it:free"),
+        ("Gemma 3 4B", "google/gemma-3-4b-it:free"),
+        ("Auto (Best Available)", "openrouter/auto")
+    ]
     
-    # Model selection
-    models = fetch_models()
     model_names = [name for name, _ in models]
     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)
     
-    # Chat history controls
+    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):
@@ -169,43 +227,102 @@ 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"]):
-        st.markdown(message["content"])
+        # 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:
+                model_name = parts[1].replace("***", "").replace("**", "")
+                st.markdown(main_content)
+                st.markdown(f"<div class='model-attribution'>Response created by: <strong>{model_name}</strong></div>", unsafe_allow_html=True)
+            else:
+                st.markdown(message["content"])
+        else:
+            st.markdown(message["content"])
 
 # 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([user_message], selected_model):
-            full_response = response
-            placeholder.markdown(full_response + "▌")
-        placeholder.markdown(full_response)
-        
-        assistant_message = {"role": "assistant", "content": full_response}
-        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]}**")