Update main.py

main.py

@@ -81,8 +81,9 @@ OPENROUTER_API_KEY = "sk-or-v1-"+os.environ['OPENROUTER_API_KEY']
 
 
 class QueryModel(BaseModel):
-    topic: str = Query(default="", description="input query to generate Report")
-    description: str = Query(default="", description="additional context for report")
+    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")
@@ -91,7 +92,7 @@ class QueryModel(BaseModel):
     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)")
 
 @cache(expire=604800)
 async def generate_report(query: QueryModel):
@@ -105,6 +106,7 @@ async def generate_report(query: QueryModel):
     optimized_search_query = ""
     all_text_with_urls = [("", "")]
     generate_charts = query.generate_charts
+    output_as_md = query.output_as_md
     
     if query.internet:
         search_query = re.sub(r'[^\w\s]', '', description).strip()
@@ -135,9 +137,8 @@ async def generate_report(query: QueryModel):
             <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_chart = extract_data_from_tag(openrouter_response(messages, model="anthropic/claude-3.5-sonnet"),"report")
-            html_report = md_to_html(md_report_chart)
-            print(html_report)
+            md_report = extract_data_from_tag(openrouter_response(messages, model="anthropic/claude-3.5-sonnet"),"report")
+            print(md_report)
         
         except Exception as e:
             print(e) 
@@ -149,8 +150,10 @@ async def generate_report(query: QueryModel):
     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)
+    
     return {
-        "report": html_report,
+        "report": final_report,
         "references": references_html,
         "search_query": optimized_search_query
     }