From 3dc2d1e647d15e9709ad5437a0c80b9ed34cfef2 Mon Sep 17 00:00:00 2001 From: vasilije Date: Sat, 14 Dec 2024 12:14:04 +0100 Subject: [PATCH] Add mcp to cognee --- cognee-mcp/mcpcognee/__init__.py | 14 ++++ cognee-mcp/mcpcognee/__main__.py | 3 + cognee-mcp/mcpcognee/server.py | 126 +++++++++++++++++++++++++++++++ cognee-mcp/pyproject.toml | 94 +++++++++++++++++++++++ 4 files changed, 237 insertions(+) create mode 100644 cognee-mcp/mcpcognee/__init__.py create mode 100644 cognee-mcp/mcpcognee/__main__.py create mode 100644 cognee-mcp/mcpcognee/server.py create mode 100644 cognee-mcp/pyproject.toml diff --git a/cognee-mcp/mcpcognee/__init__.py b/cognee-mcp/mcpcognee/__init__.py new file mode 100644 index 000000000..d353850d2 --- /dev/null +++ b/cognee-mcp/mcpcognee/__init__.py @@ -0,0 +1,14 @@ +import asyncio + +from . import server + + +def main(): + """Main entry point for the package.""" + asyncio.run(server.main()) + +# Optionally expose other important items at package level +__all__ = ['main', 'server'] + +if __name__ == "__main__": + main() diff --git a/cognee-mcp/mcpcognee/__main__.py b/cognee-mcp/mcpcognee/__main__.py new file mode 100644 index 000000000..5fe7fb5dd --- /dev/null +++ b/cognee-mcp/mcpcognee/__main__.py @@ -0,0 +1,3 @@ +from mcpcognee import main + +main() \ No newline at end of file diff --git a/cognee-mcp/mcpcognee/server.py b/cognee-mcp/mcpcognee/server.py new file mode 100644 index 000000000..ba2f85122 --- /dev/null +++ b/cognee-mcp/mcpcognee/server.py @@ -0,0 +1,126 @@ +import importlib.util +import os +from contextlib import redirect_stderr, redirect_stdout + +import cognee +import mcp.server.stdio +import mcp.types as types +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 pydantic import AnyUrl, BaseModel + +server = Server("mcpcognee") + + +def node_to_string(node): + keys_to_keep = ["chunk_index", "topological_rank", "cut_type", "id", "text"] + keyset = set(keys_to_keep) & node.keys() + return "Node(" + " ".join([key + ": " + str(node[key]) + "," for key in keyset]) + ")" + + +def retrieved_edges_to_string(search_results): + edge_strings = [] + for triplet in search_results: + node1, edge, node2 = triplet + relationship_type = edge["relationship_name"] + edge_str = f"{node_to_string(node1)} {relationship_type} {node_to_string(node2)}" + edge_strings.append(edge_str) + return "\n".join(edge_strings) + + +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 + + +@server.list_tools() +async def handle_list_tools() -> list[types.Tool]: + """ + List available tools. + Each tool specifies its arguments using JSON Schema validation. + """ + return [ + types.Tool( + name="Cognify_and_search", + description="Build knowledge graph from the input text and search in it.", + inputSchema={ + "type": "object", + "properties": { + "text": {"type": "string"}, + "search_query": {"type": "string"}, + "graph_model_file": {"type": "string"}, + "graph_model_name": {"type": "string"}, + }, + "required": ["text", "search_query"], + }, + ) + ] + + +@server.call_tool() +async def handle_call_tool( + name: str, arguments: dict | None +) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]: + """ + Handle tool execution requests. + Tools can modify server state and notify clients of changes. + """ + if name == "Cognify_and_search": + with open(os.devnull, "w") as fnull: + with redirect_stdout(fnull), redirect_stderr(fnull): + await cognee.prune.prune_data() + await cognee.prune.prune_system(metadata=True) + + if not arguments: + raise ValueError("Missing arguments") + + text = arguments.get("text") + search_query = arguments.get("search_query") + if ("graph_model_file" in arguments) and ("graph_model_name" in arguments): + model_file = arguments.get("graph_model_file") + model_name = arguments.get("graph_model_name") + graph_model = load_class(model_file, model_name) + else: + graph_model = KnowledgeGraph + + await cognee.add(text) + await cognee.cognify(graph_model=graph_model) + search_results = await cognee.search( + SearchType.INSIGHTS, query_text=search_query + ) + + results = retrieved_edges_to_string(search_results) + + return [ + types.TextContent( + type="text", + text=results, + ) + ] + else: + raise ValueError(f"Unknown tool: {name}") + + +async def main(): + # Run the server using stdin/stdout streams + async with mcp.server.stdio.stdio_server() as (read_stream, write_stream): + await server.run( + read_stream, + write_stream, + InitializationOptions( + server_name="mcpcognee", + server_version="0.1.0", + capabilities=server.get_capabilities( + notification_options=NotificationOptions(), + experimental_capabilities={}, + ), + ), + ) \ No newline at end of file diff --git a/cognee-mcp/pyproject.toml b/cognee-mcp/pyproject.toml new file mode 100644 index 000000000..4fcf86d68 --- /dev/null +++ b/cognee-mcp/pyproject.toml @@ -0,0 +1,94 @@ +[project] +name = "mcpcognee" +version = "0.1.0" +description = "A MCP server project" +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ + "mcp>=1.1.1", + "openai==1.52.0", + "pydantic==2.8.2", + "python-dotenv==1.0.1", + "fastapi>=0.109.2,<0.110.0", + "uvicorn==0.22.0", + "requests==2.32.3", + "aiohttp==3.10.10", + "typing_extensions==4.12.2", + "nest_asyncio==1.6.0", + "numpy==1.26.4", + "datasets==3.1.0", + "falkordb==1.0.9", # Optional + "boto3>=1.26.125,<2.0.0", + "botocore>=1.35.54,<2.0.0", + "gunicorn>=20.1.0,<21.0.0", + "sqlalchemy==2.0.35", + "instructor==1.5.2", + "networkx>=3.2.1,<4.0.0", + "aiosqlite>=0.20.0,<0.21.0", + "pandas==2.0.3", + "filetype>=1.2.0,<2.0.0", + "nltk>=3.8.1,<4.0.0", + "dlt[sqlalchemy]>=1.4.1,<2.0.0", + "aiofiles>=23.2.1,<24.0.0", + "qdrant-client>=1.9.0,<2.0.0", # Optional + "graphistry>=0.33.5,<0.34.0", + "tenacity>=8.4.1,<9.0.0", + "weaviate-client==4.6.7", # Optional + "scikit-learn>=1.5.0,<2.0.0", + "pypdf>=4.1.0,<5.0.0", + "neo4j>=5.20.0,<6.0.0", # Optional + "jinja2>=3.1.3,<4.0.0", + "matplotlib>=3.8.3,<4.0.0", + "tiktoken==0.7.0", + "langchain_text_splitters==0.3.2", # Optional + "langsmith==0.1.139", # Optional + "langdetect==1.0.9", + "posthog>=3.5.0,<4.0.0", # Optional + "lancedb==0.15.0", + "litellm==1.49.1", + "groq==0.8.0", # Optional + "langfuse>=2.32.0,<3.0.0", # Optional + "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 + "alembic>=1.13.3,<2.0.0", + "asyncpg==0.30.0", # Optional + "pgvector>=0.3.5,<0.4.0", # Optional + "psycopg2>=2.9.10,<3.0.0", # Optional + "llama-index-core>=0.11.22,<0.12.0", # Optional + "deepeval>=2.0.1,<3.0.0", # Optional + "transformers>=4.46.3,<5.0.0", + "pymilvus>=2.5.0,<3.0.0", # Optional + "unstructured[csv,doc,docx,epub,md,odt,org,ppt,pptx,rst,rtf,tsv,xlsx]>=0.16.10,<1.0.0", # Optional + "pytest>=7.4.0,<8.0.0", + "pytest-asyncio>=0.21.1,<0.22.0", + "coverage>=7.3.2,<8.0.0", + "mypy>=1.7.1,<2.0.0", + "deptry>=0.20.0,<0.21.0", + "debugpy==1.8.2", + "pylint>=3.0.3,<4.0.0", + "ruff>=0.2.2,<0.3.0", + "tweepy==4.14.0", + "gitpython>=3.1.43,<4.0.0", + "cognee", +] + +[[project.authors]] +name = "Rita Aleksziev" +email = "rita@topoteretes.com" + +[build-system] +requires = [ "hatchling",] +build-backend = "hatchling.build" + +[tool.uv.sources] +cognee = { path = "../../cognee" } + +[dependency-groups] +dev = [ + "cognee", +] + +[project.scripts] +mcpcognee = "mcpcognee:main"