From 6898e8f766d976cc36c4596a51ae3024d81274aa Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Wed, 2 Apr 2025 06:38:17 +0200 Subject: [PATCH] Fix codify mcp (#696) ## Description - Redirect all Cognee output to stderr for MCP ( as stdout is used to communicate between MCP Client and server ) - Add test for CODE search type - Resolve missing optional GUI dependency ## 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. --- .github/workflows/test_code_graph_example.yml | 2 +- cognee-mcp/src/server.py | 153 ++++++++++-------- examples/python/code_graph_example.py | 9 ++ poetry.lock | 76 ++++++++- pyproject.toml | 1 + 5 files changed, 167 insertions(+), 74 deletions(-) diff --git a/.github/workflows/test_code_graph_example.yml b/.github/workflows/test_code_graph_example.yml index 47c36c467..3a017de51 100644 --- a/.github/workflows/test_code_graph_example.yml +++ b/.github/workflows/test_code_graph_example.yml @@ -15,7 +15,7 @@ jobs: uses: ./.github/workflows/reusable_python_example.yml with: example-location: ./examples/python/code_graph_example.py - arguments: "--repo_path ./evals" + arguments: "--repo_path ./cognee/tasks/graph" secrets: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }} diff --git a/cognee-mcp/src/server.py b/cognee-mcp/src/server.py index 5c25c24f7..491bd699d 100755 --- a/cognee-mcp/src/server.py +++ b/cognee-mcp/src/server.py @@ -1,10 +1,11 @@ import asyncio import json import os +import sys import cognee from cognee.shared.logging_utils import get_logger, get_log_file_location import importlib.util -from contextlib import redirect_stderr, redirect_stdout +from contextlib import redirect_stdout # from PIL import Image as PILImage import mcp.types as types @@ -90,56 +91,55 @@ async def list_tools() -> list[types.Tool]: @mcp.call_tool() async def call_tools(name: str, arguments: dict) -> list[types.TextContent]: try: - with open(os.devnull, "w") as fnull: - with redirect_stdout(fnull), redirect_stderr(fnull): - log_file = get_log_file_location() + # NOTE: MCP uses stdout to communicate, we must redirect all output + # going to stdout ( like the print function ) to stderr. + with redirect_stdout(sys.stderr): + log_file = get_log_file_location() - if name == "cognify": - asyncio.create_task( - cognify( - text=arguments["text"], - graph_model_file=arguments.get("graph_model_file"), - graph_model_name=arguments.get("graph_model_name"), - ) + if name == "cognify": + asyncio.create_task( + cognify( + text=arguments["text"], + graph_model_file=arguments.get("graph_model_file"), + graph_model_name=arguments.get("graph_model_name"), ) + ) - text = ( - f"Background process launched due to MCP timeout limitations.\n" - f"Average completion time is around 4 minutes.\n" - f"For current cognify status you can check the log file at: {log_file}" + text = ( + f"Background process launched due to MCP timeout limitations.\n" + f"Average completion time is around 4 minutes.\n" + f"For current cognify status you can check the log file at: {log_file}" + ) + + return [ + types.TextContent( + type="text", + text=text, ) + ] + if name == "codify": + asyncio.create_task(codify(arguments.get("repo_path"))) - return [ - types.TextContent( - type="text", - text=text, - ) - ] - if name == "codify": - asyncio.create_task(codify(arguments.get("repo_path"))) + text = ( + f"Background process launched due to MCP timeout limitations.\n" + f"Average completion time is around 4 minutes.\n" + f"For current codify status you can check the log file at: {log_file}" + ) - text = ( - f"Background process launched due to MCP timeout limitations.\n" - f"Average completion time is around 4 minutes.\n" - f"For current codify status you can check the log file at: {log_file}" + return [ + types.TextContent( + type="text", + text=text, ) + ] + elif name == "search": + search_results = await search(arguments["search_query"], arguments["search_type"]) - return [ - types.TextContent( - type="text", - text=text, - ) - ] - elif name == "search": - search_results = await search( - arguments["search_query"], arguments["search_type"] - ) + return [types.TextContent(type="text", text=search_results)] + elif name == "prune": + await prune() - return [types.TextContent(type="text", text=search_results)] - elif name == "prune": - await prune() - - return [types.TextContent(type="text", text="Pruned")] + return [types.TextContent(type="text", text="Pruned")] except Exception as e: logger.error(f"Error calling tool '{name}': {str(e)}") return [types.TextContent(type="text", text=f"Error calling tool '{name}': {str(e)}")] @@ -147,45 +147,56 @@ async def call_tools(name: str, arguments: dict) -> list[types.TextContent]: async def cognify(text: str, graph_model_file: str = None, graph_model_name: str = None) -> str: """Build knowledge graph from the input text""" - logger.info("Cognify process starting.") - if graph_model_file and graph_model_name: - graph_model = load_class(graph_model_file, graph_model_name) - else: - graph_model = KnowledgeGraph + # NOTE: MCP uses stdout to communicate, we must redirect all output + # going to stdout ( like the print function ) to stderr. + # As cognify is an async background job the output had to be redirected again. + with redirect_stdout(sys.stderr): + logger.info("Cognify process starting.") + if graph_model_file and graph_model_name: + graph_model = load_class(graph_model_file, graph_model_name) + else: + graph_model = KnowledgeGraph - await cognee.add(text) + await cognee.add(text) - try: - await cognee.cognify(graph_model=graph_model) - logger.info("Cognify process finished.") - except Exception as e: - logger.error("Cognify process failed.") - raise ValueError(f"Failed to cognify: {str(e)}") + try: + await cognee.cognify(graph_model=graph_model) + logger.info("Cognify process finished.") + except Exception as e: + logger.error("Cognify process failed.") + raise ValueError(f"Failed to cognify: {str(e)}") async def codify(repo_path: str): - logger.info("Codify process starting.") - results = [] - async for result in run_code_graph_pipeline(repo_path, False): - results.append(result) - logger.info(result) - if all(results): - logger.info("Codify process finished succesfully.") - else: - logger.info("Codify process failed.") + # NOTE: MCP uses stdout to communicate, we must redirect all output + # going to stdout ( like the print function ) to stderr. + # As codify is an async background job the output had to be redirected again. + with redirect_stdout(sys.stderr): + logger.info("Codify process starting.") + results = [] + async for result in run_code_graph_pipeline(repo_path, False): + results.append(result) + logger.info(result) + if all(results): + logger.info("Codify process finished succesfully.") + else: + logger.info("Codify process failed.") async def search(search_query: str, search_type: str) -> str: """Search the knowledge graph""" - search_results = await cognee.search( - query_type=SearchType[search_type.upper()], query_text=search_query - ) + # NOTE: MCP uses stdout to communicate, we must redirect all output + # going to stdout ( like the print function ) to stderr. + with redirect_stdout(sys.stderr): + search_results = await cognee.search( + query_type=SearchType[search_type.upper()], query_text=search_query + ) - if search_type.upper() == "CODE": - return json.dumps(search_results, cls=JSONEncoder) - else: - results = retrieved_edges_to_string(search_results) - return results + if search_type.upper() == "CODE": + return json.dumps(search_results, cls=JSONEncoder) + else: + results = retrieved_edges_to_string(search_results) + return results async def prune(): diff --git a/examples/python/code_graph_example.py b/examples/python/code_graph_example.py index bffe95e1b..6a0745f01 100644 --- a/examples/python/code_graph_example.py +++ b/examples/python/code_graph_example.py @@ -1,5 +1,7 @@ import argparse import asyncio +import cognee +from cognee import SearchType from cognee.shared.logging_utils import get_logger, ERROR from cognee.api.v1.cognify.code_graph_pipeline import run_code_graph_pipeline @@ -10,6 +12,13 @@ async def main(repo_path, include_docs): async for run_status in run_code_graph_pipeline(repo_path, include_docs=include_docs): run_status = run_status + # Test CODE search + search_results = await cognee.search(query_type=SearchType.CODE, query_text="test") + assert len(search_results) != 0, "The search results list is empty." + print("\n\nSearch results are:\n") + for result in search_results: + print(f"{result}\n") + return run_status diff --git a/poetry.lock b/poetry.lock index 168599d5f..2582754e1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -7849,6 +7849,63 @@ files = [ [package.extras] dev = ["build", "flake8", "mypy", "pytest", "twine"] +[[package]] +name = "pyside6" +version = "6.8.3" +description = "Python bindings for the Qt cross-platform application and UI framework" +optional = true +python-versions = "<3.14,>=3.9" +groups = ["main"] +markers = "extra == \"gui\"" +files = [ + {file = "PySide6-6.8.3-cp39-abi3-macosx_12_0_universal2.whl", hash = "sha256:31f390c961b54067ae41360e5ea3b340ce0e0e5feadea2236c28226d3b37edcc"}, + {file = "PySide6-6.8.3-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8e53e2357bfbdee1fa86c48312bf637460a2c26d49e7af0b3fae2e179ccc7052"}, + {file = "PySide6-6.8.3-cp39-abi3-manylinux_2_39_aarch64.whl", hash = "sha256:5bf5153cab9484629315f57c56a9ad4b7d075b4dd275f828f7549abf712c590b"}, + {file = "PySide6-6.8.3-cp39-abi3-win_amd64.whl", hash = "sha256:722dc0061d8ef6dbb8c0b99864f21e83a5b49ece1ecb2d0b890840d969e1e461"}, +] + +[package.dependencies] +PySide6-Addons = "6.8.3" +PySide6-Essentials = "6.8.3" +shiboken6 = "6.8.3" + +[[package]] +name = "pyside6-addons" +version = "6.8.3" +description = "Python bindings for the Qt cross-platform application and UI framework (Addons)" +optional = true +python-versions = "<3.14,>=3.9" +groups = ["main"] +markers = "extra == \"gui\"" +files = [ + {file = "PySide6_Addons-6.8.3-cp39-abi3-macosx_12_0_universal2.whl", hash = "sha256:ea46649e40b9e6ab11a0da2da054d3914bff5607a5882885e9c3bc2eef200036"}, + {file = "PySide6_Addons-6.8.3-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:6983d3b01fad53637bad5360930d5923509c744cc39704f9c1190eb9934e33da"}, + {file = "PySide6_Addons-6.8.3-cp39-abi3-manylinux_2_39_aarch64.whl", hash = "sha256:7949a844a40ee10998eb2734e2c06c4c7182dfcd4c21cc4108a6b96655ebe59f"}, + {file = "PySide6_Addons-6.8.3-cp39-abi3-win_amd64.whl", hash = "sha256:67548f6db11f4e1b7e4b6efd9c3fc2e8d275188a7b2feac388961128572a6955"}, +] + +[package.dependencies] +PySide6-Essentials = "6.8.3" +shiboken6 = "6.8.3" + +[[package]] +name = "pyside6-essentials" +version = "6.8.3" +description = "Python bindings for the Qt cross-platform application and UI framework (Essentials)" +optional = true +python-versions = "<3.14,>=3.9" +groups = ["main"] +markers = "extra == \"gui\"" +files = [ + {file = "PySide6_Essentials-6.8.3-cp39-abi3-macosx_12_0_universal2.whl", hash = "sha256:aa56c135db924ecfaf50088baf32f737d28027419ca5fee67c0c7141b29184e3"}, + {file = "PySide6_Essentials-6.8.3-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:fd57fa0c886ef99b3844173322c0023ec77cc946a0c9a0cdfbc2ac5c511053c1"}, + {file = "PySide6_Essentials-6.8.3-cp39-abi3-manylinux_2_39_aarch64.whl", hash = "sha256:b4f4823f870b5bed477d6f7b6a3041839b859f70abfd703cf53208c73c2fe4cd"}, + {file = "PySide6_Essentials-6.8.3-cp39-abi3-win_amd64.whl", hash = "sha256:3c0fae5550aff69f2166f46476c36e0ef56ce73d84829eac4559770b0c034b07"}, +] + +[package.dependencies] +shiboken6 = "6.8.3" + [[package]] name = "pysocks" version = "1.7.1" @@ -9379,6 +9436,21 @@ files = [ {file = "shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de"}, ] +[[package]] +name = "shiboken6" +version = "6.8.3" +description = "Python/C++ bindings helper module" +optional = true +python-versions = "<3.14,>=3.9" +groups = ["main"] +markers = "extra == \"gui\"" +files = [ + {file = "shiboken6-6.8.3-cp39-abi3-macosx_12_0_universal2.whl", hash = "sha256:483efc7dd53c69147b8a8ade71f7619c79ffc683efcb1dc4f4cb6c40bb23d29b"}, + {file = "shiboken6-6.8.3-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:295a003466ca2cccf6660e2f2ceb5e6cef4af192a48a196a32d46b6f0c9ec5cb"}, + {file = "shiboken6-6.8.3-cp39-abi3-manylinux_2_39_aarch64.whl", hash = "sha256:2b1a41348102952d2a5fbf3630bddd4d44112e18058b5e4cf505e51f2812429d"}, + {file = "shiboken6-6.8.3-cp39-abi3-win_amd64.whl", hash = "sha256:bca3a94513ce9242f7d4bbdca902072a1631888e0aa3a8711a52cc5dbe93588f"}, +] + [[package]] name = "simplejson" version = "3.20.1" @@ -11251,7 +11323,7 @@ filesystem = ["botocore"] gemini = [] graphiti = ["graphiti-core"] groq = ["groq"] -gui = ["qasync"] +gui = ["pyside6", "qasync"] huggingface = ["transformers"] kuzu = ["kuzu"] langchain = ["langchain_text_splitters", "langsmith"] @@ -11269,4 +11341,4 @@ weaviate = ["weaviate-client"] [metadata] lock-version = "2.1" python-versions = ">=3.10,<=3.13" -content-hash = "4bda223028508503b326912854c60fa4a5f60349370d26f22dd997d0dec11e01" +content-hash = "25b759ffc908ce0b4df33344424d2043dd3126d944c6d2e9b24031bd24e1152b" diff --git a/pyproject.toml b/pyproject.toml index 704f16d74..1520a255c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,7 @@ gdown = {version = "^5.2.0", optional = true} qasync = {version = "^0.27.1", optional = true} graphiti-core = {version = "^0.7.0", optional = true} structlog = "^25.2.0" +pyside6 = {version = "^6.8.3", optional = true} [tool.poetry.extras]