add k value option

app.py

@@ -100,16 +100,18 @@ def get_docs():
             documents.extend(text_splitter.split_documents(load))
     return documents
 
-def set_chain_up(api_key, model_selector, agent):
+def set_chain_up(api_key, model_selector, k_textbox, agent):
     if api_key:
         os.environ["OPENAI_API_KEY"] = api_key
         documents = get_docs()
         embeddings = OpenAIEmbeddings()
         vectorstore = CachedChroma.from_documents_with_cache(".persisted_data", documents, embedding=embeddings, collection_name="pycbc")
-        if model_selector:
-            qa_chain = get_new_chain1(vectorstore, model_selector)
-        else:
-            qa_chain = get_new_chain1(vectorstore, "gpt-3.5-turbo")
+        if not model_selector:
+            model_selector = "gpt-3.5-turbo"
+        if not k_textbox:
+            k_textbox = 10
+        k_textbox = int(k_textbox)
+        qa_chain = get_new_chain1(vectorstore, model_selector, k_textbox)
         os.environ["OPENAI_API_KEY"] = ""
         return qa_chain
 
@@ -141,20 +143,26 @@ with block:
             type="password",
         )
         model_selector = gr.Dropdown(["gpt-3.5-turbo", "gpt-4"], label="Model", show_label=True)
-
+        k_textbox = gr.Textbox(
+            placeholder="k: Number of search results to consider",
+            label="Search Results k:"
+            show_label=True,
+            lines=1,
+        )
     chatbot = gr.Chatbot()
 
     with gr.Row():
         message = gr.Textbox(
             label="What's your question?",
-            placeholder="What's the answer to life, the universe, and everything?",
+            placeholder="What is pycbc?",
             lines=1,
         )
         submit = gr.Button(value="Send", variant="secondary").style(full_width=False)
 
     gr.Examples(
         examples=[
-            "What is pycbc?",
+            "What is PyCBC?",
+            "Where is the matched filtering done in the pycbc_live script?"
         ],
         inputs=message,
     )
@@ -176,12 +184,12 @@ with block:
 
     openai_api_key_textbox.change(
         set_chain_up,
-        inputs=[openai_api_key_textbox, model_selector, agent_state],
+        inputs=[openai_api_key_textbox, model_selector, k_textbox, agent_state],
         outputs=[agent_state],
     )
     model_selector.change(
         set_chain_up,
-        inputs=[openai_api_key_textbox, model_selector, agent_state],
+        inputs=[openai_api_key_textbox, model_selector, k_textbox, agent_state],
         outputs=[agent_state],
     )