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,
|
GraphSummaryCompletionRetriever,
|
||||||
)
|
)
|
||||||
from cognee.modules.retrieval.code_retriever import CodeRetriever
|
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.search.types import SearchType
|
||||||
from cognee.modules.storage.utils import JSONEncoder
|
from cognee.modules.storage.utils import JSONEncoder
|
||||||
from cognee.modules.users.models import User
|
from cognee.modules.users.models import User
|
||||||
|
|
@ -65,6 +66,7 @@ async def specific_search(
|
||||||
system_prompt_path=system_prompt_path
|
system_prompt_path=system_prompt_path
|
||||||
).get_completion,
|
).get_completion,
|
||||||
SearchType.CODE: CodeRetriever().get_completion,
|
SearchType.CODE: CodeRetriever().get_completion,
|
||||||
|
SearchType.CYPHER: CypherSearchRetriever().get_completion,
|
||||||
}
|
}
|
||||||
|
|
||||||
search_task = search_tasks.get(query_type)
|
search_task = search_tasks.get(query_type)
|
||||||
|
|
|
||||||
|
|
@ -9,3 +9,4 @@ class SearchType(Enum):
|
||||||
GRAPH_COMPLETION = "GRAPH_COMPLETION"
|
GRAPH_COMPLETION = "GRAPH_COMPLETION"
|
||||||
GRAPH_SUMMARY_COMPLETION = "GRAPH_SUMMARY_COMPLETION"
|
GRAPH_SUMMARY_COMPLETION = "GRAPH_SUMMARY_COMPLETION"
|
||||||
CODE = "CODE"
|
CODE = "CODE"
|
||||||
|
CYPHER = "CYPHER"
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue