Update test_web_rag.py

test_web_rag.py

@@ -15,6 +15,7 @@ from dotenv import load_dotenv
 
 load_dotenv()
 ZENROWS_KEY = os.getenv('ZENROWS_KEY')
+you_key = os.getenv("YOU_API_KEY")
 client = OpenAI()
 
 
@@ -206,6 +207,31 @@ def check_url_pdf_file(url):
         return False
 
 
+def get_ai_snippets_for_query(query, num):
+    headers = {"X-API-Key": you_key}
+    params = {"query": query}
+    return requests.get(
+        f"https://api.ydc-index.io/search?query={query}&num_web_results={num}",
+        params=params,
+        headers=headers,
+    ).json().get('hits')
+
+
+def get_web_search_you(query, num):
+    docs = get_ai_snippets_for_query(query, num)
+    markdown = ""
+    for doc in docs:
+        for key, value in doc.items():
+            if key == 'snippets':
+                markdown += f"{key}:\n"
+                for snippet in value:
+                    markdown += f"- {snippet}\n"
+            else:
+                markdown += f"{key}: {value}\n"
+        markdown += "\n"
+    return markdown
+
+
 def zenrows_scrapper(url):
     zen_client = ZenRowsClient(ZENROWS_KEY)
     params = {"js_render": "true"}
@@ -231,33 +257,39 @@ def get_new_question_from_history(pre_question, new_question, answer):
 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])
-    urls = get_google_search_url(get_google_search_query(question))[:n_web_search]
-    urls = list(set(urls))
     docs = ''
-    yield f"Scraping started for {len(urls)} urls:-\n\n"
-    for key, url in enumerate(urls):
-        if '.pdf' in url:
-            yield f"Scraping skipped pdf detected. {key + 1}/{len(urls)} - {url} ❌\n"
-            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)
-            formatted_page_source = clean_text(formatted_page_source)
-
-        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
-
-        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"
+    if strategy == 'Deep Pro':
+        docs = get_web_search_you(question, n_web_search)
+    else:
+        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"
+        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"
+                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)
+
+            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
+
+            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"
     yield {"data": docs}