Merge pull request #432 from topoteretes/COG-975

feat: Add data visualization for Anthropic

.github/pull_request_template.md

@@ -0,0 +1,7 @@
+<!-- .github/pull_request_template.md -->
+
+## Description
+<!-- Provide a clear description of the changes in this PR -->
+
+## 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

.github/workflows/approve_dco.yaml

@@ -0,0 +1,53 @@
+name: DCO Check
+
+on:
+  pull_request:
+    types: [opened, edited, reopened, synchronize, ready_for_review]
+
+jobs:
+  check-dco:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Validate Developer Certificate of Origin statement
+        uses: actions/github-script@v6
+        with:
+          # If using the built-in GITHUB_TOKEN, ensure it has 'read:org' permission.
+          # In GitHub Enterprise or private orgs, you might need a PAT (personal access token) with read:org scope.
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            const orgName = 'YOUR_ORGANIZATION_NAME'; // Replace with your org
+            const prUser = context.payload.pull_request.user.login;
+            const prBody = context.payload.pull_request.body || '';
+
+            // Exact text you require in the PR body
+            const requiredStatement = "I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin";
+
+            // 1. Check if user is in the org
+            let isOrgMember = false;
+            try {
+              // Attempt to get membership info
+              const membership = await github.rest.orgs.getMembershipForUser({
+                org: orgName,
+                username: prUser,
+              });
+              // If we get here without an error, user is in the org
+              isOrgMember = true;
+              console.log(`${prUser} is a member of ${orgName}. Skipping DCO check.`);
+            } catch (error) {
+              // If we get a 404, user is NOT an org member
+              if (error.status === 404) {
+                console.log(`${prUser} is NOT a member of ${orgName}. Enforcing DCO check.`);
+              } else {
+                // Some other error—fail the workflow or handle accordingly
+                core.setFailed(`Error checking organization membership: ${error.message}`);
+              }
+            }
+
+            // 2. If user is not in the org, enforce the DCO statement
+            if (!isOrgMember) {
+              if (!prBody.includes(requiredStatement)) {
+                core.setFailed(
+                  `DCO check failed. The PR body must include the following statement:\n\n${requiredStatement}`
+                );
+              }
+            }

.github/workflows/dockerhub.yml

@@ -1,8 +1,9 @@
-name: build | Build and Push Docker Image to DockerHub
+name: build | Build and Push Docker Image to dockerhub
 
 on:
   push:
     branches:
+      - dev
       - main
 
 jobs:
@@ -10,42 +11,38 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - name: Checkout repository
-      uses: actions/checkout@v4
+      - name: Checkout repository
+        uses: actions/checkout@v4
 
-    - name: Set up Docker Buildx
-      uses: docker/setup-buildx-action@v3
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
 
-    - name: Log in to Docker Hub
-      uses: docker/login-action@v3
-      with:
-        username: ${{ secrets.DOCKER_USERNAME }}
-        password: ${{ secrets.DOCKER_PASSWORD }}
+      - name: Log in to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_PASSWORD }}
 
-    - name: Extract Git information
-      id: git-info
-      run: |
-        echo "BRANCH_NAME=${GITHUB_REF_NAME}" >> "$GITHUB_ENV"
-        echo "COMMIT_SHA=${GITHUB_SHA::7}" >> "$GITHUB_ENV"
+      - name: Extract metadata
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: cognee/cognee
+          tags: |
+            type=ref,event=branch
+            type=sha,prefix={{branch}}-
+            type=raw,value=latest,enable={{is_default_branch}}
 
-    - name: Build and Push Docker Image
-      run: |
-        IMAGE_NAME=cognee/cognee
-        TAG_VERSION="${BRANCH_NAME}-${COMMIT_SHA}"
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          platforms: linux/amd64,linux/arm64
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=registry,ref=cognee/cognee:buildcache
+          cache-to: type=registry,ref=cognee/cognee:buildcache,mode=max
 
-        echo "Building image: ${IMAGE_NAME}:${TAG_VERSION}"
-        docker buildx build \
-          --platform linux/amd64,linux/arm64 \
-          --push \
-          --tag "${IMAGE_NAME}:${TAG_VERSION}" \
-          --tag "${IMAGE_NAME}:latest" \
-          .
-
-    - name: Verify pushed Docker images
-      run: |
-       # Verify both platform variants
-       for PLATFORM in "linux/amd64" "linux/arm64"; do
-         echo "Verifying image for $PLATFORM..."
-         docker buildx imagetools inspect "${IMAGE_NAME}:${TAG_VERSION}" --format "{{.Manifest.$PLATFORM.Digest}}"
-       done
-       echo "Successfully verified images in Docker Hub"
+      - name: Image digest
+        run: echo ${{ steps.build.outputs.digest }}

.github/workflows/test_python_3_10.yml

@@ -42,6 +42,10 @@ jobs:
 
       - name: Install dependencies
         run: poetry install --no-interaction -E docs
