Update main.py

main.py

@@ -54,45 +54,47 @@ sys_prompts = {
     },
 }
 
-class ReportParams(BaseModel):
-    topic: str = Query(default="market research", description="input query to generate report")
+class QueryModel(BaseModel):
+    topic: 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"])
+    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")
-@cache(expire=604800)  # Cache expiration set to 7 days
-async def generate_report(params: ReportParams):
-    query_str = params.topic
-    internet_status = "online" if params.internet else "offline"
-    sys_prompt_output_format = sys_prompts[internet_status][params.output_format]
+@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"
+    sys_prompt_output_format = sys_prompts[internet][query.output_format]
+    data_format = query.data_format
     optimized_search_query = ""
-    all_text_with_urls = [("","")]
+    all_text_with_urls = [("", "")]
 
-    # Combine query with user keywords
-    if params.internet:
-        search_query = params.description
+    if query.internet:
+        search_query = description
         try:
             urls, optimized_search_query = search_brave(search_query, num_results=4)
-            all_text_with_urls = fetch_and_extract_content(params.data_format, urls, query_str)
+            all_text_with_urls = fetch_and_extract_content(data_format, urls, query_str)
             additional_context = limit_tokens(str(all_text_with_urls))
-            prompt = f"#### COMPLETE THE TASK: {params.description} #### IN THE CONTEXT OF ### CONTEXT: {query_str} USING THE #### SCRAPED DATA:{additional_context}"
+            prompt = f"#### COMPLETE THE TASK: {description} #### IN THE CONTEXT OF ### CONTEXT: {query_str} USING THE #### SCRAPED DATA:{additional_context}"
         except Exception as e:
-            params.internet = False
+            query.internet = False
             print("failed to search/scrape results, falling back to LLM response")
 
-    if not params.internet:
-        prompt = f"#### COMPLETE THE TASK: {params.description} #### IN THE CONTEXT OF ### CONTEXT: {query_str}"
+    if not query.internet:
+        prompt = f"#### COMPLETE THE TASK: {description} #### IN THE CONTEXT OF ### CONTEXT: {query_str}"
 
     md_report = together_response(prompt, model=llm_default_medium, SysPrompt=sys_prompt_output_format)
 
-    if params.user_id != "test":
-        insert_data(params.user_id, query_str, params.description, str(all_text_with_urls), md_report)
-
-    references_html = dict()
+    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))
 
@@ -101,6 +103,11 @@ async def generate_report(params: ReportParams):
         "references": references_html,
         "search_query": optimized_search_query
     }
+
+@app.post("/generate_report")
+async def api_generate_report(request: Request, query: QueryModel):
+    return await generate_report(query)
+    
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],