diff --git a/cognee/modules/retrieval/cypher_search_retriever.py b/cognee/modules/retrieval/cypher_search_retriever.py new file mode 100644 index 0000000000..b07a9b63ac --- /dev/null +++ b/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 diff --git a/cognee/modules/retrieval/exceptions/__init__.py b/cognee/modules/retrieval/exceptions/__init__.py new file mode 100644 index 0000000000..1b98cddcdd --- /dev/null +++ b/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 diff --git a/cognee/modules/retrieval/exceptions/exceptions.py b/cognee/modules/retrieval/exceptions/exceptions.py new file mode 100644 index 0000000000..1b7c342516 --- /dev/null +++ b/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) diff --git a/cognee/modules/search/methods/search.py b/cognee/modules/search/methods/search.py index c0056b0480..593a5452ba 100644 --- a/cognee/modules/search/methods/search.py +++ b/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: CypherSearchRetriever().get_completion, } search_task = search_tasks.get(query_type) diff --git a/cognee/modules/search/types/SearchType.py b/cognee/modules/search/types/SearchType.py index fbab01a466..e41eadc005 100644 --- a/cognee/modules/search/types/SearchType.py +++ b/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 = "CYPHER"