gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

some refactoring for easier dev

Browse source
This commit is contained in:
Daulet Amirkhanov 2025-10-30 18:04:02 +00:00
parent 953e9917ae
commit 6569dade33
17 changed files with 4225 additions and 4083 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
cognee-mcp/pyproject.toml
View file

@ -14,6 +14,7 @@ dependencies = [
"mcp>=1.12.0,<2.0.0",
"uv>=0.6.3,<1.0.0",
"httpx>=0.27.0,<1.0.0",
"pydantic>=2.12.1",
]
authors = [
@ -32,6 +33,8 @@ packages = ["src"]
[dependency-groups]
dev = [
"debugpy>=1.8.12,<2.0.0",
"mypy>=1.18.2",
"ruff>=0.14.2",
]
[tool.hatch.metadata]

5
cognee-mcp/src/__init__.py
View file

@ -1,7 +1,4 @@
try:
from .server import main as server_main
except ImportError:
from server import main as server_main
from .server import main as server_main
import warnings
import sys

3
cognee-mcp/src/clients/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .cognee_client import CogneeClient
__all__ = ["CogneeClient"]

0
cognee-mcp/src/cognee_client.py → cognee-mcp/src/clients/cognee_client.py
View file

7
cognee-mcp/src/models/MCPServerCommandConfig.py Normal file
View file

@ -0,0 +1,7 @@
from typing import List
from pydantic import BaseModel
class MCPServerCommandConfig(BaseModel):
command: str
args: List[str]

7
cognee-mcp/src/models/MCPServerUrlConfig.py Normal file
View file

@ -0,0 +1,7 @@
from typing import Literal
from pydantic import BaseModel
class MCPServerUrlConfig(BaseModel):
url: str
transport: Literal["http", "sse"]

7
cognee-mcp/src/models/__init__.py Normal file
View file

@ -0,0 +1,7 @@
from .MCPServerCommandConfig import MCPServerCommandConfig
from .MCPServerUrlConfig import MCPServerUrlConfig
__all__ = [
"MCPServerCommandConfig",
"MCPServerUrlConfig",
]

332
cognee-mcp/src/server.py
View file

@ -9,39 +9,33 @@ Available Tools:
- list_mcp_servers: View all stored servers
- clear_registry: Clear the registry
"""
import json
import os
import sys
import argparse
import asyncio
import subprocess
from pathlib import Path
from typing import Optional
from typing import Any
from cognee.shared.logging_utils import get_logger, setup_logging, get_log_file_location
import importlib.util
from contextlib import redirect_stdout
import mcp.types as types
from cognee.shared.logging_utils import get_logger, setup_logging
from mcp.server import FastMCP
from cognee.modules.storage.utils import JSONEncoder
from starlette.responses import JSONResponse
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
import uvicorn
from fastmcp import Client as MCPClient
try:
from .cognee_client import CogneeClient
except ImportError:
from cognee_client import CogneeClient
from src.utils.context import set_cognee_client
from src.clients import CogneeClient
from src.tools import clear_registry, find_mcp_server, list_mcp_servers, remember_mcp_server
mcp = FastMCP("MCP-Registry")
mcp = FastMCP[Any]("MCP-Registry")
logger = get_logger()
cognee_client: Optional[CogneeClient] = None
mcp.tool()(clear_registry)
mcp.tool()(find_mcp_server)
mcp.tool()(list_mcp_servers)
mcp.tool()(remember_mcp_server)
async def run_sse_with_cors():
@ -91,307 +85,6 @@ async def health_check(request):
return JSONResponse({"status": "ok"})
@mcp.tool()
async def remember_mcp_server(
server_name: str,
description: str,
capabilities: str,
url: str = None,
command: str = None,
args: str = None,
installation: str = None,
repository_url: str = None,
documentation_url: str = None,
tags: str = None,
) -> list:
"""
Store information about an MCP server for future discovery.
Use this when you learn about an MCP server and want to remember its details,
capabilities, and connection information for later retrieval.
Parameters
----------
server_name : str
The name of the MCP server (e.g., "filesystem", "brave-search", "puppeteer")
description : str
A comprehensive description of what the MCP server does, its main features,
and what problems it solves. Be detailed to improve search accuracy.
capabilities : str
What the server can do. List specific capabilities, use cases, and features.
Examples: "file operations, directory listing, search files"
or "web search, real-time information, news retrieval"
url : str, optional
Server URL for HTTP/SSE connections (e.g., "http://localhost:8124/sse")
command : str, optional
Command to run for stdio-based servers (e.g., "python", "npx")
args : str, optional
Command arguments as a comma-separated string (e.g., "src/server.py, --transport, stdio")
installation : str, optional
How to install and configure the server (commands, config examples, etc.)
repository_url : str, optional
GitHub or source code repository URL
documentation_url : str, optional
Link to documentation or README
tags : str, optional
Comma-separated tags for categorization (e.g., "filesystem, tools, dev-tools")
Returns
-------
list
A TextContent object confirming the server was stored successfully.
Examples
--------
```python
await remember_mcp_server(
server_name="filesystem",
description="Provides comprehensive file system operations including reading, writing, editing files and directories",
capabilities="read files, write files, search files, list directories, create directories, move files",
installation="npx -y @modelcontextprotocol/server-filesystem /path/to/allowed/directory",
repository_url="https://github.com/modelcontextprotocol/servers",
tags="filesystem, tools, files"
)
```
"""
async def store_mcp_server() -> None:
with redirect_stdout(sys.stderr):
logger.info(f"Storing MCP server: {server_name}")
# Create structured content about the MCP server
server_content = f"""
# MCP Server: {server_name}
## Description
{description}
## Capabilities
{capabilities}
## Connection
URL: {url or 'Not provided'}
Command: {command or 'Not provided'}
Args: {args or 'Not provided'}
## Installation
{installation or 'Not provided'}
## Repository
{repository_url or 'Not provided'}
## Documentation
{documentation_url or 'Not provided'}
## Tags
{tags or 'Not provided'}
"""
try:
# Add to knowledge graph with special node set
await cognee_client.add(server_content, node_set=["mcp_servers", server_name])
# Process into knowledge graph
await cognee_client.cognify()
logger.info(f"Successfully stored MCP server: {server_name}")
except Exception as e:
logger.error(f"Failed to store MCP server {server_name}: {str(e)}")
raise ValueError(f"Failed to store MCP server: {str(e)}")
# Run as background task
asyncio.create_task(store_mcp_server())
log_file = get_log_file_location()
return [
types.TextContent(
type="text",
text=(
f"✅ Started storing MCP server '{server_name}' in background.\n"
f"Check logs at {log_file} for completion status.\n"
f"Use 'find_mcp_server' to search for it once processing is complete."
),
)
]
@mcp.tool()
async def find_mcp_server(requirements: str, max_results: int = 5) -> list:
"""
Search for MCP servers that match your requirements.
Searches through stored MCP servers and returns the ones that best match your needs
based on their capabilities and descriptions.
Parameters
----------
requirements : str
Describe what you need the MCP server to do. Be specific about your use case.
Examples:
- "I need to read and write files"
- "I want to search the web for real-time information"
- "I need to control a browser and take screenshots"
- "I want to execute code in a sandbox"
max_results : int, optional
Maximum number of MCP servers to return (default: 5)
Returns
-------
list
A TextContent object with detailed information about matching MCP servers,
including their names, descriptions, capabilities, and installation instructions.
Examples
--------
```python
# Find a server for file operations
await find_mcp_server("I need to read and modify files in my project")
# Find a server for web search
await find_mcp_server("I want to search the internet for current information")
```
"""
with redirect_stdout(sys.stderr):
try:
logger.info(f"Searching for MCP servers matching: {requirements}")
# Search using GRAPH_COMPLETION for intelligent matching
search_results = await cognee_client.search(
query_text=f"Find MCP servers that can: {requirements}. Include their capabilities, installation instructions, and documentation.",
query_type="GRAPH_COMPLETION",
)
# Format the results
if cognee_client.use_api:
if isinstance(search_results, str):
result_text = search_results
elif isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = json.dumps(search_results, cls=JSONEncoder)
else:
if isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = str(search_results)
logger.info("MCP server search completed")
return [
types.TextContent(
type="text",
text=f"🔍 MCP Servers matching your requirements:\n\n{result_text}",
)
]
except Exception as e:
error_msg = f"❌ Failed to search for MCP servers: {str(e)}"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]
@mcp.tool()
async def list_mcp_servers() -> list:
"""
List all MCP servers stored in your personal registry.
Returns detailed information about MCP servers you've previously remembered,
including their connection details (URL/command), capabilities, and documentation.
Use this information to connect to the servers with your MCP client.
Returns
-------
list
A list of all MCP servers in the registry with their connection information.
"""
with redirect_stdout(sys.stderr):
try:
logger.info("Listing all MCP servers")
# Search for all MCP servers with connection details
search_results = await cognee_client.search(
query_text="List all MCP servers with their names, descriptions, capabilities, connection information (URL, command, args), installation instructions, and documentation links",
query_type="GRAPH_COMPLETION",
)
# Format the results
if cognee_client.use_api:
if isinstance(search_results, str):
result_text = search_results
elif isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = json.dumps(search_results, cls=JSONEncoder)
else:
if isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = str(search_results)
output_text = f"📋 MCP Servers in Registry:\n\n{result_text}\n\n"
output_text += "💡 Use the connection information above (URL or command/args) to configure your MCP client."
logger.info("MCP server listing completed")
return [
types.TextContent(
type="text",
text=output_text,
)
]
except Exception as e:
error_msg = f"❌ Failed to list MCP servers: {str(e)}"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]
@mcp.tool()
async def clear_registry() -> list:
"""
Clear all stored MCP server information from the registry.
Removes all MCP servers you've stored. Use with caution as this cannot be undone.
Returns
-------
list
A TextContent object confirming the registry was cleared.
"""
with redirect_stdout(sys.stderr):
try:
await cognee_client.prune_data()
await cognee_client.prune_system(metadata=True)
logger.info("MCP server registry cleared")
return [
types.TextContent(
type="text",
text="✅ MCP server registry has been cleared. All stored servers removed.",
)
]
except NotImplementedError:
error_msg = "❌ Clear operation is not available in API mode"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]
except Exception as e:
error_msg = f"❌ Failed to clear registry: {str(e)}"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]
async def main():
global cognee_client
@ -456,6 +149,7 @@ async def main():
# Initialize the global CogneeClient
cognee_client = CogneeClient(api_url=args.api_url, api_token=args.api_token)
set_cognee_client(cognee_client)
mcp.settings.host = args.host
mcp.settings.port = args.port

11
cognee-mcp/src/tools/__init__.py Normal file
View file

@ -0,0 +1,11 @@
from .clear_registry import clear_registry
from .find_mcp_server import find_mcp_server
from .list_mcp_servers import list_mcp_servers
from .remember_mcp_server import remember_mcp_server
__all__ = [
"clear_registry",
"find_mcp_server",
"list_mcp_servers",
"remember_mcp_server",
]

40
cognee-mcp/src/tools/clear_registry.py Normal file
View file

@ -0,0 +1,40 @@
from contextlib import redirect_stdout
import sys
from src.utils.context import cognee_client
from cognee.shared.logging_utils import get_logger
import mcp.types as types
logger = get_logger()
async def clear_registry() -> list:
"""
Clear all stored MCP server information from the registry.
Removes all MCP servers you've stored. Use with caution as this cannot be undone.
Returns
-------
list
A TextContent object confirming the registry was cleared.
"""
with redirect_stdout(sys.stderr):
try:
await cognee_client.prune_data()
await cognee_client.prune_system(metadata=True)
logger.info("MCP server registry cleared")
return [
types.TextContent(
type="text",
text="✅ MCP server registry has been cleared. All stored servers removed.",
)
]
except NotImplementedError:
error_msg = "❌ Clear operation is not available in API mode"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]
except Exception as e:
error_msg = f"❌ Failed to clear registry: {str(e)}"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]

0
cognee-mcp/src/tools/context.py Normal file
View file

85
cognee-mcp/src/tools/find_mcp_server.py Normal file
View file

@ -0,0 +1,85 @@
from contextlib import redirect_stdout
import json
import sys
from cognee.shared.logging_utils import get_logger
import mcp.types as types
from src.utils.context import cognee_client
logger = get_logger()
async def find_mcp_server(requirements: str, max_results: int = 5) -> list:
"""
Search for MCP servers that match your requirements.
Searches through stored MCP servers and returns the ones that best match your needs
based on their capabilities and descriptions.
Parameters
----------
requirements : str
Describe what you need the MCP server to do. Be specific about your use case.
Examples:
- "I need to read and write files"
- "I want to search the web for real-time information"
- "I need to control a browser and take screenshots"
- "I want to execute code in a sandbox"
max_results : int, optional
Maximum number of MCP servers to return (default: 5)
Returns
-------
list
A TextContent object with detailed information about matching MCP servers,
including their names, descriptions, capabilities, and installation instructions.
Examples
--------
```python
# Find a server for file operations
await find_mcp_server("I need to read and modify files in my project")
# Find a server for web search
await find_mcp_server("I want to search the internet for current information")
```
"""
with redirect_stdout(sys.stderr):
try:
logger.info(f"Searching for MCP servers matching: {requirements}")
# Search using GRAPH_COMPLETION for intelligent matching
search_results = await cognee_client.search(
query_text=f"Find MCP servers that can: {requirements}. Include their capabilities, installation instructions, and documentation.",
query_type="GRAPH_COMPLETION",
)
# Format the results
if cognee_client.use_api:
if isinstance(search_results, str):
result_text = search_results
elif isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = json.dumps(search_results, cls=json.JSONEncoder)
else:
if isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = str(search_results)
logger.info("MCP server search completed")
return [
types.TextContent(
type="text",
text=f"🔍 MCP Servers matching your requirements:\n\n{result_text}",
)
]
except Exception as e:
error_msg = f"❌ Failed to search for MCP servers: {str(e)}"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]

65
cognee-mcp/src/tools/list_mcp_servers.py Normal file
View file

@ -0,0 +1,65 @@
from contextlib import redirect_stdout
import json
import sys
from cognee.shared.logging_utils import get_logger
import mcp.types as types
from src.utils.context import cognee_client
logger = get_logger()
async def list_mcp_servers() -> list:
"""
List all MCP servers stored in your personal registry.
Returns detailed information about MCP servers you've previously remembered,
including their connection details (URL/command), capabilities, and documentation.
Use this information to connect to the servers with your MCP client.
Returns
-------
list
A list of all MCP servers in the registry with their connection information.
"""
with redirect_stdout(sys.stderr):
try:
logger.info("Listing all MCP servers")
# Search for all MCP servers with connection details
search_results = await cognee_client.search(
query_text="List all MCP servers with their names, descriptions, capabilities, connection information (URL, command, args), installation instructions, and documentation links",
query_type="GRAPH_COMPLETION",
)
# Format the results
if cognee_client.use_api:
if isinstance(search_results, str):
result_text = search_results
elif isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = json.dumps(search_results, cls=json.JSONEncoder)
else:
if isinstance(search_results, list) and len(search_results) > 0:
result_text = str(search_results[0])
else:
result_text = str(search_results)
output_text = f"📋 MCP Servers in Registry:\n\n{result_text}\n\n"
output_text += "💡 Use the connection information above (URL or command/args) to configure your MCP client."
logger.info("MCP server listing completed")
return [
types.TextContent(
type="text",
text=output_text,
)
]
except Exception as e:
error_msg = f"❌ Failed to list MCP servers: {str(e)}"
logger.error(error_msg)
return [types.TextContent(type="text", text=error_msg)]

141
cognee-mcp/src/tools/remember_mcp_server.py Normal file
View file

@ -0,0 +1,141 @@
import asyncio
from contextlib import redirect_stdout
import sys
from cognee.shared.logging_utils import get_logger, get_log_file_location
import mcp.types as types
from src.utils.context import cognee_client
logger = get_logger()
async def remember_mcp_server(
server_name: str,
description: str,
capabilities: str,
url: str = None,
command: str = None,
args: str = None,
installation: str = None,
repository_url: str = None,
documentation_url: str = None,
tags: str = None,
) -> list:
"""
Store information about an MCP server for future discovery.
Use this when you learn about an MCP server and want to remember its details,
capabilities, and connection information for later retrieval.
Parameters
----------
server_name : str
The name of the MCP server (e.g., "filesystem", "brave-search", "puppeteer")
description : str
A comprehensive description of what the MCP server does, its main features,
and what problems it solves. Be detailed to improve search accuracy.
capabilities : str
What the server can do. List specific capabilities, use cases, and features.
Examples: "file operations, directory listing, search files"
or "web search, real-time information, news retrieval"
url : str, optional
Server URL for HTTP/SSE connections (e.g., "http://localhost:8124/sse")
command : str, optional
Command to run for stdio-based servers (e.g., "python", "npx")
args : str, optional
Command arguments as a comma-separated string (e.g., "src/server.py, --transport, stdio")
installation : str, optional
How to install and configure the server (commands, config examples, etc.)
repository_url : str, optional
GitHub or source code repository URL
documentation_url : str, optional
Link to documentation or README
tags : str, optional
Comma-separated tags for categorization (e.g., "filesystem, tools, dev-tools")
Returns
-------
list
A TextContent object confirming the server was stored successfully.
Examples
--------
```python
await remember_mcp_server(
server_name="filesystem",
description="Provides comprehensive file system operations including reading, writing, editing files and directories",
capabilities="read files, write files, search files, list directories, create directories, move files",
installation="npx -y @modelcontextprotocol/server-filesystem /path/to/allowed/directory",
repository_url="https://github.com/modelcontextprotocol/servers",
tags="filesystem, tools, files"
)
```
"""
async def store_mcp_server() -> None:
with redirect_stdout(sys.stderr):
logger.info(f"Storing MCP server: {server_name}")
# Create structured content about the MCP server
server_content = f"""
# MCP Server: {server_name}
## Description
{description}
## Capabilities
{capabilities}
## Connection
URL: {url or "Not provided"}
Command: {command or "Not provided"}
Args: {args or "Not provided"}
## Installation
{installation or "Not provided"}
## Repository
{repository_url or "Not provided"}
## Documentation
{documentation_url or "Not provided"}
## Tags
{tags or "Not provided"}
"""
try:
# Add to knowledge graph with special node set
await cognee_client.add(server_content, node_set=["mcp_servers", server_name])
# Process into knowledge graph
await cognee_client.cognify()
logger.info(f"Successfully stored MCP server: {server_name}")
except Exception as e:
logger.error(f"Failed to store MCP server {server_name}: {str(e)}")
raise ValueError(f"Failed to store MCP server: {str(e)}")
# Run as background task
asyncio.create_task(store_mcp_server())
log_file = get_log_file_location()
return [
types.TextContent(
type="text",
text=(
f"✅ Started storing MCP server '{server_name}' in background.\n"
f"Check logs at {log_file} for completion status.\n"
f"Use 'find_mcp_server' to search for it once processing is complete."
),
)
]

3
cognee-mcp/src/utils/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .context import cognee_client, set_cognee_client
__all__ = ["cognee_client", "set_cognee_client"]

12
cognee-mcp/src/utils/context.py Normal file
View file

@ -0,0 +1,12 @@
from typing import Optional, TYPE_CHECKING
if TYPE_CHECKING:
from src.clients.cognee_client import CogneeClient
cognee_client: Optional["CogneeClient"] = None
def set_cognee_client(client: "CogneeClient") -> None:
"""Set the global cognee client instance."""
global cognee_client
cognee_client = client

7587
cognee-mcp/uv.lock generated
View file

File diff suppressed because it is too large Load diff