feat: adds cypher search to retrievers module (#648)

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

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


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## 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.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
hajdul88 2025-03-19 15:01:40 +01:00 committed by GitHub
parent 24e0805f50
commit 1c65682242
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 76 additions and 0 deletions

View file

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

View file

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

View file

@ -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)

View file

@ -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)

View file

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