fix imports

cognee-mcp/pyproject.toml

@@ -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]

cognee-mcp/src/__init__.py

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

cognee-mcp/src/clients/__init__.py

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

cognee-mcp/src/models/MCPServerCommandConfig.py

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

cognee-mcp/src/models/MCPServerUrlConfig.py

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

cognee-mcp/src/models/__init__.py

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

cognee-mcp/src/tools/__init__.py

@@ -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",
+]

cognee-mcp/src/tools/clear_registry.py

@@ -0,0 +1,40 @@
+from contextlib import redirect_stdout
+import sys
+from src.utils import context
+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 context.cognee_client.prune_data()
+            await context.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)]

cognee-mcp/src/tools/find_mcp_server.py

@@ -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 import context
+
+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 context.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 context.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)]

cognee-mcp/src/tools/list_mcp_servers.py

@@ -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 import context
+
+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 context.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 context.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)]

cognee-mcp/src/tools/remember_mcp_server.py

@@ -0,0 +1,143 @@
+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 import context
+
+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 context.cognee_client.add(
+                    server_content, node_set=["mcp_servers", server_name]
+                )
+
+                # Process into knowledge graph
+                await context.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."
+            ),
+        )
+    ]

cognee-mcp/src/utils/__init__.py

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

cognee-mcp/src/utils/context.py

@@ -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