Update app.py

app.py

@@ -107,9 +107,40 @@ def reset_chat(idx, ld, state):
         gr.update(interactive=False),
     )
 
+def internet_search(ppmanager, serper_api_key, global_context, ctx_num_lconv, device="cpu"):
+    internet_search_ppm = copy.deepcopy(ppm)
+    internet_search_prompt = f"My question is '{user_msg}'. Based on the conversation history, "
+                             f"give me an appropriate query to answer my question for google search. "
+                             f"You should not say more than query. You should not say any words except the query."
+    
+    internet_search_ppm.pingpongs[-1].ping = internet_search_prompt
+    internet_search_prompt = build_prompts(internet_search_ppm, "", win_size=ctx_num_lconv)
+
+    instruction = gen_text_none_stream(internet_search_prompt, hf_model=MODEL_ID, hf_token=TOKEN)
+    ### 
+
+    searcher = SimilaritySearcher.from_pretrained(device=device)
+    iss = InternetSearchStrategy(
+        searcher, 
+        instruction=instruction, 
+        serper_api_key=serper_api_key
+    )(ppmanager)
+
+    step_ppm = None
+    while True:
+        try:
+            step_ppm, _ = next(iss)
+            yield "", step_ppm.build_uis()
+        except StopIteration:
+            break
+
+    search_prompt = build_prompts(step_ppm, global_context, ctx_num_lconv)
+    yield search_prompt, ppmanager.build_uis()
+
 async def chat_stream(
     idx, local_data, instruction_txtbox, chat_state,
-    global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv
+    global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
+    internet_option, serper_api_key
 ):
     res = [
       chat_state["ppmanager_type"].from_json(json.dumps(ppm))
@@ -121,6 +152,14 @@ async def chat_stream(
         PingPong(instruction_txtbox, "")
     )
     prompt = build_prompts(ppm, global_context, ctx_num_lconv)
+
+    #######
+    if internet_option:
+        search_prompt = None
+        for tmp_prompt, uis in internet_search(ppm, serper_api_key, global_context, ctx_num_lconv):
+            search_prompt = tmp_prompt
+            yield "", uis, prompt, str(res)
+    
     async for result in gen_text(
         prompt, hf_model=MODEL_ID, hf_token=TOKEN,
         parameters={
@@ -283,14 +322,15 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
                             elem_id="global-context"
                         )
                     
-                    # gr.Markdown("#### Internet search")
-                    # with gr.Row():
-                    #     internet_option = gr.Radio(choices=["on", "off"], value="off", label="mode")
-                    #     serper_api_key = gr.Textbox(
-                    #         value= "" if args.serper_api_key is None else args.serper_api_key,
-                    #         placeholder="Get one by visiting serper.dev", 
-                    #         label="Serper api key"
-                    #     )
+                    gr.Markdown("#### Internet search")
+                    with gr.Row():
+                        internet_option = gr.Radio(choices=["on", "off"], value="off", label="mode")
+                        serper_api_key = gr.Textbox(
+                            value= os.getenv("SERPER_API_KEY"),
+                            placeholder="Get one by visiting serper.dev", 
+                            label="Serper api key",
+                            visible=False
+                        )
                     
                     gr.Markdown("#### GenConfig for **response** text generation")
                     with gr.Row():
@@ -315,7 +355,8 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
     ).then(
         chat_stream,
         [idx, local_data, instruction_txtbox, chat_state,
-         global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv],
+         global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
+         internet_option, serper_api_key],
         [instruction_txtbox, context_inspector, chatbot, local_data, regenerate]
     ).then(
         None, local_data, None, 
@@ -346,7 +387,8 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
     regen_event = regenerate.click(
         rollback_last,
         [idx, local_data, chat_state,
-         global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv],
+         global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
+         internet_option, serper_api_key],
         [context_inspector, chatbot, local_data, regenerate]
     ).then(
         None, local_data, None,