Update main.py

main.py

@@ -1,81 +1,67 @@
-from fuzzy_json import loads
-from half_json.core import JSONFixer
-from openai import OpenAI
-from retry import retry
-import re
-from dotenv import load_dotenv
-import os
-from fastapi import FastAPI
-from fastapi import Query
+from fastapi import FastAPI, HTTPException, Request, Query
 from pydantic import BaseModel
-from fastapi.middleware.cors import CORSMiddleware
-from helper_functions_api import md_to_html
-from duckduckgo_search import DDGS
-import time
+from typing import List, Dict, Any
+from helper_functions_api import md_to_html, search_brave, fetch_and_extract_content, limit_tokens, together_response, insert_data
+import os
+from dotenv import load_dotenv, find_dotenv
 
-# Retrieve environment variables
-TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")
-GROQ_API_KEY = "gsk_"+os.getenv("GROQ_API_KEY")
+# Load environment variables from .env file
+# load_dotenv("keys.env")
+
+app = FastAPI()
+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']
+
+llm_default_small = "llama3-8b-8192"
+llm_default_medium = "llama3-70b-8192"
 
-SysPromptDefault = "You are an expert AI, complete the given task. Do not add any additional comments."    
-SysPromptList = "You are now in the role of an expert AI who can extract structured information from user request. All elements must be in double quotes. You must respond ONLY with a valid python List. Do not add any additional comments."
 SysPromptJson = "You are now in the role of an expert AI who can extract structured information from user request. Both key and value pairs must be in double quotes. You must respond ONLY with a valid JSON file. Do not add any additional comments."
+SysPromptList = "You are now in the role of an expert AI who can extract structured information from user request. All elements must be in double quotes. You must respond ONLY with a valid python List. Do not add any additional comments."
+SysPromptDefault = "You are an expert AI, complete the given task. Do not add any additional comments."
 SysPromptMd = "You are an expert AI who can create a structured report using information provided in the context from user request.The report should be in markdown format consists of markdown tables structured into subtopics. Do not add any additional comments."
-SysPromptMdOffline = "You are an expert AI who can create a structured report using your knowledge on user request.The report should be in markdown format consists of markdown tables/lists/paragraphs as needed, structured into subtopics. Do not add any additional comments."
 
-@retry(tries=3, delay=1)
-def together_response(message, model = "meta-llama/Llama-3-8b-chat-hf", SysPrompt = SysPromptDefault,temperature=0.2):
-      base_url_groq = "https://api.groq.com/openai/v1"
-      groq_model_name="llama3-8b-8192"
-      client = OpenAI(base_url= base_url_groq, api_key= GROQ_API_KEY)
-      messages=[{"role": "system", "content": SysPrompt},{"role": "user", "content": message}]
+class Query(BaseModel):
+    query: str = Query(default="market research", description="input query to generate Report")
+    description: str = Query(default="", description="additional context for report")
+    user_id: str = Query(default="", description="unique user id")
+    user_name: str = Query(default="", description="user name")
 
-      response = client.chat.completions.create(
-          model=groq_model_name,
-          messages=messages,
-          temperature=temperature,
-      )
-      return response.choices[0].message.content
+@app.post("/generate_report")
+async def generate_report(request: Request, query: Query):
+    query_str = query.query
+    description = query.description
+    user_id = query.user_id
 
-def json_from_text(text):
-    """
-    Extracts JSON from text using regex and fuzzy JSON loading.
-    """
-    match = re.search(r'\{[\s\S]*\}', text)
-    if match:
-      json_out = match.group(0)
-    else:
-      json_out = text
-    try:
-        # Using fuzzy json loader
-        return loads(json_out)
-    except Exception:
-        # Using JSON fixer/ Fixes even half json/ Remove if you need an exception
-        fix_json = JSONFixer()
-        return loads(fix_json.fix(json_out).line)
+    # Combine query with user keywords
+    search_query = query_str 
     
