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.
2025-04-02 06:38:17 +02:00 · 2025-04-02 06:38:17 +02:00 · 6898e8f766
commit 6898e8f766
parent be90fd30d6
5 changed files with 167 additions and 74 deletions
--- 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 }}
--- 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():
--- 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


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