fix: align Neptune community AOSS index

graphiti_core/driver/driver.py

@@ -52,7 +52,7 @@
 
 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')
+COMMUNITY_INDEX_NAME = os.environ.get('COMMUNITY_INDEX_NAME', 'community_name')
 ENTITY_EDGE_INDEX_NAME = os.environ.get('ENTITY_EDGE_INDEX_NAME', 'entity_edges')
 
 

graphiti_core/driver/neptune/operations/community_node_ops.py

@@ -19,7 +19,7 @@
 import logging
 from typing import TYPE_CHECKING, Any
 
-from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.driver.driver import COMMUNITY_INDEX_NAME, GraphProvider
 from graphiti_core.driver.operations.community_node_ops import CommunityNodeOperations
 from graphiti_core.driver.query_executor import QueryExecutor, Transaction
 from graphiti_core.driver.record_parsers import community_node_from_record
@@ -62,7 +62,7 @@ async def save(
 
         if self._driver is not None:
             self._driver.save_to_aoss(
-                'community_name',
+                COMMUNITY_INDEX_NAME,
                 [{'uuid': node.uuid, 'name': node.name, 'group_id': node.group_id}],
             )
 

graphiti_core/driver/neptune/operations/search_ops.py

@@ -19,7 +19,7 @@
 import logging
 from typing import TYPE_CHECKING, Any
 
-from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.driver.driver import COMMUNITY_INDEX_NAME, GraphProvider
 from graphiti_core.driver.operations.search_ops import SearchOperations
 from graphiti_core.driver.query_executor import QueryExecutor
 from graphiti_core.driver.record_parsers import (
@@ -478,7 +478,7 @@ async def community_fulltext_search(
         if self._driver is None:
             return []
         driver = self._driver
-        res = driver.run_aoss_query('community_name', query, limit=limit)
+        res = driver.run_aoss_query(COMMUNITY_INDEX_NAME, query, limit=limit)
         if not res or res.get('hits', {}).get('total', {}).get('value', 0) == 0:
             return []
 

graphiti_core/driver/neptune_driver.py

@@ -24,7 +24,12 @@
 from langchain_aws.graphs import NeptuneAnalyticsGraph, NeptuneGraph
 from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers
 
-from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+from graphiti_core.driver.driver import (
+    COMMUNITY_INDEX_NAME,
+    GraphDriver,
+    GraphDriverSession,
+    GraphProvider,
+)
 from graphiti_core.driver.neptune.operations.community_edge_ops import (
     NeptuneCommunityEdgeOperations,
 )
@@ -78,7 +83,7 @@
         },
     },
     {
-        'index_name': 'community_name',
+        'index_name': COMMUNITY_INDEX_NAME,
         'body': {
             'mappings': {
                 'properties': {

graphiti_core/nodes.py

@@ -27,6 +27,7 @@
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import (
+    COMMUNITY_INDEX_NAME,
     GraphDriver,
     GraphProvider,
 )
@@ -696,8 +697,8 @@ async def save(self, driver: GraphDriver):
                 pass
 
         if driver.provider == GraphProvider.NEPTUNE:
-            await driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
-                'communities',
+            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+                COMMUNITY_INDEX_NAME,
                 [{'name': self.name, 'uuid': self.uuid, 'group_id': self.group_id}],
             )
         result = await driver.execute_query(

graphiti_core/search/search_utils.py

@@ -24,6 +24,7 @@
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import (
+    COMMUNITY_INDEX_NAME,
     GraphDriver,
     GraphProvider,
 )
@@ -999,7 +1000,7 @@ async def community_fulltext_search(
         yield_query = 'WITH node AS c, score'
 
     if driver.provider == GraphProvider.NEPTUNE:
-        res = driver.run_aoss_query('community_name', query, limit=limit)  # pyright: ignore reportAttributeAccessIssue
+        res = driver.run_aoss_query(COMMUNITY_INDEX_NAME, query, limit=limit)  # pyright: ignore reportAttributeAccessIssue
         if res['hits']['total']['value'] > 0:
             # Calculate Cosine similarity then return the edge ids
             input_ids = []

tests/driver/test_neptune_aoss_indices.py

@@ -0,0 +1,71 @@
+from typing import Any, cast
+
+import pytest
+
+from graphiti_core.driver.driver import COMMUNITY_INDEX_NAME, GraphProvider
+from graphiti_core.driver.neptune.operations.search_ops import NeptuneSearchOperations
+from graphiti_core.driver.neptune_driver import aoss_indices
+from graphiti_core.nodes import CommunityNode
+
+
+class FakeNeptuneDriver:
+    provider = GraphProvider.NEPTUNE
+    graph_operations_interface = None
+
+    def __init__(self):
+        self.saved_to_aoss: list[tuple[str, list[dict[str, Any]]]] = []
+        self.execute_query_kwargs: dict[str, Any] | None = None
+
+    def save_to_aoss(self, name: str, data: list[dict[str, Any]]) -> int:
+        self.saved_to_aoss.append((name, data))
+        return len(data)
+
+    async def execute_query(self, *args: Any, **kwargs: Any) -> list[Any]:
+        self.execute_query_kwargs = kwargs
+        return []
+
+
+class FakeAOSSSearchDriver:
+    def __init__(self):
+        self.queries: list[tuple[str, str, int]] = []
+
+    def run_aoss_query(self, name: str, query: str, limit: int = 10) -> dict[str, Any]:
+        self.queries.append((name, query, limit))
+        return {'hits': {'total': {'value': 0}, 'hits': []}}
+
+
+def test_default_community_index_matches_registered_aoss_index():
+    registered_indices = {index['index_name'] for index in aoss_indices}
+
+    assert COMMUNITY_INDEX_NAME == 'community_name'
+    assert COMMUNITY_INDEX_NAME in registered_indices
+
+
+@pytest.mark.asyncio
+async def test_community_node_fallback_saves_to_registered_aoss_index():
+    driver = FakeNeptuneDriver()
+    node = CommunityNode(uuid='community-1', name='Community 1', group_id='group-1')
+
+    await node.save(driver)  # type: ignore[arg-type]
+
+    assert driver.saved_to_aoss == [
+        (
+            COMMUNITY_INDEX_NAME,
+            [{'name': 'Community 1', 'uuid': 'community-1', 'group_id': 'group-1'}],
+        )
+    ]
+    assert driver.execute_query_kwargs is not None
+    assert driver.execute_query_kwargs['uuid'] == 'community-1'
+
+
+@pytest.mark.asyncio
+async def test_neptune_community_search_uses_registered_aoss_index():
+    driver = FakeAOSSSearchDriver()
+    operations = NeptuneSearchOperations(cast(Any, driver))
+
+    results = await operations.community_fulltext_search(
+        cast(Any, None), 'community query', limit=3
+    )
+
+    assert results == []
+    assert driver.queries == [(COMMUNITY_INDEX_NAME, 'community query', 3)]