-def generate_topics(user_input, num_topics, previous_queries):
-    previous_context = " -> ".join(previous_queries)
-    prompt = f"""create a list of {num_topics} subtopics along with descriptions to follow for conducting {user_input} in the context of {previous_context}, RETURN A VALID PYTHON LIST"""\
-                +""" Respond in the following format: 
-        [["Subtopic","Description"],["Subtopic","Description"]]"""
-    response_topics = together_response(prompt, model="meta-llama/Llama-3-8b-chat-hf", SysPrompt=SysPromptList)
-    subtopics = json_from_text(response_topics)
-    return subtopics
+    # Search for relevant URLs
+    urls = search_brave(search_query, num_results=4)
 
-def generate_report(topic, description):
-    prompt = f"""create a detailed report on: {topic} by following the instructions: {description}"""
-    md_report = together_response(prompt, model = "meta-llama/Llama-3-70b-chat-hf", SysPrompt = SysPromptMdOffline)
-    return md_to_html(md_report)
+    # Fetch and extract content from the URLs
+    all_text_with_urls = fetch_and_extract_content(urls, query_str)
 
-def get_images(query, num_results):
-    time.sleep(0.5)
-    ddgs = DDGS() 
-    imgs = ddgs.images(keywords=query, safesearch="on", max_results=num_results)
-    return imgs
-    
-# Define the app
-app = FastAPI()
+    # Prepare the prompt for generating the report
+    additional_context = limit_tokens(str(all_text_with_urls))
+    prompt = f"#### ADDITIONAL CONTEXT:{additional_context} #### CREATE A DETAILED REPORT FOR THE QUERY:{query_str} #### IN THE CONTEXT OF ### CONTEXT: {description}"
+    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=SysPromptMd)
+
+    # Insert data into database (or other storage)
+    insert_data(user_id, query_str, description,  str(all_text_with_urls), md_report)
+    references_html = dict()
+    for text, url in all_text_with_urls:
+        references_html[url] = str(md_to_html(text))
+
+
+    # Return the generated report
+    return {
+        "report": md_to_html(md_report),
+        "references": references_html
+    }
 
 app.add_middleware(
     CORSMiddleware,
@@ -84,59 +70,3 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
-
-# Create a Pydantic model to handle the input data
-class TopicInput(BaseModel):
-    user_input: str = Query(default="market research", description="input query to generate subtopics")
-    num_topics: int = Query(default=5, description="Number of subtopics to generate (default: 5)")
-    previous_queries: list[str] = Query(default=[], description="List of previous queries for context")
-
-class imageInput(BaseModel):
-    user_input: str = Query(default="market research", description="input query to generate subtopics")
-    num_images: int = Query(default=5, description="Number of subtopics to generate (default: 5)")
-    
-class ReportInput(BaseModel):
-    topic: str = Query(default="market research",description="The main topic for the report")
-    description: str = Query(default="",description="A brief description of the topic")
-    
-class RecommendationInput(BaseModel):
-    user_input: str = Query(default="", description="Input query to generate follow-up questions")
-    num_recommendations: int = Query(default=5, description="Number of recommendations to generate")
-
-@app.get("/", tags=["Home"])
-def api_home():
-    return {'detail': 'Welcome to FastAPI Subtopics API! Visit https://pvanand-generate-subtopics.hf.space/docs to test'}
-
-@app.post("/generate_topics")
-async def create_topics(input: TopicInput):
-    topics = generate_topics(input.user_input, input.num_topics, input.previous_queries)
-    return {"topics": topics}
-
-@app.post("/generate_report")
-async def create_report(input: ReportInput):
-    report = generate_report(input.topic, input.description) 
-    return {"report": report}
-
-@app.post("/get_images")
-async def fetch_images(input: imageInput):
-    images = get_images(input.user_input, input.num_images)  
-    return {"images": images}
-
-@app.post("/get_recommendations")
-async def generate_recommendations(input: RecommendationInput):
-    
-    if input.user_input:
-        prompt = f"""create a list of {input.num_recommendations} questions that a user might ask following the question: {input.user_input}:"""
-    else:
-        prompt = f"""create a list of mixed {input.num_recommendations} questions to create a report or plan or course on any of the topics product,market,research topic """
-
-    response_topics = json_from_text(
-            together_response(
-                prompt, model="meta-llama/Llama-3-8b-chat-hf", SysPrompt=SysPromptList,temperature=1
-            )
-        )
-    return {"recommendations": response_topics}
-
-
-
-