a944871942
Adds experimental support for Apache TinkerPop Gremlin as an alternative
query language for AWS Neptune Database, alongside the existing openCypher
support. This enables users to choose their preferred query language and
opens the door for future support of other Gremlin-compatible databases.
- QueryLanguage enum (CYPHER, GREMLIN) for explicit language selection
- Dual-mode NeptuneDriver supporting both Cypher and Gremlin
- Gremlin query generation functions for common graph operations
- Graceful degradation when gremlinpython is not installed
- 100% backward compatible (defaults to CYPHER)
- graphiti_core/driver/driver.py: Added QueryLanguage enum
- graphiti_core/driver/neptune_driver.py: Dual client initialization
and query routing based on language selection
- graphiti_core/graph_queries.py: 9 new Gremlin query generation functions
- graphiti_core/utils/maintenance/graph_data_operations.py: Updated
clear_data() to support both query languages
- tests/test_neptune_gremlin_int.py: Comprehensive integration tests
- examples/quickstart/quickstart_neptune_gremlin.py: Usage example
- examples/quickstart/README.md: Updated with Gremlin instructions
- GREMLIN_FEATURE.md: Complete feature documentation
- pyproject.toml: Added gremlinpython>=3.7.0 to neptune extras
```python
from graphiti_core.driver.driver import QueryLanguage
from graphiti_core.driver.neptune_driver import NeptuneDriver
driver = NeptuneDriver(
host='neptune-db://cluster.amazonaws.com',
aoss_host='aoss-cluster.amazonaws.com',
query_language=QueryLanguage.GREMLIN
)
```
- Only Neptune Database supports Gremlin (not Neptune Analytics)
- Fulltext and vector search still use OpenSearch (AOSS) integration
- Complete search_utils.py Gremlin implementation pending (future work)
- ✅ All existing unit tests pass (103/103)
- ✅ New integration tests for Gremlin operations
- ✅ Type checking passes
- ✅ Linting passes
None. Fully backward compatible.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <noreply@anthropic.com>
130 lines
3.8 KiB
Python
130 lines
3.8 KiB
Python
"""
|
|
Copyright 2024, Zep Software, Inc.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import copy
|
|
import logging
|
|
import os
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Coroutine
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
from graphiti_core.driver.graph_operations.graph_operations import GraphOperationsInterface
|
|
from graphiti_core.driver.search_interface.search_interface import SearchInterface
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_SIZE = 10
|
|
|
|
load_dotenv()
|
|
|
|
ENTITY_INDEX_NAME = os.environ.get('ENTITY_INDEX_NAME', 'entities')
|
|
EPISODE_INDEX_NAME = os.environ.get('EPISODE_INDEX_NAME', 'episodes')
|
|
COMMUNITY_INDEX_NAME = os.environ.get('COMMUNITY_INDEX_NAME', 'communities')
|
|
ENTITY_EDGE_INDEX_NAME = os.environ.get('ENTITY_EDGE_INDEX_NAME', 'entity_edges')
|
|
|
|
|
|
class GraphProvider(Enum):
|
|
NEO4J = 'neo4j'
|
|
FALKORDB = 'falkordb'
|
|
KUZU = 'kuzu'
|
|
NEPTUNE = 'neptune'
|
|
|
|
|
|
class QueryLanguage(Enum):
|
|
CYPHER = 'cypher'
|
|
GREMLIN = 'gremlin'
|
|
|
|
|
|
class GraphDriverSession(ABC):
|
|
provider: GraphProvider
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
@abstractmethod
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
# No cleanup needed for Falkor, but method must exist
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def run(self, query: str, **kwargs: Any) -> Any:
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
async def close(self):
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
async def execute_write(self, func, *args, **kwargs):
|
|
raise NotImplementedError()
|
|
|
|
|
|
class GraphDriver(ABC):
|
|
provider: GraphProvider
|
|
query_language: QueryLanguage = QueryLanguage.CYPHER
|
|
fulltext_syntax: str = (
|
|
'' # Neo4j (default) syntax does not require a prefix for fulltext queries
|
|
)
|
|
_database: str
|
|
default_group_id: str = ''
|
|
search_interface: SearchInterface | None = None
|
|
graph_operations_interface: GraphOperationsInterface | None = None
|
|
|
|
@abstractmethod
|
|
def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
def session(self, database: str | None = None) -> GraphDriverSession:
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
def close(self):
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
def delete_all_indexes(self) -> Coroutine:
|
|
raise NotImplementedError()
|
|
|
|
def with_database(self, database: str) -> 'GraphDriver':
|
|
"""
|
|
Returns a shallow copy of this driver with a different default database.
|
|
Reuses the same connection (e.g. FalkorDB, Neo4j).
|
|
"""
|
|
cloned = copy.copy(self)
|
|
cloned._database = database
|
|
|
|
return cloned
|
|
|
|
@abstractmethod
|
|
async def build_indices_and_constraints(self, delete_existing: bool = False):
|
|
raise NotImplementedError()
|
|
|
|
def clone(self, database: str) -> 'GraphDriver':
|
|
"""Clone the driver with a different database or graph name."""
|
|
return self
|
|
|
|
def build_fulltext_query(
|
|
self, query: str, group_ids: list[str] | None = None, max_query_length: int = 128
|
|
) -> str:
|
|
"""
|
|
Specific fulltext query builder for database providers.
|
|
Only implemented by providers that need custom fulltext query building.
|
|
"""
|
|
raise NotImplementedError(f'build_fulltext_query not implemented for {self.provider}')
|