Update app.py

app.py

@@ -212,6 +212,11 @@ def process_image(image, mode, entry_type, quantity):
     if product_name:
         attributes["Product name"] = product_name
 
+    # Ensure fixed attributes are present
+    for fixed_attr in ["Stage", "H.P.", "Product name", "Model"]:
+        if fixed_attr not in attributes:
+            attributes[fixed_attr] = ""
+
     # Convert attributes to DataFrame for editing
     df = pd.DataFrame(list(attributes.items()), columns=["Attribute", "Value"])
     return f"Extracted Text:\n{extracted_text}", df, None
@@ -228,6 +233,57 @@ def export_to_salesforce(mode, entry_type, quantity, edited_df):
     except Exception as e:
         return f"❌ Error exporting to Salesforce: {str(e)}"
 
+# Function to pull structured data from Salesforce and display as a table
+def pull_data_from_salesforce(data_type):
+    try:
+        sf = Salesforce(
+            username=SALESFORCE_USERNAME,
+            password=SALESFORCE_PASSWORD,
+            security_token=SALESFORCE_SECURITY_TOKEN
+        )
+        
+        if data_type == "Inventory":
+            query = "SELECT Productname__c, Current_Stocks__c, soldstock__c FROM Inventory_Management__c LIMIT 100"
+        else:
+            query = "SELECT Productname__c, Current_Stock__c, soldstock__c FROM Un_Billable__c LIMIT 100"
+        
+        response = sf.query_all(query)
+        records = response.get("records", [])
+        
+        if not records:
+            return "No data found in Salesforce.", None, None, None
+        
+        df = pd.DataFrame(records)
+        df = df.drop(columns=['attributes'], errors='ignore')
+        
+        # Rename columns for better readability
+        df.rename(columns={
+            "Productname__c": "Product Name",
+            "Current_Stocks__c": "Current Stocks",
+            "Current_Stock__c": "Current Stocks",
+            "soldstock__c": "Sold Stock"
+        }, inplace=True)
+        
+        excel_path = "salesforce_data.xlsx"
+        df.to_excel(excel_path, index=False)
+        
+        # Generate interactive vertical bar graph using Matplotlib
+        fig, ax = plt.subplots(figsize=(12, 8))
+        df.plot(kind='bar', x="Product Name", y="Current Stocks", ax=ax, legend=False)
+        ax.set_title("Stock Distribution by Product Name")
+        ax.set_xlabel("Product Name")
+        ax.set_ylabel("Current Stocks")
+        plt.xticks(rotation=45, ha="right", fontsize=10)
+        plt.tight_layout()
+        buffer = BytesIO()
+        plt.savefig(buffer, format="png")
+        buffer.seek(0)
+        img = Image.open(buffer)
+        
+        return df, excel_path, img
+    except Exception as e:
+        return f"Error fetching data: {str(e)}", None, None, None
+
 # Gradio Interface
 def app():
     with gr.Blocks() as demo:
@@ -243,6 +299,13 @@ def app():
             ok_button = gr.Button("OK")
             result_output = gr.Text(label="🚀 Result")
 
+        with gr.Tab("📊 Salesforce Data"):
+            data_type_input = gr.Dropdown(label="Select Data Type", choices=["Inventory", "Unbilling"], value="Inventory")
+            pull_button = gr.Button("Pull Data from Salesforce")
+            salesforce_data_output = gr.Dataframe(label="📊 Salesforce Data")
+            excel_download_output = gr.File(label="📥 Download Excel")
+            graph_output = gr.Image(label="📈 Stock Distribution Graph")
+
         # Define button actions
         extract_button.click(
             fn=process_image,
@@ -254,6 +317,11 @@ def app():
             inputs=[mode_input, entry_type_input, quantity_input, editable_df_output],
             outputs=[result_output]
         )
+        pull_button.click(
+            fn=pull_data_from_salesforce,
+            inputs=[data_type_input],
+            outputs=[salesforce_data_output, excel_download_output, graph_output]
+        )
 
     return demo