From dcfd9c2d31131586cd1cbf00e7f60d690c817ec2 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Wed, 17 Sep 2025 13:50:59 -0700 Subject: [PATCH 01/11] bring monkey patches from polykg --- graphiti_core/driver/neptune_driver.py | 13 ++++-- graphiti_core/search/search_utils.py | 59 ++++++++++++++++++++++---- 2 files changed, 59 insertions(+), 13 deletions(-) diff --git a/graphiti_core/driver/neptune_driver.py b/graphiti_core/driver/neptune_driver.py index 25aa12c3f4..299d1cc2d0 100644 --- a/graphiti_core/driver/neptune_driver.py +++ b/graphiti_core/driver/neptune_driver.py @@ -26,6 +26,8 @@ from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider +from pydantic import SecretStr + logger = logging.getLogger(__name__) DEFAULT_SIZE = 10 @@ -109,7 +111,8 @@ class NeptuneDriver(GraphDriver): provider: GraphProvider = GraphProvider.NEPTUNE - def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = 443): + def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = 443, aws_access_key_id=None, + aws_secret_access_key=None, aws_region_name=None, use_https=True): """This initializes a NeptuneDriver for use with Neptune as a backend Args: @@ -124,7 +127,9 @@ def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = if host.startswith('neptune-db://'): # This is a Neptune Database Cluster endpoint = host.replace('neptune-db://', '') - self.client = NeptuneGraph(endpoint, port) + self.client = NeptuneGraph(endpoint, port, aws_access_key_id=SecretStr(aws_access_key_id), + aws_secret_access_key=SecretStr(aws_secret_access_key), + region_name=aws_region_name, use_https=use_https) logger.debug('Creating Neptune Database session for %s', host) elif host.startswith('neptune-graph://'): # This is a Neptune Analytics Graph @@ -143,7 +148,7 @@ def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = self.aoss_client = OpenSearch( hosts=[{'host': aoss_host, 'port': aoss_port}], http_auth=Urllib3AWSV4SignerAuth( - session.get_credentials(), session.region_name, 'aoss' + session.get_credentials(), aws_region_name, 'es' ), use_ssl=True, verify_certs=True, @@ -191,7 +196,7 @@ def _sanitize_parameters(self, query, params: dict): async def execute_query( self, cypher_query_, **kwargs: Any ) -> tuple[dict[str, Any], None, None]: - params = dict(kwargs) + params = kwargs.get('params', kwargs) if isinstance(cypher_query_, list): for q in cypher_query_: result, _, _ = self._run_query(q[0], q[1]) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 379662d579..081fbee3cf 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -78,6 +78,46 @@ def calculate_cosine_similarity(vector1: list[float], vector2: list[float]) -> f return dot_product / (norm_vector1 * norm_vector2) +def _lucene_sanitize(query: str) -> str: + # Escape special characters from a query before passing into Lucene + # + - && || ! ( ) { } [ ] ^ " ~ * ? : \ / + escape_map = str.maketrans( + { + '+': r'\+', + '-': r'\-', + '&': r'\&', + '|': r'\|', + '!': r'\!', + '(': r'\(', + ')': r'\)', + '{': r'\{', + '}': r'\}', + '[': r'\[', + ']': r'\]', + '^': r'\^', + '"': r'\"', + "'": r"\'", + '~': r'\~', + '*': r'\*', + '?': r'\?', + ':': r'\:', + '\\': r'\\', + '/': r'\/', + '@': r'\@', + '%': r'\%', + 'O': r'\O', + 'R': r'\R', + 'N': r'\N', + 'T': r'\T', + 'A': r'\A', + 'D': r'\D', + } + ) + + sanitized = query.translate(escape_map) + return sanitized + + def fulltext_query(query: str, group_ids: list[str] | None, driver: GraphDriver): if driver.provider == GraphProvider.KUZU: # Kuzu only supports simple queries. @@ -95,7 +135,7 @@ def fulltext_query(query: str, group_ids: list[str] | None, driver: GraphDriver) group_ids_filter += ' AND ' if group_ids_filter else '' - lucene_query = lucene_sanitize(query) + lucene_query = _lucene_sanitize(query) # If the lucene query is too long return no query if len(lucene_query.split(' ')) + len(group_ids or '') >= MAX_QUERY_LENGTH: return '' @@ -338,7 +378,7 @@ async def edge_similarity_search( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - if r['embedding']: + if 'embedding' in r and r['embedding']: score = calculate_cosine_similarity( search_vector, list(map(float, r['embedding'].split(','))) ) @@ -668,7 +708,7 @@ async def node_similarity_search( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - if r['embedding']: + if 'embedding' in r: score = calculate_cosine_similarity( search_vector, list(map(float, r['embedding'].split(','))) ) @@ -1023,7 +1063,7 @@ async def community_similarity_search( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - if r['embedding']: + if 'embedding' in r and r['embedding']: score = calculate_cosine_similarity( search_vector, list(map(float, r['embedding'].split(','))) ) @@ -1366,11 +1406,12 @@ async def get_relevant_edges( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - score = calculate_cosine_similarity( - list(map(float, r['source_embedding'].split(','))), r['target_embedding'] - ) - if score > min_score: - input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']}) + if 'source_embedding' in r and 'target_embedding' in r: + score = calculate_cosine_similarity( + list(map(float, r['source_embedding'].split(','))), r['target_embedding'] + ) + if score > min_score: + input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']}) # Match the edge ides and return the values query = """ From 1231e113b925097fff2910700bfd0ec1158ac0e7 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Wed, 17 Sep 2025 15:33:05 -0700 Subject: [PATCH 02/11] read aws env --- graphiti_core/driver/neptune_driver.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/graphiti_core/driver/neptune_driver.py b/graphiti_core/driver/neptune_driver.py index 299d1cc2d0..da88b23f5f 100644 --- a/graphiti_core/driver/neptune_driver.py +++ b/graphiti_core/driver/neptune_driver.py @@ -17,6 +17,7 @@ import asyncio import datetime import logging +import os from collections.abc import Coroutine from typing import Any @@ -111,8 +112,7 @@ class NeptuneDriver(GraphDriver): provider: GraphProvider = GraphProvider.NEPTUNE - def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = 443, aws_access_key_id=None, - aws_secret_access_key=None, aws_region_name=None, use_https=True): + def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = 443, use_https=True): """This initializes a NeptuneDriver for use with Neptune as a backend Args: @@ -127,9 +127,15 @@ def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = if host.startswith('neptune-db://'): # This is a Neptune Database Cluster endpoint = host.replace('neptune-db://', '') - self.client = NeptuneGraph(endpoint, port, aws_access_key_id=SecretStr(aws_access_key_id), - aws_secret_access_key=SecretStr(aws_secret_access_key), - region_name=aws_region_name, use_https=use_https) + aws_access_key_id = os.getenv('NEPTUNE_AWS_ACCESS_KEY_ID') + aws_secret_access_key = os.getenv('NEPTUNE_AWS_SECRET_ACCESS_KEY') + aws_region_name = os.getenv('NEPTUNE_AWS_REGION') + if aws_access_key_id and aws_secret_access_key and aws_region_name: + self.client = NeptuneGraph(endpoint, port, aws_access_key_id=SecretStr(aws_access_key_id), + aws_secret_access_key=SecretStr(aws_secret_access_key), + region_name=aws_region_name, use_https=use_https) + else: + self.client = NeptuneGraph(endpoint, port, use_https=use_https) logger.debug('Creating Neptune Database session for %s', host) elif host.startswith('neptune-graph://'): # This is a Neptune Analytics Graph @@ -148,9 +154,9 @@ def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = self.aoss_client = OpenSearch( hosts=[{'host': aoss_host, 'port': aoss_port}], http_auth=Urllib3AWSV4SignerAuth( - session.get_credentials(), aws_region_name, 'es' + session.get_credentials(), aws_region_name or 'us-east-2', 'es' ), - use_ssl=True, + use_ssl=use_https, verify_certs=True, connection_class=Urllib3HttpConnection, pool_maxsize=20, From bb49e48cbca221cd722d98bd2b466fb2b25027f1 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Thu, 18 Sep 2025 15:15:17 -0700 Subject: [PATCH 03/11] key check --- graphiti_core/search/search_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 081fbee3cf..6e31ab8e95 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1541,7 +1541,7 @@ async def get_relevant_edges( result['search_edge_uuid']: [ get_entity_edge_from_record(record, driver.provider) for record in result['matches'] ] - for result in results + for result in results if 'search_edge_uuid' in result } relevant_edges = [relevant_edges_dict.get(edge.uuid, []) for edge in edges] From 3c2f1bd226f98fc8e742012ec904dd92e0873e61 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Thu, 18 Sep 2025 15:28:50 -0700 Subject: [PATCH 04/11] fix key errors --- graphiti_core/search/search_utils.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 6e31ab8e95..a8cc214211 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1594,11 +1594,12 @@ async def get_edge_invalidation_candidates( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - score = calculate_cosine_similarity( - list(map(float, r['source_embedding'].split(','))), r['target_embedding'] - ) - if score > min_score: - input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']}) + if 'source_embedding' in r and 'target_embedding' in r: + score = calculate_cosine_similarity( + list(map(float, r['source_embedding'].split(','))), r['target_embedding'] + ) + if score > min_score: + input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) # Match the edge ides and return the values query = """ From 0573e88e9058274d6e17fa4dc9806424aadc37d8 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Thu, 18 Sep 2025 15:45:47 -0700 Subject: [PATCH 05/11] fix keyerror --- graphiti_core/search/search_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index a8cc214211..0e9707da83 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1729,7 +1729,7 @@ async def get_edge_invalidation_candidates( result['search_edge_uuid']: [ get_entity_edge_from_record(record, driver.provider) for record in result['matches'] ] - for result in results + for result in results if 'search_edge_uuid' in result } invalidation_edges = [invalidation_edges_dict.get(edge.uuid, []) for edge in edges] From a516ab89bfe1beac98f720a4bccf8d339415a5e9 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Fri, 19 Sep 2025 11:50:00 -0700 Subject: [PATCH 06/11] fix missing keys --- graphiti_core/search/search_utils.py | 30 +++++++++++++--------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 0e9707da83..aab5553dc8 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1406,12 +1406,11 @@ async def get_relevant_edges( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - if 'source_embedding' in r and 'target_embedding' in r: - score = calculate_cosine_similarity( - list(map(float, r['source_embedding'].split(','))), r['target_embedding'] - ) - if score > min_score: - input_ids.append({'id': r['id'], 'score': score, 'uuid': r['search_edge_uuid']}) + score = calculate_cosine_similarity( + list(map(float, r.get('source_embedding', '').split(','))), r.get('target_embedding') + ) + if score > min_score: + input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) # Match the edge ides and return the values query = """ @@ -1538,10 +1537,10 @@ async def get_relevant_edges( ) relevant_edges_dict: dict[str, list[EntityEdge]] = { - result['search_edge_uuid']: [ + result.get('search_edge_uuid', result.get('uuid', '')): [ get_entity_edge_from_record(record, driver.provider) for record in result['matches'] ] - for result in results if 'search_edge_uuid' in result + for result in results } relevant_edges = [relevant_edges_dict.get(edge.uuid, []) for edge in edges] @@ -1594,12 +1593,11 @@ async def get_edge_invalidation_candidates( # Calculate Cosine similarity then return the edge ids input_ids = [] for r in resp: - if 'source_embedding' in r and 'target_embedding' in r: - score = calculate_cosine_similarity( - list(map(float, r['source_embedding'].split(','))), r['target_embedding'] - ) - if score > min_score: - input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) + score = calculate_cosine_similarity( + list(map(float, r.get('source_embedding', '').split(','))), r.get('target_embedding') + ) + if score > min_score: + input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) # Match the edge ides and return the values query = """ @@ -1726,10 +1724,10 @@ async def get_edge_invalidation_candidates( **filter_params, ) invalidation_edges_dict: dict[str, list[EntityEdge]] = { - result['search_edge_uuid']: [ + result.get('search_edge_uuid', result.get('uuid', '')): [ get_entity_edge_from_record(record, driver.provider) for record in result['matches'] ] - for result in results if 'search_edge_uuid' in result + for result in results } invalidation_edges = [invalidation_edges_dict.get(edge.uuid, []) for edge in edges] From 324c85f8ccb8ea00185d06ee75e12b8c4b400662 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Fri, 19 Sep 2025 12:00:54 -0700 Subject: [PATCH 07/11] 0 default --- graphiti_core/search/search_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index aab5553dc8..1c3f1b05c4 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1407,7 +1407,7 @@ async def get_relevant_edges( input_ids = [] for r in resp: score = calculate_cosine_similarity( - list(map(float, r.get('source_embedding', '').split(','))), r.get('target_embedding') + list(map(float, r.get('source_embedding', '0').split(','))), r.get('target_embedding') ) if score > min_score: input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) @@ -1594,7 +1594,7 @@ async def get_edge_invalidation_candidates( input_ids = [] for r in resp: score = calculate_cosine_similarity( - list(map(float, r.get('source_embedding', '').split(','))), r.get('target_embedding') + list(map(float, r.get('source_embedding', '0').split(','))), r.get('target_embedding') ) if score > min_score: input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) From fee0fb15afbe96e33448ad5d18a3bf547fa66af0 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Fri, 19 Sep 2025 12:10:01 -0700 Subject: [PATCH 08/11] mend --- graphiti_core/search/search_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 1c3f1b05c4..21f03a5339 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1407,7 +1407,7 @@ async def get_relevant_edges( input_ids = [] for r in resp: score = calculate_cosine_similarity( - list(map(float, r.get('source_embedding', '0').split(','))), r.get('target_embedding') + list(map(float, r.get('source_embedding', '0').split(','))), r.get('target_embedding', 0) ) if score > min_score: input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')}) From 3e032a2e05f9397725d6e8fade7d32077df857cf Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Fri, 19 Sep 2025 14:03:56 -0700 Subject: [PATCH 09/11] fix IndexError --- graphiti_core/search/search_utils.py | 2 +- graphiti_core/utils/maintenance/node_operations.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 21f03a5339..548a0a74b9 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1538,7 +1538,7 @@ async def get_relevant_edges( relevant_edges_dict: dict[str, list[EntityEdge]] = { result.get('search_edge_uuid', result.get('uuid', '')): [ - get_entity_edge_from_record(record, driver.provider) for record in result['matches'] + get_entity_edge_from_record(record, driver.provider) for record in result.get('matches', []) ] for result in results } diff --git a/graphiti_core/utils/maintenance/node_operations.py b/graphiti_core/utils/maintenance/node_operations.py index f17706ab5c..e448837fbf 100644 --- a/graphiti_core/utils/maintenance/node_operations.py +++ b/graphiti_core/utils/maintenance/node_operations.py @@ -269,7 +269,10 @@ async def resolve_extracted_nodes( resolution_id: int = resolution.id duplicate_idx: int = resolution.duplicate_idx - extracted_node = extracted_nodes[resolution_id] + try: + extracted_node = extracted_nodes[resolution_id] + except IndexError: + continue resolved_node = ( existing_nodes[duplicate_idx] From 797d33cf273f2e28c8db725e670d13546fb4ce57 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Mon, 22 Sep 2025 09:15:05 -0700 Subject: [PATCH 10/11] fix KeyError --- graphiti_core/search/search_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 548a0a74b9..72ad942807 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1725,7 +1725,7 @@ async def get_edge_invalidation_candidates( ) invalidation_edges_dict: dict[str, list[EntityEdge]] = { result.get('search_edge_uuid', result.get('uuid', '')): [ - get_entity_edge_from_record(record, driver.provider) for record in result['matches'] + get_entity_edge_from_record(record, driver.provider) for record in result.get('matches', []) ] for result in results } From 1374bf0a87ba888f8d215c9a2ea7ac3598a28782 Mon Sep 17 00:00:00 2001 From: Michael Bradford Date: Mon, 22 Sep 2025 09:24:20 -0700 Subject: [PATCH 11/11] fix KeyError --- graphiti_core/search/search_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py index 72ad942807..8d98db5503 100644 --- a/graphiti_core/search/search_utils.py +++ b/graphiti_core/search/search_utils.py @@ -1594,7 +1594,7 @@ async def get_edge_invalidation_candidates( input_ids = [] for r in resp: score = calculate_cosine_similarity( - list(map(float, r.get('source_embedding', '0').split(','))), r.get('target_embedding') + list(map(float, r.get('source_embedding', '0').split(','))), r.get('target_embedding', 0) ) if score > min_score: input_ids.append({'id': r.get('id'), 'score': score, 'uuid': r.get('search_edge_uuid')})