From daf2d540baf5d6c7f6da92057f7bbed18d0041b6 Mon Sep 17 00:00:00 2001
From: vasilije <vas.markovic@gmail.com>
Date: Fri, 10 Jan 2025 22:28:14 +0100
Subject: [PATCH] Add data visualization for Anthropic

---
 cognee-mcp/cognee_mcp/server.py | 37 +++++++++++++++++++++++++++++++++
 cognee/shared/utils.py          | 11 +++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/cognee-mcp/cognee_mcp/server.py b/cognee-mcp/cognee_mcp/server.py
index 19dc999d7..1cc141d0f 100644
--- a/cognee-mcp/cognee_mcp/server.py
+++ b/cognee-mcp/cognee_mcp/server.py
@@ -10,6 +10,8 @@ from cognee.api.v1.search import SearchType
 from cognee.shared.data_models import KnowledgeGraph
 from mcp.server import NotificationOptions, Server
 from mcp.server.models import InitializationOptions
+from PIL import Image
+from PIL import Image as PILImage
 
 server = Server("cognee-mcp")
 
@@ -87,9 +89,36 @@ async def handle_list_tools() -> list[types.Tool]:
                 },
             },
         ),
+        types.Tool(
+            name="visualize",
+            description="Visualize the knowledge graph.",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string"},
+                },
+            },
+        ),
     ]
 
 
+def get_freshest_png(directory: str) -> Image.Image:
+    # List all files in 'directory' that end with .png
+    files = [f for f in os.listdir(directory) if f.endswith(".png")]
+    if not files:
+        raise FileNotFoundError("No PNG files found in the given directory.")
+
+    # Sort by integer value of the filename (minus the '.png')
+    # Example filename: 1673185134.png -> integer 1673185134
+    files_sorted = sorted(files, key=lambda x: int(x.replace(".png", "")))
+
+    # The "freshest" file has the largest timestamp
+    freshest_filename = files_sorted[-1]
+    freshest_path = os.path.join(directory, freshest_filename)
+
+    # Open the image with PIL and return the PIL Image object
+    return Image.open(freshest_path)
+
 @server.call_tool()
 async def handle_call_tool(
     name: str, arguments: dict | None
@@ -154,6 +183,14 @@ async def handle_call_tool(
                         text="Pruned",
                     )
                 ]
+
+    elif name == "visualize":
+        with open(os.devnull, "w") as fnull:
+            with redirect_stdout(fnull), redirect_stderr(fnull):
+                """Create a thumbnail from an image"""
+                await cognee.visualize
+                img = get_freshest_png(".")
+                return types.Image(data=img.tobytes(), format="png")
     else:
         raise ValueError(f"Unknown tool: {name}")
 
diff --git a/cognee/shared/utils.py b/cognee/shared/utils.py
index e57decde1..2dab80187 100644
--- a/cognee/shared/utils.py
+++ b/cognee/shared/utils.py
@@ -13,7 +13,7 @@ import matplotlib.pyplot as plt
 import tiktoken
 import nltk
 import base64
-
+import time
 
 import logging
 import sys
@@ -396,6 +396,7 @@ async def create_cognee_style_network_with_logo(
 
     from bokeh.embed import file_html
     from bokeh.resources import CDN
+    from bokeh.io import export_png
 
     logging.info("Converting graph to serializable format...")
     G = await convert_to_serializable_graph(G)
@@ -443,6 +444,14 @@ async def create_cognee_style_network_with_logo(
     )
     p.add_tools(hover_tool)
 
+    # Get the latest Unix timestamp as an integer
+    timestamp = int(time.time())
+
+    # Construct your filename
+    filename = f"{timestamp}.png"
+
+    export_png(p, filename=filename)
+
     logging.info(f"Saving visualization to {output_filename}...")
     html_content = file_html(p, CDN, title)
     with open(output_filename, "w") as f: