my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

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

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public