Feature: Add download visualization feature and chat without document upload

app.py

@@ -203,7 +203,7 @@ def process_text_query(query, history):
                 viz_type = vtype
                 break
     
-    # Check if we're in CSV context
+    # Check if we're in CSV context or have documents loaded
     if current_context["file_type"] == "csv" and current_context["table_name"]:
         try:
             # Connect to the database
@@ -318,8 +318,8 @@ def process_text_query(query, history):
             history[-1][1] = error_msg
             return error_msg, history
     
-    else:
-        # Handle non-CSV queries (document queries)
+    elif document_assistant.get_all_documents():
+        # Handle document queries
         try:
             response = document_assistant.process_query(query)
             history[-1][1] = response
@@ -328,6 +328,26 @@ def process_text_query(query, history):
             error_msg = f"Error processing query: {str(e)}"
             history[-1][1] = error_msg
             return error_msg, history
+    
+    else:
+        # Handle general queries with LLM when no documents are loaded
+        try:
+            # Create a general knowledge context prompt
+            general_prompt = ChatPromptTemplate.from_messages([
+                ("system", "You are a helpful assistant that provides clear, informative responses. Use your knowledge to answer the user's question concisely."),
+                ("human", "{question}")
+            ])
+            
+            # Get response from LLM
+            response = llm.invoke(general_prompt.format(question=query)).content
+            
+            # Add the response to history
+            history[-1][1] = response
+            return response, history
+        except Exception as e:
+            error_msg = f"Error processing query: {str(e)}"
+            history[-1][1] = error_msg
+            return error_msg, history
 
 def process_file_upload(files):
     """Process uploaded files and index them"""
@@ -866,8 +886,10 @@ with gr.Blocks(title="LLM Powered Database Chatbot") as demo:
                 with gr.Row():
                     clear_viz_btn = gr.Button("🗑️ Clear Visualization")
                     save_viz_btn = gr.Button("💾 Save Visualization")
+                    download_btn = gr.Button("📥 Download Visualization")
                 
                 save_status = gr.Textbox(label="Save Status", visible=False)
+                download_img = gr.Image(visible=False, type="pil", label="Download Image")
         
         # Add information about capabilities
         gr.Markdown("""
@@ -886,18 +908,50 @@ with gr.Blocks(title="LLM Powered Database Chatbot") as demo:
                 return "No visualization to save", gr.update(visible=True)
             
             try:
-                # Create a unique filename
-                timestamp = time.strftime("%Y%m%d_%H%M%S")
-                filename = f"visualization_{timestamp}.html"
-                filepath = os.path.join(DATA_DIR, filename)
+                # Extract the base64 image data from the HTML
+                img_data_match = re.search(r'src=\'data:image/png;base64,([^\']+)\'', viz_html)
                 
-                # Save the visualization
-                with open(filepath, "w") as f:
-                    f.write(viz_html)
+                if img_data_match:
+                    # Get the base64 data
+                    img_data = img_data_match.group(1)
+                    
+                    # Create a downloadable file
+                    timestamp = time.strftime("%Y%m%d_%H%M%S")
+                    filename = f"visualization_{timestamp}.png"
+                    
+                    # Return success message with download link
+                    return f"Visualization ready for download", gr.update(visible=True)
+                else:
+                    return "Could not extract image data", gr.update(visible=True)
+            except Exception as e:
+                return f"Error preparing visualization for download: {str(e)}", gr.update(visible=True)
+
+        def download_visualization(viz_html):
+            if not viz_html:
+                return None
+            
+            try:
+                # Extract the base64 image data from the HTML
+                img_data_match = re.search(r'src=\'data:image/png;base64,([^\']+)\'', viz_html)
                 
-                return f"Visualization saved as {filename}", gr.update(visible=True)
+                if img_data_match:
+                    # Get the base64 data
+                    img_data = img_data_match.group(1)
+                    
+                    # Convert base64 to image
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+                    
+                    image_data = base64.b64decode(img_data)
+                    image = Image.open(BytesIO(image_data))
+                    
+                    return image, gr.update(visible=True)
+                else:
+                    return None, gr.update(visible=False)
             except Exception as e:
-                return f"Error saving visualization: {str(e)}", gr.update(visible=True)
+                print(f"Error downloading visualization: {str(e)}")
+                return None, gr.update(visible=False)
         
         clear_viz_btn.click(
             clear_visualization,
@@ -909,6 +963,12 @@ with gr.Blocks(title="LLM Powered Database Chatbot") as demo:
             inputs=[current_visualization],
             outputs=[save_status, save_status]
         )
+        
+        download_btn.click(
+            download_visualization,
+            inputs=[current_visualization],
+            outputs=[download_img, download_img]
+        )
     
     # Update the process_text_query function to handle visualizations
     def process_text_query_with_visualization(query, history, current_viz):