Update main.py

main.py

@@ -1,172 +1,182 @@
-from fastapi import FastAPI, HTTPException, Request, Query
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from typing import List, Dict, Any
-from helper_functions_api import has_tables, extract_data_from_tag, openrouter_response,md_to_html, search_brave, fetch_and_extract_content, limit_tokens, together_response, insert_data
 import os
-from dotenv import load_dotenv, find_dotenv
+from typing import List, Dict, Any
 from datetime import datetime, timedelta
+import re
+from functools import lru_cache
+
+from fastapi import FastAPI, HTTPException, Request, Query, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
 from fastapi_cache import FastAPICache
 from fastapi_cache.backends.inmemory import InMemoryBackend
 from fastapi_cache.decorator import cache
-import asyncio
-import re
-# Load environment variables from .env file
-#load_dotenv("keys.env")
+from dotenv import load_dotenv
+
+from helper_functions_api import (
+    has_tables, extract_data_from_tag, openrouter_response, md_to_html,
+    search_brave, fetch_and_extract_content, limit_tokens, together_response, insert_data
+)
+
+# Load environment variables
+load_dotenv()
+
+# Constants
+LLM_MODELS = {
+    "default": {
+        "small": "llama3-8b-8192",
+        "medium": "llama3-70b-8192"
+    },
+    "fallback": {
+        "small": "meta-llama/Llama-3-8b-chat-hf",
+        "medium": "meta-llama/Llama-3-70b-chat-hf"
+    }
+}
+
+SYSTEM_PROMPTS = {
+    "json": "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.",
+    "list": "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.",
+    "default": "You are an expert AI, complete the given task. Do not add any additional comments.",
+    "md": "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.",
+    "online": """You are an expert AI who can create a detailed structured report using internet search results.
+                1. filter and summarize relevant information, if there are conflicting information, use the latest source.
+                2. use it to construct a clear and factual answer.
+                Your response should be structured and properly formatted using markdown headings, subheadings, tables, use as necessary. Ignore Links and references""",
+    "offline": "You are an expert AI who can create detailed answers. Your response should be properly formatted and well readable using markdown formatting."
+}
+
+# Prompt templates
+PROMPT_TEMPLATES = {
+    "online": {
+        "chat": "Write a well thought out, detailed and structured answer to the query:: {description} #### , refer the provided internet search results reference:{reference}",
+        "report": "Write a well thought out, detailed and structured Report to the query:: {description} #### , refer the provided internet search results reference:{reference}, The report should be well formatted using markdown format structured into subtopics as necessary",
+        "report_table": "Write a well thought out Report to the query:: {description},#### , refer the provided internet search results reference:{reference}. The report should be well formatted using markdown format, structured into subtopics, include tables or lists as needed to make it well readable"
+    },
+    "offline": {
+        "chat": "Write a well thought out, detailed and structured answer to the query:: {description}",
+        "report": "Write a well thought out, detailed and structured Report to the query:: {description}. The report should be well formatted using markdown format, structured into subtopics",
+        "report_table": "Write a detailed and structured Report to the query:: {description}, The report should be well formatted using markdown format, structured into subtopics, include tables or lists as needed to make it well readable"
+    }
+}
 
+# FastAPI app setup
 app = FastAPI()
 
 @app.on_event("startup")
 async def startup():
     FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")
 
