my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import logging
#2
#3	from .base import NeptuneBase
#4
#5	try:
#6	from langchain_aws import NeptuneAnalyticsGraph
#7	from botocore.config import Config
#8	except ImportError:
#9	raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.")
#10
#11	logger = logging.getLogger(__name__)
#12
#13
#14	class MemoryGraph(NeptuneBase):
#15	def __init__(self, config):
#16	self.config = config
#17
#18	self.graph = None
#19	endpoint = self.config.graph_store.config.endpoint
#20	app_id = self.config.graph_store.config.app_id
#21	if endpoint and endpoint.startswith("neptune-graph://"):
#22	graph_identifier = endpoint.replace("neptune-graph://", "")
#23	self.graph = NeptuneAnalyticsGraph(graph_identifier = graph_identifier,
#24	config = Config(user_agent_appid=app_id))
#25
#26	if not self.graph:
#27	raise ValueError("Unable to create a Neptune client: missing 'endpoint' in config")
#28
#29	self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
#30
#31	self.embedding_model = NeptuneBase._create_embedding_model(self.config)
#32
#33	# Default to openai if no specific provider is configured
#34	self.llm_provider = "openai"
#35	if self.config.llm.provider:
#36	self.llm_provider = self.config.llm.provider
#37	if self.config.graph_store.llm:
#38	self.llm_provider = self.config.graph_store.llm.provider
#39
#40	self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
#41	self.user_id = None
#42	# Use threshold from graph_store config, default to 0.7 for backward compatibility
#43	self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, 'threshold') else 0.7
#44
#45	def _delete_entities_cypher(self, source, destination, relationship, user_id):
#46	"""
#47	Returns the OpenCypher query and parameters for deleting entities in the graph DB
#48
#49	:param source: source node
#50	:param destination: destination node
#51	:param relationship: relationship label
#52	:param user_id: user_id to use
#53	:return: str, dict
#54	"""
#55
#56	cypher = f"""
#57	MATCH (n {self.node_label} {{name: $source_name, user_id: $user_id}})
#58	-[r:{relationship}]->
#59	(m {self.node_label} {{name: $dest_name, user_id: $user_id}})
#60	DELETE r
#61	RETURN
#62	n.name AS source,
#63	m.name AS target,
#64	type(r) AS relationship
#65	"""
#66	params = {
#67	"source_name": source,
#68	"dest_name": destination,
#69	"user_id": user_id,
#70	}
#71	logger.debug(f"_delete_entities\n query={cypher}")
#72	return cypher, params
#73
#74	def _add_entities_by_source_cypher(
#75	self,
#76	source_node_list,
#77	destination,
#78	dest_embedding,
#79	destination_type,
#80	relationship,
#81	user_id,
#82	):
#83	"""
#84	Returns the OpenCypher query and parameters for adding entities in the graph DB
#85
#86	:param source_node_list: list of source nodes
#87	:param destination: destination name
#88	:param dest_embedding: destination embedding
#89	:param destination_type: destination node label
#90	:param relationship: relationship label
#91	:param user_id: user id to use
#92	:return: str, dict
#93	"""
#94
#95	destination_label = self.node_label if self.node_label else f":`{destination_type}`"
#96	destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
#97
#98	cypher = f"""
#99	MATCH (source {{user_id: $user_id}})
#100	WHERE id(source) = $source_id
#101	SET source.mentions = coalesce(source.mentions, 0) + 1
#102	WITH source
#103	MERGE (destination {destination_label} {{name: $destination_name, user_id: $user_id}})
#104	ON CREATE SET
#105	destination.created = timestamp(),
#106	destination.updated = timestamp(),
#107	destination.mentions = 1
#108	{destination_extra_set}
#109	ON MATCH SET
#110	destination.mentions = coalesce(destination.mentions, 0) + 1,
#111	destination.updated = timestamp()
#112	WITH source, destination, $dest_embedding as dest_embedding
#113	CALL neptune.algo.vectors.upsert(destination, dest_embedding)
#114	WITH source, destination
#115	MERGE (source)-[r:{relationship}]->(destination)
#116	ON CREATE SET
#117	r.created = timestamp(),
#118	r.updated = timestamp(),
#119	r.mentions = 1
#120	ON MATCH SET
#121	r.mentions = coalesce(r.mentions, 0) + 1,
#122	r.updated = timestamp()
#123	RETURN source.name AS source, type(r) AS relationship, destination.name AS target
#124	"""
#125
#126	params = {
#127	"source_id": source_node_list[0]["id(source_candidate)"],
#128	"destination_name": destination,
#129	"dest_embedding": dest_embedding,
#130	"user_id": user_id,
#131	}
#132	logger.debug(
#133	f"_add_entities:\n source_node_search_result={source_node_list[0]}\n query={cypher}"
#134	)
#135	return cypher, params
#136
#137	def _add_entities_by_destination_cypher(
#138	self,
#139	source,
#140	source_embedding,
#141	source_type,
#142	destination_node_list,
#143	relationship,
#144	user_id,
#145	):
#146	"""
#147	Returns the OpenCypher query and parameters for adding entities in the graph DB
#148
#149	:param source: source node name
#150	:param source_embedding: source node embedding
#151	:param source_type: source node label
#152	:param destination_node_list: list of dest nodes
#153	:param relationship: relationship label
#154	:param user_id: user id to use
#155	:return: str, dict
#156	"""
#157
#158	source_label = self.node_label if self.node_label else f":`{source_type}`"
#159	source_extra_set = f", source:`{source_type}`" if self.node_label else ""
#160
#161	cypher = f"""
#162	MATCH (destination {{user_id: $user_id}})
#163	WHERE id(destination) = $destination_id
#164	SET
#165	destination.mentions = coalesce(destination.mentions, 0) + 1,
#166	destination.updated = timestamp()
#167	WITH destination
#168	MERGE (source {source_label} {{name: $source_name, user_id: $user_id}})
#169	ON CREATE SET
#170	source.created = timestamp(),
#171	source.updated = timestamp(),
#172	source.mentions = 1
#173	{source_extra_set}
#174	ON MATCH SET
#175	source.mentions = coalesce(source.mentions, 0) + 1,
#176	source.updated = timestamp()
#177	WITH source, destination, $source_embedding as source_embedding
#178	CALL neptune.algo.vectors.upsert(source, source_embedding)
#179	WITH source, destination
#180	MERGE (source)-[r:{relationship}]->(destination)
#181	ON CREATE SET
#182	r.created = timestamp(),
#183	r.updated = timestamp(),
#184	r.mentions = 1
#185	ON MATCH SET
#186	r.mentions = coalesce(r.mentions, 0) + 1,
#187	r.updated = timestamp()
#188	RETURN source.name AS source, type(r) AS relationship, destination.name AS target
#189	"""
#190
#191	params = {
#192	"destination_id": destination_node_list[0]["id(destination_candidate)"],
#193	"source_name": source,
#194	"source_embedding": source_embedding,
#195	"user_id": user_id,
#196	}
#197	logger.debug(
#198	f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n query={cypher}"
#199	)
#200	return cypher, params
#201
#202	def _add_relationship_entities_cypher(
#203	self,
#204	source_node_list,
#205	destination_node_list,
#206	relationship,
#207	user_id,
#208	):
#209	"""
#210	Returns the OpenCypher query and parameters for adding entities in the graph DB
#211
#212	:param source_node_list: list of source node ids
#213	:param destination_node_list: list of dest node ids
#214	:param relationship: relationship label
#215	:param user_id: user id to use
#216	:return: str, dict
#217	"""
#218
#219	cypher = f"""
#220	MATCH (source {{user_id: $user_id}})
#221	WHERE id(source) = $source_id
#222	SET
#223	source.mentions = coalesce(source.mentions, 0) + 1,
#224	source.updated = timestamp()
#225	WITH source
#226	MATCH (destination {{user_id: $user_id}})
#227	WHERE id(destination) = $destination_id
#228	SET
#229	destination.mentions = coalesce(destination.mentions) + 1,
#230	destination.updated = timestamp()
#231	MERGE (source)-[r:{relationship}]->(destination)
#232	ON CREATE SET
#233	r.created_at = timestamp(),
#234	r.updated_at = timestamp(),
#235	r.mentions = 1
#236	ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
#237	RETURN source.name AS source, type(r) AS relationship, destination.name AS target
#238	"""
#239	params = {
#240	"source_id": source_node_list[0]["id(source_candidate)"],
#241	"destination_id": destination_node_list[0]["id(destination_candidate)"],
#242	"user_id": user_id,
#243	}
#244	logger.debug(
#245	f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n source_node_search_result={source_node_list[0]}\n query={cypher}"
#246	)
#247	return cypher, params
#248
#249	def _add_new_entities_cypher(
#250	self,
#251	source,
#252	source_embedding,
#253	source_type,
#254	destination,
#255	dest_embedding,
#256	destination_type,
#257	relationship,
#258	user_id,
#259	):
#260	"""
#261	Returns the OpenCypher query and parameters for adding entities in the graph DB
#262
#263	:param source: source node name
#264	:param source_embedding: source node embedding
#265	:param source_type: source node label
#266	:param destination: destination name
#267	:param dest_embedding: destination embedding
#268	:param destination_type: destination node label
#269	:param relationship: relationship label
#270	:param user_id: user id to use
#271	:return: str, dict
#272	"""
#273
#274	source_label = self.node_label if self.node_label else f":`{source_type}`"
#275	source_extra_set = f", source:`{source_type}`" if self.node_label else ""
#276	destination_label = self.node_label if self.node_label else f":`{destination_type}`"
#277	destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
#278
#279	cypher = f"""
#280	MERGE (n {source_label} {{name: $source_name, user_id: $user_id}})
#281	ON CREATE SET n.created = timestamp(),
#282	n.updated = timestamp(),
#283	n.mentions = 1
#284	{source_extra_set}
#285	ON MATCH SET
#286	n.mentions = coalesce(n.mentions, 0) + 1,
#287	n.updated = timestamp()
#288	WITH n, $source_embedding as source_embedding
#289	CALL neptune.algo.vectors.upsert(n, source_embedding)
#290	WITH n
#291	MERGE (m {destination_label} {{name: $dest_name, user_id: $user_id}})
#292	ON CREATE SET
#293	m.created = timestamp(),
#294	m.updated = timestamp(),
#295	m.mentions = 1
#296	{destination_extra_set}
#297	ON MATCH SET
#298	m.updated = timestamp(),
#299	m.mentions = coalesce(m.mentions, 0) + 1
#300	WITH n, m, $dest_embedding as dest_embedding
#301	CALL neptune.algo.vectors.upsert(m, dest_embedding)
#302	WITH n, m
#303	MERGE (n)-[rel:{relationship}]->(m)
#304	ON CREATE SET
#305	rel.created = timestamp(),
#306	rel.updated = timestamp(),
#307	rel.mentions = 1
#308	ON MATCH SET
#309	rel.updated = timestamp(),
#310	rel.mentions = coalesce(rel.mentions, 0) + 1
#311	RETURN n.name AS source, type(rel) AS relationship, m.name AS target
#312	"""
#313	params = {
#314	"source_name": source,
#315	"dest_name": destination,
#316	"source_embedding": source_embedding,
#317	"dest_embedding": dest_embedding,
#318	"user_id": user_id,
#319	}
#320	logger.debug(
#321	f"_add_new_entities_cypher:\n query={cypher}"
#322	)
#323	return cypher, params
#324
#325	def _search_source_node_cypher(self, source_embedding, user_id, threshold):
#326	"""
#327	Returns the OpenCypher query and parameters to search for source nodes
#328
#329	:param source_embedding: source vector
#330	:param user_id: user_id to use
#331	:param threshold: the threshold for similarity
#332	:return: str, dict
#333	"""
#334	cypher = f"""
#335	MATCH (source_candidate {self.node_label})
#336	WHERE source_candidate.user_id = $user_id
#337
#338	WITH source_candidate, $source_embedding as v_embedding
#339	CALL neptune.algo.vectors.distanceByEmbedding(
#340	v_embedding,
#341	source_candidate,
#342	{{metric:"CosineSimilarity"}}
#343	) YIELD distance
#344	WITH source_candidate, distance AS cosine_similarity
#345	WHERE cosine_similarity >= $threshold
#346
#347	WITH source_candidate, cosine_similarity
#348	ORDER BY cosine_similarity DESC
#349	LIMIT 1
#350
#351	RETURN id(source_candidate), cosine_similarity
#352	"""
#353
#354	params = {
#355	"source_embedding": source_embedding,
#356	"user_id": user_id,
#357	"threshold": threshold,
#358	}
#359	logger.debug(f"_search_source_node\n query={cypher}")
#360	return cypher, params
#361
#362	def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
#363	"""
#364	Returns the OpenCypher query and parameters to search for destination nodes
#365
#366	:param source_embedding: source vector
#367	:param user_id: user_id to use
#368	:param threshold: the threshold for similarity
#369	:return: str, dict
#370	"""
#371	cypher = f"""
#372	MATCH (destination_candidate {self.node_label})
#373	WHERE destination_candidate.user_id = $user_id
#374
#375	WITH destination_candidate, $destination_embedding as v_embedding
#376	CALL neptune.algo.vectors.distanceByEmbedding(
#377	v_embedding,
#378	destination_candidate,
#379	{{metric:"CosineSimilarity"}}
#380	) YIELD distance
#381	WITH destination_candidate, distance AS cosine_similarity
#382	WHERE cosine_similarity >= $threshold
#383
#384	WITH destination_candidate, cosine_similarity
#385	ORDER BY cosine_similarity DESC
#386	LIMIT 1
#387
#388	RETURN id(destination_candidate), cosine_similarity
#389	"""
#390	params = {
#391	"destination_embedding": destination_embedding,
#392	"user_id": user_id,
#393	"threshold": threshold,
#394	}
#395
#396	logger.debug(f"_search_destination_node\n query={cypher}")
#397	return cypher, params
#398
#399	def _delete_all_cypher(self, filters):
#400	"""
#401	Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
#402
#403	:param filters: search filters
#404	:return: str, dict
#405	"""
#406	cypher = f"""
#407	MATCH (n {self.node_label} {{user_id: $user_id}})
#408	DETACH DELETE n
#409	"""
#410	params = {"user_id": filters["user_id"]}
#411
#412	logger.debug(f"delete_all query={cypher}")
#413	return cypher, params
#414
#415	def _get_all_cypher(self, filters, limit):
#416	"""
#417	Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
#418
#419	:param filters: search filters
#420	:param limit: return limit
#421	:return: str, dict
#422	"""
#423
#424	cypher = f"""
#425	MATCH (n {self.node_label} {{user_id: $user_id}})-[r]->(m {self.node_label} {{user_id: $user_id}})
#426	RETURN n.name AS source, type(r) AS relationship, m.name AS target
#427	LIMIT $limit
#428	"""
#429	params = {"user_id": filters["user_id"], "limit": limit}
#430	return cypher, params
#431
#432	def _search_graph_db_cypher(self, n_embedding, filters, limit):
#433	"""
#434	Returns the OpenCypher query and parameters to search for similar nodes in the memory store
#435
#436	:param n_embedding: node vector
#437	:param filters: search filters
#438	:param limit: return limit
#439	:return: str, dict
#440	"""
#441
#442	cypher_query = f"""
#443	MATCH (n {self.node_label})
#444	WHERE n.user_id = $user_id
#445	WITH n, $n_embedding as n_embedding
#446	CALL neptune.algo.vectors.distanceByEmbedding(
#447	n_embedding,
#448	n,
#449	{{metric:"CosineSimilarity"}}
#450	) YIELD distance
#451	WITH n, distance as similarity
#452	WHERE similarity >= $threshold
#453	CALL {{
#454	WITH n
#455	MATCH (n)-[r]->(m)
#456	RETURN n.name AS source, id(n) AS source_id, type(r) AS relationship, id(r) AS relation_id, m.name AS destination, id(m) AS destination_id
#457	UNION ALL
#458	WITH n
#459	MATCH (m)-[r]->(n)
#460	RETURN m.name AS source, id(m) AS source_id, type(r) AS relationship, id(r) AS relation_id, n.name AS destination, id(n) AS destination_id
#461	}}
#462	WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
#463	RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
#464	ORDER BY similarity DESC
#465	LIMIT $limit
#466	"""
#467	params = {
#468	"n_embedding": n_embedding,
#469	"threshold": self.threshold,
#470	"user_id": filters["user_id"],
#471	"limit": limit,
#472	}
#473	logger.debug(f"_search_graph_db\n query={cypher_query}")
#474
#475	return cypher_query, params
#476

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public