From f5cd39c09d77067d791b6c8e3fc41933f8a22367 Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Mon, 7 Apr 2025 18:13:15 +0200 Subject: [PATCH] fix: Resolve failing test for RAG_COMPLETION, add RAG_COMPLETION to MCP (#706) ## Description Resolve failing test for RAG_COMPLETION, add RAG_COMPLETION to MCP ## 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. Co-authored-by: Boris --- cognee-mcp/src/server.py | 2 +- notebooks/graphrag_vs_rag.ipynb | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cognee-mcp/src/server.py b/cognee-mcp/src/server.py index 1d1b90143..4b658d1e3 100755 --- a/cognee-mcp/src/server.py +++ b/cognee-mcp/src/server.py @@ -194,7 +194,7 @@ async def search(search_query: str, search_type: str) -> str: if search_type.upper() == "CODE": return json.dumps(search_results, cls=JSONEncoder) - elif search_type.upper() == "GRAPH_COMPLETION" or search_type.upper() == "COMPLETION": + elif search_type.upper() == "GRAPH_COMPLETION" or search_type.upper() == "RAG_COMPLETION": return search_results[0] else: results = retrieved_edges_to_string(search_results) diff --git a/notebooks/graphrag_vs_rag.ipynb b/notebooks/graphrag_vs_rag.ipynb index 29d8db019..8cfb49529 100644 --- a/notebooks/graphrag_vs_rag.ipynb +++ b/notebooks/graphrag_vs_rag.ipynb @@ -160,7 +160,7 @@ "metadata": {}, "cell_type": "code", "source": [ - "search_results = await cognee.search(SearchType.RAG_COMPLETION, \"Tell me who are the people mentioned?\")\n", + "search_results = await cognee.search(SearchType.COMPLETION, \"Tell me who are the people mentioned?\")\n", "\n", "print(\"\\n\\nAnswer based on RAG:\\n\")\n", "for result in search_results:\n",