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:
parent
24e0805f50
commit
1c65682242
5 changed files with 76 additions and 0 deletions
44
cognee/modules/retrieval/cypher_search_retriever.py
Normal file
44
cognee/modules/retrieval/cypher_search_retriever.py
Normal 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
|
||||
7
cognee/modules/retrieval/exceptions/__init__.py
Normal file
7
cognee/modules/retrieval/exceptions/__init__.py
Normal 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
|
||||
22
cognee/modules/retrieval/exceptions/exceptions.py
Normal file
22
cognee/modules/retrieval/exceptions/exceptions.py
Normal 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)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -9,3 +9,4 @@ class SearchType(Enum):
|
|||
GRAPH_COMPLETION = "GRAPH_COMPLETION"
|
||||
GRAPH_SUMMARY_COMPLETION = "GRAPH_SUMMARY_COMPLETION"
|
||||
CODE = "CODE"
|
||||
CYPHER = "CYPHER"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue