Update main.py

main.py

@@ -7,7 +7,7 @@ import os
 from dotenv import load_dotenv, find_dotenv
 
 # Load environment variables from .env file
-# load_dotenv("keys.env")
+#load_dotenv("keys.env")
 
 app = FastAPI()
 TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
@@ -25,18 +25,39 @@ SysPromptList = "You are now in the role of an expert AI who can extract structu
 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."
 
-class Query(BaseModel):
+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 provided in the context 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 using information provided in the context 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 using information provided in the context from user request. The report should be in markdown format consists of only markdown tables. Do not add any additional comments.",
+    },
+}
+
+class QueryModel(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")
+    internet: bool = Query(default=True, description="Enable Internet search")
+    output_format: str = Query(default="Tabular Report", description="Output format for the report", enum=["Chat", "Full Text Report", "Tabular Report", "Tables only"])
+    data_format: str = Query(default="Structured data", description="Type of data to extract from the internet", enum=["No presets", "Structured data", "Quantitative data"])
 
 @app.post("/generate_report")
-async def generate_report(request: Request, query: Query):
+async def generate_report(request: Request, query: QueryModel):
     query_str = query.query
     description = query.description
     user_id = query.user_id
-
+    internet = "online" if query.internet else "offline"
+    output_format = sys_prompts[internet][query.output_format]
+    data_format = query.data_format
+    #return {output_format,data_format}
     # Combine query with user keywords
     search_query = query_str 
     
@@ -44,7 +65,7 @@ async def generate_report(request: Request, query: Query):
     urls = search_brave(search_query, num_results=4)
 
     # Fetch and extract content from the URLs
-    all_text_with_urls = fetch_and_extract_content(urls, query_str)
+    all_text_with_urls = fetch_and_extract_content(data_format, urls, query_str)
 
     # Prepare the prompt for generating the report
     additional_context = limit_tokens(str(all_text_with_urls))
@@ -52,7 +73,7 @@ async def generate_report(request: Request, query: Query):
     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)
+    #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))
@@ -69,5 +90,4 @@ app.add_middleware(
     allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
-    allow_headers=["*"],
-)
+    allow_headers=["*"],)