Update helper_functions_api.py

helper_functions_api.py

@@ -1,4 +1,3 @@
-# !pip install mistune
 import mistune
 from mistune.plugins.table import table
 from jinja2 import Template
@@ -9,97 +8,18 @@ import markdown
 from bs4 import BeautifulSoup
 from lxml import etree
 import markdown
-
-def md_to_html(md_text):
-    html_content = markdown.markdown(md_text,extensions=["extra"])
-    return html_content.replace('\n', '')
-
-def has_tables(html_string):
-    try:
-        # Use BeautifulSoup with lxml parser
-        soup = BeautifulSoup(html_string, 'lxml')
-        
-        # First, try BeautifulSoup's find_all method
-        if soup.find_all('table'):
-            return True
-        
-        # If no tables found, try a more aggressive search using lxml's XPath
-        tree = etree.HTML(str(soup))
-        return len(tree.xpath('//table')) > 0
-    
-    except Exception as e:
-        # Log the exception if needed
-        print(f"An error occurred: {str(e)}")
-        return False
-
-def extract_data_from_tag(input_string, tag):
-    # Create the regex pattern
-    pattern = f'<{tag}.*?>(.*?)</{tag}>'
-    
-    # Find all matches
-    matches = re.findall(pattern, input_string, re.DOTALL)
-    
-    # If matches are found, return them joined by newlines
-    if matches:
-        out = '\n'.join(match.strip() for match in matches)
-        # Check for incorrect tagging
-        if len(out) > 0.8*len(input_string):
-            return out
-        else:
-            return input_string
-    
-    # If no matches are found, return the original string
-    return input_string
-
-####------------------------------ OPTIONAL--> User id and persistant data storage-------------------------------------####
+import logging
 from datetime import datetime
 import psycopg2
-
-from dotenv import load_dotenv, find_dotenv
-
-# Load environment variables from .env file
-load_dotenv("keys.env")
-
-TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
-BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-HELICON_API_KEY = os.getenv("HELICON_API_KEY")
-SUPABASE_USER = os.environ['SUPABASE_USER']
-SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']
-OPENROUTER_API_KEY = "sk-or-v1-"+os.environ['OPENROUTER_API_KEY']
-
-def insert_data(user_id, user_query, subtopic_query, response, html_report):
-    # Connect to your database
-    conn = psycopg2.connect(
-    dbname="postgres",
-    user=SUPABASE_USER,
-    password=SUPABASE_PASSWORD,
-    host="aws-0-us-west-1.pooler.supabase.com",
-    port="5432"
-)
-    cur = conn.cursor()
-    insert_query = """
-    INSERT INTO research_pro_chat_v2 (user_id, user_query, subtopic_query, response, html_report, created_at)
-    VALUES (%s, %s, %s, %s, %s, %s);
-    """
-    cur.execute(insert_query, (user_id,user_query, subtopic_query, response, html_report, datetime.now()))
-    conn.commit()
-    cur.close()
-    conn.close()
-
-####-----------------------------------------------------END----------------------------------------------------------####
-
-
+from dotenv import load_dotenv
 import ast
 from fpdf import FPDF
-import re
 import pandas as pd
 import nltk
 import requests
 import json
 from retry import retry
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from bs4 import BeautifulSoup
 from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
 from brave import Brave
@@ -109,21 +29,28 @@ from openai import OpenAI
 from together import Together
 from urllib.parse import urlparse
 import trafilatura
+import tiktoken
 
