From 1c65682242e9ae8226dbb9d8e726f6d3304e84ae Mon Sep 17 00:00:00 2001 From: hajdul88 <52442977+hajdul88@users.noreply.github.com> Date: Wed, 19 Mar 2025 15:01:40 +0100 Subject: [PATCH] feat: adds cypher search to retrievers module (#648) ## Description Exposes the query method of the adapter in the search interface for Kuzu and Neo4j (cypher compatible adapters) ## 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 ## Summary by CodeRabbit - **New Features** - Introduced a new cypher-based search option that expands the app's search functionality. - Enabled asynchronous processing for advanced query execution. - Enhanced error messaging for unsupported search types and query execution issues. - Added a new enumeration value for `CYPHER` to support the new search type. --- .../retrieval/cypher_search_retriever.py | 44 +++++++++++++++++++ .../modules/retrieval/exceptions/__init__.py | 7 +++ .../retrieval/exceptions/exceptions.py | 22 ++++++++++ cognee/modules/search/methods/search.py | 2 + cognee/modules/search/types/SearchType.py | 1 + 5 files changed, 76 insertions(+) create mode 100644 cognee/modules/retrieval/cypher_search_retriever.py create mode 100644 cognee/modules/retrieval/exceptions/__init__.py create mode 100644 cognee/modules/retrieval/exceptions/exceptions.py diff --git a/cognee/modules/retrieval/cypher_search_retriever.py b/cognee/modules/retrieval/cypher_search_retriever.py new file mode 100644 index 000000000..b07a9b63a --- /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 000000000..1b98cddcd --- /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 000000000..1b7c34251 --- /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 c0056b048..593a5452b 100644 --- a/cognee/modules/search/methods/search.py +++ b/cognee/modules/search/methods/search.py @@ -12,6 +12,7 @@ from cognee.modules.retrieval.graph_summary_completion_retriever import ( 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 fbab01a46..e41eadc00 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"