-# 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"
-
-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."
-
-prompt_user = {}
-prompt_system = {}
-prompt_user["online"] = {}
-prompt_user["offline"] = {}
-prompt_user["online"]["chat"] = "Write a well thought out, detailed and structured answer to the query:: {description} #### , refer the provided internet search results reference:{reference}"
-prompt_user["online"]["report"] = "Write a well thought out, detailed and structured Report to the query:: {description} #### , refer the provided internet search results reference:{reference}, The report should be well formatted using markdown format structured into subtopics as necessory"
-prompt_user["online"]["report_table"] = "Write a well thought out Report to the query:: {description},#### , refer the provided internet search results reference:{reference}. The report should be well formatted using markdown format,  structured into subtopics, include tables or lists as needed to make it well readable"
-
-prompt_user["offline"]["chat"] = "Write a well thought out, detailed and structured answer to the query:: {description}"
-prompt_user["offline"]["report"] = "Write a well thought out, detailed and structured Report to the query:: {description}. The report should be well formatted using markdown format,  structured into subtopics"
-prompt_user["offline"]["report_table"] = "Write a detailed and structured Report to the query:: {description}, The report should be well formatted using markdown format,  structured into subtopics, include tables or lists as needed to make it well readable"
-
-prompt_system["online"] = """You are an expert AI who can create a detailed structured report using internet search results.
-
-                                1 filter and summarize relevant information, if there are conflicting information, use the latest source.
-                                2. use it to construct a clear and factual answer.
-                                Your response should be structured and properly formatted using markdown headings, subheadings, tables, use as necessory. Ignore Links and references"""
-
-prompt_system["offline"] = """You are an expert AI who can create detailed answers. Your response should be properly formatted and well readable using markdown formatting."""
-
-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']
-
-# sys_prompts = {
-#     "offline": {
-#         "Chat": "You are an expert AI, complete the given task. Do not add any additional comments.",
-#         "Full Text Report": "You are an expert AI who can create a detailed report from user request. The report should be in markdown format. Do not add any additional comments.",
-#         "Tabular Report": "You are an expert AI who can create a structured report from user request.The report should be in markdown format structured into subtopics/tables/lists. Do not add any additional comments.",
-#         "Tables only": "You are an expert AI who can create a structured tabular report from user request.The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
-#     },
-#     "online": {
-#         "Chat": "You are an expert AI, complete the given task using the provided context. Do not add any additional comments.",
-#         "Full Text Report": "You are an expert AI who can create a detailed report using information scraped from the internet. You should decide which information is relevant to the given task and use it to create a report. The report should be in markdown format. Do not add any additional comments.",
-#         "Tabular Report": """You are an expert AI who can provide answers using internet search results.
-#                                 1 filter and summarize relevant information, if there are conflicting information, use the latest source.
-#                                 2. use it to construct a clear and factual answer.
-#                                 Your response should be properly formatted and well readable using markdown formatting. """,
-#         "Tables only": "You are an expert AI who can create a structured tabular report using information scraped from the internet. You should decide which information is relevant to the given task. The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
-#     },
-# }
-
-
+# Pydantic model for query parameters
 class QueryModel(BaseModel):
-    user_query: str = Query(default="", description="Initial user query")
-    topic: str = Query(default="", description="Topic name to generate Report")
-    description: str = Query(description="Description/prompt for report (REQUIRED)")
-    user_id: str = Query(default="", description="unique user id")
-    user_name: str = Query(default="", description="user name")
-    internet: bool = Query(default=True, description="Enable Internet search")
-    output_format: str = Query(default="report_table", description="Output format for the report",
-                               enum=["chat", "report", "report_table"])
-    data_format: str = Query(default="Structured data", description="Type of data to extract from the internet",
-                             enum=["No presets", "Structured data", "Quantitative data"])
-    generate_charts: bool = Query(default=False, description="Include generated charts")
-    output_as_md: bool = Query(default=False, description="Output report in markdown (default output in HTML)")
+    user_query: str = Field(default="", description="Initial user query")
+    topic: str = Field(default="", description="Topic name to generate Report")
+    description: str = Field(..., description="Description/prompt for report (REQUIRED)")
+    user_id: str = Field(default="", description="unique user id")
+    user_name: str = Field(default="", description="user name")
+    internet: bool = Field(default=True, description="Enable Internet search")
+    output_format: str = Field(default="report_table", description="Output format for the report")
+    data_format: str = Field(default="Structured data", description="Type of data to extract from the internet")
+    generate_charts: bool = Field(default=False, description="Include generated charts")
+    output_as_md: bool = Field(default=False, description="Output report in markdown (default output in HTML)")
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "user_query": "How does climate change affect biodiversity?",
+                "topic": "Climate Change and Biodiversity",
+                "description": "Provide a detailed report on the impacts of climate change on global biodiversity",
+                "user_id": "user123",
+                "user_name": "John Doe",
+                "internet": True,
+                "output_format": "report_table",
+                "data_format": "Structured data",
+                "generate_charts": True,
+                "output_as_md": False
+            }
+        }
+
+@lru_cache()
+def get_api_keys():
+    return {
+        "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": f"sk-or-v1-{os.environ['OPENROUTER_API_KEY']}"
+    }
+
+def get_internet_data(description: str, data_format: str):
+    search_query = re.sub(r'[^\w\s]', '', description).strip()
+    urls, optimized_search_query, full_search_object = search_brave(search_query, num_results=8)
+    all_text_with_urls = fetch_and_extract_content(data_format, urls, optimized_search_query)
+    reference = limit_tokens(str(all_text_with_urls), token_limit=5000)
+    return all_text_with_urls, optimized_search_query, full_search_object, reference
+
+def generate_charts(md_report: str):
+    chart_prompt = (
+        "Convert the numerical data tables in the given content to embedded html plotly.js charts if appropriate, "
+        "use appropriate colors. Output format: <report>output the full content without any other changes in md "
+        f"format enclosed in tags like this</report> using the following: {md_report}"
+    )
+    messages = [{"role": 'user', "content": chart_prompt}]
+    return extract_data_from_tag(openrouter_response(messages, model="anthropic/claude-3.5-sonnet"), "report")
 
 @cache(expire=604800)