-llm_default_small = "meta-llama/Llama-3-8b-chat-hf"
-llm_default_medium = "meta-llama/Llama-3-70b-chat-hf"
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# SysPromptData = """You are expert in information extraction from the given context.
-#                     Steps to follow:
-#                     1. Check if relevant factual data regarding <USER QUERY> is present in the <SCRAPED DATA>.
-#                        - IF YES, extract the maximum relevant factual information related to <USER QUERY> from the <SCRAPED DATA>.
-#                        - IF NO, then return "N/A"
-                    
-#                     Rules to follow:
-#                     - Return N/A if information is not present in the scraped data.
-#                     - FORGET EVERYTHING YOU KNOW, Only output information that is present in the scraped data, DO NOT MAKE UP INFORMATION
-#                 """
-SysPromptData = """
+# Load environment variables
+load_dotenv("keys.env")
+TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
+BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+HELICON_API_KEY = os.getenv("HELICON_API_KEY")
+SUPABASE_USER = os.environ['SUPABASE_USER']
+SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']
+OPENROUTER_API_KEY = "sk-or-v1-" + os.environ['OPENROUTER_API_KEY']
+
+# Define constants
+LLM_DEFAULT_SMALL = "llama3-8b-8192"
+LLM_DEFAULT_MEDIUM = "llama3-70b-8192"
+LLM_FALLBACK_SMALL = "meta-llama/Llama-3-8b-chat-hf"
+LLM_FALLBACK_MEDIUM = "meta-llama/Llama-3-70b-chat-hf"
+
+SYS_PROMPT_DATA = """
 You are an AI assistant tasked with extracting relevant information from scraped website data based on a given query. Your goal is to provide accurate and concise information that directly relates to the query, using only the data provided.
 Guidelines for extraction:
 1. Only use information present in the scraped data.
@@ -131,66 +58,108 @@ Guidelines for extraction:
 3. If there is no relevant information in the scraped data, state that clearly.
 4. Do not make assumptions or add information not present in the data.
 5. If the query is ambiguous, interpret it in the most reasonable way based on the available data.
-                """
+"""
 
-SysPromptDefault = "You are an expert AI, complete the given task. Do not add any additional comments."
-SysPromptSearch = """You are a search query generator, create a concise Google search query, focusing only on the main topic and omitting additional redundant details, include year if necessory, 2024, Do not add any additional comments. OUTPUT ONLY THE SEARCH QUERY
-                        #Additional instructions:
-                        ##Use the following search operator if necessory
-                        OR #to cover multiple topics"""
+SYS_PROMPT_DEFAULT = "You are an expert AI, complete the given task. Do not add any additional comments."
+SYS_PROMPT_SEARCH = """You are a search query generator, create a concise Google search query, focusing only on the main topic and omitting additional redundant details, include year if necessary, 2024, Do not add any additional comments. OUTPUT ONLY THE SEARCH QUERY
+#Additional instructions:
+##Use the following search operator if necessary
+OR #to cover multiple topics"""
 
-import tiktoken # Used to limit tokens
-encoding = tiktoken.encoding_for_model("gpt-3.5-turbo") # Instead of Llama3 using available option/ replace if found anything better
-
-def limit_tokens(input_string, token_limit=7500):
-    """
-    Limit tokens sent to the model
-    """
-    return encoding.decode(encoding.encode(input_string)[:token_limit])
+# Initialize API clients
+encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
 
 together_client = OpenAI(
-        api_key=TOGETHER_API_KEY, 
-        base_url="https://together.hconeai.com/v1", 
-        default_headers={ "Helicone-Auth": f"Bearer {HELICON_API_KEY}"})
+    api_key=TOGETHER_API_KEY, 
+    base_url="https://together.hconeai.com/v1", 
+    default_headers={"Helicone-Auth": f"Bearer {HELICON_API_KEY}"})
 
 groq_client = OpenAI(
-        api_key=GROQ_API_KEY, 
-        base_url="https://groq.hconeai.com/openai/v1", 
-        default_headers={ "Helicone-Auth": f"Bearer {HELICON_API_KEY}"})
+    api_key=GROQ_API_KEY, 
+    base_url="https://groq.hconeai.com/openai/v1", 
+    default_headers={"Helicone-Auth": f"Bearer {HELICON_API_KEY}"})
 
 or_client = OpenAI(
     base_url="https://openrouter.ai/api/v1",
     api_key=OPENROUTER_API_KEY)
 
