From 690d028928e0c088c9f89920ec6e37e6409ec9f9 Mon Sep 17 00:00:00 2001
From: Hande <159312713+hande-k@users.noreply.github.com>
Date: Tue, 4 Feb 2025 19:32:04 +0100
Subject: [PATCH] fix: mcp CODE search error (#495)

<!-- .github/pull_request_template.md -->

## Description
<!-- Provide a clear description of the changes in this PR -->

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


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Enhanced search functionality now returns JSON formatted results for
code-based queries for clearer response presentation.

- **Refactor**
- Updated the search logic to differentiate result handling based on
query type, maintaining string output for non-code queries.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---
 cognee-mcp/src/server.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/cognee-mcp/src/server.py b/cognee-mcp/src/server.py
index ec0a6564b..093bcc632 100755
--- a/cognee-mcp/src/server.py
+++ b/cognee-mcp/src/server.py
@@ -1,5 +1,5 @@
+import json
 import os
-
 import cognee
 import logging
 import importlib.util
@@ -12,6 +12,7 @@ from mcp.server.models import InitializationOptions
 from cognee.api.v1.cognify.code_graph_pipeline import run_code_graph_pipeline
 from cognee.modules.search.types import SearchType
 from cognee.shared.data_models import KnowledgeGraph
+from cognee.modules.storage.utils import JSONEncoder
 
 mcp = Server("cognee")
 
@@ -143,9 +144,11 @@ async def search(search_query: str, search_type: str) -> str:
         query_type=SearchType[search_type.upper()], query_text=search_query
     )
 
-    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():