Update main.py

main.py

@@ -2,7 +2,7 @@ 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 md_to_html, search_brave, fetch_and_extract_content, limit_tokens, together_response, insert_data
+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 datetime import datetime, timedelta
@@ -119,7 +119,25 @@ async def generate_report(query: QueryModel):
         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 has_tables(html_report):
+        print("tables found, creating charts")
+        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_chart = extract_data_from_tag(openrouter_response(messages, model="anthropic/claude-3.5-sonnet"),"report")
+            html_report = md_to_html(md_report_chart)
+        
+        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 = {}