-# Groq model names
-llm_default_small = "llama3-8b-8192"
-llm_default_medium = "llama3-70b-8192"
+def md_to_html(md_text):
+    try:
+        html_content = markdown.markdown(md_text, extensions=["extra"])
+        return html_content.replace('\n', '')
+    except Exception as e:
+        logging.error(f"Error converting markdown to HTML: {e}")
+        return md_text
+
+def has_tables(html_string):
+    try:
+        soup = BeautifulSoup(html_string, 'lxml')
+        if soup.find_all('table'):
+            return True
+        tree = etree.HTML(str(soup))
+        return len(tree.xpath('//table')) > 0
+    except Exception as e:
+        logging.error(f"Error checking for tables: {e}")
+        return False
+
+def extract_data_from_tag(input_string, tag):
+    try:
+        pattern = f'<{tag}.*?>(.*?)</{tag}>'
+        matches = re.findall(pattern, input_string, re.DOTALL)
+        if matches:
+            out = '\n'.join(match.strip() for match in matches)
+            return out if len(out) <= 0.8 * len(input_string) else input_string
+        return input_string
+    except Exception as e:
+        logging.error(f"Error extracting data from tag: {e}")
+        return input_string
 
-# Together Model names (fallback)
-llm_fallback_small = "meta-llama/Llama-3-8b-chat-hf"
-llm_fallback_medium = "meta-llama/Llama-3-70b-chat-hf"
+def insert_data(user_id, user_query, subtopic_query, response, html_report):
+    try:
+        with psycopg2.connect(
+            dbname="postgres",
+            user=SUPABASE_USER,
+            password=SUPABASE_PASSWORD,
+            host="aws-0-us-west-1.pooler.supabase.com",
+            port="5432"
+        ) as conn:
+            with conn.cursor() as cur:
+                insert_query = """
+                INSERT INTO research_pro_chat_v2 (user_id, user_query, subtopic_query, response, html_report, created_at)
+                VALUES (%s, %s, %s, %s, %s, %s);
+                """
+                cur.execute(insert_query, (user_id, user_query, subtopic_query, response, html_report, datetime.now()))
+    except Exception as e:
+        logging.error(f"Error inserting data into database: {e}")
 
-### ------END OF LLM CONFIG-------- ###
+def limit_tokens(input_string, token_limit=7500):
+    try:
+        return encoding.decode(encoding.encode(input_string)[:token_limit])
+    except Exception as e:
+        logging.error(f"Error limiting tokens: {e}")
+        return input_string[:token_limit]  # Fallback to simple string slicing
 
-def together_response(message, model = llm_default_small, SysPrompt = SysPromptDefault, temperature=0.2, frequency_penalty =0.1, max_tokens= 2000):
-    
-    messages=[{"role": "system", "content": SysPrompt},{"role": "user", "content": message}]
+def together_response(message, model=LLM_DEFAULT_SMALL, SysPrompt=SYS_PROMPT_DEFAULT, temperature=0.2, frequency_penalty=0.1, max_tokens=2000):
+    messages = [{"role": "system", "content": SysPrompt}, {"role": "user", "content": message}]
     params = {
-      "model": model,
-      "messages": messages,
-      "temperature": temperature,
-      "frequency_penalty": frequency_penalty,
-      "max_tokens": max_tokens
+        "model": model,
+        "messages": messages,
+        "temperature": temperature,
+        "frequency_penalty": frequency_penalty,
+        "max_tokens": max_tokens
     }
     try:
-      response = groq_client.chat.completions.create(**params)
-      return response.choices[0].message.content
-    
+        response = groq_client.chat.completions.create(**params)
+        return response.choices[0].message.content
     except Exception as e:
