Spaces:
Running
Running
# !pip install mistune | |
import mistune | |
from mistune.plugins.table import table | |
from jinja2 import Template | |
import re | |
import os | |
import hrequests | |
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-------------------------------------#### | |
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 = 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----------------------------------------------------------#### | |
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 | |
from fuzzy_json import loads | |
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" | |
# 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 = """ | |
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. | |
2. Focus on extracting facts, tables, and direct quotes that are relevant to the query. | |
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""" | |
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]) | |
together_client = OpenAI( | |
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}"}) | |
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" | |
# Together Model names (fallback) | |
llm_fallback_small = "meta-llama/Llama-3-8b-chat-hf" | |
llm_fallback_medium = "meta-llama/Llama-3-70b-chat-hf" | |
### ------END OF LLM CONFIG-------- ### | |
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}] | |
params = { | |
"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 | |
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 | |
def openrouter_response(messages, model="meta-llama/llama-3-70b-instruct:nitro"): | |
try: | |
response = or_client.chat.completions.create( | |
model=model, | |
messages=messages, | |
max_tokens=4096, | |
) | |
response_message = response.choices[0].message.content | |
return response_message | |
except Exception as e: | |
print(f"An error occurred: {str(e)}") | |
return None | |
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) | |
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) | |
def rephrase_content(data_format, content, query): | |
try: | |
if data_format == "Structured data": | |
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, | |
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, | |
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, | |
max_tokens=500, | |
) | |
except Exception as e: | |
print(f"An error occurred: {str(e)}") | |
return limit_tokens(content,token_limit=500) | |
def fetch_content(url): | |
try: | |
response = hrequests.get(url) | |
if response.status_code == 200: | |
return response.text | |
except Exception as e: | |
print(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: | |
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 | |
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 | |
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 | |
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 | |