diff --git a/graphiti_core/edges.py b/graphiti_core/edges.py index f7d26b6f..165dee53 100644 --- a/graphiti_core/edges.py +++ b/graphiti_core/edges.py @@ -79,7 +79,9 @@ class Edge(BaseModel, ABC): if driver.aoss_client: await driver.aoss_client.delete( - index=ENTITY_EDGE_INDEX_NAME, id=self.uuid, routing=self.group_id + index=ENTITY_EDGE_INDEX_NAME, + id=self.uuid, + params={'routing': self.group_id}, ) logger.debug(f'Deleted Edge: {self.uuid}') @@ -273,7 +275,7 @@ class EntityEdge(Edge): 'size': 1, }, index=ENTITY_EDGE_INDEX_NAME, - routing=self.group_id, + params={'routing': self.group_id}, ) if resp['hits']['hits']: diff --git a/graphiti_core/nodes.py b/graphiti_core/nodes.py index 45599b01..7fafbe4f 100644 --- a/graphiti_core/nodes.py +++ b/graphiti_core/nodes.py @@ -119,7 +119,9 @@ class Node(BaseModel, ABC): # Delete the node from OpenSearch indices for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME): await driver.aoss_client.delete( - index=index, id=self.uuid, routing=self.group_id + index=index, + id=self.uuid, + params={'routing': self.group_id}, ) # Bulk delete the detached edges @@ -198,25 +200,25 @@ class Node(BaseModel, ABC): await driver.aoss_client.delete_by_query( index=EPISODE_INDEX_NAME, body={'query': {'term': {'group_id': group_id}}}, - routing=group_id, + params={'routing': group_id}, ) await driver.aoss_client.delete_by_query( index=ENTITY_INDEX_NAME, body={'query': {'term': {'group_id': group_id}}}, - routing=group_id, + params={'routing': group_id}, ) await driver.aoss_client.delete_by_query( index=COMMUNITY_INDEX_NAME, body={'query': {'term': {'group_id': group_id}}}, - routing=group_id, + params={'routing': group_id}, ) await driver.aoss_client.delete_by_query( index=ENTITY_EDGE_INDEX_NAME, body={'query': {'term': {'group_id': group_id}}}, - routing=group_id, + params={'routing': group_id}, ) case GraphProvider.KUZU: @@ -520,7 +522,7 @@ class EntityNode(Node): 'size': 1, }, index=ENTITY_INDEX_NAME, - routing=self.group_id, + params={'routing': self.group_id}, ) if resp['hits']['hits']: diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 2af8dd10..55c438da 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -215,11 +215,11 @@ async def edge_fulltext_search( # Match the edge ids and return the values query = ( """ - UNWIND $ids as id - MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity) - WHERE e.group_id IN $group_ids - AND id(e)=id - """ + UNWIND $ids as id + MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity) + WHERE e.group_id IN $group_ids + AND id(e)=id + """ + filter_query + """ AND id(e)=id @@ -256,15 +256,16 @@ async def edge_fulltext_search( filters = build_aoss_edge_filters(group_ids or [], search_filter) res = await driver.aoss_client.search( index=ENTITY_EDGE_INDEX_NAME, - routing=route, - _source=['uuid'], + params={'routing': route}, body={ + 'size': limit, + '_source': ['uuid'], 'query': { 'bool': { 'filter': filters, 'must': [{'match': {'fact': {'query': query, 'operator': 'or'}}}], } - } + }, }, ) @@ -353,8 +354,8 @@ async def edge_similarity_search( if driver.provider == GraphProvider.NEPTUNE: query = ( """ - MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity) - """ + MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity) + """ + filter_query + """ RETURN DISTINCT id(e) as id, e.fact_embedding as embedding @@ -417,10 +418,10 @@ async def edge_similarity_search( filters = build_aoss_edge_filters(group_ids or [], search_filter) res = await driver.aoss_client.search( index=ENTITY_EDGE_INDEX_NAME, - routing=route, - _source=['uuid'], - size=limit, + params={'routing': route}, body={ + 'size': limit, + '_source': ['uuid'], 'query': { 'knn': { 'fact_embedding': { @@ -429,7 +430,7 @@ async def edge_similarity_search( 'filter': {'bool': {'filter': filters}}, } } - } + }, }, ) @@ -637,11 +638,11 @@ async def node_fulltext_search( # Match the edge ides and return the values query = ( """ - UNWIND $ids as i - MATCH (n:Entity) - WHERE n.uuid=i.id - RETURN - """ + UNWIND $ids as i + MATCH (n:Entity) + WHERE n.uuid=i.id + RETURN + """ + get_entity_node_return_query(driver.provider) + """ ORDER BY i.score DESC @@ -663,10 +664,10 @@ async def node_fulltext_search( filters = build_aoss_node_filters(group_ids or [], search_filter) res = await driver.aoss_client.search( index=ENTITY_INDEX_NAME, - routing=route, - _source=['uuid'], - size=limit, + params={'routing': route}, body={ + '_source': ['uuid'], + 'size': limit, 'query': { 'bool': { 'filter': filters, @@ -680,7 +681,7 @@ async def node_fulltext_search( } ], } - } + }, }, ) @@ -751,8 +752,8 @@ async def node_similarity_search( if driver.provider == GraphProvider.NEPTUNE: query = ( """ - MATCH (n:Entity) - """ + MATCH (n:Entity) + """ + filter_query + """ RETURN DISTINCT id(n) as id, n.name_embedding as embedding @@ -781,11 +782,11 @@ async def node_similarity_search( # Match the edge ides and return the values query = ( """ - UNWIND $ids as i - MATCH (n:Entity) - WHERE id(n)=i.id - RETURN - """ + UNWIND $ids as i + MATCH (n:Entity) + WHERE id(n)=i.id + RETURN + """ + get_entity_node_return_query(driver.provider) + """ ORDER BY i.score DESC @@ -809,9 +810,9 @@ async def node_similarity_search( res = await driver.aoss_client.search( index=ENTITY_INDEX_NAME, params={'routing': route}, - _source=['uuid'], - size=limit, body={ + 'size': limit, + '_source': ['uuid'], 'query': { 'knn': { 'name_embedding': { @@ -820,7 +821,7 @@ async def node_similarity_search( 'filter': {'bool': {'filter': filters}}, } } - } + }, }, ) @@ -837,8 +838,8 @@ async def node_similarity_search( else: query = ( """ - MATCH (n:Entity) - """ + MATCH (n:Entity) + """ + filter_query + """ WITH n, """ @@ -1013,9 +1014,10 @@ async def episode_fulltext_search( route = group_ids[0] if group_ids else None res = await driver.aoss_client.search( EPISODE_INDEX_NAME, - routing=route, - _source=['uuid'], + params={'routing': route}, query={ + 'size': limit, + '_source': ['uuid'], 'bool': { 'filter': {'terms': group_ids}, 'must': [ @@ -1027,7 +1029,7 @@ async def episode_fulltext_search( } } ], - } + }, }, limit=limit, ) @@ -1170,8 +1172,8 @@ async def community_similarity_search( if driver.provider == GraphProvider.NEPTUNE: query = ( """ - MATCH (n:Community) - """ + MATCH (n:Community) + """ + group_filter_query + """ RETURN DISTINCT id(n) as id, n.name_embedding as embedding @@ -1230,8 +1232,8 @@ async def community_similarity_search( query = ( """ - MATCH (c:Community) - """ + MATCH (c:Community) + """ + group_filter_query + """ WITH c, @@ -1373,9 +1375,9 @@ async def get_relevant_nodes( # FIXME: Kuzu currently does not support using variables such as `node.fulltext_query` as an input to FTS, which means `get_relevant_nodes()` won't work with Kuzu as the graph driver. query = ( """ - UNWIND $nodes AS node - MATCH (n:Entity {group_id: $group_id}) - """ + UNWIND $nodes AS node + MATCH (n:Entity {group_id: $group_id}) + """ + filter_query + """ WITH node, n, """ @@ -1420,9 +1422,9 @@ async def get_relevant_nodes( else: query = ( """ - UNWIND $nodes AS node - MATCH (n:Entity {group_id: $group_id}) - """ + UNWIND $nodes AS node + MATCH (n:Entity {group_id: $group_id}) + """ + filter_query + """ WITH node, n, """ @@ -1511,9 +1513,9 @@ async def get_relevant_edges( if driver.provider == GraphProvider.NEPTUNE: query = ( """ - UNWIND $edges AS edge - MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid}) - """ + UNWIND $edges AS edge + MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid}) + """ + filter_query + """ WITH e, edge @@ -1583,9 +1585,9 @@ async def get_relevant_edges( query = ( """ - UNWIND $edges AS edge - MATCH (n:Entity {uuid: edge.source_node_uuid})-[:RELATES_TO]-(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]-(m:Entity {uuid: edge.target_node_uuid}) - """ + UNWIND $edges AS edge + MATCH (n:Entity {uuid: edge.source_node_uuid})-[:RELATES_TO]-(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]-(m:Entity {uuid: edge.target_node_uuid}) + """ + filter_query + """ WITH e, edge, n, m, """ @@ -1621,9 +1623,9 @@ async def get_relevant_edges( else: query = ( """ - UNWIND $edges AS edge - MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid}) - """ + UNWIND $edges AS edge + MATCH (n:Entity {uuid: edge.source_node_uuid})-[e:RELATES_TO {group_id: edge.group_id}]-(m:Entity {uuid: edge.target_node_uuid}) + """ + filter_query + """ WITH e, edge, """ @@ -1696,10 +1698,10 @@ async def get_edge_invalidation_candidates( if driver.provider == GraphProvider.NEPTUNE: query = ( """ - UNWIND $edges AS edge - MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity) - WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid] - """ + UNWIND $edges AS edge + MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity) + WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid] + """ + filter_query + """ WITH e, edge @@ -1769,10 +1771,10 @@ async def get_edge_invalidation_candidates( query = ( """ - UNWIND $edges AS edge - MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]->(m:Entity) - WHERE (n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]) - """ + UNWIND $edges AS edge + MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_ {group_id: edge.group_id})-[:RELATES_TO]->(m:Entity) + WHERE (n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid]) + """ + filter_query + """ WITH edge, e, n, m, """ @@ -1808,10 +1810,10 @@ async def get_edge_invalidation_candidates( else: query = ( """ - UNWIND $edges AS edge - MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity) - WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid] - """ + UNWIND $edges AS edge + MATCH (n:Entity)-[e:RELATES_TO {group_id: edge.group_id}]->(m:Entity) + WHERE n.uuid IN [edge.source_node_uuid, edge.target_node_uuid] OR m.uuid IN [edge.target_node_uuid, edge.source_node_uuid] + """ + filter_query + """ WITH edge, e, """