Update test_web_rag.py

test_web_rag.py

@@ -1,3 +1,4 @@
+import time
 import urllib.request
 from urllib.parse import quote
 from seleniumbase import SB
@@ -12,6 +13,7 @@ from zenrows import ZenRowsClient
 import requests
 import os
 from dotenv import load_dotenv
+from threading import Thread
 
 load_dotenv()
 ZENROWS_KEY = os.getenv('ZENROWS_KEY')
@@ -254,6 +256,25 @@ def get_new_question_from_history(pre_question, new_question, answer):
     return call_open_ai(system_prompt=system_prompt, max_tokens=50)
 
 
+def scraping_job(strategy, question, url, results, key):
+    if strategy == 'Deep':
+        # page_source = get_page_source_selenium_base(url)
+        page_source = zenrows_scrapper(url)
+        formatted_page_source = format_output(page_source)
+        formatted_page_source = clean_text(formatted_page_source)
+    else:
+        page_source = get_fast_url_source(url)
+        formatted_page_source = format_output(page_source)
+        formatted_page_source = clean_text(formatted_page_source)
+
+    tokens = token_counter(formatted_page_source, 'gpt-3.5-turbo')
+    if tokens >= 15585:
+        results[key] = ''
+    else:
+        summary = url_summary(formatted_page_source, question)
+        results[key] = summary
+
+
 def get_docs_from_web(question, history, n_web_search, strategy):
     if history:
         question = get_new_question_from_history(history[0][0], question, history[0][1])
@@ -264,32 +285,26 @@ def get_docs_from_web(question, history, n_web_search, strategy):
         urls = get_google_search_url(get_google_search_query(question))[:n_web_search]
         urls = list(set(urls))
         yield f"Scraping started for {len(urls)} urls:-\n\n"
+
+        threads = [None] * len(urls)
+        results = [None] * len(urls)
+
         for key, url in enumerate(urls):
             if '.pdf' in url or '.PDF' in url:
                 yield f"Scraping skipped pdf detected. {key + 1}/{len(urls)} - {url} ❌\n"
+                results[key] = ''
                 continue
 
-            if strategy == 'Deep':
-                # page_source = get_page_source_selenium_base(url)
-                page_source = zenrows_scrapper(url)
-                formatted_page_source = format_output(page_source)
-                formatted_page_source = clean_text(formatted_page_source)
-            else:
-                page_source = get_fast_url_source(url)
-                formatted_page_source = format_output(page_source)
-                if formatted_page_source == '':
-                    yield f"Scraping failed. {key + 1}/{len(urls)} - {url} ❌\n"
-                    continue
-                formatted_page_source = clean_text(formatted_page_source)
+            threads[key] = Thread(target=scraping_job, args=(strategy, question, url, results, key))
+            threads[key].start()
 
-            tokens = token_counter(formatted_page_source, 'gpt-3.5-turbo')
-
-            if tokens >= 15585:
-                yield f"Scraping skipped as token limit exceeded. {key + 1}/{len(urls)} - {url} ❌\n"
-                continue
+        for i in range(len(threads)):
+            if threads[i] is not None:
+                threads[i].join()
 
-            summary = url_summary(formatted_page_source, question)
-            docs += summary
-            docs += '\n Source:-' + url + '\n\n'
-            yield f"Scraping Done {key + 1}/{len(urls)} - {url} ✅\n"
+        for key, result in enumerate(results):
+            if result is not None and result != '':
+                docs += result
+                docs += '\n Source:-' + urls[key] + '\n\n'
+                yield f"Scraping Done {key + 1}/{len(urls)} - {urls[key]} ✅\n"
     yield {"data": docs}