Update main.py

main.py

@@ -40,6 +40,24 @@ 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."
 
+prompt_user = {}
+prompt_user["online"] = {}
+prompt_user["offline"] = {}
+prompt_user["online"]["chat"] = "Write a well thought out answer to the query:: {description} #### , refer the provided internet search results reference:{reference}
+prompt_user["online"]["report"] = "Write a well thought out Report to the query:: {description} #### , refer the provided internet search results reference:{reference}
+prompt_user["online"]["report_table"] = "Write a well thought out Report to the query:: {description}, include any relevant tables if present in the source #### , refer the provided internet search results reference:{reference}
+
+prompt_user["offline"]["chat"] = "Write a well thought out answer to the query:: {description}"
+prompt_user["offline"]["report"] = "Write a well thought out Report to the query:: {description}"
+prompt_user["offline"]["report_table"] = "Write a well thought out Report to the query:: {description}, include any relevant tables as markdown"
+
+prompt_system["online"] = """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. """
+
+prompt_system["offline"] = """Your response should be properly formatted and well readable using markdown formatting."""
+
 sys_prompts = {
     "offline": {
         "Chat": "You are an expert AI, complete the given task. Do not add any additional comments.",
@@ -47,7 +65,8 @@ sys_prompts = {
         "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": {
+    "online": 
+    "user":{
         "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 create a structured 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 structured into subtopics/tables/lists. Do not add any additional comments.",
@@ -61,8 +80,8 @@ class QueryModel(BaseModel):
     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"])
+    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"])
 
@@ -72,7 +91,8 @@ async def generate_report(query: QueryModel):
     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]
+    user_prompt_final = prompt_user[internet][query.output_format]
+    system_prompt_final = prompt_system[internet]
     data_format = query.data_format
     optimized_search_query = ""
     all_text_with_urls = [("", "")]
@@ -82,8 +102,8 @@ async def generate_report(query: QueryModel):
         try:
             urls, optimized_search_query = search_brave(search_query, num_results=8)
             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: {description} #### IN THE CONTEXT OF ### CONTEXT: {query_str} USING THE #### SCRAPED DATA:{additional_context}"
+            reference = limit_tokens(str(all_text_with_urls))
+            prompt = user_prompt_final.format(description,reference)
         except Exception as e:
             print(e)
             query.internet = False
@@ -91,8 +111,9 @@ async def generate_report(query: QueryModel):
 
     if not query.internet:
         prompt = f"#### COMPLETE THE TASK: {description} #### IN THE CONTEXT OF ### CONTEXT: {query_str}"
+        system_prompt_final = prompt_system["offline"]
 
-    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=sys_prompt_output_format)
+    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=system_prompt_final)
 
     if user_id != "test":
         insert_data(user_id, query_str, description, str(all_text_with_urls), md_report)