Update main.py

main.py

@@ -55,25 +55,31 @@ async def generate_report(request: Request, query: QueryModel):
     description = query.description
     user_id = query.user_id
     internet = "online" if query.internet else "offline"
-    output_format = sys_prompts[internet][query.output_format]
+    sys_prompt_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 
-    
-    # Search for relevant URLs
-    urls = search_brave(search_query, num_results=4)
+    if query.internet:
+        search_query = query_str 
+        
+        # Search for relevant URLs
+        urls = search_brave(search_query, num_results=4)
 
-    # Fetch and extract content from the URLs
-    all_text_with_urls = fetch_and_extract_content(data_format, urls, query_str)
+        # Fetch and extract content from the URLs
+        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))
-    prompt = f"#### ADDITIONAL CONTEXT:{additional_context} #### CREATE A DETAILED REPORT FOR THE QUERY:{query_str} #### IN THE CONTEXT OF ### CONTEXT: {description}"
-    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=SysPromptMd)
+        # Prepare the prompt for generating the report
+        additional_context = limit_tokens(str(all_text_with_urls))
+        prompt = f"#### COMPLETE THE TASK: {query_str} #### IN THE CONTEXT OF ### CONTEXT: {description} #### ADDITIONAL CONTEXT:{additional_context}"
 
+    else:
+        prompt = f"#### COMPLETE THE TASK: {query_str} #### IN THE CONTEXT OF ### CONTEXT: {description}"
+        all_text_with_urls = [("","")]
+
+    md_report = together_response(prompt, model=llm_default_medium, SysPrompt=sys_prompt_output_format)
+    
     # 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))