Update app.py

app.py

@@ -49,8 +49,8 @@ def format_elapsed_time(time):
     # Format the elapsed time to two decimal places
     return "{:.2f}".format(time)
 
-def search_knowledgebase(query):
-    results = retriever(query)
+def search_knowledgebase(query, k):
+    results = retriever(query, k)
     return results
 
 def process_search_results(search_results):
@@ -109,7 +109,7 @@ def lookup_related_chunks(df_chunks, chunk_id):
     return df_chunks[(df_chunks['Title'] == title) & (df_chunks['PageNumber'].isin(page_range))]
 
 
-def search_and_reconstruct(query, df_chunks):
+def search_and_reconstruct(query, df_chunks, k):
     """
     Combines search, lookup of related chunks, and text reconstruction.
 
@@ -119,7 +119,7 @@ def search_and_reconstruct(query, df_chunks):
     :param top_k: Number of top search results to retrieve.
     :return: A list of dictionaries with document title, page number, and reconstructed text.
     """
-    search_results = search_knowledgebase(query)
+    search_results = search_knowledgebase(query, k)
     processed_results = process_search_results(search_results)
 
     reconstructed_results = []
@@ -171,7 +171,7 @@ persona1SystemMessage = st.sidebar.text_area("Query Designer System Message", va
 llm1 = st.sidebar.selectbox("Model Selection", ['GPT-4', 'GPT3.5'], key='persona1_size')
 temp1 = st.sidebar.slider("Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.6, key='persona1_temp')
 tokens1 = st.sidebar.slider("Tokens", min_value=1, max_value=4000, step=100, value=500, key='persona1_tokens')
-k = st.sidebar.slider("Tokens", min_value=1, max_value=10, step=1, value=3, key='k')
+k = st.sidebar.slider("Returned Docs", min_value=1, max_value=10, step=1, value=3, key='k')
 
 st.sidebar.caption(f"Session ID: {genuuid()}")