Update app.py

app.py

@@ -1,12 +1,7 @@
 import os
-import traceback
-
-# Try using a different directory path where you should have permissions
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/model_cache'
-os.environ['HF_HOME'] = '/tmp/model_cache'
-os.makedirs('/tmp/model_cache', exist_ok=True)
-
-from flask import Flask, render_template, jsonify, request
+import streamlit as st
+import json
+from datetime import datetime, timedelta
 from src.helper import download_hugging_face_embeddings
 from langchain_community.vectorstores import Pinecone
 from langchain_openai import OpenAI
@@ -14,149 +9,197 @@ from langchain.chains import create_retrieval_chain
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate
 from dotenv import load_dotenv
-from src.prompt import *
+from src.prompt import system_prompt
+
+# Set up cache directories
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/model_cache'
+os.environ['HF_HOME'] = '/tmp/model_cache'
+os.makedirs('/tmp/model_cache', exist_ok=True)
+
+# Load environment variables
+load_dotenv()
 
-app = Flask(__name__)
+# Rate limiting configuration
+RATE_LIMIT_FILE = "/tmp/rate_limits.json"
+MAX_REQUESTS_PER_DAY = 5
 
-# Load environment variables - these will be set in Hugging Face Space secrets
-load_dotenv()  # Still useful for local development
+# Initialize rate limiting storage
+def init_rate_limiting():
+    if not os.path.exists(RATE_LIMIT_FILE):
+        with open(RATE_LIMIT_FILE, 'w') as f:
+            json.dump({}, f)
 
-print("Starting application initialization")
-print(f"Python version: {os.sys.version}")
+# Check if a user has exceeded their daily limit
+def check_rate_limit(user_id):
+    today = datetime.now().strftime('%Y-%m-%d')
+    
+    try:
+        with open(RATE_LIMIT_FILE, 'r') as f:
+            rate_limits = json.load(f)
+    except (json.JSONDecodeError, FileNotFoundError):
+        rate_limits = {}
+    
+    # Clean up old entries
+    yesterday = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')
+    users_to_remove = []
+    for uid in rate_limits:
+        if yesterday in rate_limits[uid]:
+            del rate_limits[uid][yesterday]
+            if not rate_limits[uid]:  # If user has no other days, remove them
+                users_to_remove.append(uid)
+    
+    for uid in users_to_remove:
+        del rate_limits[uid]
+    
+    # Check and update current user's limit
+    if user_id not in rate_limits:
+        rate_limits[user_id] = {}
+    
+    if today not in rate_limits[user_id]:
+        rate_limits[user_id][today] = 0
+    
+    # Check if limit exceeded
+    if rate_limits[user_id][today] >= MAX_REQUESTS_PER_DAY:
+        return False, rate_limits[user_id][today]
+    
+    # Increment count and save
+    rate_limits[user_id][today] += 1
+    with open(RATE_LIMIT_FILE, 'w') as f:
+        json.dump(rate_limits, f)
+    
+    return True, rate_limits[user_id][today]
 
-# Add debugging endpoints
-@app.route("/test")
-def test():
-    return "Flask app is working. This is a test endpoint."
+def get_user_id():
+    # For Streamlit, we'll use session_id as user identifier
+    if not hasattr(st.session_state, 'user_id'):
+        st.session_state.user_id = str(hash(datetime.now().strftime("%Y%m%d%H%M%S")))
+    return st.session_state.user_id
 
-@app.route("/check-env")
-def check_env():
-    has_pinecone = "Yes" if os.environ.get("PINECONE_API_KEY") else "No"
-    has_openai = "Yes" if os.environ.get("OPENAI_API_KEY") else "No"
+def get_remaining_queries(user_id):
+    today = datetime.now().strftime('%Y-%m-%d')
     
-    # Check if keys appear valid (without revealing them)
-    pinecone_valid = len(os.environ.get("PINECONE_API_KEY", "")) > 10 if has_pinecone == "Yes" else "N/A"
-    openai_valid = os.environ.get("OPENAI_API_KEY", "").startswith("sk-") if has_openai == "Yes" else "N/A"
+    try:
+        with open(RATE_LIMIT_FILE, 'r') as f:
+            rate_limits = json.load(f)
+    except (json.JSONDecodeError, FileNotFoundError):
+        return MAX_REQUESTS_PER_DAY
     
-    return f"Pinecone key present: {has_pinecone} (appears valid: {pinecone_valid})<br>OpenAI key present: {has_openai} (appears valid: {openai_valid})"
+    count = rate_limits.get(user_id, {}).get(today, 0)
+    return MAX_REQUESTS_PER_DAY - count
+
+# Set up page configuration
+st.set_page_config(
+    page_title="Medical Assistant RAG Chatbot",
+    page_icon="🩺",
+    layout="centered"
+)
+
+# Initialize session state for chat history
+if 'messages' not in st.session_state:
+    st.session_state.messages = []
+
+# Initialize rate limiting
+init_rate_limiting()
+
+# Display remaining queries
+user_id = get_user_id()
+remaining_queries = get_remaining_queries(user_id)
+st.sidebar.write(f"Remaining queries today: {remaining_queries}/{MAX_REQUESTS_PER_DAY}")
 
-print("Checking environment variables...")
+# Check for API keys
 PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 
-if not PINECONE_API_KEY:
-    print("WARNING: Missing PINECONE_API_KEY")
-if not OPENAI_API_KEY:
-    print("WARNING: Missing OPENAI_API_KEY")
-
 if not PINECONE_API_KEY or not OPENAI_API_KEY:
