Update app.py

app.py

@@ -161,59 +161,67 @@ if "messages" not in st.session_state:
     # return res['result']
 
 def generate_response(prompt_input):
-    # Retrieve vector database context using ONLY the current user input
-    retriever = st.session_state.chain.retriever
-    relevant_context = retriever.get_relevant_documents(prompt_input)  # Retrieve context only for the current prompt
-    
-    # Prepare full conversation history for the LLM
-    conversation_history = ""
-    for message in st.session_state.messages:
-        conversation_history += f"{message['role']}: {message['content']}\n"
-    
-    # Append the current user prompt to the conversation history
-    conversation_history += f"user: {prompt_input}\n"
-
-    # Format the input for the chain with the retrieved context
-    formatted_input = (
-        f"Context:\n"
-        f"{' '.join([doc.page_content for doc in relevant_context])}\n\n"
-        f"Conversation:\n{conversation_history}"
-    )
-
-    # Invoke the RetrievalQA chain directly with the formatted input
-    res = st.session_state.chain.invoke({"query": formatted_input})
-    
-    # Process the response text
-    result_text = res['result']
-    
-    # Clean up prefixing phrases and capitalize the first letter
-    if result_text.startswith('According to the provided context, '):
-        result_text = result_text[35:].strip()
-    elif result_text.startswith('Based on the provided context, '):
-        result_text = result_text[31:].strip()
-    elif result_text.startswith('According to the provided text, '):
-        result_text = result_text[34:].strip()
-    elif result_text.startswith('According to the context, '):
-        result_text = result_text[26:].strip()
-    
-    # Ensure the first letter is uppercase
-    result_text = result_text[0].upper() + result_text[1:] if result_text else result_text
-
-    # Extract and format sources (if available)
-    sources = []
-    for doc in relevant_context:
-        source_path = doc.metadata.get('source', '')
-        formatted_source = source_path[122:-4] if source_path else "Unknown source"
-        sources.append(formatted_source)
+    try:
+        # Retrieve vector database context using ONLY the current user input
+        retriever = st.session_state.chain.retriever
+        relevant_context = retriever.get_relevant_documents(prompt_input)  # Retrieve context only for the current prompt
+        
+        # Prepare full conversation history for the LLM
+        conversation_history = ""
+        for message in st.session_state.messages:
+            conversation_history += f"{message['role']}: {message['content']}\n"
+        
+        # Append the current user prompt to the conversation history
+        conversation_history += f"user: {prompt_input}\n"
     
-    # Remove duplicates and combine into a single string
-    unique_sources = list(set(sources))
-    source_list = ", ".join(unique_sources)
+        # Format the input for the chain with the retrieved context
+        formatted_input = (
+            f"Context:\n"
+            f"{' '.join([doc.page_content for doc in relevant_context])}\n\n"
+            f"Conversation:\n{conversation_history}"
+        )
     
-    # # Combine response text with sources
-    # result_text += f"\n\n**Sources:** {source_list}" if source_list else "\n\n**Sources:** None"
+        # Invoke the RetrievalQA chain directly with the formatted input
+        res = st.session_state.chain.invoke({"query": formatted_input})
+        
+        # Process the response text
+        result_text = res['result']
+        
+        # Clean up prefixing phrases and capitalize the first letter
+        if result_text.startswith('According to the provided context, '):
+            result_text = result_text[35:].strip()
+        elif result_text.startswith('Based on the provided context, '):
+            result_text = result_text[31:].strip()
+        elif result_text.startswith('According to the provided text, '):
+            result_text = result_text[34:].strip()
+        elif result_text.startswith('According to the context, '):
+            result_text = result_text[26:].strip()
+        
+        # Ensure the first letter is uppercase
+        result_text = result_text[0].upper() + result_text[1:] if result_text else result_text
     
-    return result_text
+        # Extract and format sources (if available)
+        sources = []
+        for doc in relevant_context:
+            source_path = doc.metadata.get('source', '')
+            formatted_source = source_path[122:-4] if source_path else "Unknown source"
+            sources.append(formatted_source)
+        
+        # Remove duplicates and combine into a single string
+        unique_sources = list(set(sources))
+        source_list = ", ".join(unique_sources)
+        
+        # # Combine response text with sources
+        # result_text += f"\n\n**Sources:** {source_list}" if source_list else "\n\n**Sources:** None"
+        
+        return result_text
+
+    except Exception as e:
+        # Handle rate limit or other errors gracefully
+        if "rate_limit_exceeded" in str(e).lower():
+            return "⚠️ Rate limit exceeded. Please clear the chat history and try again."
+        else:
+            return f"❌ An error occurred: {str(e)}"
 
 # Display chat messages
 for message in st.session_state.messages: