feat: Add configurable embedding similarity threshold for graph store node matching (#3593)

This commit is contained in:
Ronak Bhalgami
2025-10-18 13:03:30 -07:00
committed by GitHub
parent 9ef644b95e
commit de224dd26d
9 changed files with 237 additions and 13 deletions
+9
View File
@@ -89,6 +89,15 @@ class GraphStoreConfig(BaseModel):
custom_prompt: Optional[str] = Field(
description="Custom prompt to fetch entities from the given text", default=None
)
threshold: float = Field(
description="Threshold for embedding similarity when matching nodes during graph ingestion. "
"Range: 0.0 to 1.0. Higher values require closer matches. "
"Use lower values (e.g., 0.5-0.7) for distinct entities with similar embeddings. "
"Use higher values (e.g., 0.9+) when you want stricter matching.",
default=0.7,
ge=0.0,
le=1.0,
)
@field_validator("config")
def validate_config(cls, v, values):
+2 -2
View File
@@ -234,8 +234,8 @@ class NeptuneBase(ABC):
dest_embedding = self.embedding_model.embed(destination)
# search for the nodes with the closest embeddings
source_node_search_result = self._search_source_node(source_embedding, user_id, threshold=0.9)
destination_node_search_result = self._search_destination_node(dest_embedding, user_id, threshold=0.9)
source_node_search_result = self._search_source_node(source_embedding, user_id, threshold=self.threshold)
destination_node_search_result = self._search_destination_node(dest_embedding, user_id, threshold=self.threshold)
cypher, params = self._add_entities_cypher(
source_node_search_result,
+2 -1
View File
@@ -56,7 +56,8 @@ class MemoryGraph(NeptuneBase):
self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
self.user_id = None
self.threshold = 0.7
# Use threshold from graph_store config, default to 0.7 for backward compatibility
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
self.vector_store_limit=5
def _delete_entities_cypher(self, source, destination, relationship, user_id):
+2 -1
View File
@@ -39,7 +39,8 @@ class MemoryGraph(NeptuneBase):
self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
self.user_id = None
self.threshold = 0.7
# Use threshold from graph_store config, default to 0.7 for backward compatibility
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
def _delete_entities_cypher(self, source, destination, relationship, user_id):
"""
+4 -3
View File
@@ -70,7 +70,8 @@ class MemoryGraph:
llm_config = self.config.llm.config
self.llm = LlmFactory.create(self.llm_provider, llm_config)
self.user_id = None
self.threshold = 0.7
# Use threshold from graph_store config, default to 0.7 for backward compatibility
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
def add(self, data, filters):
"""
@@ -434,8 +435,8 @@ class MemoryGraph:
dest_embedding = self.embedding_model.embed(destination)
# search for the nodes with the closest embeddings
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=0.9)
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=0.9)
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
# TODO: Create a cypher query and common params for all the cases
if not destination_node_search_result and source_node_search_result:
+4 -3
View File
@@ -62,7 +62,8 @@ class MemoryGraph:
self.llm = LlmFactory.create(self.llm_provider, llm_config)
self.user_id = None
self.threshold = 0.7
# Use threshold from graph_store config, default to 0.7 for backward compatibility
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
def kuzu_create_schema(self):
self.kuzu_execute(
@@ -452,8 +453,8 @@ class MemoryGraph:
dest_embedding = self.embedding_model.embed(destination)
# search for the nodes with the closest embeddings
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=0.9)
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=0.9)
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
if not destination_node_search_result and source_node_search_result:
params = {
+4 -3
View File
@@ -55,7 +55,8 @@ class MemoryGraph:
llm_config = self.config.llm.config
self.llm = LlmFactory.create(self.llm_provider, llm_config)
self.user_id = None
self.threshold = 0.7
# Use threshold from graph_store config, default to 0.7 for backward compatibility
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
# Setup Memgraph:
# 1. Create vector index (created Entity label on all nodes)
@@ -422,8 +423,8 @@ class MemoryGraph:
dest_embedding = self.embedding_model.embed(destination)
# search for the nodes with the closest embeddings
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=0.9)
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=0.9)
source_node_search_result = self._search_source_node(source_embedding, filters, threshold=self.threshold)
destination_node_search_result = self._search_destination_node(dest_embedding, filters, threshold=self.threshold)
# Prepare agent_id for node creation
agent_id_clause = ""