-      print(f"Error calling GROQ API: {e}")
-      params["model"] = llm_fallback_small if model == llm_default_small else llm_fallback_medium 
-      response = together_client.chat.completions.create(**params)
-      return response.choices[0].message.content
+        logging.error(f"Error calling GROQ API: {e}")
+        try:
+            params["model"] = LLM_FALLBACK_SMALL if model == LLM_DEFAULT_SMALL else LLM_FALLBACK_MEDIUM 
+            response = together_client.chat.completions.create(**params)
+            return response.choices[0].message.content
+        except Exception as e:
+            logging.error(f"Error calling Together API: {e}")
+            return "An error occurred while processing your request."
 
 def openrouter_response(messages, model="meta-llama/llama-3-70b-instruct:nitro"):
     try:
@@ -199,45 +168,47 @@ def openrouter_response(messages, model="meta-llama/llama-3-70b-instruct:nitro")
             messages=messages,
             max_tokens=4096,
         )
-        
-        response_message = response.choices[0].message.content
-        return response_message
+        return response.choices[0].message.content
     except Exception as e:
-        print(f"An error occurred: {str(e)}")
+        logging.error(f"Error calling OpenRouter API: {e}")
         return None
 
 def openrouter_response_stream(messages, model="meta-llama/llama-3-70b-instruct:nitro"):
-    response = or_client.chat.completions.create(
-        model=model,
-        messages=messages,
-        max_tokens=4096,
-        stream=True
-    )
-    
-    for chunk in response:
-        if chunk.choices[0].delta.content is not None:
-            yield chunk.choices[0].delta.content
+    try:
+        response = or_client.chat.completions.create(
+            model=model,
+            messages=messages,
+            max_tokens=4096,
+            stream=True
+        )
+        for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                yield chunk.choices[0].delta.content
+    except Exception as e:
+        logging.error(f"Error streaming response from OpenRouter API: {e}")
+        yield "An error occurred while streaming the response."
 
 def json_from_text(text):
-    """
-    Extracts JSON from text using regex and fuzzy JSON loading.
-    """
     try:
-      return json.loads(text)
-    except:
-      match = re.search(r'\{[\s\S]*\}', text)
-      if match:
-        json_out = match.group(0)
-      else:
-        json_out = text
-      # Use Fuzzy JSON loading
-      return loads(json_out)
+        return json.loads(text)
+    except json.JSONDecodeError:
+        try:
+            match = re.search(r'\{[\s\S]*\}', text)
+            json_out = match.group(0) if match else text
+            return loads(json_out)
+        except Exception as e:
+            logging.error(f"Error parsing JSON from text: {e}")
+            return {}
 
 def remove_stopwords(text):
-    stop_words = set(stopwords.words('english'))
-    words = word_tokenize(text)
-    filtered_text = [word for word in words if word.lower() not in stop_words]
-    return ' '.join(filtered_text)
+    try:
+        stop_words = set(stopwords.words('english'))
+        words = word_tokenize(text)
+        filtered_text = [word for word in words if word.lower() not in stop_words]
+        return ' '.join(filtered_text)
+    except Exception as e:
+        logging.error(f"Error removing stopwords: {e}")
+        return text
 
 def rephrase_content(data_format, content, query):
     try:
@@ -245,101 +216,109 @@ def rephrase_content(data_format, content, query):
             return together_response(
                 f"""return only the relevant information regarding the query: {{{query}}}. Output should be concise chunks of \
                 paragraphs or tables or both, extracted from the following scraped context {{{limit_tokens(content,token_limit=2000)}}}""",
-                SysPrompt=SysPromptData,
+                SysPrompt=SYS_PROMPT_DATA,
                 max_tokens=900,
             )
         elif data_format == "Quantitative data":
             return together_response(
                 f"return only the numerical or quantitative data regarding the query: {{{query}}} structured into .md tables, using the scraped context:{{{limit_tokens(content,token_limit=2000)}}}",
-                SysPrompt=SysPromptData,
+                SysPrompt=SYS_PROMPT_DATA,
                 max_tokens=500,
             )
         else:
             return together_response(
                 f"return only the relevant information regarding the query: {{{query}}} using the scraped context:{{{limit_tokens(content,token_limit=2000)}}}",
-                SysPrompt=SysPromptData,
+                SysPrompt=SYS_PROMPT_DATA,
                 max_tokens=500,
             )
     except Exception as e:
