feat: adds cypher search to retrievers module

cognee/modules/retrieval/cypher_search_retriever.py

@@ -0,0 +1,44 @@
+from typing import Any, Optional
+import logging
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter
+from cognee.modules.retrieval.base_retriever import BaseRetriever
+from cognee.modules.retrieval.utils.completion import generate_completion
+from cognee.modules.retrieval.exceptions import SearchTypeNotSupported, CypherSearchError
+
+logger = logging.getLogger("CypherSearchRetriever")
+
+
+class CypherSearchRetriever(BaseRetriever):
+    """Retriever for handling cypher-based search"""
+
+    def __init__(
+        self,
+        user_prompt_path: str = "context_for_question.txt",
+        system_prompt_path: str = "answer_simple_question.txt",
+    ):
+        """Initialize retriever with optional custom prompt paths."""
+        self.user_prompt_path = user_prompt_path
+        self.system_prompt_path = system_prompt_path
+
+    async def get_context(self, query: str) -> Any:
+        """Retrieves relevant context using a cypher query."""
+        try:
+            graph_engine = await get_graph_engine()
+
+            if isinstance(graph_engine, NetworkXAdapter):
+                raise SearchTypeNotSupported(
+                    "CYPHER search type not supported for NetworkXAdapter."
+                )
+
+            result = await graph_engine.query(query)
+        except Exception as e:
+            logger.error("Failed to execture cypher search retrieval: %s", str(e))
+            raise CypherSearchError() from e
+        return result
+
+    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
+        """Returns the graph connections context."""
+        if context is None:
+            context = await self.get_context(query)
+        return context

cognee/modules/retrieval/exceptions/__init__.py

@@ -0,0 +1,7 @@
+"""
+Custom exceptions for the Cognee API.
+
+This module defines a set of exceptions for handling various data errors
+"""
+
+from .exceptions import SearchTypeNotSupported, CypherSearchError

cognee/modules/retrieval/exceptions/exceptions.py

@@ -0,0 +1,22 @@
+from cognee.exceptions import CogneeApiError
+from fastapi import status
+
+
+class SearchTypeNotSupported(CogneeApiError):
+    def __init__(
+        self,
+        message: str = "CYPHER search type not supported by the adapter.",
+        name: str = "SearchTypeNotSupported",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)
+
+
+class CypherSearchError(CogneeApiError):
+    def __init__(
+        self,
+        message: str = "An error occurred during the execution of the Cypher query.",
+        name: str = "CypherSearchError",
+        status_code: int = status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)

cognee/modules/search/methods/search.py

@@ -12,6 +12,7 @@
     GraphSummaryCompletionRetriever,
 )
 from cognee.modules.retrieval.code_retriever import CodeRetriever
+from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
 from cognee.modules.search.types import SearchType
 from cognee.modules.storage.utils import JSONEncoder
 from cognee.modules.users.models import User
@@ -65,6 +66,7 @@ async def specific_search(
             system_prompt_path=system_prompt_path
         ).get_completion,
         SearchType.CODE: CodeRetriever().get_completion,
+        SearchType.CYPHER_SEARCH: CypherSearchRetriever().get_completion,
     }
 
     search_task = search_tasks.get(query_type)

cognee/modules/search/types/SearchType.py

@@ -9,3 +9,4 @@ class SearchType(Enum):
     GRAPH_COMPLETION = "GRAPH_COMPLETION"
     GRAPH_SUMMARY_COMPLETION = "GRAPH_SUMMARY_COMPLETION"
     CODE = "CODE"
+    CYPHER_SEARCH = "CYPHER_SEARCH"