From ecf6a19ab8478aaf2c68e0e6fc6b9405b4fbb53f Mon Sep 17 00:00:00 2001
From: Igor Ilic <30923996+dexters1@users.noreply.github.com>
Date: Sat, 15 Feb 2025 03:08:26 +0100
Subject: [PATCH] Cognee gui add visualziation[COG-1334] (#538)

<!-- .github/pull_request_template.md -->

## Description
Add visualization to Cognee GUI

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Introduced an interactive data visualization option that opens
insights in your web browser.
- Enhanced the interface layout for a more streamlined file upload and
visualization experience.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 cognee-gui.py | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/cognee-gui.py b/cognee-gui.py
index f5cf55210..2a158349b 100644
--- a/cognee-gui.py
+++ b/cognee-gui.py
@@ -10,6 +10,7 @@ try:
         QLineEdit,
         QFileDialog,
         QVBoxLayout,
+        QHBoxLayout,
         QLabel,
         QMessageBox,
         QTextEdit,
@@ -38,9 +39,18 @@ class FileSearchApp(QWidget):
         self.init_ui()
 
     def init_ui(self):
+        # Horizontal layout for file upload and visualization buttons
+        button_layout = QHBoxLayout()
+
         # Button to open file dialog
         self.file_button = QPushButton("Upload File to Cognee", parent=self)
         self.file_button.clicked.connect(self.open_file_dialog)
+        button_layout.addWidget(self.file_button)
+
+        # Button to visualize data
+        self.visualize_button = QPushButton("Visualize Data", parent=self)
+        self.visualize_button.clicked.connect(lambda: asyncio.ensure_future(self.visualize_data()))
+        button_layout.addWidget(self.visualize_button)
 
         # Label to display selected file path
         self.file_label = QLabel("No file selected", parent=self)
@@ -61,13 +71,12 @@ class FileSearchApp(QWidget):
         # Progress dialog
         self.progress_dialog = QProgressDialog("Processing..", None, 0, 0, parent=self)
         self.progress_dialog.setWindowModality(Qt.WindowModal)
-        # self.progress_dialog.setAttribute(Qt.WA_DeleteOnClose)
         self.progress_dialog.setCancelButton(None)  # Remove the cancel button
         self.progress_dialog.close()
 
         # Layout setup
         layout = QVBoxLayout()
-        layout.addWidget(self.file_button)
+        layout.addLayout(button_layout)
         layout.addWidget(self.file_label)
         layout.addWidget(self.search_input)
         layout.addWidget(self.search_button)
@@ -119,6 +128,17 @@ class FileSearchApp(QWidget):
         self.setEnabled(True)
         self.progress_dialog.close()
 
+    async def visualize_data(self):
+        """Async slot for handling visualize data button press."""
+        import webbrowser
+        from cognee.api.v1.visualize.visualize import visualize_graph
+        import os
+        import pathlib
+
+        html_file = os.path.join(pathlib.Path(__file__).parent, ".data", "graph_visualization.html")
+        await visualize_graph(html_file)
+        webbrowser.open(f"file://{html_file}")
+
 
 if __name__ == "__main__":
     app = QApplication(sys.argv)