-        print(f"An error occurred: {str(e)}")
-        return limit_tokens(content,token_limit=500)
+        logging.error(f"Error rephrasing content: {e}")
+        return limit_tokens(content, token_limit=500)
 
 def fetch_content(url):
     try:
         response = hrequests.get(url, timeout=5)
         if response.status_code == 200:
             return response.text
+        else:
+            logging.warning(f"Failed to fetch content from {url}. Status code: {response.status_code}")
     except Exception as e:
-        print(f"Error fetching page content for {url}: {e}")
+        logging.error(f"Error fetching page content for {url}: {e}")
     return None
 
 def extract_main_content(html):
-    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:
+    try:
+        extracted = trafilatura.extract(
+            html,
+            output_format="markdown",
+            target_language="en",
+            include_tables=True,
+            include_images=False,
+            include_links=False,
+            deduplicate=True,
+        )
+        return trafilatura.utils.sanitize(extracted) if extracted else ""
+    except Exception as e:
+        logging.error(f"Error extracting main content: {e}")
         return ""
 
 def process_content(data_format, url, query):
-    html_content = 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=4000),
-                query=query,
-            )
-            return rephrased_content, url
+    try:
+        html_content = 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=4000),
+                    query=query,
+                )
+                return rephrased_content, url
+    except Exception as e:
+        logging.error(f"Error processing content for {url}: {e}")
     return "", 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
+    try:
+        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
+    except Exception as e:
+        logging.error(f"Error fetching and extracting content: {e}")
+        return []
 
 def search_brave(query, num_results=5):
-    """Fetch search results from Brave's API."""
-    
-    cleaned_query = 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 = re.sub(r'[^\w\s]', '', search_query).strip() #re.sub(r'[^a-zA-Z0-9*]+', '', search_query)
-    
-    url = "https://api.search.brave.com/res/v1/web/search"
-    headers = {
-        "Accept": "application/json",
-        "Accept-Encoding": "gzip",
-        "X-Subscription-Token": BRAVE_API_KEY 
-    }
-    params = {"q": cleaned_search_query}
-    
-    response = requests.get(url, headers=headers, params=params)
-    
-    if response.status_code == 200:
-        result = response.json()  # Return the JSON response if successful
-        return [item["url"] for item in result["web"]["results"]][:num_results],cleaned_search_query, result
-    else:
-        return [],cleaned_search_query  # Return error code if not successful
-
-# #@retry(tries=3, delay=0.25)
-# def search_brave(query, num_results=5):
-#     cleaned_query = 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 = re.sub(r'[^\w\s]', '', search_query).strip() #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
+    try:
+        cleaned_query = query
+        search_query = together_response(cleaned_query, model=LLM_DEFAULT_SMALL, SysPrompt=SYS_PROMPT_SEARCH, max_tokens=25).strip()
+        cleaned_search_query = re.sub(r'[^\w\s]', '', search_query).strip()
+        
+        url = "https://api.search.brave.com/res/v1/web/search"
+        headers = {
+            "Accept": "application/json",
+            "Accept-Encoding": "gzip",
+            "X-Subscription-Token": BRAVE_API_KEY 
+        }
+        params = {"q": cleaned_search_query}
+        
+        response = requests.get(url, headers=headers, params=params)
+        
+        if response.status_code == 200:
+            result = response.json()
+            return [item["url"] for item in result["web"]["results"]][:num_results], cleaned_search_query, result
+        else:
+            logging.warning(f"Brave search API returned status code {response.status_code}")
+            return [], cleaned_search_query, None
+    except Exception as e:
+        logging.error(f"Error in Brave search: {e}")
+        return [], query, None
+
+# Main execution
+if __name__ == "__main__":
+    logging.info("Script started")
+    # Add your main execution logic here
+    logging.info("Script completed")