+      - name: Download NLTK tokenizer data
+        run: |
+          poetry run python -m nltk.downloader punkt_tab averaged_perceptron_tagger_eng
+
 
       - name: Run unit tests
         run: poetry run pytest cognee/tests/unit/

.github/workflows/test_python_3_11.yml

@@ -44,6 +44,11 @@ jobs:
       - name: Install dependencies
         run: poetry install --no-interaction -E docs
 
+      - name: Download NLTK tokenizer data
+        run: |
+          poetry run python -m nltk.downloader punkt_tab averaged_perceptron_tagger_eng
+
+
       - name: Run unit tests
         run: poetry run pytest cognee/tests/unit/
 

.github/workflows/test_python_3_12.yml

@@ -43,6 +43,9 @@ jobs:
 
       - name: Install dependencies
         run: poetry install --no-interaction -E docs
+      - name: Download NLTK tokenizer data
+        run: |
+          poetry run python -m nltk.downloader punkt_tab averaged_perceptron_tagger_eng
 
       - name: Run unit tests
         run: poetry run pytest cognee/tests/unit/

cognee-mcp/README.md

@@ -37,7 +37,15 @@ source .venv/bin/activate
 4. Add the new server to your Claude config:
 
 The file should be located here: ~/Library/Application\ Support/Claude/
+```
+cd ~/Library/Application\ Support/Claude/
+```
 You need to create claude_desktop_config.json in this folder if it doesn't exist
