feat: Add data visualization for Anthropic

cognee-mcp/README.md

@@ -83,3 +83,12 @@ npx -y @smithery/cli install cognee --client claude
 
 Define cognify tool in server.py
 Restart your Claude desktop.
+
+
+To use debugger, run:
+```bash
+npx @modelcontextprotocol/inspector uv --directory /Users/name/folder run cognee
+```
+
+To reset the installation
+uv sync --dev --all-extras --reinstall 

cognee-mcp/cognee_mcp/server.py

@@ -3,13 +3,17 @@
 import asyncio
 from contextlib import redirect_stderr, redirect_stdout
 
+from sqlalchemy.testing.plugin.plugin_base import logging
+
 import cognee
 import mcp.server.stdio
 import mcp.types as types
 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
+
 
 server = Server("cognee-mcp")
 
@@ -87,9 +91,45 @@ 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:
+    if not os.path.exists(directory):
+        raise FileNotFoundError(f"Directory {directory} does not exist")
+
+    # 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
+    try:
+        files_sorted = sorted(files, key=lambda x: int(x.replace(".png", "")))
+    except ValueError as e:
+        raise ValueError("Invalid PNG filename format. Expected timestamp format.") from e
+
+    # 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
+    try:
+        return Image.open(freshest_path)
+    except (IOError, OSError) as e:
+        raise IOError(f"Failed to open PNG file {freshest_path}") from e
+
 @server.call_tool()
 async def handle_call_tool(
     name: str, arguments: dict | None
@@ -154,6 +194,22 @@ async def handle_call_tool(
                         text="Pruned",
                     )
                 ]
+
+    elif name == "visualize":
+        with open(os.devnull, "w") as fnull:
+            with redirect_stdout(fnull), redirect_stderr(fnull):
+                try:
+                    result = await cognee.visualize_graph()
+                    results = retrieved_edges_to_string(result)
+
+                    return [
+                        types.TextContent(
+                            type="text",
+                            text=results,
+                        )
+                    ]
+                except (FileNotFoundError, IOError, ValueError) as e:
+                    raise ValueError(f"Failed to create visualization: {str(e)}")
     else:
         raise ValueError(f"Unknown tool: {name}")
 

cognee-mcp/pyproject.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 description = "A MCP server project"
 readme = "README.md"
 requires-python = ">=3.10"
+
 dependencies = [
     "mcp>=1.1.1",
     "openai==1.59.4",
@@ -51,7 +52,7 @@ dependencies = [
     "pydantic-settings>=2.2.1,<3.0.0",
     "anthropic>=0.26.1,<1.0.0",
     "sentry-sdk[fastapi]>=2.9.0,<3.0.0",
-    "fastapi-users[sqlalchemy]", # Optional
+    "fastapi-users[sqlalchemy]>=14.0.0", # Optional
     "alembic>=1.13.3,<2.0.0",
     "asyncpg==0.30.0", # Optional
     "pgvector>=0.3.5,<0.4.0", # Optional
@@ -91,4 +92,4 @@ dev = [
 ]
 
 [project.scripts]
-cognee = "cognee_mcp:main"
+cognee = "cognee_mcp:main"

cognee/__init__.py

@@ -4,7 +4,7 @@
 from .api.v1.datasets.datasets import datasets
 from .api.v1.prune import prune
 from .api.v1.search import SearchType, get_search_history, search
-from .api.v1.visualize import visualize
+from .api.v1.visualize import visualize_graph
 from .shared.utils import create_cognee_style_network_with_logo
 
 # Pipelines

cognee/shared/utils.py

@@ -13,7 +13,7 @@
 import tiktoken
 import nltk
 import base64
-
+import time
 
 import logging
 import sys
@@ -30,6 +30,34 @@
 proxy_url = "https://test.prometh.ai"
 
 
+
+def get_entities(tagged_tokens):
+    nltk.download("maxent_ne_chunker", quiet=True)
+    from nltk.chunk import ne_chunk
+
+    return ne_chunk(tagged_tokens)
+
+
+def extract_pos_tags(sentence):
+    """Extract Part-of-Speech (POS) tags for words in a sentence."""
+
+    # Ensure that the necessary NLTK resources are downloaded
+    nltk.download("words", quiet=True)
+    nltk.download("punkt", quiet=True)
+    nltk.download("averaged_perceptron_tagger", quiet=True)
+
+    from nltk.tag import pos_tag
+    from nltk.tokenize import word_tokenize
+
+    # Tokenize the sentence into words
+    tokens = word_tokenize(sentence)
+
+    # Tag each word with its corresponding POS tag
+    pos_tags = pos_tag(tokens)
+
+    return pos_tags
+
+
 def get_anonymous_id():
     """Creates or reads a anonymous user id"""
     home_dir = str(pathlib.Path(pathlib.Path(__file__).parent.parent.parent.resolve()))
@@ -243,31 +271,9 @@ async def render_graph(
 #     return df.replace([np.inf, -np.inf, np.nan], None)
 
 
-def get_entities(tagged_tokens):
-    nltk.download("maxent_ne_chunker", quiet=True)
-    from nltk.chunk import ne_chunk
-
-    return ne_chunk(tagged_tokens)
-
-
-def extract_pos_tags(sentence):
-    """Extract Part-of-Speech (POS) tags for words in a sentence."""
-
-    # Ensure that the necessary NLTK resources are downloaded
-    nltk.download("words", quiet=True)
-    nltk.download("punkt", quiet=True)
-    nltk.download("averaged_perceptron_tagger", quiet=True)
-
-    from nltk.tag import pos_tag
-    from nltk.tokenize import word_tokenize
 
-    # Tokenize the sentence into words
-    tokens = word_tokenize(sentence)
 
-    # Tag each word with its corresponding POS tag
-    pos_tags = pos_tag(tokens)
 
-    return pos_tags
 
 
 logging.basicConfig(level=logging.INFO)
@@ -396,6 +402,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,15 +450,19 @@ 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"
+
+
+
     logging.info(f"Saving visualization to {output_filename}...")
     html_content = file_html(p, CDN, title)
     with open(output_filename, "w") as f:
         f.write(html_content)
 
-    logging.info("Visualization complete.")
-
-    if bokeh_object:
-        return p
     return html_content
 
 
@@ -512,7 +523,7 @@ def setup_logging(log_level=logging.INFO):
             G,
             output_filename="example_network.html",
             title="Example Cognee Network",
-            node_attribute="group",  # Attribute to use for coloring nodes
+            label="group",  # Attribute to use for coloring nodes
             layout_func=nx.spring_layout,  # Layout function
             layout_scale=3.0,  # Scale for the layout
             logo_alpha=0.2,