Update app.py

app.py

@@ -1,7 +1,9 @@
 import os
 from collections.abc import Iterator
 from threading import Thread
-
+import requests
+from bs4 import BeautifulSoup
+from readability import Document
 import gradio as gr
 import spaces
 import torch
@@ -27,6 +29,38 @@ model = AutoModelForCausalLM.from_pretrained(
 model.config.sliding_window = 4096
 model.eval()
 
+def extract_text_from_webpage(html_content):
+    doc = Document(html_content)
+    return doc.summary()
+
+def search(query):
+    term = query
+    all_results = []
+    max_chars_per_page = 8000
+    with requests.Session() as session:
+        resp = session.get(
+            url="https://www.google.com/search",
+            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
+            params={"q": term, "num": 4, "udm": 14},
+            timeout=5,
+            verify=None,
+        )
+        resp.raise_for_status()
+        soup = BeautifulSoup(resp.text, "html.parser")
+        result_block = soup.find_all("div", attrs={"class": "g"})
+        for result in result_block:
+            link = result.find("a", href=True)
+            link = link["href"]
+            try:
+                webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"}, timeout=5, verify=False)
+                webpage.raise_for_status()
+                visible_text = extract_text_from_webpage(webpage.text)
+                if len(visible_text) > max_chars_per_page:
+                    visible_text = visible_text[:max_chars_per_page]
+                all_results.append({"link": link, "text": visible_text})
+            except requests.exceptions.RequestException:
+                all_results.append({"link": link, "text": None})
+    return all_results
 
 @spaces.GPU(duration=120)
 def generate(
@@ -41,6 +75,16 @@ def generate(
     conversation = chat_history.copy()
     conversation.append({"role": "user", "content": message})
 
+    # Check if the message requires a web search
+    if "search" in message.lower() or "find" in message.lower():
+        search_query = message
+        search_results = search(search_query)
+        if search_results:
+            search_context = "\n".join([result["text"] for result in search_results if result["text"]])
+            conversation.append({"role": "assistant", "content": f"Here are some search results:\n{search_context}"})
+        else:
+            conversation.append({"role": "assistant", "content": "I couldn't find any relevant information."})
+
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -67,7 +111,6 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
-
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
@@ -121,6 +164,5 @@ demo = gr.ChatInterface(
     fill_height=True,
 )
 
-
 if __name__ == "__main__":
     demo.queue(max_size=20).launch()