-    print("CRITICAL ERROR: Missing API keys")
-    # We'll continue anyway to allow debugging, but the app won't work properly
+    st.error("Missing API keys. Please set PINECONE_API_KEY and OPENAI_API_KEY environment variables.")
+    st.stop()
 
 os.environ["PINECONE_API_KEY"] = PINECONE_API_KEY
 os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 
-# Initialize embeddings and chain at startup
-embeddings = None
-rag_chain = None
-
-def initialize_chain():
-    global embeddings, rag_chain
+# Cache the RAG chain initialization
+@st.cache_resource
+def initialize_rag_chain():
     try:
-        print("Step 1: Starting to download embeddings")
+        st.sidebar.write("Loading embeddings model...")
         embeddings = download_hugging_face_embeddings()
-        print("Step 2: Successfully downloaded embeddings")
         
+        st.sidebar.write("Connecting to Pinecone...")
         index_name = "medprep"
-        print(f"Step 3: Connecting to Pinecone index: {index_name}")
-        
-        try:
-            from pinecone import Pinecone as PineconeClient
-            pc = PineconeClient(api_key=PINECONE_API_KEY)
-            # List available indexes to verify connection
-            indexes = pc.list_indexes()
-            print(f"Available Pinecone indexes: {indexes}")
-            
-            if index_name not in [idx.name for idx in indexes]:
-                print(f"WARNING: Index '{index_name}' not found in your Pinecone account!")
-        except Exception as e:
-            print(f"Failed to connect to Pinecone API: {e}")
-        
         docsearch = Pinecone.from_existing_index(
             index_name=index_name,
             embedding=embeddings
         )
-        print("Step 4: Successfully connected to Pinecone")
         
-        retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k":3})
-        print("Step 5: Created retriever")
+        retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k": 3})
         
-        print("Step 6: Initializing OpenAI")
+        st.sidebar.write("Initializing OpenAI...")
         llm = OpenAI(temperature=0.4, max_tokens=500)
-        print("Step 7: OpenAI initialized")
         
-        print("Step 8: Creating prompt template")
-        prompt = ChatPromptTemplate.from_messages(
-            [
-                ("system", system_prompt),
-                ("human", "{input}"),
-            ]
-        )
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", system_prompt),
+            ("human", "{input}")
+        ])
         
-        print("Step 9: Creating QA chain")
         question_answer_chain = create_stuff_documents_chain(llm, prompt)
-        
-        print("Step 10: Creating RAG chain")
         rag_chain = create_retrieval_chain(retriever, question_answer_chain)
-        print("Step 11: RAG chain initialized successfully")
-        return True
+        
+        st.sidebar.success("✅ System initialized successfully!")
+        return rag_chain
     except Exception as e:
-        print(f"Failed to initialize RAG chain: {e}")
-        print(f"Error type: {type(e)}")
-        traceback.print_exc()
-        return False
-
-# Initialize the chain when the application starts
-print("Starting chain initialization...")
-initialization_result = initialize_chain()
-print(f"Chain initialization result: {initialization_result}")
-
-@app.route("/")
-def index():
-    return render_template('chat.html')
-
-@app.route("/get", methods=["GET", "POST"])
-def chat():
-    global rag_chain
+        st.sidebar.error(f"Error initializing system: {str(e)}")
+        import traceback
+        st.sidebar.text(traceback.format_exc())
+        return None
+
+# Main app title
+st.title("Medical Assistant Chatbot")
+st.write("Ask me any medical question, and I'll try to help!")
+
+# Initialize the RAG chain
+rag_chain = initialize_rag_chain()
+
+if rag_chain is None:
+    st.error("Failed to initialize the system. Please check the sidebar for error details.")
+    st.stop()
+
+# Display chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Get user input
+if prompt := st.chat_input("Ask a question..."):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
     
-    # Make sure chain is initialized
-    if rag_chain is None:
-        print("RAG chain not initialized, attempting to initialize again...")
-        if not initialize_chain():
-            return "Error: System not initialized properly. Please check the logs."
+    # Display user message
+    with st.chat_message("user"):
+        st.markdown(prompt)
     
-    msg = request.form["msg"]
-    try:
-        print(f"Processing message: {msg[:30]}...")  # Log only first 30 chars for privacy
-        response = rag_chain.invoke({"input": msg})
-        print("Successfully generated response")
-        return str(response["answer"])
-    except Exception as e:
-        error_msg = f"Error processing request: {e}"
-        print(error_msg)
-        traceback.print_exc()
-        return f"Error: {str(e)}"
-
-# Health check endpoint for monitoring
-@app.route("/health")
-def health_check():
-    is_initialized = rag_chain is not None
-    return jsonify({
-        "status": "healthy", 
-        "rag_chain_initialized": is_initialized,
-        "embeddings_loaded": embeddings is not None
-    })
-
-if __name__ == '__main__':
-    port = int(os.environ.get("PORT", 7860))
-    app.run(host="0.0.0.0", port=port, debug=False)
+    # Check rate limit
+    user_id = get_user_id()
+    allowed, count = check_rate_limit(user_id)
+    
+    if not allowed:
+        response = f"⚠️ Daily limit reached. You've used {count} queries today. Please try again tomorrow."
+    else:
+        # Process the query with the RAG chain
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                try:
+                    result = rag_chain.invoke({"input": prompt})
+                    response = result.get("answer", "Sorry, I couldn't find an answer to that.")
+                    remaining = MAX_REQUESTS_PER_DAY - count
+                    response += f"\n\n\n_You have {remaining} queries remaining today._"
+                except Exception as e:
+                    response = f"Error processing your request: {str(e)}"
+            
+            st.markdown(response)
+    
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+# Footer
+st.markdown("---")
+st.markdown("*This is a RAG-based medical assistant chatbot. It retrieves information from a medical knowledge base to answer your questions.*")