Update main.py

main.py

@@ -577,5 +577,45 @@ async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
         content={"detail": "Too many requests. Please try again later."}
     )
 
+# ===== ADD THIS AT THE BOTTOM OF main.py =====
 if __name__ == "__main__":
+    import uvicorn
+    from fastapi.testclient import TestClient
+    from io import BytesIO
+    import base64
+    from PIL import Image
+    import matplotlib.pyplot as plt
+
+    # 1. Start the app (or connect to a running instance)
+    client = TestClient(app)
+
+    # 2. Test the visualization endpoint
+    test_file = "test.xlsx"  # Replace with your test file
+    test_prompt = "Show me a bar chart of sales by region"
+
+    # 3. Send request to your own API
+    with open(test_file, "rb") as f:
+        response = client.post(
+            "/visualize/natural",
+            files={"file": ("test.xlsx", f, "application/vnd.ms-excel")},
+            data={"prompt": test_prompt}
+        )
+
+    # 4. Check if successful
+    if response.status_code == 200:
+        result = response.json()
+        print("Visualization generated successfully!")
+
+        # 5. Decode and display the image
+        image_data = result["image"].split(",")[1]  # Remove header
+        image_bytes = base64.b64decode(image_data)
+        image = Image.open(BytesIO(image_bytes))
+
+        plt.imshow(image)
+        plt.axis("off")
+        plt.show()
+    else:
+        print(f"Error: {response.status_code}\n{response.text}")
+
+    # 6. Optional: Run the server (if not already running)
     uvicorn.run(app, host="0.0.0.0", port=7860)