-async def generate_report(query: QueryModel):
-    query_str = query.topic
-    description = query.description
-    user_id = query.user_id
-    internet = "online" if query.internet else "offline"
-    user_prompt_final = prompt_user[internet][query.output_format]
-    system_prompt_final = prompt_system[internet]
-    data_format = query.data_format
+async def generate_report(query: QueryModel, api_keys: Dict[str, str] = Depends(get_api_keys)):
+    internet_mode = "online" if query.internet else "offline"
+    user_prompt = PROMPT_TEMPLATES[internet_mode][query.output_format]
+    system_prompt = SYSTEM_PROMPTS[internet_mode]
+    
+    all_text_with_urls = []
     optimized_search_query = ""
-    all_text_with_urls = [("", "")]
     full_search_object = {}
-    generate_charts = query.generate_charts
-    output_as_md = query.output_as_md
-    
+
     if query.internet:
-        search_query = re.sub(r'[^\w\s]', '', description).strip()
         try:
-            urls, optimized_search_query, full_search_object = search_brave(search_query, num_results=8)
-            all_text_with_urls = fetch_and_extract_content(data_format, urls, optimized_search_query)
-            reference = limit_tokens(str(all_text_with_urls),token_limit=5000)
-            user_prompt_final = user_prompt_final.format(description=description, reference=reference)
+            all_text_with_urls, optimized_search_query, full_search_object, reference = get_internet_data(query.description, query.data_format)
+            user_prompt = user_prompt.format(description=query.description, reference=reference)
         except Exception as e:
-            print(e)
-            query.internet = False
-            print("failed to search/scrape results, falling back to LLM response")
-
-    if not query.internet:
-        user_prompt_final = prompt_user["offline"][query.output_format].format(description=description)
-        system_prompt_final = prompt_system["offline"]
-
-    md_report = together_response(user_prompt_final, model=llm_default_medium, SysPrompt=system_prompt_final)
-    html_report = md_to_html(md_report)
-
-    # Render Charts
-    if generate_charts and has_tables(html_report):
-        print("tables found, creating charts")
+            print(f"Failed to search/scrape results: {e}")
+            internet_mode = "offline"
+            user_prompt = PROMPT_TEMPLATES[internet_mode][query.output_format].format(description=query.description)
+            system_prompt = SYSTEM_PROMPTS[internet_mode]
+    else:
+        user_prompt = user_prompt.format(description=query.description)
+
+    md_report = together_response(user_prompt, model=LLM_MODELS["default"]["medium"], SysPrompt=system_prompt)
+    
+    if query.generate_charts and has_tables(md_to_html(md_report)):
         try:
-            
-            prompt = "convert the numerical data tables in the given content to embedded html plotly.js charts if appropriate, use appropriate colors, \
-            output format:\
-            <report>output the full content without any other changes in md format enclosed in tags like this</report> using the following:" + str(md_report)
-                        
-            messages = [{"role": 'user', "content": prompt}]
-            md_report = extract_data_from_tag(openrouter_response(messages, model="anthropic/claude-3.5-sonnet"),"report")
-            print(md_report)
-        
+            md_report = generate_charts(md_report)
         except Exception as e:
-            print(e) 
-            print("failed to generate charts, falling back to original report")
-        
-    if user_id != "test":
-        insert_data(user_id, query_str, description, str(all_text_with_urls), md_report)
-    references_html = {}
-    for text, url in all_text_with_urls:
-        references_html[url] = str(md_to_html(text))
-
-    final_report = md_report if output_as_md else md_to_html(md_report)
+            print(f"Failed to generate charts: {e}")
+
+    if query.user_id != "test":
+        insert_data(query.user_id, query.topic, query.description, str(all_text_with_urls), md_report)
+
+    references_html = {url: str(md_to_html(text)) for text, url in all_text_with_urls}
+    final_report = md_report if query.output_as_md else md_to_html(md_report)
     
     return {
         "report": final_report,
         "references": references_html,
         "search_query": optimized_search_query,
-        "search_data_full":full_search_object
+        "search_data_full": full_search_object
     }
 
-@app.post("/generate_report")
-async def api_generate_report(request: Request, query: QueryModel):
-    return await generate_report(query)
-    
+@app.post("/generate_report", response_model=Dict[str, Any])
+async def api_generate_report(query: QueryModel, api_keys: Dict[str, str] = Depends(get_api_keys)):
+    try:
+        return await generate_report(query, api_keys)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# CORS middleware setup
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
-    allow_headers=["*"],)
+    allow_headers=["*"],
+)