feat: Add data visualization for Anthropic #432
Conversation
WalkthroughThe pull request introduces enhanced image processing and visualization capabilities for the Cognee knowledge graph system. The changes span two files: Changes
Sequence DiagramsequenceDiagram
participant User
participant Server
participant Cognee
participant FileSystem
User->>Server: Request graph visualization
Server->>Cognee: Call visualize method
Cognee->>FileSystem: Generate PNG visualization
FileSystem-->>Cognee: Save PNG with timestamp
Cognee-->>Server: Return visualization image
Server-->>User: Display visualization
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
Vasilije1990
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
cognee-mcp/cognee_mcp/server.py(3 hunks)cognee/shared/utils.py(3 hunks)
🧰 Additional context used
🪛 GitHub Actions: ruff format
cognee-mcp/cognee_mcp/server.py
[warning] File requires formatting. Code does not conform to Ruff formatting standards.
🔇 Additional comments (1)
cognee-mcp/cognee_mcp/server.py (1)
92-101: LGTM: Tool definition looks goodThe "visualize" tool is properly defined with a clear description and appropriate schema.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff formatting standards.
| from PIL import Image | ||
| from PIL import Image as PILImage |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove duplicate PIL Image import
The Image module from PIL is imported twice with different aliases. This is redundant and could cause confusion.
-from PIL import Image
-from PIL import Image as PILImage
+from PIL import Image📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| from PIL import Image | |
| from PIL import Image as PILImage | |
| from PIL import Image |
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff formatting standards.
| # Get the latest Unix timestamp as an integer | ||
| timestamp = int(time.time()) | ||
|
|
||
| # Construct your filename | ||
| filename = f"{timestamp}.png" | ||
|
|
||
| export_png(p, filename=filename) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add error handling for PNG export
The PNG export functionality lacks error handling and cleanup of old files.
# Get the latest Unix timestamp as an integer
timestamp = int(time.time())
# Construct your filename
filename = f"{timestamp}.png"
- export_png(p, filename=filename)
+ try:
+ # Cleanup old PNG files to prevent disk space issues
+ cleanup_old_pngs(directory=".", keep_latest=5)
+
+ # Export the new PNG
+ export_png(p, filename=filename)
+ except Exception as e:
+ logging.error(f"Failed to export PNG: {str(e)}")
+ raiseConsider adding a helper function to cleanup old PNG files:
def cleanup_old_pngs(directory: str, keep_latest: int = 5):
"""Cleanup old PNG files, keeping only the N latest files."""
png_files = [f for f in os.listdir(directory) if f.endswith('.png')]
if len(png_files) <= keep_latest:
return
# Sort by timestamp in filename
sorted_files = sorted(png_files, key=lambda x: int(x.replace(".png", "")))
# Remove older files
for f in sorted_files[:-keep_latest]:
try:
os.remove(os.path.join(directory, f))
except OSError as e:
logging.warning(f"Failed to remove old PNG file {f}: {str(e)}")
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🔭 Outside diff range comments (1)
cognee-mcp/cognee_mcp/server.py (1)
Fix import statements organization
The file has the following formatting issues:
- Remove duplicate PIL import:
from PIL import Image as PILImageis redundant sinceImageis already imported- Organize imports according to PEP 8 style guide:
- Standard library imports
- Third-party imports
- Local application imports
🔗 Analysis chain
Line range hint
1-230: Fix formatting issuesThe file requires formatting according to Ruff's standards. Please run the formatter:
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
ruff format cognee-mcp/cognee_mcp/server.pyLength of output: 140
Script:
#!/bin/bash # Check for basic formatting patterns echo "=== Checking import statements ===" rg "^import|^from.*import" cognee-mcp/cognee_mcp/server.py echo -e "\n=== Checking blank lines between functions ===" rg "^def|^async def" -A 1 cognee-mcp/cognee_mcp/server.py echo -e "\n=== Checking indentation ===" rg "^[[:space:]]+[^[:space:]]" cognee-mcp/cognee_mcp/server.pyLength of output: 8158
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
🧹 Nitpick comments (1)
cognee-mcp/cognee_mcp/server.py (1)
105-130: Add docstring to document timestamp format requirementThe function assumes PNG files are named with timestamps but this requirement isn't documented.
def get_freshest_png(directory: str) -> Image.Image: + """ + Retrieve the most recently created PNG file from a directory. + + Args: + directory (str): Path to directory containing PNG files + + Returns: + Image.Image: PIL Image object of the most recent PNG + + Raises: + FileNotFoundError: If directory doesn't exist or no PNG files found + ValueError: If PNG filenames don't follow timestamp format (e.g., '1673185134.png') + IOError: If PNG file can't be opened + + Note: + PNG files must be named with Unix timestamps (e.g., '1673185134.png') + """ if not os.path.exists(directory):🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
cognee-mcp/cognee_mcp/server.py(3 hunks)
🧰 Additional context used
🪛 GitHub Actions: ruff format
cognee-mcp/cognee_mcp/server.py
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: run_notebook_test / test
- GitHub Check: run_notebook_test / test
- GitHub Check: run_notebook_test / test
- GitHub Check: docker-compose-test
🔇 Additional comments (1)
cognee-mcp/cognee_mcp/server.py (1)
13-14: Remove duplicate PIL Image importThe Image module from PIL is imported twice with different aliases. This is redundant and could cause confusion.
-from PIL import Image -from PIL import Image as PILImage +from PIL import Image🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
| with open(os.devnull, "w") as fnull: | ||
| with redirect_stdout(fnull), redirect_stderr(fnull): | ||
| try: | ||
| await cognee.visualize |
There was a problem hiding this comment.
Fix suspicious await on property access
The line await cognee.visualize looks incorrect. Properties typically don't need await. This might be a typo where the parentheses for method invocation are missing.
- await cognee.visualize
+ await cognee.visualize()📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| await cognee.visualize | |
| await cognee.visualize() |
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
| with redirect_stdout(fnull), redirect_stderr(fnull): | ||
| try: | ||
| await cognee.visualize | ||
| img = get_freshest_png(".") |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use configured output directory instead of current directory
Using "." as the directory path is brittle and depends on the current working directory when the server starts. Consider using a configured output directory.
- img = get_freshest_png(".")
+ img = get_freshest_png(cognee.config.visualization_output_dir)You'll need to ensure the output directory is properly configured in the Cognee configuration.
Committable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
| types.Tool( | ||
| name="visualize", | ||
| description="Visualize the knowledge graph.", | ||
| inputSchema={ | ||
| "type": "object", | ||
| "properties": { | ||
| "query": {"type": "string"}, | ||
| }, | ||
| }, | ||
| ), |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Clarify query parameter requirement in visualization tool schema
The tool's schema defines a "query" property but doesn't mark it as required. Either:
- Mark it as required if visualization needs a query parameter
- Document that it's optional if visualization can work without a query
types.Tool(
name="visualize",
description="Visualize the knowledge graph.",
inputSchema={
"type": "object",
"properties": {
"query": {"type": "string"},
},
+ "required": ["query"], # Add this if query is required
},
),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| types.Tool( | |
| name="visualize", | |
| description="Visualize the knowledge graph.", | |
| inputSchema={ | |
| "type": "object", | |
| "properties": { | |
| "query": {"type": "string"}, | |
| }, | |
| }, | |
| ), | |
| types.Tool( | |
| name="visualize", | |
| description="Visualize the knowledge graph.", | |
| inputSchema={ | |
| "type": "object", | |
| "properties": { | |
| "query": {"type": "string"}, | |
| }, | |
| "required": ["query"], # Add this if query is required | |
| }, | |
| ), |
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting. Code does not conform to Ruff's formatting standards.
Summary by CodeRabbit
New Features
Improvements