feat: use fastmcp for mcp server #465
Conversation
WalkthroughThis pull request introduces significant restructuring of the Cognee MCP (Model Control Protocol) project. The changes involve relocating server functionality from Changes
Suggested labels
Suggested reviewers
Possibly related PRs
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 (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (8)
cognee-mcp/src/server.py (1)
1-118: Code formatting requiredThe pipeline failure indicates that this file requires formatting using the Ruff formatter. Please format the code according to the project's style guidelines to resolve the pipeline issue.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
cognee-mcp/src/__init__.py (1)
1-5: Code formatting requiredThe pipeline failure indicates that this file requires formatting using the Ruff formatter. Please format the code to adhere to the project's style guidelines and resolve the pipeline issue.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
cognee-mcp/pyproject.toml (1)
10-10: Consider using version range for mcp dependency.While pinning to exact version ensures reproducible builds, using
>=1.2.0,<1.3.0would allow for patch updates while maintaining compatibility.- "mcp==1.2.0", + "mcp>=1.2.0,<1.3.0",cognee-mcp/README.md (2)
11-11: Consider platform-specific installation instructions.Using
brew installlimits the instructions to macOS users. Consider adding alternative installation methods for other platforms.# macOS brew install uv # Other platforms pip install uv
54-54: Replace hard tabs with spaces.Use spaces instead of hard tabs for consistent formatting across different editors.
🧰 Tools
🪛 Markdownlint (0.37.0)
54-54: Column: 1
Hard tabs(MD010, no-hard-tabs)
cognee-mcp/src/client.py (3)
6-10: Consider making server parameters configurable.The server parameters are currently hard-coded. Consider making them configurable through environment variables or a config file.
-server_params = StdioServerParameters( - command="mcp", # Executable - args=["run", "src/server.py"], # Optional command line arguments - env=None # Optional environment variables -) +server_params = StdioServerParameters( + command=os.getenv("MCP_COMMAND", "mcp"), + args=os.getenv("MCP_ARGS", "run src/server.py").split(), + env=json.loads(os.getenv("MCP_ENV", "{}")) +)🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
12-28: Move example text to a separate file.Long string literals should be moved to separate resource files for better maintainability.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
30-39: Consider configurable session timeout.The session timeout is hard-coded to 3 minutes. Consider making it configurable for different use cases.
- async with ClientSession(read, write, timedelta(minutes=3)) as session: + timeout_minutes = int(os.getenv("SESSION_TIMEOUT_MINUTES", "3")) + async with ClientSession(read, write, timedelta(minutes=timeout_minutes)) as session:🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (2)
cognee-mcp/uv.lockis excluded by!**/*.lockpoetry.lockis excluded by!**/*.lock
📒 Files selected for processing (9)
cognee-mcp/README.md(3 hunks)cognee-mcp/cognee_mcp/__init__.py(0 hunks)cognee-mcp/cognee_mcp/server.py(0 hunks)cognee-mcp/pyproject.toml(1 hunks)cognee-mcp/src/__init__.py(1 hunks)cognee-mcp/src/client.py(1 hunks)cognee-mcp/src/server.py(1 hunks)cognee/infrastructure/databases/graph/networkx/adapter.py(2 hunks)pyproject.toml(1 hunks)
💤 Files with no reviewable changes (2)
- cognee-mcp/cognee_mcp/init.py
- cognee-mcp/cognee_mcp/server.py
🧰 Additional context used
🪛 GitHub Actions: ruff format
cognee-mcp/src/__init__.py
[warning] File requires formatting using Ruff formatter
cognee-mcp/src/client.py
[warning] File requires formatting using Ruff formatter
cognee-mcp/src/server.py
[warning] File requires formatting using Ruff formatter
🪛 Markdownlint (0.37.0)
cognee-mcp/README.md
54-54: Column: 1
Hard tabs
(MD010, no-hard-tabs)
85-85: null
Spaces inside code span elements
(MD038, no-space-in-code)
78-78: null
Fenced code blocks should have a language specified
(MD040, fenced-code-language)
⏰ Context from checks skipped due to timeout of 90000ms (6)
- GitHub Check: test
- GitHub Check: run_notebook_test / test
- GitHub Check: test
- GitHub Check: test
- GitHub Check: windows-latest
- GitHub Check: docker-compose-test
🔇 Additional comments (7)
cognee/infrastructure/databases/graph/networkx/adapter.py (1)
250-250: Verify compatibility ofedges="links"parameter in graph serializationThe addition of the
edges="links"parameter innode_link_dataandnode_link_graphchanges the key used for edges in the serialized graph data. This might impact compatibility with existing graph files or other parts of the codebase that expect the default key.Please ensure that this change does not introduce issues when saving and loading graphs, especially if there are existing graph data files serialized with the default
edgeskey. Consider testing the serialization and deserialization process to confirm that graphs are correctly persisted and restored.Run the following script to verify that serialization and deserialization work as expected:
Also applies to: 313-313
✅ Verification successful
Parameter usage is consistent and safe
The
edges="links"parameter is used consistently in both serialization (node_link_data) and deserialization (node_link_graph) operations within the same adapter class. No existing graph data files were found in the codebase that could be affected by this change.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Test graph serialization and deserialization with edges="links" # Create a temporary Python script to perform the test cat << 'EOF' > test_graph_serialization.py import networkx as nx from networkx.readwrite import json_graph import json # Create a sample graph G = nx.MultiDiGraph() G.add_node('node1') G.add_edge('node1', 'node2', key='relation') # Serialize graph with edges='links' graph_data = json_graph.node_link_data(G, edges='links') with open('test_graph.json', 'w') as f: json.dump(graph_data, f) # Deserialize graph with open('test_graph.json', 'r') as f: graph_data_loaded = json.load(f) G_loaded = json_graph.node_link_graph(graph_data_loaded, edges='links') # Verify that original and loaded graphs are isomorphic if nx.is_isomorphic(G, G_loaded): print('Graphs are serialized and deserialized successfully.') else: print('Graphs are not isomorphic after deserialization.') EOF # Run the Python script python test_graph_serialization.py # Clean up rm test_graph_serialization.py test_graph.jsonLength of output: 336
Script:
#!/bin/bash # Search for node_link_data and node_link_graph usage echo "=== Searching for node_link_data and node_link_graph usage ===" rg -A 2 "node_link_(data|graph)" --type py echo -e "\n=== Looking for JSON graph data files ===" fd -e json -x head -n 5 {} \; -x echo "---" \;Length of output: 4300
cognee-mcp/pyproject.toml (2)
21-22: LGTM! Good packaging practice.The wheel configuration properly specifies the src directory layout, which is a Python best practice for package organization.
28-28: LGTM! Entry point update.The entry point change aligns with the project restructuring to use the src directory.
cognee-mcp/README.md (2)
75-79: LGTM! Simplified debugger instructions.The new debugger command is more straightforward and the timeout configuration is clearly documented.
🧰 Tools
🪛 Markdownlint (0.37.0)
78-78: null
Fenced code blocks should have a language specified(MD040, fenced-code-language)
84-86: LGTM! Clear development workflow.The step-by-step instructions for applying changes during development are well-organized and clear.
🧰 Tools
🪛 Markdownlint (0.37.0)
85-85: null
Spaces inside code span elements(MD038, no-space-in-code)
cognee-mcp/src/client.py (1)
1-42: Format code using Ruff.The pipeline indicates that the file requires formatting using Ruff formatter.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
pyproject.toml (1)
23-23: LGTM! Dependency updates.The updates to pydantic (2.10.5) and uvicorn (0.34.0) are to stable versions with no breaking changes.
Also applies to: 26-26
| def load_class(model_file, model_name): | ||
| model_file = os.path.abspath(model_file) | ||
| spec = importlib.util.spec_from_file_location("graph_model", model_file) | ||
| module = importlib.util.module_from_spec(spec) | ||
| spec.loader.exec_module(module) | ||
|
|
||
| model_class = getattr(module, model_name) | ||
|
|
||
| return model_class |
There was a problem hiding this comment.
Potential security risk with dynamic module loading in load_class function
The load_class function uses user-supplied inputs graph_model_file and graph_model_name to dynamically load and execute code. This can introduce security vulnerabilities if untrusted inputs are provided.
Consider implementing input validation to ensure that only trusted files and classes can be loaded. Restrict the directories from which files can be loaded, and verify that the specified class exists within the module. Additionally, you might want to handle exceptions that could arise during the dynamic loading process.
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
| async def search(search_query: str) -> str: | ||
| """Search the knowledge graph""" | ||
| search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query) | ||
|
|
||
| results = retrieved_edges_to_string(search_results) | ||
|
|
||
| return results |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add exception handling to search function
The search function lacks exception handling, which means any errors during the search operation could cause the server to crash or return unintended responses.
Consider adding a try-except block to handle potential exceptions from cognee.search and to provide meaningful error messages.
Apply this diff to enhance error handling:
def search(search_query: str) -> str:
"""Search the knowledge graph"""
+ try:
search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query)
results = retrieved_edges_to_string(search_results)
return results
+ except Exception as e:
+ raise ValueError(f"Search failed: {str(e)}")📝 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.
| async def search(search_query: str) -> str: | |
| """Search the knowledge graph""" | |
| search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query) | |
| results = retrieved_edges_to_string(search_results) | |
| return results | |
| async def search(search_query: str) -> str: | |
| """Search the knowledge graph""" | |
| try: | |
| search_results = await cognee.search(SearchType.INSIGHTS, query_text=search_query) | |
| results = retrieved_edges_to_string(search_results) | |
| return results | |
| except Exception as e: | |
| raise ValueError(f"Search failed: {str(e)}") |
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
| async def prune() -> str: | ||
| """Reset the knowledge graph""" | ||
| await cognee.prune.prune_data() | ||
| await cognee.prune.prune_system(metadata=True) | ||
|
|
||
| return "Pruned" |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add exception handling to prune function
Similarly, the prune function does not handle exceptions that may occur during the pruning process, which could lead to unhandled errors.
Consider wrapping the pruning operations in a try-except block to handle any exceptions and inform the user appropriately.
Apply this diff to improve error handling:
def prune() -> str:
"""Reset the knowledge graph"""
+ try:
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
return "Pruned"
+ except Exception as e:
+ raise ValueError(f"Prune operation failed: {str(e)}")📝 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.
| async def prune() -> str: | |
| """Reset the knowledge graph""" | |
| await cognee.prune.prune_data() | |
| await cognee.prune.prune_system(metadata=True) | |
| return "Pruned" | |
| async def prune() -> str: | |
| """Reset the knowledge graph""" | |
| try: | |
| await cognee.prune.prune_data() | |
| await cognee.prune.prune_system(metadata=True) | |
| return "Pruned" | |
| except Exception as e: | |
| raise ValueError(f"Prune operation failed: {str(e)}") |
🧰 Tools
🪛 GitHub Actions: ruff format
[warning] File requires formatting using Ruff formatter
Description
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
Summary by CodeRabbit
Documentation
Dependency Updates
mcpdependency to version 1.2.0pydanticto version 2.10.5uvicornto version 0.34.0Project Restructuring
Bug Fixes