Update app.py

app.py

@@ -1,8 +1,10 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 import requests
 import os
+import toml
+import pathlib
+import json
+import time
 
 # Page configuration
 st.set_page_config(
@@ -11,29 +13,63 @@ st.set_page_config(
     layout="centered"
 )
 
-# Check for API usage vs local model
-use_api = st.sidebar.checkbox("Use Hugging Face API (recommended)", value=True)
+# Try to load secrets from .streamlit/secrets.toml
+def load_secrets():
+    try:
+        # First, try streamlit's built-in secrets
+        if hasattr(st, 'secrets'):
+            try:
+                return st.secrets.get("HF_API_TOKEN", "")
+            except:
+                pass
+                
+        # Second, try to find secrets.toml in the .streamlit directory
+        secrets_path = pathlib.Path(".streamlit/secrets.toml")
+        if secrets_path.exists():
+            secrets = toml.load(secrets_path)
+            return secrets.get("HF_API_TOKEN", "")
+        
+        # Third, check in user's home directory
+        home_secrets_path = pathlib.Path.home() / ".streamlit/secrets.toml"
+        if home_secrets_path.exists():
+            secrets = toml.load(home_secrets_path)
+            return secrets.get("HF_API_TOKEN", "")
+            
+        # Last, check environment variables
+        return os.environ.get("HF_API_TOKEN", "")
+    except Exception as e:
+        st.sidebar.error(f"Error loading secrets: {e}")
+        return ""
 
-if use_api:
-    # Hugging Face API setup
-    API_URL = "https://api-inference.huggingface.co/models/meta-llama/Llama-2-7b-chat-hf"
-    API_TOKEN = os.environ.get("HF_API_TOKEN", st.secrets.get("HF_API_TOKEN", ""))
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-    
-    # Show warning if no API token is provided
-    if not API_TOKEN:
-        st.sidebar.warning("⚠️ No API token found. Please add one in your secrets.")
+# Model selection
+MODEL_OPTIONS = {
+    "Llama-3-8B-Instruct": "meta-llama/Meta-Llama-3-8B-Instruct",  # More up-to-date model
+    "Llama-2-7B-Chat": "meta-llama/Llama-2-7b-chat-hf",
+    "Mistral-7B-Instruct": "mistralai/Mistral-7B-Instruct-v0.2",
+    "Falcon-7B-Instruct": "tiiuae/falcon-7b-instruct",
+    "OpenAssistant": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+    "Flan-T5-Large": "google/flan-t5-large",  # Open access model
+    "GPT2": "gpt2",  # Fully open model
+    "Rule-Based (No API)": "local"  # Completely local option
+}
+
+# Allow model selection and initialize with default
+if "selected_model" not in st.session_state:
+    st.session_state.selected_model = "Rule-Based (No API)"  # Default to local option for reliability
+
+# Hugging Face API setup
+def get_api_url():
+    model_id = MODEL_OPTIONS[st.session_state.selected_model]
+    return f"https://api-inference.huggingface.co/models/{model_id}"
+
+API_TOKEN = load_secrets()
+headers = {"Authorization": f"Bearer {API_TOKEN}"}
+
+# Show warning if no API token is provided
+if not API_TOKEN:
+    st.sidebar.warning("⚠️ No API token found. Using fallback responses. Add your API token in environment variables or create a .streamlit/secrets.toml file.")
 else:
-    # Load model and tokenizer locally (fallback option)
-    @st.cache_resource
-    def load_model():
-        model_name = "facebook/blenderbot-400M-distill"
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
-        return tokenizer, model
-    
-    with st.spinner("Loading model... This might take a minute."):
-        tokenizer, model = load_model()
+    st.sidebar.success("✅ API token loaded successfully")
 
 # Customer support context and guidelines
 SYSTEM_PROMPT = """You are an AI-powered customer support assistant integrated into a company website. 
@@ -46,6 +82,7 @@ Rules:
 - Direct users to relevant resources when needed
 - Be available 24/7 and never mention that you're an AI unless explicitly asked
 - Avoid hallucinating facts - if you don't know something, politely let the user know and offer to connect them with a human representative
+- Keep responses concise and to the point
 
 Example questions you should be able to handle:
 - "I need help tracking my order"
@@ -65,86 +102,205 @@ COMPANY_INFO = {
     "contact": "For complex issues, you can reach our human support team at support@example.com or call 1-800-123-4567."
 }
 
-def format_user_prompt(user_input, history):
-    """Format the input with history for local model"""
-    prompt = SYSTEM_PROMPT + "\n\n"
-    for user_msg, bot_msg in history:
-        prompt += f"Customer: {user_msg}\nSupport: {bot_msg}\n\n"
-    prompt += f"Customer: {user_input}\nSupport:"
-    return prompt
+# Fallback responses for common queries when API is not available
+FALLBACK_RESPONSES = {
+    "greeting": "Hello! Welcome to our customer support. How can I assist you today?",
+    "return": f"Our return policy allows returns within 30 days of purchase with original receipt. Items must be unused and in original packaging.",
+    "track": "You can track your order by logging into your account or using the tracking number provided in your confirmation email.",
+    "login": "For account login issues, try resetting your password. If problems persist, please provide your email address so we can investigate.",
+    "subscription": "Yes, you can cancel your subscription at any time through your account settings. There are no cancellation fees.",
+    "hours": "Our customer support team is available 24/7 via this chat. For phone support, our hours are 9 AM - 8 PM Monday through Friday, and 10 AM - 6 PM on weekends.",
+    "contact": "For complex issues, you can reach our human support team at support@example.com or call 1-800-123-4567.",
+    "default": "Thank you for your inquiry. I'd be happy to help with that. Could you please provide more details so I can assist you better?"
+}
 
-def format_prompt_for_llama2(user_input, history):
-    """Format the prompt for Llama-2 chat model"""
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+def format_prompt(user_input, history):
+    """Format the prompt for the selected model"""
+    if st.session_state.selected_model in ["Llama-3-8B-Instruct", "Llama-2-7B-Chat"]:
+        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+        
+        # Add chat history
+        for user_msg, bot_msg in history:
+            messages.append({"role": "user", "content": user_msg})
+            messages.append({"role": "assistant", "content": bot_msg})
+        
+        # Add current user message
+        messages.append({"role": "user", "content": user_input})
+        
+        return {"inputs": messages}
+    elif st.session_state.selected_model == "Flan-T5-Large":
+        # Special format for Flan-T5
+        # Flan-T5 works better with direct questions/instructions
+        prompt = f"Answer this customer support question in a helpful way: {user_input}"
+        return {"inputs": prompt}
+    elif st.session_state.selected_model == "GPT2":
+        # Special format for GPT2
+        prompt = f"Customer: {user_input}\nSupport agent:"
+        return {"inputs": prompt}
+    else:
+        # For other models that use text completion format
+        prompt = SYSTEM_PROMPT + "\n\n"
+        for user_msg, bot_msg in history:
+            prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n\n"
+        prompt += f"User: {user_input}\nAssistant:"
+        
+        return {"inputs": prompt}
+
+def query_huggingface_api(payload, user_input):
+    """Send request to Hugging Face API with retry logic"""
+    api_url = get_api_url()
+    max_retries = 3
+    retry_delay = 2
     
-    # Add chat history
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": bot_msg})
+    # Try alternative TGI API endpoint if selected
+    use_tgi_api = False
+    if "use_tgi_api" in st.session_state and st.session_state.use_tgi_api:
+        use_tgi_api = True
+        # The TGI API has a different format
+        if st.session_state.selected_model in ["Llama-3-8B-Instruct", "Llama-2-7B-Chat", "Mistral-7B-Instruct"]:
+            # Extract just the messages for chat models
+            if "inputs" in payload and isinstance(payload["inputs"], list):
+                tgi_payload = {
+                    "inputs": payload["inputs"][-1]["content"],
+                    "parameters": {
+                        "max_new_tokens": 256,
+                        "temperature": 0.7,
+                        "top_p": 0.95,
+                        "do_sample": True
+                    }
+                }
+                payload = tgi_payload
+                api_url = f"https://api-inference.huggingface.co/models/{MODEL_OPTIONS[st.session_state.selected_model]}"
     
-    # Add current user message
-    messages.append({"role": "user", "content": user_input})
+    original_payload = payload.copy()
+    tried_simple_string = False
     
-    return messages
+    for attempt in range(max_retries):
+        try:
+            response = requests.post(api_url, headers=headers, json=payload, timeout=90)
+            
+            # Check if model is still loading
+            if response.status_code == 503:
+                st.warning(f"Model is loading. Retrying in {retry_delay} seconds... (Attempt {attempt+1}/{max_retries})")
+                time.sleep(retry_delay)
+                retry_delay *= 2  # Exponential backoff
+                continue
+            
+            # Handle 422 errors specifically - try simplifying the input format
+            if response.status_code == 422 and not tried_simple_string:
+                st.warning("Model expecting different input format. Trying simpler format...")
+                # Convert to simple string input
+                if isinstance(payload["inputs"], dict) or isinstance(payload["inputs"], list):
+                    # Extract just the user's query for simplicity
+                    payload = {"inputs": user_input}
+                    tried_simple_string = True
+                    continue
+                    
+            # Handle other errors
+            if response.status_code != 200:
+                st.error(f"API error: {response.status_code} - {response.text}")
+                return {"error": f"API error: {response.status_code} - {response.text}"}
+                
+            return response.json()
+            
+        except Exception as e:
+            st.error(f"Error querying API (attempt {attempt+1}/{max_retries}): {e}")
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
+                retry_delay *= 2
+            else:
+                return {"error": str(e)}
 
-def query_huggingface_api(payload):
-    """Send request to Hugging Face API"""
+def extract_response(api_response):
+    """Extract the response text from different model response formats"""
     try:
-        response = requests.post(API_URL, headers=headers, json=payload)
-        return response.json()
+        if "debug_mode" in st.session_state and st.session_state.debug_mode:
+            st.sidebar.subheader("Raw Response:")
+            st.sidebar.json(api_response)
+            
+        # Handle list format responses
+        if isinstance(api_response, list):
+            if len(api_response) > 0:
+                # Standard response format for some models
+                if "generated_text" in api_response[0]:
+                    return api_response[0]["generated_text"]
+                # TGI API direct text response
+                elif isinstance(api_response[0], str):
+                    return api_response[0]
+                
+        # Handle dictionary format responses
+        if isinstance(api_response, dict):
+            # For text generation models
+            if "generated_text" in api_response:
+                return api_response["generated_text"]
+            
+            # For simple string response in a dict
+            if "text" in api_response:
+                return api_response["text"]
+                
+            # For chat models that return as list array
+            if "conversation" in api_response:
+                return api_response["conversation"]["messages"][-1]["content"]
+                
+            # For Llama-3/Mistral chat format
+            if "outputs" in api_response and len(api_response["outputs"]) > 0:
+                return api_response["outputs"][0]["text"]
+                
+        # Fallback: attempt to extract any text or return the response as string
+        if isinstance(api_response, str):
+            return api_response
+            
+        return str(api_response)
+        
     except Exception as e:
-        st.error(f"Error querying API: {e}")
-        return {"error": str(e)}
-
-def get_local_model_response(user_input, history):
-    """Generate response using the local model"""
-    prompt = format_user_prompt(user_input, history)
-    
-    # Generate response
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-    with torch.no_grad():
-        output = model.generate(
-            inputs["input_ids"],
-            max_length=200,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    # Extract just the assistant's response (after Support:)
-    try:
-        response = response.split("Support:")[-1].strip()
-    except:
-        response = "I apologize, but I'm having trouble generating a response. Please try rephrasing your question, or I'd be happy to connect you with a human representative."
+        st.error(f"Error extracting response: {e}")
+        return None
+
+def get_fallback_response(user_input):
+    """Get a fallback response based on keywords in the user input"""
+    lower_input = user_input.lower()
     
-    return response
+    if any(word in lower_input for word in ["hi", "hello", "hey"]):
+        return FALLBACK_RESPONSES["greeting"]
+    elif "return" in lower_input:
+        return FALLBACK_RESPONSES["return"]
+    elif any(word in lower_input for word in ["track", "order", "package", "shipping"]):
+        return FALLBACK_RESPONSES["track"]
+    elif any(word in lower_input for word in ["login", "sign in", "account", "password"]):
+        return FALLBACK_RESPONSES["login"]
+    elif "subscription" in lower_input:
+        return FALLBACK_RESPONSES["subscription"]
+    elif any(word in lower_input for word in ["hours", "available", "weekend"]):
+        return FALLBACK_RESPONSES["hours"]
+    elif any(word in lower_input for word in ["human", "person", "agent", "representative"]):
+        return FALLBACK_RESPONSES["contact"]
+    else:
+        return FALLBACK_RESPONSES["default"]
 
-def get_api_response(user_input, history):
-    """Get response from the LLM via Hugging Face API"""
-    # Format the messages for the API
-    messages = format_prompt_for_llama2(user_input, history)
+def get_response(user_input, history):
+    """Get response from the LLM via Hugging Face API or fall back to simple responses"""
+    # Use rule-based responses if selected or if API token is missing
+    if st.session_state.selected_model == "Rule-Based (No API)" or not API_TOKEN:
+        return get_fallback_response(user_input)
+    
+    # Format the messages for the API based on selected model
+    payload = format_prompt(user_input, history)
     
     # Query the API
-    response = query_huggingface_api({"inputs": messages})
+    api_response = query_huggingface_api(payload, user_input)
+    
+    # Debug mode: show raw API response
+    if "debug_mode" in st.session_state and st.session_state.debug_mode:
+        st.sidebar.json(api_response)
     
     # Extract the assistant's response
-    try:
-        if isinstance(response, list) and len(response) > 0:
-            assistant_response = response[0].get("generated_text", "")
-        elif isinstance(response, dict) and "generated_text" in response:
-            assistant_response = response["generated_text"]
-        else:
-            st.warning(f"Unexpected API response format")
-            assistant_response = "I apologize, but I'm having trouble generating a response. Please try rephrasing your question, or I'd be happy to connect you with a human representative."
-    except Exception as e:
-        st.error(f"Error processing API response: {e}")
-        assistant_response = "I apologize, but I'm having trouble generating a response. Please try rephrasing your question, or I'd be happy to connect you with a human representative."
+    if api_response and "error" not in api_response:
+        assistant_response = extract_response(api_response)
+        if assistant_response:
+            return assistant_response
     
-    return assistant_response
+    # Fall back to keyword-based responses if API fails
+    return get_fallback_response(user_input)
 
 def enhance_response(user_input, base_response):
     """Enhance response with specific company information"""
@@ -165,23 +321,19 @@ def enhance_response(user_input, base_response):
     
     return base_response
 
-def get_response(user_input, history):
-    """Get chatbot response based on the selected mode"""
-    if use_api:
-        base_response = get_api_response(user_input, history)
-    else:
-        base_response = get_local_model_response(user_input, history)
-        
-    # Enhance with specific company information
-    return enhance_response(user_input, base_response)
-
 # App title and intro
 st.title("Customer Support Assistant")
 st.markdown("Welcome to our customer support chat! How can I help you today?")
 
-# Initialize session state for chat history
+# Initialize session state for chat history and API token
 if "chat_history" not in st.session_state:
     st.session_state.chat_history = []
+    
+if "api_token" not in st.session_state:
+    st.session_state.api_token = API_TOKEN
+    
+if "debug_mode" not in st.session_state:
+    st.session_state.debug_mode = False
 
 # Display chat history
 for user_msg, bot_msg in st.session_state.chat_history:
@@ -200,21 +352,85 @@ if user_input:
     
     # Get bot response with a spinner
     with st.chat_message("assistant", avatar="🧑‍💼"):
-        with st.spinner("Thinking..."):
-            bot_response = get_response(user_input, st.session_state.chat_history)
-        st.write(bot_response)
-    
-    # Add to chat history
-    st.session_state.chat_history.append((user_input, bot_response))
+        with st.spinner(f"Thinking using {st.session_state.selected_model}..."):
+            try:
+                base_response = get_response(user_input, st.session_state.chat_history)
+                bot_response = enhance_response(user_input, base_response)
+                st.write(bot_response)
+                # Add to chat history
+                st.session_state.chat_history.append((user_input, bot_response))
+            except Exception as e:
+                st.error(f"Sorry, I encountered an error processing your request: {str(e)}")
 
 # Sidebar options
 with st.sidebar:
     st.title("Options")
     
+    # Model selection
+    st.markdown("---")
+    st.markdown("### Model Settings")
+    selected_model = st.selectbox(
+        "Select AI Model",
+        list(MODEL_OPTIONS.keys()),
+        index=list(MODEL_OPTIONS.keys()).index(st.session_state.selected_model)
+    )
+    
+    if selected_model != st.session_state.selected_model:
+        st.session_state.selected_model = selected_model
+        st.success(f"Model changed to {selected_model}")
+    
+    # API mode toggle
+    st.checkbox("Use TGI API (try if regular API fails)", key="use_tgi_api", 
+                help="Use Text Generation Inference API format which works better for some models")
+    
+    # API token input
+    st.markdown("---")
+    st.markdown("### API Settings")
+    manually_entered_token = st.text_input("Enter Hugging Face API Token", type="password", help="Your API token will not be stored permanently")
+    if manually_entered_token:
+        st.session_state.api_token = manually_entered_token
+        API_TOKEN = manually_entered_token
+        headers = {"Authorization": f"Bearer {API_TOKEN}"}
+        st.success("✅ API Token set for this session")
+
+    # Test API button
+    if st.button("Test API Connection"):
+        with st.spinner("Testing API connection..."):
+            test_model = "gpt2"  # Use a simple model that everyone has access to
+            test_url = f"https://api-inference.huggingface.co/models/{test_model}"
+            test_payload = {"inputs": "Hello, I'm testing the API connection."}
+            try:
+                response = requests.post(test_url, headers=headers, json=test_payload, timeout=10)
+                if response.status_code == 200:
+                    st.success("✅ API connection successful! Your token is working correctly.")
+                    # Update the API token in session
+                    if not API_TOKEN:
+                        st.session_state.api_token = manually_entered_token
+                        API_TOKEN = manually_entered_token
+                        headers = {"Authorization": f"Bearer {API_TOKEN}"}
+                else:
+                    st.error(f"❌ API Error: {response.status_code} - {response.text}")
+                    st.info("Please check your API token and try again. Make sure you're using a valid Hugging Face API token.")
+            except Exception as e:
+                st.error(f"❌ Connection Error: {str(e)}")
+                st.info("Please check your internet connection and try again.")
+    
+    # Debug toggle
+    st.checkbox("Debug Mode", key="debug_mode", help="Show raw API responses in the sidebar")
+    
     if st.button("Clear Conversation"):
         st.session_state.chat_history = []
         st.experimental_rerun()
     
+    # API status indicator
+    st.markdown("---")
+    if API_TOKEN:
+        st.success("✅ API Connected")
+        st.info(f"Using model: {st.session_state.selected_model}")
+    else:
+        st.error("❌ API Not Connected")
+        st.info("Add your Hugging Face API token in the field above, in environment variables, or create a .streamlit/secrets.toml file")
+    
     st.markdown("---")
     st.markdown("### About")
     st.markdown("This customer support chatbot is powered by AI and provides assistance for common customer inquiries.")