Update app.py

app.py

@@ -1193,19 +1193,25 @@ update_message_counter()
  
 
 # Define search functions
-def search_knowledge_base(query):
-    """
-    Searches the FAISS index using the provided query.
-    Returns the most relevant documents based on the query.
-    """
-    if "faiss_db" not in st.session_state:
-        st.error("FAISS database is not initialized.")
-        return []
+def search_knowledge_base(query, k=3):
+    """Optimized FAISS search for main knowledge base and user-specific knowledge base."""
+    results = []
+
+    # Search in the main FAISS index
+    if "main_faiss_db" in st.session_state and st.session_state["main_faiss_db"] is not None:
+        main_results = st.session_state["main_faiss_db"].similarity_search_with_score(query, k=5)  # Fetch extra results for better ranking
+        results.extend(main_results)
+
+    # Search in the selected document's FAISS index
+    if "faiss_db" in st.session_state and st.session_state["faiss_db"] is not None:
+        user_results = st.session_state["faiss_db"].similarity_search_with_score(query, k=5)  # Fetch extra results for better ranking
+        results.extend(user_results)
 
-    # Retrieve the top 5 most relevant documents
-    retrieved_docs = st.session_state["faiss_db"].similarity_search(query, k=7)
-    return retrieved_docs
+    # Sort results by similarity score (higher score = more relevant)
+    sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
 
+    # Return only the top `k` most relevant results
+    return [result[0] for result in sorted_results[:k]]
 
 def google_search(query):
     """