Update app.py

app.py

@@ -118,17 +118,82 @@ def convert_history_to_token(history: List[Tuple[str, str]]):
         input_token = tok(text, return_tensors="pt", **tokenizer_kwargs).input_ids
     return input_token
 
+# Initialize the search tool
+search = DuckDuckGoSearchRun()
+
+# Function to retrieve and format search results based on user input
+def fetch_search_results(query: str) -> str:
+    search_results = search.invoke(query)
+    # Displaying search results for debugging
+    print("Search results: ", search_results)
+    return f"Relevant and recent information:\n{search_results}"
+
+# Function to decide if a search is needed based on the user query
+def should_use_search(query: str) -> bool:
+    # Simple heuristic, can be extended with more advanced intent analysis
+    search_keywords = ["latest", "news", "update", "which" "who", "what", "when", "why","how", "recent", "result", "tell", "explain", 
+                       "announcement", "bulletin", "report", "brief", "insight", "disclosure", "update", 
+                        "release", "memo", "headline", "current", "ongoing", "fresh", "upcoming", "immediate", 
+                        "recently", "new", "now", "in-progress", "inquiry", "query", "ask", "investigate", 
+                        "explore", "seek", "clarify", "confirm", "discover", "learn", "describe", "define", 
+                        "illustrate", "outline", "interpret", "expound", "detail", "summarize", "elucidate", 
+                        "break down", "outcome", "effect", "consequence", "finding", "achievement", "conclusion", 
+                        "product", "performance", "resolution"
+                        ]
+    return any(keyword in query.lower() for keyword in search_keywords)
+
+# Generate prompt for model with optional search context
+def construct_model_prompt(user_query: str, search_context: str, history: List[Tuple[str, str]]) -> str:
+    # Simple instruction for the model to prioritize search information if available
+    instructions = (
+        "If relevant information is provided below, use it to give an accurate and concise answer. If there is no relevant information available, please rely on your general knowledge and indicate that no recent or specific information is available to answer."
+    )
+    
+    # Build the prompt with instructions, search context, and user query
+    prompt = f"{instructions}\n\n"
+    if search_context:
+        prompt += f"{search_context}\n\n"  # Include search context prominently at the top
+    
+    # Add the user's query
+    prompt += f"{user_query} ?\n\n"
+    
+    # Optionally add recent history for context, without labels
+    # if history:
+    #     prompt += "Recent conversation:\n"
+    #     for user_msg, assistant_msg in history[:-1]:  # Exclude the last message to prevent duplication
+    #         prompt += f"{user_msg}\n{assistant_msg}\n"
+    
+    return prompt
+
+
 def bot(history, temperature, top_p, top_k, repetition_penalty, conversation_id):
-    # Callback function for running chatbot on submit button click
-    input_ids = convert_history_to_token(history)
+    """
+    Main callback function for running chatbot on submit button click.
+    """
+    user_query = history[-1][0]
+    search_context = ""
+
+    # Decide if search is required based on the user query
+    if should_use_search(user_query):
+        search_context = fetch_search_results(user_query)
+        prompt = construct_model_prompt(user_query, search_context, history)
+        input_ids = tok(prompt, return_tensors="pt", truncation=True, max_length=2500).input_ids
+    else:
+        # If no search context, use the original logic with tokenization
+        prompt = construct_model_prompt(user_query, "", history)
+        input_ids = convert_history_to_token(history)
+
+    # Ensure input length does not exceed a threshold (e.g., 2000 tokens)
     if input_ids.shape[1] > 2000:
+        # If input exceeds the limit, only use the most recent conversation
         history = [history[-1]]
-        input_ids = convert_history_to_token(history)
 
-    streamer = TextIteratorStreamer(tok, timeout=3600.0, skip_prompt=True, skip_special_tokens=True)
+    # Streamer for model response generation
+    streamer = TextIteratorStreamer(tok, timeout=4600.0, skip_prompt=True, skip_special_tokens=True)
+
     generate_kwargs = dict(
         input_ids=input_ids,
-        max_new_tokens=256,
+        max_new_tokens=256,  # Adjust this as needed
         temperature=temperature,
         do_sample=temperature > 0.0,
         top_p=top_p,
@@ -136,23 +201,32 @@ def bot(history, temperature, top_p, top_k, repetition_penalty, conversation_id)
         repetition_penalty=repetition_penalty,
         streamer=streamer,
     )
+
+    if stop_tokens is not None:
+        generate_kwargs["stopping_criteria"] = StoppingCriteriaList(stop_tokens)
+
+    # Event to signal when streaming is complete
     stream_complete = Event()
 
     def generate_and_signal_complete():
         ov_model.generate(**generate_kwargs)
         stream_complete.set()
 
-    Thread(target=generate_and_signal_complete).start()
+    t1 = Thread(target=generate_and_signal_complete)
+    t1.start()
+
+    # Initialize an empty string to store the generated text
     partial_text = ""
     for new_text in streamer:
-        partial_text += new_text
-        history[-1][1] = partial_text
+        partial_text = text_processor(partial_text, new_text)
+        # Update the last entry in the original history with the response
+        history[-1] = (user_query, partial_text)
         yield history
 
 def request_cancel():
     ov_model.request.cancel()
 
 # Gradio setup and launch
-demo = make_demo(run_fn=bot, stop_fn=request_cancel, title=f"OpenVINO {model_id_value} Chatbot", language=model_language_value)
+demo = make_demo(run_fn=bot, stop_fn=request_cancel, title=f"OpenVINO Search & Reasoning Chatbot", language=model_language_value)
 if __name__ == "__main__":
     demo.launch(debug=True, share=True, server_name="0.0.0.0", server_port=7860)