Update app.py

app.py

@@ -597,7 +597,7 @@ import traceback
 def generate_answer(message, choice, retrieval_mode, selected_model):
     logging.debug(f"generate_answer called with choice: {choice}, retrieval_mode: {retrieval_mode}, and selected_model: {selected_model}")
 
-    # Logic for disabling options for Phi-3.5
+    # Logic for disabling options for Phi-3.5 (LM-2)
     if selected_model == phi_pipe:
         choice = None
         retrieval_mode = None
@@ -611,10 +611,10 @@ def generate_answer(message, choice, retrieval_mode, selected_model):
         else:
             prompt_template = QA_CHAIN_PROMPT_1  # Fallback to template1
 
+        # VDB retrieval mode
         if retrieval_mode == "VDB":
             logging.debug("Using VDB retrieval mode")
-            # Adjust this block to handle both LM-1 and LM-3
-            if selected_model in [gpt_model, gpt_mini_model]:  # Both LM-1 and LM-3 should use the same logic
+            if selected_model in [gpt_model, gpt_mini_model]:  # Handle both LM-1 and LM-3
                 logging.debug(f"Selected model: {'LM-1' if selected_model == gpt_model else 'LM-3'}")
                 retriever = gpt_retriever
                 context = retriever.get_relevant_documents(message)
@@ -659,11 +659,13 @@ def generate_answer(message, choice, retrieval_mode, selected_model):
                     logging.error("LM-2 did not return any response.")
                     return "No response generated.", []
 
+        # KGF retrieval mode
         elif retrieval_mode == "KGF":
             logging.debug("Using KGF retrieval mode")
             response = chain_neo4j.invoke({"question": message})
             logging.debug(f"KGF response: {response}")
             return response, extract_addresses(response)
+
         else:
             logging.error("Invalid retrieval mode selected.")
             return "Invalid retrieval mode selected.", []