Update helper_functions_api.py

helper_functions_api.py

@@ -67,6 +67,7 @@ from half_json.core import JSONFixer
 from openai import OpenAI
 from together import Together
 from urllib.parse import urlparse
+import trafilatura
 
 llm_default_small = "meta-llama/Llama-3-8b-chat-hf"
 llm_default_medium = "meta-llama/Llama-3-70b-chat-hf"
@@ -197,13 +198,20 @@ class Scraper:
         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 ""
+    extracted = trafilatura.extract(
+        html,
+        output_format="markdown",
+        target_language="en",
+        include_tables=True,
+        include_images=False,
+        include_links=False,
+        deduplicate=True,
+    )
+    
+    if extracted:
+        return trafilatura.utils.sanitize(extracted)
+    else:
+        return ""
 
 def process_content(data_format, url, query):
     scraper = Scraper()
@@ -213,7 +221,7 @@ def process_content(data_format, url, query):
         if content:
             rephrased_content = rephrase_content(
                 data_format=data_format,
-                content=limit_tokens(remove_stopwords(content), token_limit=1000),
+                content=limit_tokens(remove_stopwords(content), token_limit=4000),
                 query=query,
             )
             return rephrased_content, url