Update helper_functions_api.py

helper_functions_api.py

@@ -4,8 +4,6 @@ from mistune.plugins.table import table
 from jinja2 import Template
 import re
 import os
-from urllib.parse import urlparse
-from typing import Dict, Any, List, Tuple
 
 def md_to_html(md_text):
     renderer = mistune.HTMLRenderer()
@@ -182,77 +180,60 @@ def rephrase_content(data_format, content, query):
             max_tokens=500,
         )
         
-def extract_main_content(url):
-    if url:
-        try:
-            result = urlparse(url)
-            if all([result.scheme, result.netloc]):
-                # Prepare query parameters
-                params = {
-                    "url": url,
-                    "favor_precision": False,
-                    "favor_recall": False,
-                    "output_format": "markdown",
-                    "target_language": "en",
-                    "include_tables": True,
-                    "include_images": False,
-                    "include_links": False,
-                    "deduplicate": True,
-                }
-
-                # Make request to FastAPI endpoint
-                response = requests.get("https://pvanand-web-scraping.hf.space/extract-article", params=params)
+class Scraper:
+    def __init__(self, user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"):
+        self.session = requests.Session()
+        self.session.headers.update({"User-Agent": user_agent})
 
-                if response.status_code == 200:
-                    return response.json()["article"]
-                else:
-                    return ""
-        except:
-            return ""
+    @retry(tries=3, delay=1)
+    def fetch_content(self, url):
+        try:
+            response = self.session.get(url, timeout=2)
+            if response.status_code == 200:
+                return response.text
+        except requests.exceptions.RequestException as e:
+            print(f"Error fetching page content for {url}: {e}")
+        return None
+
+def extract_main_content(html):
+    if html:
+        plain_text = ""
+        soup = BeautifulSoup(html, 'lxml')
+        for element in soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'table']):
+            plain_text += element.get_text(separator=" ", strip=True) + "\n"
+        return plain_text
     return ""
 
 def process_content(data_format, url, query):
-    content = extract_main_content(url)
-    if content:
-        rephrased_content = rephrase_content(
-            data_format=data_format,
-            content=limit_tokens(content, token_limit=4000),
-            query=query,
-        )
-        return rephrased_content, url
+    scraper = Scraper()
+    html_content = scraper.fetch_content(url)
+    if html_content:
+        content = extract_main_content(html_content)
+        if content:
+            rephrased_content = rephrase_content(
+                data_format=data_format,
+                content=limit_tokens(remove_stopwords(content), token_limit=1000),
+                query=query,
+            )
+            return rephrased_content, url
     return "", url
 
-def fetch_and_extract_content(
-    data_format: str, query: str, urls: List[str], num_refrences: int = 6
-) -> List[Tuple[str | None, str]]:
-    """
-    Asynchronously makeing request to urls and doing further process
-    """
-    all_text_with_urls = []
-    start_url = 0
-    while (len(all_text_with_urls) != num_refrences) and (start_url < len(urls)):
-        end_url = start_url + (num_refrences - len(all_text_with_urls))
-        urls_subset = urls[start_url:end_url]
-        with ThreadPoolExecutor(max_workers=len(urls_subset)) as executor:
-            future_to_url = {
-                executor.submit(process_content, data_format, url, query): url
-                for url in urls_subset
-            }
-            all_text_with_urls += [
-                future.result()
-                for future in as_completed(future_to_url)
-                if future.result()[0] != ""
-            ]
-        start_url = end_url
+def fetch_and_extract_content(data_format, urls, query):
+    with ThreadPoolExecutor(max_workers=len(urls)) as executor:
+        future_to_url = {
+            executor.submit(process_content, data_format, url, query): url
+            for url in urls
+        }
+        all_text_with_urls = [future.result() for future in as_completed(future_to_url)]
 
     return all_text_with_urls
 
 
 @retry(tries=3, delay=0.25)
 def search_brave(query, num_results=5):
-    cleaned_query = query #re.sub(r'[^a-zA-Z0-9]+', '', query)
+    cleaned_query = re.sub(r'[^a-zA-Z0-9]+', '', query)
     search_query = together_response(cleaned_query, model=llm_default_small, SysPrompt=SysPromptSearch, max_tokens = 25).strip()
-    cleaned_search_query = search_query #re.sub(r'[^a-zA-Z0-9*]+', '', search_query)
+    cleaned_search_query = re.sub(r'[^a-zA-Z0-9*]+', '', search_query)
     brave = Brave(BRAVE_API_KEY)
     search_results = brave.search(q=cleaned_search_query, count=num_results)
     return [url.__str__() for url in search_results.urls],cleaned_search_query