+Make sure to add your paths and LLM API key to the file bellow
+Use your editor of choice, for example Nano:
+```
+nano claude_desktop_config.json
+```
 
 ```
 
@@ -83,3 +91,17 @@ 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 apply new changes while development you do:
+
+1. Poetry lock in cognee folder
+2. uv sync --dev --all-extras --reinstall 
+3. npx @modelcontextprotocol/inspector uv --directory /Users/vasilije/cognee/cognee-mcp run cognee
+
+

cognee-mcp/cognee_mcp/server.py

@@ -3,6 +3,8 @@ import os
 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
@@ -10,6 +12,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
+
 
 server = Server("cognee-mcp")
 
@@ -87,9 +91,46 @@ 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 +195,20 @@ 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:
+                    results = await cognee.visualize_graph()
+                    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-mcp/uv.lock

@@ -561,7 +561,7 @@ name = "click"
 version = "8.1.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "platform_system == 'Windows'" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 }
 wheels = [
@@ -570,7 +570,7 @@ wheels = [
 
 [[package]]
 name = "cognee"
-version = "0.1.21"
+version = "0.1.22"
 source = { directory = "../" }
 dependencies = [
     { name = "aiofiles" },
@@ -633,7 +633,7 @@ requires-dist = [
     { name = "dlt", extras = ["sqlalchemy"], specifier = ">=1.4.1,<2.0.0" },
     { name = "falkordb", marker = "extra == 'falkordb'", specifier = "==1.0.9" },
     { name = "fastapi", specifier = ">=0.109.2,<0.116.0" },
-    { name = "fastapi-users", extras = ["sqlalchemy"] },
+    { name = "fastapi-users", extras = ["sqlalchemy"], specifier = "==14.0.0" },
     { name = "filetype", specifier = ">=1.2.0,<2.0.0" },
     { name = "graphistry", specifier = ">=0.33.5,<0.34.0" },
     { name = "groq", marker = "extra == 'groq'", specifier = "==0.8.0" },
@@ -647,12 +647,12 @@ requires-dist = [
     { name = "langfuse", specifier = ">=2.32.0,<3.0.0" },
     { name = "langsmith", marker = "extra == 'langchain'", specifier = "==0.2.3" },
     { name = "litellm", specifier = "==1.57.2" },
-    { name = "llama-index-core", marker = "extra == 'llama-index'", specifier = ">=0.12.10.post1,<0.13.0" },
+    { name = "llama-index-core", marker = "extra == 'llama-index'", specifier = ">=0.12.11,<0.13.0" },
     { name = "matplotlib", specifier = ">=3.8.3,<4.0.0" },
     { name = "neo4j", marker = "extra == 'neo4j'", specifier = ">=5.20.0,<6.0.0" },
     { name = "nest-asyncio", specifier = "==1.6.0" },
     { name = "networkx", specifier = ">=3.2.1,<4.0.0" },
-    { name = "nltk", specifier = ">=3.8.1,<4.0.0" },
+    { name = "nltk", specifier = "==3.9.1" },
     { name = "numpy", specifier = "==1.26.4" },
     { name = "openai", specifier = "==1.59.4" },
     { name = "pandas", specifier = "==2.2.3" },
@@ -674,7 +674,7 @@ requires-dist = [
     { name = "tiktoken", specifier = "==0.7.0" },
     { name = "transformers", specifier = ">=4.46.3,<5.0.0" },
     { name = "typing-extensions", specifier = "==4.12.2" },
-    { name = "unstructured", extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], marker = "extra == 'docs'", specifier = ">=0.16.10,<0.17.0" },
+    { name = "unstructured", extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], marker = "extra == 'docs'", specifier = ">=0.16.13,<0.17.0" },
     { name = "uvicorn", specifier = "==0.22.0" },
     { name = "weaviate-client", marker = "extra == 'weaviate'", specifier = "==4.9.6" },
 ]
@@ -777,7 +777,7 @@ requires-dist = [
     { name = "dlt", extras = ["sqlalchemy"], specifier = ">=1.4.1,<2.0.0" },
     { name = "falkordb", specifier = "==1.0.9" },
     { name = "fastapi", specifier = ">=0.109.2,<0.110.0" },
-    { name = "fastapi-users", extras = ["sqlalchemy"] },
+    { name = "fastapi-users", extras = ["sqlalchemy"], specifier = ">=14.0.0" },
     { name = "filetype", specifier = ">=1.2.0,<2.0.0" },
     { name = "gitpython", specifier = ">=3.1.43,<4.0.0" },
     { name = "graphistry", specifier = ">=0.33.5,<0.34.0" },
@@ -3359,7 +3359,7 @@ name = "portalocker"
 version = "2.10.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "pywin32", marker = "platform_system == 'Windows'" },
+    { name = "pywin32", marker = "sys_platform == 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/ed/d3/c6c64067759e87af98cc668c1cc75171347d0f1577fab7ca3749134e3cd4/portalocker-2.10.1.tar.gz", hash = "sha256:ef1bf844e878ab08aee7e40184156e1151f228f103aa5c6bd0724cc330960f8f", size = 40891 }
 wheels = [
@@ -4954,7 +4954,7 @@ name = "tqdm"
 version = "4.67.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "platform_system == 'Windows'" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
 wheels = [

cognee/__init__.py

@@ -4,7 +4,7 @@ from .api.v1.config.config import config
 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/api/v1/visualize/visualize.py

@@ -10,5 +10,6 @@ async def visualize_graph(label: str = "name"):
     logging.info(graph_data)
 
     graph = await create_cognee_style_network_with_logo(graph_data, label=label)
+    logging.info("The HTML file has been stored on your home directory! Navigate there with cd ~")
 
     return graph

cognee/shared/utils.py

@@ -11,9 +11,7 @@ import networkx as nx
 import pandas as pd
 import matplotlib.pyplot as plt
 import tiktoken
-import nltk
-import base64
-
+import time
 
 import logging
 import sys
@@ -23,13 +21,40 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 
 from uuid import uuid4
 import pathlib
-
+import nltk
 from cognee.shared.exceptions import IngestionError
 
 # Analytics Proxy Url, currently hosted by Vercel
 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,33 +268,6 @@ 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 +394,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)
@@ -445,13 +444,14 @@ async def create_cognee_style_network_with_logo(
 
     logging.info(f"Saving visualization to {output_filename}...")
     html_content = file_html(p, CDN, title)
-    with open(output_filename, "w") as f:
+
+    home_dir = os.path.expanduser("~")
+
+    # Construct the final output file path
+    output_filepath = os.path.join(home_dir, output_filename)
+    with open(output_filepath, "w") as f:
         f.write(html_content)
 
-    logging.info("Visualization complete.")
-
-    if bokeh_object:
-        return p
     return html_content
 
 
@@ -512,7 +512,7 @@ if __name__ == "__main__":
             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,

pyproject.toml

@@ -40,7 +40,6 @@ networkx = "^3.2.1"
 aiosqlite = "^0.20.0"
 pandas = "2.2.3"
 filetype = "^1.2.0"
-nltk = "^3.8.1"
 dlt = {extras = ["sqlalchemy"], version = "^1.4.1"}
 aiofiles = "^23.2.1"
 qdrant-client = {version = "^1.9.0", optional = true}
@@ -64,19 +63,20 @@ langfuse = "^2.32.0"
 pydantic-settings = "^2.2.1"
 anthropic = "^0.26.1"
 sentry-sdk = {extras = ["fastapi"], version = "^2.9.0"}
-fastapi-users = {version = "*", extras = ["sqlalchemy"]}
+fastapi-users = {version = "14.0.0", extras = ["sqlalchemy"]}
 alembic = "^1.13.3"
 asyncpg =  {version = "0.30.0", optional = true}
 pgvector = {version = "^0.3.5", optional = true}
 psycopg2 = {version = "^2.9.10", optional = true}
-llama-index-core = {version = "^0.12.10.post1", optional = true}
+llama-index-core = {version = "^0.12.11", optional = true}
 deepeval = {version = "^2.0.1", optional = true}
 transformers = "^4.46.3"
 pymilvus = {version = "^2.5.0", optional = true}
-unstructured = { extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], version = "^0.16.10", optional = true }
+unstructured = { extras = ["csv", "doc", "docx", "epub", "md", "odt", "org", "ppt", "pptx", "rst", "rtf", "tsv", "xlsx"], version = "^0.16.13", optional = true }
 pre-commit = "^4.0.1"
 httpx = "0.27.0"
 bokeh="^3.6.2"
+nltk = "3.9.1"