feat: TS SDK v3 port, graph store removal, client v3 API migration
TypeScript SDK v3 pipeline (full parity with Python): - Single-pass additive extraction with ADDITIVE_EXTRACTION_PROMPT - Hybrid search (semantic + BM25 + entity boost) with additive scoring - 8-phase batch pipeline (batch embed, persist, entity linking) - New utils: scoring.ts, lemmatization.ts (natural), entity_extraction.ts (compromise) - keywordSearch() on 8 vector stores (3 full: PGVector, Memory, Azure AI Search) - Message persistence in SQLiteManager (rolling window of 10) - Entity store as second vector collection - MAX_BATCH=100 chunking guard on OpenAI/Azure embedBatch - Updated default LLM model to gpt-4.1-nano-2025-04-14 - compromise + natural added as peer dependencies Graph store removal (Python + TypeScript): - Removed Neo4j, Memgraph, Kuzu, Neptune, Apache AGE integrations - Deleted 18 graph-related files across both SDKs - Removed GraphStoreFactory, GraphStoreConfig, graph_store config field - Removed "relations" key from all API responses - Removed graph optional dependency group from pyproject.toml - Removed neo4j-driver from TS peerDependencies - Simplified add/search/delete/reset (no more parallel graph operations) Client SDK v3 API migration: - add() endpoint: /v1/memories/ -> /v3/memories/ (async response) - search() endpoint: /v2/memories/search/ -> /v3/memories/search/ - Removed output_format injection and v1.1 unwrapping logic - Applied to both Python (sync + async) and TypeScript clients Review feedback fixes: - Removed deprecated custom_update_memory_prompt from MemoryConfig - Added MAX_BATCH=100 chunking to Python + TS embed_batch - Moved all inline imports to top level in main.py - Cleaned up GraphStoreError dead code from exceptions.py Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Soumil Rathi
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
3402c67805
commit
59880a6f8f
+6
-22
@@ -162,12 +162,9 @@ class MemoryClient:
|
||||
elif not isinstance(messages, list):
|
||||
raise ValueError(f"messages must be str, dict, or list[dict], got {type(messages).__name__}")
|
||||
|
||||
# Force v1.1 format for all add operations
|
||||
kwargs["output_format"] = "v1.1"
|
||||
|
||||
kwargs = self._prepare_params(kwargs)
|
||||
payload = self._prepare_payload(messages, kwargs)
|
||||
response = self.client.post("/v1/memories/", json=payload)
|
||||
response = self.client.post("/v3/memories/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -271,7 +268,7 @@ class MemoryClient:
|
||||
params = self._prepare_params(kwargs)
|
||||
payload = {"query": query, **params}
|
||||
|
||||
response = self.client.post("/v2/memories/search/", json=payload)
|
||||
response = self.client.post("/v3/memories/search/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -283,12 +280,7 @@ class MemoryClient:
|
||||
"sync_type": "sync",
|
||||
},
|
||||
)
|
||||
result = response.json()
|
||||
|
||||
# Ensure v1.1 format (wrap raw list if needed)
|
||||
if isinstance(result, list):
|
||||
return {"results": result}
|
||||
return result
|
||||
return response.json()
|
||||
|
||||
@api_error_handler
|
||||
def update(
|
||||
@@ -1083,12 +1075,9 @@ class AsyncMemoryClient:
|
||||
elif not isinstance(messages, list):
|
||||
raise ValueError(f"messages must be str, dict, or list[dict], got {type(messages).__name__}")
|
||||
|
||||
# Force v1.1 format for all add operations
|
||||
kwargs["output_format"] = "v1.1"
|
||||
|
||||
kwargs = self._prepare_params(kwargs)
|
||||
payload = self._prepare_payload(messages, kwargs)
|
||||
response = await self.async_client.post("/v1/memories/", json=payload)
|
||||
response = await self.async_client.post("/v3/memories/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -1175,7 +1164,7 @@ class AsyncMemoryClient:
|
||||
params = self._prepare_params(kwargs)
|
||||
payload = {"query": query, **params}
|
||||
|
||||
response = await self.async_client.post("/v2/memories/search/", json=payload)
|
||||
response = await self.async_client.post("/v3/memories/search/", json=payload)
|
||||
response.raise_for_status()
|
||||
if "metadata" in kwargs:
|
||||
del kwargs["metadata"]
|
||||
@@ -1187,12 +1176,7 @@ class AsyncMemoryClient:
|
||||
"sync_type": "async",
|
||||
},
|
||||
)
|
||||
result = response.json()
|
||||
|
||||
# Ensure v1.1 format (wrap raw list if needed)
|
||||
if isinstance(result, list):
|
||||
return {"results": result}
|
||||
return result
|
||||
return response.json()
|
||||
|
||||
@api_error_handler
|
||||
async def update(
|
||||
|
||||
@@ -5,7 +5,6 @@ from pydantic import BaseModel, Field
|
||||
|
||||
from mem0.configs.rerankers.config import RerankerConfig
|
||||
from mem0.embeddings.configs import EmbedderConfig
|
||||
from mem0.graphs.configs import GraphStoreConfig
|
||||
from mem0.llms.configs import LlmConfig
|
||||
from mem0.vector_stores.configs import VectorStoreConfig
|
||||
|
||||
@@ -44,10 +43,6 @@ class MemoryConfig(BaseModel):
|
||||
description="Path to the history database",
|
||||
default=os.path.join(mem0_dir, "history.db"),
|
||||
)
|
||||
graph_store: GraphStoreConfig = Field(
|
||||
description="Configuration for the graph",
|
||||
default_factory=GraphStoreConfig,
|
||||
)
|
||||
reranker: Optional[RerankerConfig] = Field(
|
||||
description="Configuration for the reranker",
|
||||
default=None,
|
||||
|
||||
@@ -321,24 +321,6 @@ class VectorStoreError(MemoryError):
|
||||
super().__init__(message, error_code, details, suggestion, debug_info)
|
||||
|
||||
|
||||
class GraphStoreError(MemoryError):
|
||||
"""Raised when graph store operations fail.
|
||||
|
||||
This exception is raised when graph store operations fail,
|
||||
such as relationship creation, entity management, or graph queries.
|
||||
|
||||
Example:
|
||||
raise GraphStoreError(
|
||||
message="Graph store operation failed",
|
||||
error_code="GRAPH_001",
|
||||
details={"operation": "create_relationship", "entity": "user_123"},
|
||||
suggestion="Please check your graph store configuration and connection"
|
||||
)
|
||||
"""
|
||||
def __init__(self, message: str, error_code: str = "GRAPH_001", details: dict = None,
|
||||
suggestion: str = "Please check your graph store configuration and connection",
|
||||
debug_info: dict = None):
|
||||
super().__init__(message, error_code, details, suggestion, debug_info)
|
||||
|
||||
|
||||
class EmbeddingError(MemoryError):
|
||||
|
||||
@@ -1,136 +0,0 @@
|
||||
from typing import Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
from mem0.llms.configs import LlmConfig
|
||||
|
||||
|
||||
class Neo4jConfig(BaseModel):
|
||||
url: Optional[str] = Field(None, description="Host address for the graph database")
|
||||
username: Optional[str] = Field(None, description="Username for the graph database")
|
||||
password: Optional[str] = Field(None, description="Password for the graph database")
|
||||
database: Optional[str] = Field(None, description="Database for the graph database")
|
||||
base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
url, username, password = (
|
||||
values.get("url"),
|
||||
values.get("username"),
|
||||
values.get("password"),
|
||||
)
|
||||
if not url or not username or not password:
|
||||
raise ValueError("Please provide 'url', 'username' and 'password'.")
|
||||
return values
|
||||
|
||||
|
||||
class MemgraphConfig(BaseModel):
|
||||
url: Optional[str] = Field(None, description="Host address for the graph database")
|
||||
username: Optional[str] = Field(None, description="Username for the graph database")
|
||||
password: Optional[str] = Field(None, description="Password for the graph database")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
url, username, password = (
|
||||
values.get("url"),
|
||||
values.get("username"),
|
||||
values.get("password"),
|
||||
)
|
||||
if not url or not username or not password:
|
||||
raise ValueError("Please provide 'url', 'username' and 'password'.")
|
||||
return values
|
||||
|
||||
|
||||
class NeptuneConfig(BaseModel):
|
||||
app_id: Optional[str] = Field("Mem0", description="APP_ID for the connection")
|
||||
endpoint: Optional[str] = (
|
||||
Field(
|
||||
None,
|
||||
description="Endpoint to connect to a Neptune-DB Cluster as 'neptune-db://<host>' or Neptune Analytics Server as 'neptune-graph://<graphid>'",
|
||||
),
|
||||
)
|
||||
base_label: Optional[bool] = Field(None, description="Whether to use base node label __Entity__ for all entities")
|
||||
collection_name: Optional[str] = Field(None, description="vector_store collection name to store vectors when using Neptune-DB Clusters")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_host_port_or_path(cls, values):
|
||||
endpoint = values.get("endpoint")
|
||||
if not endpoint:
|
||||
raise ValueError("Please provide 'endpoint' with the format as 'neptune-db://<endpoint>' or 'neptune-graph://<graphid>'.")
|
||||
if endpoint.startswith("neptune-db://"):
|
||||
# This is a Neptune DB Graph
|
||||
return values
|
||||
elif endpoint.startswith("neptune-graph://"):
|
||||
# This is a Neptune Analytics Graph
|
||||
graph_identifier = endpoint.replace("neptune-graph://", "")
|
||||
if not graph_identifier.startswith("g-"):
|
||||
raise ValueError("Provide a valid 'graph_identifier'.")
|
||||
values["graph_identifier"] = graph_identifier
|
||||
return values
|
||||
else:
|
||||
raise ValueError(
|
||||
"You must provide an endpoint to create a NeptuneServer as either neptune-db://<endpoint> or neptune-graph://<graphid>"
|
||||
)
|
||||
|
||||
|
||||
class KuzuConfig(BaseModel):
|
||||
db: Optional[str] = Field(":memory:", description="Path to a Kuzu database file")
|
||||
|
||||
|
||||
class ApacheAgeConfig(BaseModel):
|
||||
host: Optional[str] = Field("localhost", description="PostgreSQL server hostname")
|
||||
port: Optional[int] = Field(5432, description="PostgreSQL server port")
|
||||
database: Optional[str] = Field(None, description="PostgreSQL database name")
|
||||
username: Optional[str] = Field(None, description="PostgreSQL username")
|
||||
password: Optional[str] = Field(None, description="PostgreSQL password")
|
||||
graph_name: Optional[str] = Field("mem0_graph", description="Name of the Apache AGE graph")
|
||||
|
||||
@model_validator(mode="before")
|
||||
def check_required_fields(cls, values):
|
||||
database, username, password = (
|
||||
values.get("database"),
|
||||
values.get("username"),
|
||||
values.get("password"),
|
||||
)
|
||||
if not database or not username or not password:
|
||||
raise ValueError("Please provide 'database', 'username' and 'password'.")
|
||||
return values
|
||||
|
||||
|
||||
class GraphStoreConfig(BaseModel):
|
||||
provider: str = Field(
|
||||
description="Provider of the data store (e.g., 'neo4j', 'memgraph', 'neptune', 'kuzu', 'apache_age')",
|
||||
default="neo4j",
|
||||
)
|
||||
config: Union[Neo4jConfig, MemgraphConfig, NeptuneConfig, KuzuConfig, ApacheAgeConfig] = Field(
|
||||
description="Configuration for the specific data store", default=None
|
||||
)
|
||||
llm: Optional[LlmConfig] = Field(description="LLM configuration for querying the graph store", default=None)
|
||||
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):
|
||||
provider = values.data.get("provider")
|
||||
if provider == "neo4j":
|
||||
return Neo4jConfig(**v.model_dump())
|
||||
elif provider == "memgraph":
|
||||
return MemgraphConfig(**v.model_dump())
|
||||
elif provider == "neptune" or provider == "neptunedb":
|
||||
return NeptuneConfig(**v.model_dump())
|
||||
elif provider == "kuzu":
|
||||
return KuzuConfig(**v.model_dump())
|
||||
elif provider == "apache_age":
|
||||
return ApacheAgeConfig(**v.model_dump())
|
||||
else:
|
||||
raise ValueError(f"Unsupported graph store provider: {provider}")
|
||||
@@ -1,515 +0,0 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory, VectorStoreFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NeptuneBase(ABC):
|
||||
"""
|
||||
Abstract base class for neptune (neptune analytics and neptune db) calls using OpenCypher
|
||||
to store/retrieve data
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _create_embedding_model(config):
|
||||
"""
|
||||
:return: the Embedder model used for memory store
|
||||
"""
|
||||
return EmbedderFactory.create(
|
||||
config.embedder.provider,
|
||||
config.embedder.config,
|
||||
{"enable_embeddings": True},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _create_llm(config, llm_provider):
|
||||
"""
|
||||
:return: the llm model used for memory store
|
||||
"""
|
||||
return LlmFactory.create(llm_provider, config.llm.config)
|
||||
|
||||
@staticmethod
|
||||
def _create_vector_store(vector_store_provider, config):
|
||||
"""
|
||||
:param vector_store_provider: name of vector store
|
||||
:param config: the vector_store configuration
|
||||
:return:
|
||||
"""
|
||||
return VectorStoreFactory.create(vector_store_provider, config.vector_store.config)
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters["user_id"])
|
||||
added_entities = self._add_entities(to_be_added, filters["user_id"], entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""
|
||||
Extract all entities mentioned in the query.
|
||||
"""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""
|
||||
Establish relations among the extracted nodes.
|
||||
"""
|
||||
if self.config.graph_store.custom_prompt:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]).replace(
|
||||
"CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}"
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}",
|
||||
},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities["tool_calls"]:
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=False)
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""
|
||||
Get the entities to be deleted from the search output.
|
||||
"""
|
||||
|
||||
search_output_string = format_entities(search_output)
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, filters["user_id"])
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates["tool_calls"]:
|
||||
if item["name"] == "delete_graph_memory":
|
||||
to_be_deleted.append(item["arguments"])
|
||||
# in case if it is not in the correct format
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, user_id):
|
||||
"""
|
||||
Delete the entities from the graph.
|
||||
"""
|
||||
|
||||
results = []
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# Delete the specific relationship between nodes
|
||||
cypher, params = self._delete_entities_cypher(source, destination, relationship, user_id)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
@abstractmethod
|
||||
def _delete_entities_cypher(self, source, destination, relationship, user_id):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for deleting entities in the graph DB
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
def _add_entities(self, to_be_added, user_id, entity_type_map):
|
||||
"""
|
||||
Add the new entities to the graph. Merge the nodes if they already exist.
|
||||
"""
|
||||
|
||||
results = []
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# types
|
||||
source_type = entity_type_map.get(source, "__User__")
|
||||
destination_type = entity_type_map.get(destination, "__User__")
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
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=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,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_search_result,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _add_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
"""
|
||||
if not destination_node_list and source_node_list:
|
||||
return self._add_entities_by_source_cypher(
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id)
|
||||
elif destination_node_list and not source_node_list:
|
||||
return self._add_entities_by_destination_cypher(
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id)
|
||||
elif source_node_list and destination_node_list:
|
||||
return self._add_relationship_entities_cypher(
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id)
|
||||
# else source_node_list and destination_node_list are empty
|
||||
return self._add_new_entities_cypher(
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id)
|
||||
|
||||
@abstractmethod
|
||||
def _add_entities_by_source_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _add_entities_by_destination_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _add_relationship_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def _add_new_entities_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
pass
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
return search_results
|
||||
|
||||
def _search_source_node(self, source_embedding, user_id, threshold=0.9):
|
||||
cypher, params = self._search_source_node_cypher(source_embedding, user_id, threshold)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
@abstractmethod
|
||||
def _search_source_node_cypher(self, source_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for source nodes
|
||||
"""
|
||||
pass
|
||||
|
||||
def _search_destination_node(self, destination_embedding, user_id, threshold=0.9):
|
||||
cypher, params = self._search_destination_node_cypher(destination_embedding, user_id, threshold)
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
@abstractmethod
|
||||
def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for destination nodes
|
||||
"""
|
||||
pass
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters["user_id"])
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
cypher, params = self._delete_all_cypher(filters)
|
||||
self.graph.query(cypher, params=params)
|
||||
|
||||
@abstractmethod
|
||||
def _delete_all_cypher(self, filters):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on filtering criteria.
|
||||
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'contexts': The base data store response for each memory.
|
||||
- 'entities': A list of strings representing the nodes and relationships
|
||||
"""
|
||||
|
||||
# return all nodes and relationships
|
||||
query, params = self._get_all_cypher(filters, top_k)
|
||||
results = self.graph.query(query, params=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.debug(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
@abstractmethod
|
||||
def _get_all_cypher(self, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
|
||||
"""
|
||||
pass
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""
|
||||
Search similar nodes among and their respective incoming and outgoing relations.
|
||||
"""
|
||||
result_relations = []
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
cypher_query, params = self._search_graph_db_cypher(n_embedding, filters, top_k)
|
||||
ans = self.graph.query(cypher_query, params=params)
|
||||
result_relations.extend(ans)
|
||||
|
||||
return result_relations
|
||||
|
||||
@abstractmethod
|
||||
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
|
||||
"""
|
||||
pass
|
||||
|
||||
# Reset is not defined in base.py
|
||||
def reset(self):
|
||||
"""
|
||||
Reset the graph by clearing all nodes and relationships.
|
||||
|
||||
link: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/neptune-graph/client/reset_graph.html
|
||||
"""
|
||||
|
||||
logger.warning("Clearing graph...")
|
||||
graph_id = self.graph.graph_identifier
|
||||
self.graph.client.reset_graph(
|
||||
graphIdentifier=graph_id,
|
||||
skipSnapshot=True,
|
||||
)
|
||||
waiter = self.graph.client.get_waiter("graph_available")
|
||||
waiter.wait(graphIdentifier=graph_id, WaiterConfig={"Delay": 10, "MaxAttempts": 60})
|
||||
@@ -1,511 +0,0 @@
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .base import NeptuneBase
|
||||
|
||||
try:
|
||||
from langchain_aws import NeptuneGraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class MemoryGraph(NeptuneBase):
|
||||
def __init__(self, config):
|
||||
"""
|
||||
Initialize the Neptune DB memory store.
|
||||
"""
|
||||
|
||||
self.config = config
|
||||
|
||||
self.graph = None
|
||||
endpoint = self.config.graph_store.config.endpoint
|
||||
if endpoint and endpoint.startswith("neptune-db://"):
|
||||
host = endpoint.replace("neptune-db://", "")
|
||||
port = 8182
|
||||
self.graph = NeptuneGraph(host, port)
|
||||
|
||||
if not self.graph:
|
||||
raise ValueError("Unable to create a Neptune-DB client: missing 'endpoint' in config")
|
||||
|
||||
self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
|
||||
|
||||
self.embedding_model = NeptuneBase._create_embedding_model(self.config)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.graph_store.llm:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
elif self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
|
||||
# fetch the vector store as a provider
|
||||
self.vector_store_provider = self.config.vector_store.provider
|
||||
if self.config.graph_store.config.collection_name:
|
||||
vector_store_collection_name = self.config.graph_store.config.collection_name
|
||||
else:
|
||||
vector_store_config = self.config.vector_store.config
|
||||
if vector_store_config.collection_name:
|
||||
vector_store_collection_name = vector_store_config.collection_name + "_neptune_vector_store"
|
||||
else:
|
||||
vector_store_collection_name = "mem0_neptune_vector_store"
|
||||
self.config.vector_store.config.collection_name = vector_store_collection_name
|
||||
self.vector_store = NeptuneBase._create_vector_store(self.vector_store_provider, self.config)
|
||||
|
||||
self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
|
||||
self.user_id = None
|
||||
# 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):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for deleting entities in the graph DB
|
||||
|
||||
:param source: source node
|
||||
:param destination: destination node
|
||||
:param relationship: relationship label
|
||||
:param user_id: user_id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{name: $source_name, user_id: $user_id}})
|
||||
-[r:{relationship}]->
|
||||
(m {self.node_label} {{name: $dest_name, user_id: $user_id}})
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(f"_delete_entities\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_source_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source nodes
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
destination_id = str(uuid.uuid4())
|
||||
destination_payload = {
|
||||
"name": destination,
|
||||
"type": destination_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
self.vector_store.insert(
|
||||
vectors=[dest_embedding],
|
||||
payloads=[destination_payload],
|
||||
ids=[destination_id],
|
||||
)
|
||||
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{`~id`: $destination_id, name: $destination_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.updated = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target, id(destination) AS destination_id
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_id": destination_id,
|
||||
"destination_name": destination,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
f"_add_entities:\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_destination_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination_node_list: list of dest nodes
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
source_id = str(uuid.uuid4())
|
||||
source_payload = {
|
||||
"name": source,
|
||||
"type": source_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
self.vector_store.insert(
|
||||
vectors=[source_embedding],
|
||||
payloads=[source_payload],
|
||||
ids=[source_id],
|
||||
)
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{`~id`: $source_id, name: $source_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.updated = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"source_id": source_id,
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_relationship_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source node ids
|
||||
:param destination_node_list: list of dest node ids
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions) + 1,
|
||||
destination.updated = timestamp()
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_new_entities_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
source_id = str(uuid.uuid4())
|
||||
source_payload = {
|
||||
"name": source,
|
||||
"type": source_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
destination_id = str(uuid.uuid4())
|
||||
destination_payload = {
|
||||
"name": destination,
|
||||
"type": destination_type,
|
||||
"user_id": user_id,
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
self.vector_store.insert(
|
||||
vectors=[source_embedding, dest_embedding],
|
||||
payloads=[source_payload, destination_payload],
|
||||
ids=[source_id, destination_id],
|
||||
)
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MERGE (n {source_label} {{name: $source_name, user_id: $user_id, `~id`: $source_id}})
|
||||
ON CREATE SET n.created = timestamp(),
|
||||
n.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET n.mentions = coalesce(n.mentions, 0) + 1
|
||||
WITH n
|
||||
MERGE (m {destination_label} {{name: $dest_name, user_id: $user_id, `~id`: $dest_id}})
|
||||
ON CREATE SET m.created = timestamp(),
|
||||
m.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET m.mentions = coalesce(m.mentions, 0) + 1
|
||||
WITH n, m
|
||||
MERGE (n)-[rel:{relationship}]->(m)
|
||||
ON CREATE SET rel.created = timestamp(), rel.mentions = 1
|
||||
ON MATCH SET rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||
RETURN n.name AS source, type(rel) AS relationship, m.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_id,
|
||||
"dest_id": destination_id,
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_new_entities_cypher:\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _search_source_node_cypher(self, source_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for source nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
source_nodes = self.vector_store.search(
|
||||
query="",
|
||||
vectors=source_embedding,
|
||||
top_k=self.vector_store_limit,
|
||||
filters={"user_id": user_id},
|
||||
)
|
||||
|
||||
ids = [n.id for n in filter(lambda s: s.score > threshold, source_nodes)]
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE source_candidate.user_id = $user_id AND id(source_candidate) IN $ids
|
||||
RETURN id(source_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"ids": ids,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
logger.debug(f"_search_source_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for destination nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
destination_nodes = self.vector_store.search(
|
||||
query="",
|
||||
vectors=destination_embedding,
|
||||
top_k=self.vector_store_limit,
|
||||
filters={"user_id": user_id},
|
||||
)
|
||||
|
||||
ids = [n.id for n in filter(lambda d: d.score > threshold, destination_nodes)]
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE destination_candidate.user_id = $user_id AND id(destination_candidate) IN $ids
|
||||
RETURN id(destination_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"ids": ids,
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
logger.debug(f"_search_destination_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _delete_all_cypher(self, filters):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
# remove the vector store index
|
||||
self.vector_store.reset()
|
||||
|
||||
# create a query that: deletes the nodes of the graph_store
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
|
||||
logger.debug(f"delete_all query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _get_all_cypher(self, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})-[r]->(m {self.node_label} {{user_id: $user_id}})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
return cypher, params
|
||||
|
||||
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
|
||||
|
||||
:param n_embedding: node vector
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
# search vector store for applicable nodes using cosine similarity
|
||||
search_nodes = self.vector_store.search(
|
||||
query="",
|
||||
vectors=n_embedding,
|
||||
top_k=self.vector_store_limit,
|
||||
filters=filters,
|
||||
)
|
||||
|
||||
ids = [n.id for n in search_nodes]
|
||||
|
||||
cypher_query = f"""
|
||||
MATCH (n {self.node_label})-[r]->(m)
|
||||
WHERE n.user_id = $user_id AND id(n) IN $n_ids
|
||||
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
|
||||
UNION
|
||||
MATCH (m)-[r]->(n {self.node_label})
|
||||
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
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {
|
||||
"n_ids": ids,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
logger.debug(f"_search_graph_db\n query={cypher_query}")
|
||||
|
||||
return cypher_query, params
|
||||
@@ -1,475 +0,0 @@
|
||||
import logging
|
||||
|
||||
from .base import NeptuneBase
|
||||
|
||||
try:
|
||||
from botocore.config import Config
|
||||
from langchain_aws import NeptuneAnalyticsGraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_aws is not installed. Please install it using 'make install_all'.")
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph(NeptuneBase):
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
self.graph = None
|
||||
endpoint = self.config.graph_store.config.endpoint
|
||||
app_id = self.config.graph_store.config.app_id
|
||||
if endpoint and endpoint.startswith("neptune-graph://"):
|
||||
graph_identifier = endpoint.replace("neptune-graph://", "")
|
||||
self.graph = NeptuneAnalyticsGraph(graph_identifier = graph_identifier,
|
||||
config = Config(user_agent_appid=app_id))
|
||||
|
||||
if not self.graph:
|
||||
raise ValueError("Unable to create a Neptune client: missing 'endpoint' in config")
|
||||
|
||||
self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
|
||||
|
||||
self.embedding_model = NeptuneBase._create_embedding_model(self.config)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store.llm:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
self.llm = NeptuneBase._create_llm(self.config, self.llm_provider)
|
||||
self.user_id = None
|
||||
# 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):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for deleting entities in the graph DB
|
||||
|
||||
:param source: source node
|
||||
:param destination: destination node
|
||||
:param relationship: relationship label
|
||||
:param user_id: user_id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{name: $source_name, user_id: $user_id}})
|
||||
-[r:{relationship}]->
|
||||
(m {self.node_label} {{name: $dest_name, user_id: $user_id}})
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(f"_delete_entities\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_source_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source nodes
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{name: $destination_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.updated = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH source, destination, $dest_embedding as dest_embedding
|
||||
CALL neptune.algo.vectors.upsert(destination, dest_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_name": destination,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_entities_by_destination_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination_node_list: list of dest nodes
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.updated = timestamp()
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{name: $source_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.updated = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source, destination, $source_embedding as source_embedding
|
||||
CALL neptune.algo.vectors.upsert(source, source_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp(),
|
||||
r.updated = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.updated = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_relationship_entities_cypher(
|
||||
self,
|
||||
source_node_list,
|
||||
destination_node_list,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source_node_list: list of source node ids
|
||||
:param destination_node_list: list of dest node ids
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source {{user_id: $user_id}})
|
||||
WHERE id(source) = $source_id
|
||||
SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.updated = timestamp()
|
||||
WITH source
|
||||
MATCH (destination {{user_id: $user_id}})
|
||||
WHERE id(destination) = $destination_id
|
||||
SET
|
||||
destination.mentions = coalesce(destination.mentions) + 1,
|
||||
destination.updated = timestamp()
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_node_list[0]["id(source_candidate)"],
|
||||
"destination_id": destination_node_list[0]["id(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_entities:\n destination_node_search_result={destination_node_list[0]}\n source_node_search_result={source_node_list[0]}\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _add_new_entities_cypher(
|
||||
self,
|
||||
source,
|
||||
source_embedding,
|
||||
source_type,
|
||||
destination,
|
||||
dest_embedding,
|
||||
destination_type,
|
||||
relationship,
|
||||
user_id,
|
||||
):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters for adding entities in the graph DB
|
||||
|
||||
:param source: source node name
|
||||
:param source_embedding: source node embedding
|
||||
:param source_type: source node label
|
||||
:param destination: destination name
|
||||
:param dest_embedding: destination embedding
|
||||
:param destination_type: destination node label
|
||||
:param relationship: relationship label
|
||||
:param user_id: user id to use
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
cypher = f"""
|
||||
MERGE (n {source_label} {{name: $source_name, user_id: $user_id}})
|
||||
ON CREATE SET n.created = timestamp(),
|
||||
n.updated = timestamp(),
|
||||
n.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
n.mentions = coalesce(n.mentions, 0) + 1,
|
||||
n.updated = timestamp()
|
||||
WITH n, $source_embedding as source_embedding
|
||||
CALL neptune.algo.vectors.upsert(n, source_embedding)
|
||||
WITH n
|
||||
MERGE (m {destination_label} {{name: $dest_name, user_id: $user_id}})
|
||||
ON CREATE SET
|
||||
m.created = timestamp(),
|
||||
m.updated = timestamp(),
|
||||
m.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
m.updated = timestamp(),
|
||||
m.mentions = coalesce(m.mentions, 0) + 1
|
||||
WITH n, m, $dest_embedding as dest_embedding
|
||||
CALL neptune.algo.vectors.upsert(m, dest_embedding)
|
||||
WITH n, m
|
||||
MERGE (n)-[rel:{relationship}]->(m)
|
||||
ON CREATE SET
|
||||
rel.created = timestamp(),
|
||||
rel.updated = timestamp(),
|
||||
rel.mentions = 1
|
||||
ON MATCH SET
|
||||
rel.updated = timestamp(),
|
||||
rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||
RETURN n.name AS source, type(rel) AS relationship, m.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
logger.debug(
|
||||
f"_add_new_entities_cypher:\n query={cypher}"
|
||||
)
|
||||
return cypher, params
|
||||
|
||||
def _search_source_node_cypher(self, source_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for source nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE source_candidate.user_id = $user_id
|
||||
|
||||
WITH source_candidate, $source_embedding as v_embedding
|
||||
CALL neptune.algo.vectors.distanceByEmbedding(
|
||||
v_embedding,
|
||||
source_candidate,
|
||||
{{metric:"CosineSimilarity"}}
|
||||
) YIELD distance
|
||||
WITH source_candidate, distance AS cosine_similarity
|
||||
WHERE cosine_similarity >= $threshold
|
||||
|
||||
WITH source_candidate, cosine_similarity
|
||||
ORDER BY cosine_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN id(source_candidate), cosine_similarity
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
logger.debug(f"_search_source_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _search_destination_node_cypher(self, destination_embedding, user_id, threshold):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for destination nodes
|
||||
|
||||
:param source_embedding: source vector
|
||||
:param user_id: user_id to use
|
||||
:param threshold: the threshold for similarity
|
||||
:return: str, dict
|
||||
"""
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE destination_candidate.user_id = $user_id
|
||||
|
||||
WITH destination_candidate, $destination_embedding as v_embedding
|
||||
CALL neptune.algo.vectors.distanceByEmbedding(
|
||||
v_embedding,
|
||||
destination_candidate,
|
||||
{{metric:"CosineSimilarity"}}
|
||||
) YIELD distance
|
||||
WITH destination_candidate, distance AS cosine_similarity
|
||||
WHERE cosine_similarity >= $threshold
|
||||
|
||||
WITH destination_candidate, cosine_similarity
|
||||
ORDER BY cosine_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN id(destination_candidate), cosine_similarity
|
||||
"""
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
|
||||
logger.debug(f"_search_destination_node\n query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _delete_all_cypher(self, filters):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to delete all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:return: str, dict
|
||||
"""
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
|
||||
logger.debug(f"delete_all query={cypher}")
|
||||
return cypher, params
|
||||
|
||||
def _get_all_cypher(self, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to get all edges/nodes in the memory store
|
||||
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{user_id: $user_id}})-[r]->(m {self.node_label} {{user_id: $user_id}})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
return cypher, params
|
||||
|
||||
def _search_graph_db_cypher(self, n_embedding, filters, top_k):
|
||||
"""
|
||||
Returns the OpenCypher query and parameters to search for similar nodes in the memory store
|
||||
|
||||
:param n_embedding: node vector
|
||||
:param filters: search filters
|
||||
:param top_k: return limit
|
||||
:return: str, dict
|
||||
"""
|
||||
|
||||
cypher_query = f"""
|
||||
MATCH (n {self.node_label})
|
||||
WHERE n.user_id = $user_id
|
||||
WITH n, $n_embedding as n_embedding
|
||||
CALL neptune.algo.vectors.distanceByEmbedding(
|
||||
n_embedding,
|
||||
n,
|
||||
{{metric:"CosineSimilarity"}}
|
||||
) YIELD distance
|
||||
WITH n, distance as similarity
|
||||
WHERE similarity >= $threshold
|
||||
CALL {{
|
||||
WITH n
|
||||
MATCH (n)-[r]->(m)
|
||||
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
|
||||
UNION ALL
|
||||
WITH n
|
||||
MATCH (m)-[r]->(n)
|
||||
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
|
||||
}}
|
||||
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
logger.debug(f"_search_graph_db\n query={cypher_query}")
|
||||
|
||||
return cypher_query, params
|
||||
@@ -1,371 +0,0 @@
|
||||
UPDATE_MEMORY_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "update_graph_memory",
|
||||
"description": "Update the relationship key of an existing graph memory based on new information. This function should be called when there's a need to modify an existing relationship in the knowledge graph. The update should only be performed if the new information is more recent, more accurate, or provides additional context compared to the existing information. The source and destination nodes of the relationship must remain the same as in the existing graph memory; only the relationship itself can be updated.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The new or updated relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
},
|
||||
"required": ["source", "destination", "relationship"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ADD_MEMORY_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add_graph_memory",
|
||||
"description": "Add a new graph memory to the knowledge graph. This function creates a new relationship between two nodes, potentially creating new nodes if they don't exist.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The type of relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
"source_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the source node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
"destination_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the destination node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"destination",
|
||||
"relationship",
|
||||
"source_type",
|
||||
"destination_type",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
NOOP_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "No operation should be performed to the graph entities. This function is called when the system determines that no changes or additions are necessary based on the current input or context. It serves as a placeholder action when no other actions are required, ensuring that the system can explicitly acknowledge situations where no modifications to the graph are needed.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
RELATIONS_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "establish_relationships",
|
||||
"description": "Establish relationships among the entities based on the provided text.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {"type": "string", "description": "The source entity of the relationship."},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The relationship between the source and destination entities.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The destination entity of the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
EXTRACT_ENTITIES_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "extract_entities",
|
||||
"description": "Extract entities and their types from the text.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {"type": "string", "description": "The name or identifier of the entity."},
|
||||
"entity_type": {"type": "string", "description": "The type or category of the entity."},
|
||||
},
|
||||
"required": ["entity", "entity_type"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"description": "An array of entities with their types.",
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
UPDATE_MEMORY_STRUCT_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "update_graph_memory",
|
||||
"description": "Update the relationship key of an existing graph memory based on new information. This function should be called when there's a need to modify an existing relationship in the knowledge graph. The update should only be performed if the new information is more recent, more accurate, or provides additional context compared to the existing information. The source and destination nodes of the relationship must remain the same as in the existing graph memory; only the relationship itself can be updated.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship to be updated. This should match an existing node in the graph.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The new or updated relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
},
|
||||
"required": ["source", "destination", "relationship"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ADD_MEMORY_STRUCT_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "add_graph_memory",
|
||||
"description": "Add a new graph memory to the knowledge graph. This function creates a new relationship between two nodes, potentially creating new nodes if they don't exist.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the new relationship. This can be an existing node or a new node to be created.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The type of relationship between the source and destination nodes. This should be a concise, clear description of how the two nodes are connected.",
|
||||
},
|
||||
"source_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the source node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
"destination_type": {
|
||||
"type": "string",
|
||||
"description": "The type or category of the destination node. This helps in classifying and organizing nodes in the graph.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"destination",
|
||||
"relationship",
|
||||
"source_type",
|
||||
"destination_type",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
NOOP_STRUCT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "No operation should be performed to the graph entities. This function is called when the system determines that no changes or additions are necessary based on the current input or context. It serves as a placeholder action when no other actions are required, ensuring that the system can explicitly acknowledge situations where no modifications to the graph are needed.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
RELATIONS_STRUCT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "establish_relations",
|
||||
"description": "Establish relationships among the entities based on the provided text.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The source entity of the relationship.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The relationship between the source and destination entities.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The destination entity of the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "extract_entities",
|
||||
"description": "Extract entities and their types from the text.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {"type": "string", "description": "The name or identifier of the entity."},
|
||||
"entity_type": {"type": "string", "description": "The type or category of the entity."},
|
||||
},
|
||||
"required": ["entity", "entity_type"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"description": "An array of entities with their types.",
|
||||
}
|
||||
},
|
||||
"required": ["entities"],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "delete_graph_memory",
|
||||
"description": "Delete the relationship between two nodes. This function deletes the existing relationship.",
|
||||
"strict": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The existing relationship between the source and destination nodes that needs to be deleted.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
DELETE_MEMORY_TOOL_GRAPH = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "delete_graph_memory",
|
||||
"description": "Delete the relationship between two nodes. This function deletes the existing relationship.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the source node in the relationship.",
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "The existing relationship between the source and destination nodes that needs to be deleted.",
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "The identifier of the destination node in the relationship.",
|
||||
},
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination",
|
||||
],
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -1,97 +0,0 @@
|
||||
UPDATE_GRAPH_PROMPT = """
|
||||
You are an AI expert specializing in graph memory management and optimization. Your task is to analyze existing graph memories alongside new information, and update the relationships in the memory list to ensure the most accurate, current, and coherent representation of knowledge.
|
||||
|
||||
Input:
|
||||
1. Existing Graph Memories: A list of current graph memories, each containing source, target, and relationship information.
|
||||
2. New Graph Memory: Fresh information to be integrated into the existing graph structure.
|
||||
|
||||
Guidelines:
|
||||
1. Identification: Use the source and target as primary identifiers when matching existing memories with new information.
|
||||
2. Conflict Resolution:
|
||||
- If new information contradicts an existing memory:
|
||||
a) For matching source and target but differing content, update the relationship of the existing memory.
|
||||
b) If the new memory provides more recent or accurate information, update the existing memory accordingly.
|
||||
3. Comprehensive Review: Thoroughly examine each existing graph memory against the new information, updating relationships as necessary. Multiple updates may be required.
|
||||
4. Consistency: Maintain a uniform and clear style across all memories. Each entry should be concise yet comprehensive.
|
||||
5. Semantic Coherence: Ensure that updates maintain or improve the overall semantic structure of the graph.
|
||||
6. Temporal Awareness: If timestamps are available, consider the recency of information when making updates.
|
||||
7. Relationship Refinement: Look for opportunities to refine relationship descriptions for greater precision or clarity.
|
||||
8. Redundancy Elimination: Identify and merge any redundant or highly similar relationships that may result from the update.
|
||||
|
||||
Memory Format:
|
||||
source -- RELATIONSHIP -- destination
|
||||
|
||||
Task Details:
|
||||
======= Existing Graph Memories:=======
|
||||
{existing_memories}
|
||||
|
||||
======= New Graph Memory:=======
|
||||
{new_memories}
|
||||
|
||||
Output:
|
||||
Provide a list of update instructions, each specifying the source, target, and the new relationship to be set. Only include memories that require updates.
|
||||
"""
|
||||
|
||||
EXTRACT_RELATIONS_PROMPT = """
|
||||
|
||||
You are an advanced algorithm designed to extract structured information from text to construct knowledge graphs. Your goal is to capture comprehensive and accurate information. Follow these key principles:
|
||||
|
||||
1. Extract only explicitly stated information from the text.
|
||||
2. Establish relationships among the entities provided.
|
||||
3. Use "USER_ID" as the source entity for any self-references (e.g., "I," "me," "my," etc.) in user messages.
|
||||
CUSTOM_PROMPT
|
||||
|
||||
Relationships:
|
||||
- Use consistent, general, and timeless relationship types.
|
||||
- Example: Prefer "professor" over "became_professor."
|
||||
- Relationships should only be established among the entities explicitly mentioned in the user message.
|
||||
|
||||
Entity Consistency:
|
||||
- Ensure that relationships are coherent and logically align with the context of the message.
|
||||
- Maintain consistent naming for entities across the extracted data.
|
||||
|
||||
Strive to construct a coherent and easily understandable knowledge graph by establishing all the relationships among the entities and adherence to the user’s context.
|
||||
|
||||
Adhere strictly to these guidelines to ensure high-quality knowledge graph extraction."""
|
||||
|
||||
DELETE_RELATIONS_SYSTEM_PROMPT = """
|
||||
You are a graph memory manager specializing in identifying, managing, and optimizing relationships within graph-based memories. Your primary task is to analyze a list of existing relationships and determine which ones should be deleted based on the new information provided.
|
||||
Input:
|
||||
1. Existing Graph Memories: A list of current graph memories, each containing source, relationship, and destination information.
|
||||
2. New Text: The new information to be integrated into the existing graph structure.
|
||||
3. Use "USER_ID" as node for any self-references (e.g., "I," "me," "my," etc.) in user messages.
|
||||
|
||||
Guidelines:
|
||||
1. Identification: Use the new information to evaluate existing relationships in the memory graph.
|
||||
2. Deletion Criteria: Delete a relationship only if it meets at least one of these conditions:
|
||||
- Outdated or Inaccurate: The new information is more recent or accurate.
|
||||
- Contradictory: The new information conflicts with or negates the existing information.
|
||||
3. DO NOT DELETE if their is a possibility of same type of relationship but different destination nodes.
|
||||
4. Comprehensive Analysis:
|
||||
- Thoroughly examine each existing relationship against the new information and delete as necessary.
|
||||
- Multiple deletions may be required based on the new information.
|
||||
5. Semantic Integrity:
|
||||
- Ensure that deletions maintain or improve the overall semantic structure of the graph.
|
||||
- Avoid deleting relationships that are NOT contradictory/outdated to the new information.
|
||||
6. Temporal Awareness: Prioritize recency when timestamps are available.
|
||||
7. Necessity Principle: Only DELETE relationships that must be deleted and are contradictory/outdated to the new information to maintain an accurate and coherent memory graph.
|
||||
|
||||
Note: DO NOT DELETE if their is a possibility of same type of relationship but different destination nodes.
|
||||
|
||||
For example:
|
||||
Existing Memory: alice -- loves_to_eat -- pizza
|
||||
New Information: Alice also loves to eat burger.
|
||||
|
||||
Do not delete in the above example because there is a possibility that Alice loves to eat both pizza and burger.
|
||||
|
||||
Memory Format:
|
||||
source -- relationship -- destination
|
||||
|
||||
Provide a list of deletion instructions, each specifying the relationship to be deleted.
|
||||
"""
|
||||
|
||||
|
||||
def get_delete_messages(existing_memories_string, data, user_id):
|
||||
return DELETE_RELATIONS_SYSTEM_PROMPT.replace(
|
||||
"USER_ID", user_id
|
||||
), f"Here are the existing memories: {existing_memories_string} \n\n New Information: {data}"
|
||||
@@ -1,595 +0,0 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
|
||||
from mem0.memory.utils import format_entities, sanitize_relationship_for_cypher
|
||||
|
||||
try:
|
||||
import age
|
||||
except ImportError:
|
||||
raise ImportError("apache-age-python is not installed. Please install it using pip install apache-age-python")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _cosine_similarity(vec1, vec2):
|
||||
"""Compute cosine similarity between two vectors without numpy."""
|
||||
dot = sum(a * b for a, b in zip(vec1, vec2))
|
||||
norm1 = sum(a * a for a in vec1) ** 0.5
|
||||
norm2 = sum(b * b for b in vec2) ** 0.5
|
||||
if norm1 == 0 or norm2 == 0:
|
||||
return 0.0
|
||||
return dot / (norm1 * norm2)
|
||||
|
||||
|
||||
def _get_similar_nodes(nodes, query_embedding, filters, threshold):
|
||||
"""Find nodes above the similarity threshold from a fetched node list.
|
||||
|
||||
Shared by ``_find_similar_node`` and ``_search_graph_db`` to avoid
|
||||
duplicating the client-side cosine similarity logic.
|
||||
"""
|
||||
matches = []
|
||||
for node in nodes:
|
||||
props = node if isinstance(node, dict) else {}
|
||||
stored_emb = props.get("embedding")
|
||||
if not stored_emb:
|
||||
continue
|
||||
if isinstance(stored_emb, str):
|
||||
stored_emb = json.loads(stored_emb)
|
||||
|
||||
if filters.get("agent_id") and props.get("agent_id") != filters["agent_id"]:
|
||||
continue
|
||||
if filters.get("run_id") and props.get("run_id") != filters["run_id"]:
|
||||
continue
|
||||
|
||||
sim = _cosine_similarity(query_embedding, stored_emb)
|
||||
if sim >= threshold:
|
||||
matches.append({"name": props.get("name"), "similarity": sim, "props": props})
|
||||
|
||||
matches.sort(key=lambda x: x["similarity"], reverse=True)
|
||||
return matches
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
graph_cfg = self.config.graph_store.config
|
||||
self.graph_name = graph_cfg.graph_name
|
||||
|
||||
# Connect using the Apache AGE Python driver (psycopg2-based)
|
||||
self.ag = age.connect(
|
||||
graph=self.graph_name,
|
||||
host=graph_cfg.host,
|
||||
port=graph_cfg.port,
|
||||
dbname=graph_cfg.database,
|
||||
user=graph_cfg.username,
|
||||
password=graph_cfg.password,
|
||||
)
|
||||
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider, self.config.embedder.config, self.config.vector_store.config
|
||||
)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
self.user_id = None
|
||||
self.threshold = self.config.graph_store.threshold if hasattr(self.config.graph_store, "threshold") else 0.7
|
||||
|
||||
# -- helpers ---------------------------------------------------------------
|
||||
|
||||
def _exec_cypher(self, cypher_stmt, cols=None, params=None):
|
||||
"""Execute a Cypher query via the AGE driver and return results.
|
||||
|
||||
Uses ``ag.execCypher`` which delegates to psycopg2's safe parameter
|
||||
substitution (``%s`` placeholders). The *cols* argument specifies the
|
||||
column names in the ``AS (…)`` clause — when ``None`` the driver
|
||||
defaults to a single ``v agtype`` column.
|
||||
|
||||
When *cols* are provided, returns a list of dicts keyed by column name.
|
||||
When *cols* is ``None``, returns vertex/edge property dicts or raw values.
|
||||
"""
|
||||
cursor = self.ag.execCypher(cypher_stmt, cols=cols, params=params)
|
||||
rows = cursor.fetchall()
|
||||
if not rows:
|
||||
return []
|
||||
|
||||
col_names = [desc[0] for desc in cursor.description] if cursor.description else None
|
||||
results = []
|
||||
for row in rows:
|
||||
if col_names and len(col_names) > 1:
|
||||
record = {}
|
||||
for i, col_name in enumerate(col_names):
|
||||
val = row[i]
|
||||
if hasattr(val, "properties"):
|
||||
record[col_name] = val.properties
|
||||
else:
|
||||
record[col_name] = val
|
||||
results.append(record)
|
||||
else:
|
||||
val = row[0] if len(row) == 1 else row
|
||||
if hasattr(val, "properties"):
|
||||
results.append(val.properties)
|
||||
else:
|
||||
results.append(val)
|
||||
return results
|
||||
|
||||
def _fetch_user_nodes_with_embeddings(self, user_id):
|
||||
"""Fetch all nodes with embeddings for a given user_id."""
|
||||
return self._exec_cypher(
|
||||
"MATCH (n {user_id: %s}) WHERE n.embedding IS NOT NULL RETURN n",
|
||||
params=(user_id,),
|
||||
)
|
||||
|
||||
def _find_similar_node(self, embedding, filters, threshold=0.9):
|
||||
"""Find the most similar existing node by cosine similarity.
|
||||
|
||||
Apache AGE does not have a built-in vector index, so we fetch all
|
||||
node embeddings matching the filters and compute cosine similarity
|
||||
on the client side. This is adequate for moderate graph sizes; for
|
||||
very large graphs consider pairing AGE with pgvector.
|
||||
"""
|
||||
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
|
||||
matches = _get_similar_nodes(nodes, embedding, filters, threshold)
|
||||
return matches[0]["props"] if matches else None
|
||||
|
||||
def _merge_node(self, user_id, name, embedding, agent_id=None, run_id=None):
|
||||
"""Create a node if it doesn't exist, or update mentions if it does.
|
||||
|
||||
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so
|
||||
we use ``MERGE … SET`` which always applies the SET clause. Embeddings
|
||||
and optional filter properties are set in a single query.
|
||||
"""
|
||||
set_parts = [
|
||||
"n.embedding = %s",
|
||||
"n.mentions = coalesce(n.mentions, 0) + 1",
|
||||
"n.created = coalesce(n.created, %s)",
|
||||
]
|
||||
params = [user_id, name, json.dumps(embedding), int(time.time() * 1000)]
|
||||
|
||||
if agent_id:
|
||||
set_parts.append("n.agent_id = %s")
|
||||
params.append(agent_id)
|
||||
if run_id:
|
||||
set_parts.append("n.run_id = %s")
|
||||
params.append(run_id)
|
||||
|
||||
set_clause = ", ".join(set_parts)
|
||||
self._exec_cypher(
|
||||
f"MERGE (n {{user_id: %s, name: %s}}) SET {set_clause}",
|
||||
params=tuple(params),
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""Close the underlying database connection."""
|
||||
if self.ag:
|
||||
self.ag.close()
|
||||
|
||||
# -- public API ------------------------------------------------------------
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
list: A list of dicts with keys "source", "relationship", "destination".
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
"""Delete all nodes and relationships for a user or specific agent."""
|
||||
where_parts = ["n.user_id = %s"]
|
||||
params = [filters["user_id"]]
|
||||
if filters.get("agent_id"):
|
||||
where_parts.append("n.agent_id = %s")
|
||||
params.append(filters["agent_id"])
|
||||
if filters.get("run_id"):
|
||||
where_parts.append("n.run_id = %s")
|
||||
params.append(filters["run_id"])
|
||||
where_clause = " AND ".join(where_parts)
|
||||
|
||||
self._exec_cypher(
|
||||
f"MATCH (n) WHERE {where_clause} DETACH DELETE n",
|
||||
params=tuple(params),
|
||||
)
|
||||
self.ag.commit()
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'source': The source node name.
|
||||
- 'relationship': The relationship type.
|
||||
- 'target': The target node name.
|
||||
"""
|
||||
where_parts = ["n.user_id = %s", "m.user_id = %s"]
|
||||
params = [filters["user_id"], filters["user_id"]]
|
||||
if filters.get("agent_id"):
|
||||
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
|
||||
params.extend([filters["agent_id"], filters["agent_id"]])
|
||||
if filters.get("run_id"):
|
||||
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
|
||||
params.extend([filters["run_id"], filters["run_id"]])
|
||||
where_clause = " AND ".join(where_parts)
|
||||
params.append(top_k)
|
||||
|
||||
results = self._exec_cypher(
|
||||
f"MATCH (n)-[r]->(m) WHERE {where_clause} "
|
||||
f"RETURN n.name, type(r), m.name LIMIT %s",
|
||||
cols=["source", "relationship", "target"],
|
||||
params=tuple(params),
|
||||
)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
return final_results
|
||||
|
||||
# -- LLM-driven extraction -------------------------------------------------
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
if "entity" in item and "entity_type" in item:
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Establish relations among the extracted nodes."""
|
||||
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
if self.config.graph_store.custom_prompt:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities and extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
# -- graph DB operations ---------------------------------------------------
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""Search similar nodes and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
|
||||
nodes = self._fetch_user_nodes_with_embeddings(filters["user_id"])
|
||||
similar_nodes = _get_similar_nodes(nodes, n_embedding, filters, self.threshold)
|
||||
|
||||
# Build WHERE clause for relationship target filtering
|
||||
rel_where_parts = ["m.user_id = %s"]
|
||||
rel_params_suffix = [filters["user_id"]]
|
||||
if filters.get("agent_id"):
|
||||
rel_where_parts.append("m.agent_id = %s")
|
||||
rel_params_suffix.append(filters["agent_id"])
|
||||
if filters.get("run_id"):
|
||||
rel_where_parts.append("m.run_id = %s")
|
||||
rel_params_suffix.append(filters["run_id"])
|
||||
rel_where = " AND ".join(rel_where_parts)
|
||||
|
||||
# For each similar node, fetch its relationships
|
||||
for sn in similar_nodes[:top_k]:
|
||||
node_name = sn["name"]
|
||||
similarity = sn["similarity"]
|
||||
|
||||
out_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
|
||||
out_results = self._exec_cypher(
|
||||
f"MATCH (n {{user_id: %s, name: %s}})-[r]->(m) "
|
||||
f"WHERE {rel_where} "
|
||||
f"RETURN n.name, type(r), m.name",
|
||||
cols=["source", "relationship", "destination"],
|
||||
params=out_params,
|
||||
)
|
||||
|
||||
in_params = (filters["user_id"], node_name) + tuple(rel_params_suffix)
|
||||
in_results = self._exec_cypher(
|
||||
f"MATCH (n {{user_id: %s, name: %s}})<-[r]-(m) "
|
||||
f"WHERE {rel_where} "
|
||||
f"RETURN m.name, type(r), n.name",
|
||||
cols=["source", "relationship", "destination"],
|
||||
params=in_params,
|
||||
)
|
||||
|
||||
for rel in out_results + in_results:
|
||||
rel["similarity"] = similarity
|
||||
result_relations.append(rel)
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [DELETE_MEMORY_STRUCT_TOOL_GRAPH]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates.get("tool_calls", []):
|
||||
if item.get("name") == "delete_graph_memory":
|
||||
to_be_deleted.append(item.get("arguments"))
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id")
|
||||
run_id = filters.get("run_id")
|
||||
results = []
|
||||
|
||||
try:
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
where_parts = [
|
||||
"n.user_id = %s", "n.name = %s",
|
||||
"m.user_id = %s", "m.name = %s",
|
||||
]
|
||||
params = [user_id, source, user_id, destination]
|
||||
if agent_id:
|
||||
where_parts.extend(["n.agent_id = %s", "m.agent_id = %s"])
|
||||
params.extend([agent_id, agent_id])
|
||||
if run_id:
|
||||
where_parts.extend(["n.run_id = %s", "m.run_id = %s"])
|
||||
params.extend([run_id, run_id])
|
||||
where_clause = " AND ".join(where_parts)
|
||||
|
||||
result = self._exec_cypher(
|
||||
f"MATCH (n)-[r:{relationship}]->(m) "
|
||||
f"WHERE {where_clause} "
|
||||
f"DELETE r "
|
||||
f"RETURN n.name, type(r), m.name",
|
||||
cols=["source", "relationship", "target"],
|
||||
params=tuple(params),
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
self.ag.commit()
|
||||
except Exception:
|
||||
self.ag.rollback()
|
||||
raise
|
||||
|
||||
return results
|
||||
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add new entities to the graph. Merge nodes if they already exist.
|
||||
|
||||
Apache AGE does not support ``ON CREATE SET`` / ``ON MATCH SET``, so we
|
||||
use ``MERGE … SET`` which always applies. The ``coalesce`` pattern
|
||||
ensures ``created`` is only set on the first merge.
|
||||
"""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id")
|
||||
run_id = filters.get("run_id")
|
||||
results = []
|
||||
|
||||
try:
|
||||
for item in to_be_added:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
dest_embedding = self.embedding_model.embed(destination)
|
||||
|
||||
source_match = self._find_similar_node(source_embedding, filters, threshold=self.threshold)
|
||||
dest_match = self._find_similar_node(dest_embedding, filters, threshold=self.threshold)
|
||||
|
||||
effective_source = source_match["name"] if source_match else source
|
||||
effective_dest = dest_match["name"] if dest_match else destination
|
||||
|
||||
# Merge source and destination nodes
|
||||
self._merge_node(user_id, effective_source, source_embedding, agent_id, run_id)
|
||||
self._merge_node(user_id, effective_dest, dest_embedding, agent_id, run_id)
|
||||
|
||||
# Merge relationship
|
||||
result = self._exec_cypher(
|
||||
f"MATCH (s {{user_id: %s, name: %s}}), (d {{user_id: %s, name: %s}}) "
|
||||
f"MERGE (s)-[r:{relationship}]->(d) "
|
||||
f"RETURN s.name, type(r), d.name",
|
||||
cols=["source", "relationship", "target"],
|
||||
params=(user_id, effective_source, user_id, effective_dest),
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
self.ag.commit()
|
||||
except Exception:
|
||||
self.ag.rollback()
|
||||
raise
|
||||
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
for item in entity_list:
|
||||
item["source"] = item["source"].lower().replace(" ", "_")
|
||||
item["relationship"] = sanitize_relationship_for_cypher(item["relationship"].lower().replace(" ", "_"))
|
||||
item["destination"] = item["destination"].lower().replace(" ", "_")
|
||||
return entity_list
|
||||
|
||||
def reset(self):
|
||||
"""Reset the graph by clearing all nodes and relationships."""
|
||||
logger.warning("Clearing graph...")
|
||||
self._exec_cypher("MATCH (n) DETACH DELETE n")
|
||||
self.ag.commit()
|
||||
@@ -1,744 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
from langchain_neo4j import Neo4jGraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_neo4j is not installed. Please install it using pip install langchain-neo4j")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.graph = Neo4jGraph(
|
||||
url=self.config.graph_store.config.url,
|
||||
username=self.config.graph_store.config.username,
|
||||
password=self.config.graph_store.config.password,
|
||||
database=self.config.graph_store.config.database,
|
||||
refresh_schema=False,
|
||||
driver_config={"notifications_min_severity": "OFF"},
|
||||
)
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider, self.config.embedder.config, self.config.vector_store.config
|
||||
)
|
||||
self.node_label = ":`__Entity__`" if self.config.graph_store.config.base_label else ""
|
||||
|
||||
if self.config.graph_store.config.base_label:
|
||||
# Safely add user_id index
|
||||
try:
|
||||
self.graph.query(f"CREATE INDEX entity_single IF NOT EXISTS FOR (n {self.node_label}) ON (n.user_id)")
|
||||
except Exception:
|
||||
pass
|
||||
try: # Safely try to add composite index (Enterprise only)
|
||||
self.graph.query(
|
||||
f"CREATE INDEX entity_composite IF NOT EXISTS FOR (n {self.node_label}) ON (n.name, n.user_id)"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
self.user_id = None
|
||||
# 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):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
# TODO: Batch queries with APOC plugin
|
||||
# TODO: Add more filter support
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then soft-deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
self.graph.query(cypher, params=params)
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'contexts': The base data store response for each memory.
|
||||
- 'entities': A list of strings representing the nodes and relationships
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
|
||||
# Build node properties based on filters
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
query = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
results = self.graph.query(query, params=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Establish relations among the extracted nodes."""
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
if self.config.graph_store.custom_prompt:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
# Add the custom prompt line if configured
|
||||
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""Search similar nodes among and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
|
||||
cypher_query = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
WHERE n.embedding IS NOT NULL
|
||||
WITH n, round(2 * vector.similarity.cosine(n.embedding, $n_embedding) - 1, 4) AS similarity // denormalize for backward compatibility
|
||||
WHERE similarity >= $threshold
|
||||
CALL {{
|
||||
WITH n
|
||||
MATCH (n)-[r]->(m {self.node_label} {{{node_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
RETURN n.name AS source, elementId(n) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, m.name AS destination, elementId(m) AS destination_id
|
||||
UNION
|
||||
WITH n
|
||||
MATCH (n)<-[r]-(m {self.node_label} {{{node_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
RETURN m.name AS source, elementId(m) AS source_id, type(r) AS relationship, elementId(r) AS relation_id, n.name AS destination, elementId(n) AS destination_id
|
||||
}}
|
||||
WITH distinct source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
RETURN source, source_id, relationship, relation_id, destination, destination_id, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit
|
||||
"""
|
||||
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
|
||||
ans = self.graph.query(cypher_query, params=params)
|
||||
result_relations.extend(ans)
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates.get("tool_calls", []):
|
||||
if item.get("name") == "delete_graph_memory":
|
||||
to_be_deleted.append(item.get("arguments"))
|
||||
# Clean entities formatting
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# Build the agent filter for the query
|
||||
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
# Build node properties for filtering
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
# Soft-delete: mark relationship as invalid instead of removing it,
|
||||
# enabling temporal reasoning over historical graph state.
|
||||
# See: https://github.com/mem0ai/mem0/issues/4187
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{source_props_str}}})
|
||||
-[r:{relationship}]->
|
||||
(m {self.node_label} {{{dest_props_str}}})
|
||||
WHERE r.valid IS NULL OR r.valid = true
|
||||
SET r.valid = false, r.invalidated_at = datetime()
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add the new entities to the graph. Merge the nodes if they already exist."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# types
|
||||
source_type = entity_type_map.get(source, "__User__")
|
||||
source_label = self.node_label if self.node_label else f":`{source_type}`"
|
||||
source_extra_set = f", source:`{source_type}`" if self.node_label else ""
|
||||
destination_type = entity_type_map.get(destination, "__User__")
|
||||
destination_label = self.node_label if self.node_label else f":`{destination_type}`"
|
||||
destination_extra_set = f", destination:`{destination_type}`" if self.node_label else ""
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
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=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:
|
||||
# Build destination MERGE properties
|
||||
merge_props = ["name: $destination_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE elementId(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH source, destination
|
||||
CALL db.create.setNodeVectorProperty(destination, 'embedding', $destination_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["elementId(source_candidate)"],
|
||||
"destination_name": destination,
|
||||
"destination_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
elif destination_node_search_result and not source_node_search_result:
|
||||
# Build source MERGE properties
|
||||
merge_props = ["name: $source_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination)
|
||||
WHERE elementId(destination) = $destination_id
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source, destination
|
||||
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_search_result[0]["elementId(destination_candidate)"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
elif source_node_search_result and destination_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE elementId(source) = $source_id
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MATCH (destination)
|
||||
WHERE elementId(destination) = $destination_id
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["elementId(source_candidate)"],
|
||||
"destination_id": destination_node_search_result[0]["elementId(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
|
||||
else:
|
||||
# Build dynamic MERGE props for both source and destination
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
cypher = f"""
|
||||
MERGE (source {source_label} {{{source_props_str}}})
|
||||
ON CREATE SET source.created = timestamp(),
|
||||
source.mentions = 1
|
||||
{source_extra_set}
|
||||
ON MATCH SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
CALL db.create.setNodeVectorProperty(source, 'embedding', $source_embedding)
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{dest_props_str}}})
|
||||
ON CREATE SET destination.created = timestamp(),
|
||||
destination.mentions = 1
|
||||
{destination_extra_set}
|
||||
ON MATCH SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH source, destination
|
||||
CALL db.create.setNodeVectorProperty(destination, 'embedding', $dest_embedding)
|
||||
WITH source, destination
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp(),
|
||||
r.mentions = 1,
|
||||
r.valid = true
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1,
|
||||
r.valid = true,
|
||||
r.updated_at = timestamp(),
|
||||
r.invalidated_at = null
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
params["run_id"] = run_id
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=True)
|
||||
|
||||
def _search_source_node(self, source_embedding, filters, threshold=0.9):
|
||||
# Build WHERE conditions
|
||||
where_conditions = ["source_candidate.embedding IS NOT NULL", "source_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("source_candidate.agent_id = $agent_id")
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("source_candidate.run_id = $run_id")
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH source_candidate,
|
||||
round(2 * vector.similarity.cosine(source_candidate.embedding, $source_embedding) - 1, 4) AS source_similarity // denormalize for backward compatibility
|
||||
WHERE source_similarity >= $threshold
|
||||
|
||||
WITH source_candidate, source_similarity
|
||||
ORDER BY source_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN elementId(source_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
|
||||
# Build WHERE conditions
|
||||
where_conditions = ["destination_candidate.embedding IS NOT NULL", "destination_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("destination_candidate.agent_id = $agent_id")
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("destination_candidate.run_id = $run_id")
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH destination_candidate,
|
||||
round(2 * vector.similarity.cosine(destination_candidate.embedding, $destination_embedding) - 1, 4) AS destination_similarity // denormalize for backward compatibility
|
||||
|
||||
WHERE destination_similarity >= $threshold
|
||||
|
||||
WITH destination_candidate, destination_similarity
|
||||
ORDER BY destination_similarity DESC
|
||||
LIMIT 1
|
||||
|
||||
RETURN elementId(destination_candidate)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
# Reset is not defined in base.py
|
||||
def reset(self):
|
||||
"""Reset the graph by clearing all nodes and relationships."""
|
||||
logger.warning("Clearing graph...")
|
||||
cypher_query = """
|
||||
MATCH (n) DETACH DELETE n
|
||||
"""
|
||||
return self.graph.query(cypher_query)
|
||||
@@ -1,732 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
import kuzu
|
||||
except ImportError:
|
||||
raise ImportError("kuzu is not installed. Please install it using pip install kuzu")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider,
|
||||
self.config.embedder.config,
|
||||
self.config.vector_store.config,
|
||||
)
|
||||
self.embedding_dims = self.embedding_model.config.embedding_dims
|
||||
|
||||
if self.embedding_dims is None or self.embedding_dims <= 0:
|
||||
raise ValueError(f"embedding_dims must be a positive integer. Given: {self.embedding_dims}")
|
||||
|
||||
self.db = kuzu.Database(self.config.graph_store.config.db)
|
||||
self.graph = kuzu.Connection(self.db)
|
||||
|
||||
self.node_label = ":Entity"
|
||||
self.rel_label = ":CONNECTED_TO"
|
||||
self.kuzu_create_schema()
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
|
||||
self.user_id = None
|
||||
# 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(
|
||||
"""
|
||||
CREATE NODE TABLE IF NOT EXISTS Entity(
|
||||
id SERIAL PRIMARY KEY,
|
||||
user_id STRING,
|
||||
agent_id STRING,
|
||||
run_id STRING,
|
||||
name STRING,
|
||||
mentions INT64,
|
||||
created TIMESTAMP,
|
||||
embedding FLOAT[]);
|
||||
"""
|
||||
)
|
||||
self.kuzu_execute(
|
||||
"""
|
||||
CREATE REL TABLE IF NOT EXISTS CONNECTED_TO(
|
||||
FROM Entity TO Entity,
|
||||
name STRING,
|
||||
mentions INT64,
|
||||
created TIMESTAMP,
|
||||
updated TIMESTAMP
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def kuzu_execute(self, query, parameters=None):
|
||||
results = self.graph.execute(query, parameters)
|
||||
return list(results.rows_as_dict())
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=5):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=top_k)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id, run_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
if filters.get("agent_id"):
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
params["run_id"] = filters["run_id"]
|
||||
self.kuzu_execute(cypher, parameters=params)
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'contexts': The base data store response for each memory.
|
||||
- 'entities': A list of strings representing the nodes and relationships
|
||||
"""
|
||||
|
||||
params = {
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
# Build node properties based on filters
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
query = f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})-[r]->(m {self.node_label} {{{node_props_str}}})
|
||||
RETURN
|
||||
n.name AS source,
|
||||
r.name AS relationship,
|
||||
m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
results = self.kuzu_execute(query, parameters=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Establish relations among the extracted nodes."""
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
if self.config.graph_store.custom_prompt:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
# Add the custom prompt line if configured
|
||||
system_content = system_content.replace("CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}")
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
system_content = EXTRACT_RELATIONS_PROMPT.replace("USER_ID", user_identity)
|
||||
messages = [
|
||||
{"role": "system", "content": system_content},
|
||||
{"role": "user", "content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}"},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100, threshold=None):
|
||||
"""Search similar nodes among and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
params = {
|
||||
"threshold": threshold if threshold else self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
# Build node properties for filtering
|
||||
node_props = ["user_id: $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
node_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
node_props.append("run_id: $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
node_props_str = ", ".join(node_props)
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
params["n_embedding"] = n_embedding
|
||||
|
||||
results = []
|
||||
for match_fragment in [
|
||||
f"(n)-[r]->(m {self.node_label} {{{node_props_str}}}) WITH n as src, r, m as dst, similarity",
|
||||
f"(m {self.node_label} {{{node_props_str}}})-[r]->(n) WITH m as src, r, n as dst, similarity"
|
||||
]:
|
||||
results.extend(self.kuzu_execute(
|
||||
f"""
|
||||
MATCH (n {self.node_label} {{{node_props_str}}})
|
||||
WHERE n.embedding IS NOT NULL
|
||||
WITH n, array_cosine_similarity(n.embedding, CAST($n_embedding,'FLOAT[{self.embedding_dims}]')) AS similarity
|
||||
WHERE similarity >= CAST($threshold, 'DOUBLE')
|
||||
MATCH {match_fragment}
|
||||
RETURN
|
||||
src.name AS source,
|
||||
id(src) AS source_id,
|
||||
r.name AS relationship,
|
||||
id(r) AS relation_id,
|
||||
dst.name AS destination,
|
||||
id(dst) AS destination_id,
|
||||
similarity
|
||||
LIMIT $limit
|
||||
""",
|
||||
parameters=params))
|
||||
|
||||
# Kuzu does not support sort/limit over unions. Do it manually for now.
|
||||
result_relations.extend(sorted(results, key=lambda x: x["similarity"], reverse=True)[:top_k])
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
|
||||
# Compose user identification string for prompt
|
||||
user_identity = f"user_id: {filters['user_id']}"
|
||||
if filters.get("agent_id"):
|
||||
user_identity += f", agent_id: {filters['agent_id']}"
|
||||
if filters.get("run_id"):
|
||||
user_identity += f", run_id: {filters['run_id']}"
|
||||
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, user_identity)
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
to_be_deleted = []
|
||||
for item in memory_updates.get("tool_calls", []):
|
||||
if item.get("name") == "delete_graph_memory":
|
||||
to_be_deleted.append(item.get("arguments"))
|
||||
# Clean entities formatting
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
"relationship_name": relationship,
|
||||
}
|
||||
# Build node properties for filtering
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
# Delete the specific relationship between nodes
|
||||
cypher = f"""
|
||||
MATCH (n {self.node_label} {{{source_props_str}}})
|
||||
-[r {self.rel_label} {{name: $relationship_name}}]->
|
||||
(m {self.node_label} {{{dest_props_str}}})
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
r.name AS relationship,
|
||||
m.name AS target
|
||||
"""
|
||||
|
||||
result = self.kuzu_execute(cypher, parameters=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add the new entities to the graph. Merge the nodes if they already exist."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
run_id = filters.get("run_id", None)
|
||||
results = []
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
source_label = self.node_label
|
||||
|
||||
destination = item["destination"]
|
||||
destination_label = self.node_label
|
||||
|
||||
relationship = item["relationship"]
|
||||
relationship_label = self.rel_label
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
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=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 = {
|
||||
"table_id": source_node_search_result[0]["id"]["table"],
|
||||
"offset_id": source_node_search_result[0]["id"]["offset"],
|
||||
"destination_name": destination,
|
||||
"destination_embedding": dest_embedding,
|
||||
"relationship_name": relationship,
|
||||
"user_id": user_id,
|
||||
}
|
||||
# Build source MERGE properties
|
||||
merge_props = ["name: $destination_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE id(source) = internal_id($table_id, $offset_id)
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
destination.created = current_timestamp(),
|
||||
destination.mentions = 1,
|
||||
destination.embedding = CAST($destination_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.embedding = CAST($destination_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source, destination
|
||||
MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = current_timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
r.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
elif destination_node_search_result and not source_node_search_result:
|
||||
params = {
|
||||
"table_id": destination_node_search_result[0]["id"]["table"],
|
||||
"offset_id": destination_node_search_result[0]["id"]["offset"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"relationship_name": relationship,
|
||||
}
|
||||
# Build source MERGE properties
|
||||
merge_props = ["name: $source_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
merge_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
merge_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
merge_props_str = ", ".join(merge_props)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination)
|
||||
WHERE id(destination) = internal_id($table_id, $offset_id)
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
WITH destination
|
||||
MERGE (source {source_label} {{{merge_props_str}}})
|
||||
ON CREATE SET
|
||||
source.created = current_timestamp(),
|
||||
source.mentions = 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source, destination
|
||||
MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = current_timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET
|
||||
r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
r.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
elif source_node_search_result and destination_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source)
|
||||
WHERE id(source) = internal_id($src_table, $src_offset)
|
||||
SET source.mentions = coalesce(source.mentions, 0) + 1
|
||||
WITH source
|
||||
MATCH (destination)
|
||||
WHERE id(destination) = internal_id($dst_table, $dst_offset)
|
||||
SET destination.mentions = coalesce(destination.mentions, 0) + 1
|
||||
MERGE (source)-[r {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = current_timestamp(),
|
||||
r.updated = current_timestamp(),
|
||||
r.mentions = 1
|
||||
ON MATCH SET r.mentions = coalesce(r.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
r.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"src_table": source_node_search_result[0]["id"]["table"],
|
||||
"src_offset": source_node_search_result[0]["id"]["offset"],
|
||||
"dst_table": destination_node_search_result[0]["id"]["table"],
|
||||
"dst_offset": destination_node_search_result[0]["id"]["offset"],
|
||||
"relationship_name": relationship,
|
||||
}
|
||||
else:
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"relationship_name": relationship,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
# Build dynamic MERGE props for both source and destination
|
||||
source_props = ["name: $source_name", "user_id: $user_id"]
|
||||
dest_props = ["name: $dest_name", "user_id: $user_id"]
|
||||
if agent_id:
|
||||
source_props.append("agent_id: $agent_id")
|
||||
dest_props.append("agent_id: $agent_id")
|
||||
params["agent_id"] = agent_id
|
||||
if run_id:
|
||||
source_props.append("run_id: $run_id")
|
||||
dest_props.append("run_id: $run_id")
|
||||
params["run_id"] = run_id
|
||||
source_props_str = ", ".join(source_props)
|
||||
dest_props_str = ", ".join(dest_props)
|
||||
|
||||
cypher = f"""
|
||||
MERGE (source {source_label} {{{source_props_str}}})
|
||||
ON CREATE SET
|
||||
source.created = current_timestamp(),
|
||||
source.mentions = 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
source.mentions = coalesce(source.mentions, 0) + 1,
|
||||
source.embedding = CAST($source_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source
|
||||
MERGE (destination {destination_label} {{{dest_props_str}}})
|
||||
ON CREATE SET
|
||||
destination.created = current_timestamp(),
|
||||
destination.mentions = 1,
|
||||
destination.embedding = CAST($dest_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
ON MATCH SET
|
||||
destination.mentions = coalesce(destination.mentions, 0) + 1,
|
||||
destination.embedding = CAST($dest_embedding,'FLOAT[{self.embedding_dims}]')
|
||||
WITH source, destination
|
||||
MERGE (source)-[rel {relationship_label} {{name: $relationship_name}}]->(destination)
|
||||
ON CREATE SET
|
||||
rel.created = current_timestamp(),
|
||||
rel.mentions = 1
|
||||
ON MATCH SET
|
||||
rel.mentions = coalesce(rel.mentions, 0) + 1
|
||||
RETURN
|
||||
source.name AS source,
|
||||
rel.name AS relationship,
|
||||
destination.name AS target
|
||||
"""
|
||||
|
||||
result = self.kuzu_execute(cypher, parameters=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=False)
|
||||
|
||||
def _search_source_node(self, source_embedding, filters, threshold=0.9):
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
where_conditions = ["source_candidate.embedding IS NOT NULL", "source_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("source_candidate.agent_id = $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("source_candidate.run_id = $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (source_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH source_candidate,
|
||||
array_cosine_similarity(source_candidate.embedding, CAST($source_embedding,'FLOAT[{self.embedding_dims}]')) AS source_similarity
|
||||
|
||||
WHERE source_similarity >= $threshold
|
||||
|
||||
WITH source_candidate, source_similarity
|
||||
ORDER BY source_similarity DESC
|
||||
LIMIT 2
|
||||
|
||||
RETURN id(source_candidate) as id, source_similarity
|
||||
"""
|
||||
|
||||
return self.kuzu_execute(cypher, parameters=params)
|
||||
|
||||
def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": filters["user_id"],
|
||||
"threshold": threshold,
|
||||
}
|
||||
where_conditions = ["destination_candidate.embedding IS NOT NULL", "destination_candidate.user_id = $user_id"]
|
||||
if filters.get("agent_id"):
|
||||
where_conditions.append("destination_candidate.agent_id = $agent_id")
|
||||
params["agent_id"] = filters["agent_id"]
|
||||
if filters.get("run_id"):
|
||||
where_conditions.append("destination_candidate.run_id = $run_id")
|
||||
params["run_id"] = filters["run_id"]
|
||||
where_clause = " AND ".join(where_conditions)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (destination_candidate {self.node_label})
|
||||
WHERE {where_clause}
|
||||
|
||||
WITH destination_candidate,
|
||||
array_cosine_similarity(destination_candidate.embedding, CAST($destination_embedding,'FLOAT[{self.embedding_dims}]')) AS destination_similarity
|
||||
|
||||
WHERE destination_similarity >= $threshold
|
||||
|
||||
WITH destination_candidate, destination_similarity
|
||||
ORDER BY destination_similarity DESC
|
||||
LIMIT 2
|
||||
|
||||
RETURN id(destination_candidate) as id, destination_similarity
|
||||
"""
|
||||
|
||||
return self.kuzu_execute(cypher, parameters=params)
|
||||
|
||||
# Reset is not defined in base.py
|
||||
def reset(self):
|
||||
"""Reset the graph by clearing all nodes and relationships."""
|
||||
logger.warning("Clearing graph...")
|
||||
cypher_query = """
|
||||
MATCH (n) DETACH DELETE n
|
||||
"""
|
||||
return self.kuzu_execute(cypher_query)
|
||||
+18
-211
@@ -1,5 +1,4 @@
|
||||
import asyncio
|
||||
import concurrent
|
||||
import gc
|
||||
import hashlib
|
||||
import json
|
||||
@@ -41,7 +40,6 @@ from mem0.memory.utils import (
|
||||
)
|
||||
from mem0.utils.factory import (
|
||||
EmbedderFactory,
|
||||
GraphStoreFactory,
|
||||
LlmFactory,
|
||||
VectorStoreFactory,
|
||||
RerankerFactory,
|
||||
@@ -217,11 +215,6 @@ class Memory(MemoryBase):
|
||||
# Entity store is initialized lazily on first use
|
||||
self._entity_store = None
|
||||
|
||||
if self.config.graph_store.config:
|
||||
provider = self.config.graph_store.provider
|
||||
self.graph = GraphStoreFactory.create(provider, self.config)
|
||||
else:
|
||||
self.graph = None
|
||||
if MEM0_TELEMETRY:
|
||||
# Create telemetry config manually to avoid deepcopy issues with thread locks
|
||||
telemetry_config_dict = {}
|
||||
@@ -322,13 +315,6 @@ class Memory(MemoryBase):
|
||||
|
||||
@staticmethod
|
||||
def _process_config(config_dict: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if "graph_store" in config_dict:
|
||||
if "vector_store" not in config_dict and "embedder" in config_dict:
|
||||
config_dict["vector_store"] = {}
|
||||
config_dict["vector_store"]["config"] = {}
|
||||
config_dict["vector_store"]["config"]["embedding_model_dims"] = config_dict["embedder"]["config"][
|
||||
"embedding_dims"
|
||||
]
|
||||
try:
|
||||
return config_dict
|
||||
except ValidationError as e:
|
||||
@@ -339,20 +325,20 @@ class Memory(MemoryBase):
|
||||
"""Determine whether to use agent memory extraction based on the logic:
|
||||
- If agent_id is present and messages contain assistant role -> True
|
||||
- Otherwise -> False
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries
|
||||
metadata: Metadata containing user_id, agent_id, etc.
|
||||
|
||||
|
||||
Returns:
|
||||
bool: True if should use agent memory extraction, False for user memory extraction
|
||||
"""
|
||||
# Check if agent_id is present in metadata
|
||||
has_agent_id = metadata.get("agent_id") is not None
|
||||
|
||||
|
||||
# Check if there are assistant role messages
|
||||
has_assistant_messages = any(msg.get("role") == "assistant" for msg in messages)
|
||||
|
||||
|
||||
# Use agent memory extraction if agent_id is present and there are assistant messages
|
||||
return has_agent_id and has_assistant_messages
|
||||
|
||||
@@ -393,14 +379,12 @@ class Memory(MemoryBase):
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the result of the memory addition operation, typically
|
||||
including a list of memory items affected (added, updated) under a "results" key,
|
||||
and potentially "relations" if graph store is enabled.
|
||||
including a list of memory items affected (added, updated) under a "results" key.
|
||||
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", "event": "ADD"}]}`
|
||||
|
||||
Raises:
|
||||
Mem0ValidationError: If input validation fails (invalid memory_type, messages format, etc.).
|
||||
VectorStoreError: If vector store operations fail.
|
||||
GraphStoreError: If graph store operations fail.
|
||||
EmbeddingError: If embedding generation fails.
|
||||
LLMError: If LLM operations fail.
|
||||
DatabaseError: If database operations fail.
|
||||
@@ -444,21 +428,6 @@ class Memory(MemoryBase):
|
||||
else:
|
||||
messages = parse_vision_messages(messages)
|
||||
|
||||
if self.graph:
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future1 = executor.submit(self._add_to_vector_store, messages, processed_metadata, effective_filters, infer)
|
||||
future2 = executor.submit(self._add_to_graph, messages, effective_filters)
|
||||
|
||||
concurrent.futures.wait([future1, future2])
|
||||
|
||||
vector_store_result = future1.result()
|
||||
graph_result = future2.result()
|
||||
|
||||
return {
|
||||
"results": vector_store_result,
|
||||
"relations": graph_result,
|
||||
}
|
||||
|
||||
vector_store_result = self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
|
||||
return {"results": vector_store_result}
|
||||
|
||||
@@ -773,17 +742,6 @@ class Memory(MemoryBase):
|
||||
)
|
||||
return returned_memories
|
||||
|
||||
def _add_to_graph(self, messages, filters):
|
||||
added_entities = []
|
||||
if self.graph:
|
||||
if filters.get("user_id") is None:
|
||||
filters["user_id"] = "user"
|
||||
|
||||
data = "\n".join([msg["content"] for msg in messages if "content" in msg and msg["role"] != "system"])
|
||||
added_entities = self.graph.add(data, filters)
|
||||
|
||||
return added_entities
|
||||
|
||||
def get(self, memory_id):
|
||||
"""
|
||||
Retrieve a memory by ID.
|
||||
@@ -849,9 +807,7 @@ class Memory(MemoryBase):
|
||||
limit (int, optional): The maximum number of memories to return. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing a list of memories under the "results" key,
|
||||
and potentially "relations" if graph store is enabled. For API v1.0,
|
||||
it might return a direct list (see deprecation warning).
|
||||
dict: A dictionary containing a list of memories under the "results" key.
|
||||
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", ...}]}`
|
||||
"""
|
||||
|
||||
@@ -867,21 +823,7 @@ class Memory(MemoryBase):
|
||||
"mem0.get_all", self, {"limit": limit, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "sync"}
|
||||
)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future_memories = executor.submit(self._get_all_from_vector_store, effective_filters, limit)
|
||||
future_graph_entities = (
|
||||
executor.submit(self.graph.get_all, effective_filters, limit) if self.graph else None
|
||||
)
|
||||
|
||||
concurrent.futures.wait(
|
||||
[future_memories, future_graph_entities] if future_graph_entities else [future_memories]
|
||||
)
|
||||
|
||||
all_memories_result = future_memories.result()
|
||||
graph_entities_result = future_graph_entities.result() if future_graph_entities else None
|
||||
|
||||
if self.graph:
|
||||
return {"results": all_memories_result, "relations": graph_entities_result}
|
||||
all_memories_result = self._get_all_from_vector_store(effective_filters, limit)
|
||||
|
||||
return {"results": all_memories_result}
|
||||
|
||||
@@ -972,8 +914,7 @@ class Memory(MemoryBase):
|
||||
- {"NOT": [filter1]} - logical NOT
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the search results, typically under a "results" key,
|
||||
and potentially "relations" if graph store is enabled.
|
||||
dict: A dictionary containing the search results under a "results" key.
|
||||
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", "score": 0.8, ...}]}`
|
||||
"""
|
||||
_, effective_filters = _build_filters_and_metadata(
|
||||
@@ -1013,18 +954,7 @@ class Memory(MemoryBase):
|
||||
},
|
||||
)
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
future_memories = executor.submit(self._search_vector_store, query, effective_filters, limit, threshold)
|
||||
future_graph_entities = (
|
||||
executor.submit(self.graph.search, query, effective_filters, limit) if self.graph else None
|
||||
)
|
||||
|
||||
concurrent.futures.wait(
|
||||
[future_memories, future_graph_entities] if future_graph_entities else [future_memories]
|
||||
)
|
||||
|
||||
original_memories = future_memories.result()
|
||||
graph_entities = future_graph_entities.result() if future_graph_entities else None
|
||||
original_memories = self._search_vector_store(query, effective_filters, limit, threshold)
|
||||
|
||||
# Apply reranking if enabled and reranker is available
|
||||
if rerank and self.reranker and original_memories:
|
||||
@@ -1034,9 +964,6 @@ class Memory(MemoryBase):
|
||||
except Exception as e:
|
||||
logger.warning(f"Reranking failed, using original results: {e}")
|
||||
|
||||
if self.graph:
|
||||
return {"results": original_memories, "relations": graph_entities}
|
||||
|
||||
return {"results": original_memories}
|
||||
|
||||
def _process_metadata_filters(self, metadata_filters: Dict[str, Any]) -> Dict[str, Any]:
|
||||
@@ -1334,21 +1261,6 @@ class Memory(MemoryBase):
|
||||
if existing_memory is None:
|
||||
raise ValueError(f"Memory with id {memory_id} not found")
|
||||
|
||||
# Clean up graph entities before deleting from vector store
|
||||
if self.graph:
|
||||
try:
|
||||
memory_text = existing_memory.payload.get("data", "")
|
||||
if memory_text:
|
||||
filters = {}
|
||||
for key in ("user_id", "agent_id", "run_id"):
|
||||
val = existing_memory.payload.get(key)
|
||||
if val:
|
||||
filters[key] = val
|
||||
if filters.get("user_id"):
|
||||
self.graph.delete(memory_text, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up graph for memory {memory_id}: {e}")
|
||||
|
||||
self._delete_memory(memory_id, existing_memory)
|
||||
return {"message": "Memory deleted successfully!"}
|
||||
|
||||
@@ -1383,9 +1295,6 @@ class Memory(MemoryBase):
|
||||
|
||||
logger.info(f"Deleted {len(memories)} memories")
|
||||
|
||||
if self.graph:
|
||||
self.graph.delete_all(filters)
|
||||
|
||||
return {"message": "Memories deleted successfully!"}
|
||||
|
||||
def history(self, memory_id):
|
||||
@@ -1584,12 +1493,6 @@ class Memory(MemoryBase):
|
||||
logger.warning(f"Failed to reset entity store: {e}")
|
||||
self._entity_store = None
|
||||
|
||||
if self.graph:
|
||||
try:
|
||||
self.graph.reset()
|
||||
except Exception:
|
||||
logger.warning("Failed to reset graph store, continuing with reset")
|
||||
|
||||
capture_event("mem0.reset", self, {"sync_type": "sync"})
|
||||
|
||||
def close(self):
|
||||
@@ -1629,12 +1532,6 @@ class AsyncMemory(MemoryBase):
|
||||
config.reranker.config
|
||||
)
|
||||
|
||||
if self.config.graph_store.config:
|
||||
provider = self.config.graph_store.provider
|
||||
self.graph = GraphStoreFactory.create(provider, self.config)
|
||||
else:
|
||||
self.graph = None
|
||||
|
||||
if MEM0_TELEMETRY:
|
||||
telemetry_config = _safe_deepcopy_config(self.config.vector_store.config)
|
||||
telemetry_config.collection_name = "mem0migrations"
|
||||
@@ -1673,13 +1570,6 @@ class AsyncMemory(MemoryBase):
|
||||
|
||||
@staticmethod
|
||||
def _process_config(config_dict: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if "graph_store" in config_dict:
|
||||
if "vector_store" not in config_dict and "embedder" in config_dict:
|
||||
config_dict["vector_store"] = {}
|
||||
config_dict["vector_store"]["config"] = {}
|
||||
config_dict["vector_store"]["config"]["embedding_model_dims"] = config_dict["embedder"]["config"][
|
||||
"embedding_dims"
|
||||
]
|
||||
try:
|
||||
return config_dict
|
||||
except ValidationError as e:
|
||||
@@ -1690,20 +1580,20 @@ class AsyncMemory(MemoryBase):
|
||||
"""Determine whether to use agent memory extraction based on the logic:
|
||||
- If agent_id is present and messages contain assistant role -> True
|
||||
- Otherwise -> False
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries
|
||||
metadata: Metadata containing user_id, agent_id, etc.
|
||||
|
||||
|
||||
Returns:
|
||||
bool: True if should use agent memory extraction, False for user memory extraction
|
||||
"""
|
||||
# Check if agent_id is present in metadata
|
||||
has_agent_id = metadata.get("agent_id") is not None
|
||||
|
||||
|
||||
# Check if there are assistant role messages
|
||||
has_assistant_messages = any(msg.get("role") == "assistant" for msg in messages)
|
||||
|
||||
|
||||
# Use agent memory extraction if agent_id is present and there are assistant messages
|
||||
return has_agent_id and has_assistant_messages
|
||||
|
||||
@@ -1771,19 +1661,6 @@ class AsyncMemory(MemoryBase):
|
||||
else:
|
||||
messages = parse_vision_messages(messages)
|
||||
|
||||
if self.graph:
|
||||
vector_store_task = asyncio.create_task(
|
||||
self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
|
||||
)
|
||||
graph_task = asyncio.create_task(self._add_to_graph(messages, effective_filters))
|
||||
|
||||
vector_store_result, graph_result = await asyncio.gather(vector_store_task, graph_task)
|
||||
|
||||
return {
|
||||
"results": vector_store_result,
|
||||
"relations": graph_result,
|
||||
}
|
||||
|
||||
vector_store_result = await self._add_to_vector_store(messages, processed_metadata, effective_filters, infer)
|
||||
return {"results": vector_store_result}
|
||||
|
||||
@@ -2104,17 +1981,6 @@ class AsyncMemory(MemoryBase):
|
||||
)
|
||||
return returned_memories
|
||||
|
||||
async def _add_to_graph(self, messages, filters):
|
||||
added_entities = []
|
||||
if self.graph:
|
||||
if filters.get("user_id") is None:
|
||||
filters["user_id"] = "user"
|
||||
|
||||
data = "\n".join([msg["content"] for msg in messages if "content" in msg and msg["role"] != "system"])
|
||||
added_entities = await asyncio.to_thread(self.graph.add, data, filters)
|
||||
|
||||
return added_entities
|
||||
|
||||
async def get(self, memory_id):
|
||||
"""
|
||||
Retrieve a memory by ID asynchronously.
|
||||
@@ -2180,9 +2046,7 @@ class AsyncMemory(MemoryBase):
|
||||
limit (int, optional): The maximum number of memories to return. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing a list of memories under the "results" key,
|
||||
and potentially "relations" if graph store is enabled. For API v1.0,
|
||||
it might return a direct list (see deprecation warning).
|
||||
dict: A dictionary containing a list of memories under the "results" key.
|
||||
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", ...}]}`
|
||||
"""
|
||||
|
||||
@@ -2201,25 +2065,9 @@ class AsyncMemory(MemoryBase):
|
||||
"mem0.get_all", self, {"limit": limit, "keys": keys, "encoded_ids": encoded_ids, "sync_type": "async"}
|
||||
)
|
||||
|
||||
vector_store_task = asyncio.create_task(self._get_all_from_vector_store(effective_filters, limit))
|
||||
all_memories_result = await self._get_all_from_vector_store(effective_filters, limit)
|
||||
|
||||
graph_task = None
|
||||
if self.graph:
|
||||
graph_get_all = getattr(self.graph, "get_all", None)
|
||||
if callable(graph_get_all):
|
||||
if asyncio.iscoroutinefunction(graph_get_all):
|
||||
graph_task = asyncio.create_task(graph_get_all(effective_filters, limit))
|
||||
else:
|
||||
graph_task = asyncio.create_task(asyncio.to_thread(graph_get_all, effective_filters, limit))
|
||||
|
||||
results_dict = {}
|
||||
if graph_task:
|
||||
vector_store_result, graph_entities_result = await asyncio.gather(vector_store_task, graph_task)
|
||||
results_dict.update({"results": vector_store_result, "relations": graph_entities_result})
|
||||
else:
|
||||
results_dict.update({"results": await vector_store_task})
|
||||
|
||||
return results_dict
|
||||
return {"results": all_memories_result}
|
||||
|
||||
async def _get_all_from_vector_store(self, filters, limit):
|
||||
memories_result = await asyncio.to_thread(self.vector_store.list, filters=filters, top_k=limit)
|
||||
@@ -2309,8 +2157,7 @@ class AsyncMemory(MemoryBase):
|
||||
- {"NOT": [filter1]} - logical NOT
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the search results, typically under a "results" key,
|
||||
and potentially "relations" if graph store is enabled.
|
||||
dict: A dictionary containing the search results under a "results" key.
|
||||
Example for v1.1+: `{"results": [{"id": "...", "memory": "...", "score": 0.8, ...}]}`
|
||||
"""
|
||||
|
||||
@@ -2351,20 +2198,7 @@ class AsyncMemory(MemoryBase):
|
||||
},
|
||||
)
|
||||
|
||||
vector_store_task = asyncio.create_task(self._search_vector_store(query, effective_filters, limit, threshold))
|
||||
|
||||
graph_task = None
|
||||
if self.graph:
|
||||
if hasattr(self.graph.search, "__await__"): # Check if graph search is async
|
||||
graph_task = asyncio.create_task(self.graph.search(query, effective_filters, limit))
|
||||
else:
|
||||
graph_task = asyncio.create_task(asyncio.to_thread(self.graph.search, query, effective_filters, limit))
|
||||
|
||||
if graph_task:
|
||||
original_memories, graph_entities = await asyncio.gather(vector_store_task, graph_task)
|
||||
else:
|
||||
original_memories = await vector_store_task
|
||||
graph_entities = None
|
||||
original_memories = await self._search_vector_store(query, effective_filters, limit, threshold)
|
||||
|
||||
# Apply reranking if enabled and reranker is available
|
||||
if rerank and self.reranker and original_memories:
|
||||
@@ -2377,9 +2211,6 @@ class AsyncMemory(MemoryBase):
|
||||
except Exception as e:
|
||||
logger.warning(f"Reranking failed, using original results: {e}")
|
||||
|
||||
if self.graph:
|
||||
return {"results": original_memories, "relations": graph_entities}
|
||||
|
||||
return {"results": original_memories}
|
||||
|
||||
def _process_metadata_filters(self, metadata_filters: Dict[str, Any]) -> Dict[str, Any]:
|
||||
@@ -2663,21 +2494,6 @@ class AsyncMemory(MemoryBase):
|
||||
if existing_memory is None:
|
||||
raise ValueError(f"Memory with id {memory_id} not found")
|
||||
|
||||
# Clean up graph entities before deleting from vector store
|
||||
if self.graph:
|
||||
try:
|
||||
memory_text = existing_memory.payload.get("data", "")
|
||||
if memory_text:
|
||||
filters = {}
|
||||
for key in ("user_id", "agent_id", "run_id"):
|
||||
val = existing_memory.payload.get(key)
|
||||
if val:
|
||||
filters[key] = val
|
||||
if filters.get("user_id"):
|
||||
await asyncio.to_thread(self.graph.delete, memory_text, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up graph for memory {memory_id}: {e}")
|
||||
|
||||
await self._delete_memory(memory_id, existing_memory)
|
||||
return {"message": "Memory deleted successfully!"}
|
||||
|
||||
@@ -2715,9 +2531,6 @@ class AsyncMemory(MemoryBase):
|
||||
|
||||
logger.info(f"Deleted {len(memories[0])} memories")
|
||||
|
||||
if self.graph:
|
||||
await asyncio.to_thread(self.graph.delete_all, filters)
|
||||
|
||||
return {"message": "Memories deleted successfully!"}
|
||||
|
||||
async def history(self, memory_id):
|
||||
@@ -2934,12 +2747,6 @@ class AsyncMemory(MemoryBase):
|
||||
self.config.vector_store.provider, self.config.vector_store.config
|
||||
)
|
||||
|
||||
if self.graph:
|
||||
try:
|
||||
await asyncio.to_thread(self.graph.reset)
|
||||
except Exception:
|
||||
logger.warning("Failed to reset graph store, continuing with reset")
|
||||
|
||||
capture_event("mem0.reset", self, {"sync_type": "async"})
|
||||
|
||||
def close(self):
|
||||
|
||||
@@ -1,708 +0,0 @@
|
||||
import logging
|
||||
|
||||
from mem0.memory.utils import format_entities, remove_spaces_from_entities
|
||||
|
||||
try:
|
||||
from langchain_memgraph.graphs.memgraph import Memgraph
|
||||
except ImportError:
|
||||
raise ImportError("langchain_memgraph is not installed. Please install it using pip install langchain-memgraph")
|
||||
|
||||
try:
|
||||
from rank_bm25 import BM25Okapi
|
||||
except ImportError:
|
||||
raise ImportError("rank_bm25 is not installed. Please install it using pip install rank-bm25")
|
||||
|
||||
from mem0.graphs.tools import (
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
DELETE_MEMORY_TOOL_GRAPH,
|
||||
EXTRACT_ENTITIES_STRUCT_TOOL,
|
||||
EXTRACT_ENTITIES_TOOL,
|
||||
RELATIONS_STRUCT_TOOL,
|
||||
RELATIONS_TOOL,
|
||||
)
|
||||
from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_delete_messages
|
||||
from mem0.utils.factory import EmbedderFactory, LlmFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MemoryGraph:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self.graph = Memgraph(
|
||||
self.config.graph_store.config.url,
|
||||
self.config.graph_store.config.username,
|
||||
self.config.graph_store.config.password,
|
||||
)
|
||||
self.embedding_model = EmbedderFactory.create(
|
||||
self.config.embedder.provider,
|
||||
self.config.embedder.config,
|
||||
{"enable_embeddings": True},
|
||||
)
|
||||
|
||||
# Default to openai if no specific provider is configured
|
||||
self.llm_provider = "openai"
|
||||
if self.config.llm and self.config.llm.provider:
|
||||
self.llm_provider = self.config.llm.provider
|
||||
if self.config.graph_store and self.config.graph_store.llm and self.config.graph_store.llm.provider:
|
||||
self.llm_provider = self.config.graph_store.llm.provider
|
||||
|
||||
# Get LLM config with proper null checks
|
||||
llm_config = None
|
||||
if self.config.graph_store and self.config.graph_store.llm and hasattr(self.config.graph_store.llm, "config"):
|
||||
llm_config = self.config.graph_store.llm.config
|
||||
elif hasattr(self.config.llm, "config"):
|
||||
llm_config = self.config.llm.config
|
||||
self.llm = LlmFactory.create(self.llm_provider, llm_config)
|
||||
self.user_id = None
|
||||
# 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)
|
||||
# 2. Create label property index for performance optimizations
|
||||
embedding_dims = self.config.embedder.config["embedding_dims"]
|
||||
index_info = self._fetch_existing_indexes()
|
||||
|
||||
# Create vector index if not exists
|
||||
if not self._vector_index_exists(index_info, "memzero"):
|
||||
self.graph.query(
|
||||
f"CREATE VECTOR INDEX memzero ON :Entity(embedding) WITH CONFIG {{'dimension': {embedding_dims}, 'capacity': 1000, 'metric': 'cos'}};"
|
||||
)
|
||||
|
||||
# Create label+property index if not exists
|
||||
if not self._label_property_index_exists(index_info, "Entity", "user_id"):
|
||||
self.graph.query("CREATE INDEX ON :Entity(user_id);")
|
||||
|
||||
# Create label index if not exists
|
||||
if not self._label_index_exists(index_info, "Entity"):
|
||||
self.graph.query("CREATE INDEX ON :Entity;")
|
||||
|
||||
def add(self, data, filters):
|
||||
"""
|
||||
Adds data to the graph.
|
||||
|
||||
Args:
|
||||
data (str): The data to add to the graph.
|
||||
filters (dict): A dictionary containing filters to be applied during the addition.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
to_be_added = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
to_be_deleted = self._get_delete_entities_from_search_output(search_output, data, filters)
|
||||
|
||||
# TODO: Batch queries with APOC plugin
|
||||
# TODO: Add more filter support
|
||||
deleted_entities = self._delete_entities(to_be_deleted, filters)
|
||||
added_entities = self._add_entities(to_be_added, filters, entity_type_map)
|
||||
|
||||
return {"deleted_entities": deleted_entities, "added_entities": added_entities}
|
||||
|
||||
def search(self, query, filters, top_k=100):
|
||||
"""
|
||||
Search for memories and related graph data.
|
||||
|
||||
Args:
|
||||
query (str): Query to search for.
|
||||
filters (dict): A dictionary containing filters to be applied during the search.
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing:
|
||||
- "contexts": List of search results from the base data store.
|
||||
- "entities": List of related graph data based on the query.
|
||||
"""
|
||||
entity_type_map = self._retrieve_nodes_from_data(query, filters)
|
||||
search_output = self._search_graph_db(node_list=list(entity_type_map.keys()), filters=filters)
|
||||
|
||||
if not search_output:
|
||||
return []
|
||||
|
||||
search_outputs_sequence = [
|
||||
[item["source"], item["relationship"], item["destination"]] for item in search_output
|
||||
]
|
||||
bm25 = BM25Okapi(search_outputs_sequence)
|
||||
|
||||
tokenized_query = query.split(" ")
|
||||
reranked_results = bm25.get_top_n(tokenized_query, search_outputs_sequence, n=5)
|
||||
|
||||
search_results = []
|
||||
for item in reranked_results:
|
||||
search_results.append({"source": item[0], "relationship": item[1], "destination": item[2]})
|
||||
|
||||
logger.info(f"Returned {len(search_results)} search results")
|
||||
|
||||
return search_results
|
||||
|
||||
def delete(self, data, filters):
|
||||
"""
|
||||
Delete graph entities associated with the given memory text.
|
||||
|
||||
Extracts entities and relationships from the memory text using the same
|
||||
pipeline as add(), then deletes the matching relationships in the graph.
|
||||
|
||||
Args:
|
||||
data (str): The memory text whose graph entities should be removed.
|
||||
filters (dict): Scope filters (user_id, agent_id).
|
||||
"""
|
||||
try:
|
||||
entity_type_map = self._retrieve_nodes_from_data(data, filters)
|
||||
if not entity_type_map:
|
||||
logger.debug("No entities found in memory text, skipping graph cleanup")
|
||||
return
|
||||
to_be_deleted = self._establish_nodes_relations_from_data(data, filters, entity_type_map)
|
||||
if to_be_deleted:
|
||||
self._delete_entities(to_be_deleted, filters)
|
||||
except Exception as e:
|
||||
logger.error(f"Error during graph cleanup for memory delete: {e}")
|
||||
|
||||
def delete_all(self, filters):
|
||||
"""Delete all nodes and relationships for a user or specific agent."""
|
||||
if filters.get("agent_id"):
|
||||
cypher = """
|
||||
MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "agent_id": filters["agent_id"]}
|
||||
else:
|
||||
cypher = """
|
||||
MATCH (n:Entity {user_id: $user_id})
|
||||
DETACH DELETE n
|
||||
"""
|
||||
params = {"user_id": filters["user_id"]}
|
||||
self.graph.query(cypher, params=params)
|
||||
|
||||
def get_all(self, filters, top_k=100):
|
||||
"""
|
||||
Retrieves all nodes and relationships from the graph database based on optional filtering criteria.
|
||||
|
||||
Args:
|
||||
filters (dict): A dictionary containing filters to be applied during the retrieval.
|
||||
Supports 'user_id' (required) and 'agent_id' (optional).
|
||||
top_k (int): The maximum number of nodes and relationships to retrieve. Defaults to 100.
|
||||
Returns:
|
||||
list: A list of dictionaries, each containing:
|
||||
- 'source': The source node name.
|
||||
- 'relationship': The relationship type.
|
||||
- 'target': The target node name.
|
||||
"""
|
||||
# Build query based on whether agent_id is provided
|
||||
if filters.get("agent_id"):
|
||||
query = """
|
||||
MATCH (n:Entity {user_id: $user_id, agent_id: $agent_id})-[r]->(m:Entity {user_id: $user_id, agent_id: $agent_id})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "agent_id": filters["agent_id"], "limit": top_k}
|
||||
else:
|
||||
query = """
|
||||
MATCH (n:Entity {user_id: $user_id})-[r]->(m:Entity {user_id: $user_id})
|
||||
RETURN n.name AS source, type(r) AS relationship, m.name AS target
|
||||
LIMIT $limit
|
||||
"""
|
||||
params = {"user_id": filters["user_id"], "limit": top_k}
|
||||
|
||||
results = self.graph.query(query, params=params)
|
||||
|
||||
final_results = []
|
||||
for result in results:
|
||||
final_results.append(
|
||||
{
|
||||
"source": result["source"],
|
||||
"relationship": result["relationship"],
|
||||
"target": result["target"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f"Retrieved {len(final_results)} relationships")
|
||||
|
||||
return final_results
|
||||
|
||||
def _retrieve_nodes_from_data(self, data, filters):
|
||||
"""Extracts all the entities mentioned in the query."""
|
||||
_tools = [EXTRACT_ENTITIES_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
|
||||
search_results = self.llm.generate_response(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": f"You are a smart assistant who understands entities and their types in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question.",
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entity_type_map = {}
|
||||
|
||||
try:
|
||||
for tool_call in search_results["tool_calls"]:
|
||||
if tool_call["name"] != "extract_entities":
|
||||
continue
|
||||
for item in tool_call.get("arguments", {}).get("entities", []):
|
||||
if "entity" in item and "entity_type" in item:
|
||||
entity_type_map[item["entity"]] = item["entity_type"]
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
f"Error in search tool: {e}, llm_provider={self.llm_provider}, search_results={search_results}"
|
||||
)
|
||||
|
||||
entity_type_map = {k.lower().replace(" ", "_"): v.lower().replace(" ", "_") for k, v in entity_type_map.items()}
|
||||
logger.debug(f"Entity type map: {entity_type_map}\n search_results={search_results}")
|
||||
return entity_type_map
|
||||
|
||||
def _establish_nodes_relations_from_data(self, data, filters, entity_type_map):
|
||||
"""Eshtablish relations among the extracted nodes."""
|
||||
if self.config.graph_store.custom_prompt:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]).replace(
|
||||
"CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}"
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": data},
|
||||
]
|
||||
else:
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["user_id"]),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"List of entities: {list(entity_type_map.keys())}. \n\nText: {data}",
|
||||
},
|
||||
]
|
||||
|
||||
_tools = [RELATIONS_TOOL]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [RELATIONS_STRUCT_TOOL]
|
||||
|
||||
extracted_entities = self.llm.generate_response(
|
||||
messages=messages,
|
||||
tools=_tools,
|
||||
)
|
||||
|
||||
entities = []
|
||||
if extracted_entities and extracted_entities.get("tool_calls"):
|
||||
entities = extracted_entities["tool_calls"][0].get("arguments", {}).get("entities", [])
|
||||
|
||||
entities = self._remove_spaces_from_entities(entities)
|
||||
logger.debug(f"Extracted entities: {entities}")
|
||||
return entities
|
||||
|
||||
def _search_graph_db(self, node_list, filters, top_k=100):
|
||||
"""Search similar nodes among and their respective incoming and outgoing relations."""
|
||||
result_relations = []
|
||||
|
||||
for node in node_list:
|
||||
n_embedding = self.embedding_model.embed(node)
|
||||
|
||||
# Build query based on whether agent_id is provided
|
||||
if filters.get("agent_id"):
|
||||
cypher_query = """
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (n)-[r]->(m:Entity)
|
||||
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, similarity
|
||||
UNION
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.agent_id = $agent_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (m:Entity)-[r]->(n)
|
||||
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, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit;
|
||||
"""
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"agent_id": filters["agent_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
else:
|
||||
cypher_query = """
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (n)-[r]->(m:Entity)
|
||||
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, similarity
|
||||
UNION
|
||||
CALL vector_search.search("memzero", $limit, $n_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS n, similarity
|
||||
WHERE n:Entity AND n.user_id = $user_id AND n.embedding IS NOT NULL AND similarity >= $threshold
|
||||
MATCH (m:Entity)-[r]->(n)
|
||||
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, similarity
|
||||
ORDER BY similarity DESC
|
||||
LIMIT $limit;
|
||||
"""
|
||||
params = {
|
||||
"n_embedding": n_embedding,
|
||||
"threshold": self.threshold,
|
||||
"user_id": filters["user_id"],
|
||||
"limit": top_k,
|
||||
}
|
||||
|
||||
ans = self.graph.query(cypher_query, params=params)
|
||||
result_relations.extend(ans)
|
||||
|
||||
return result_relations
|
||||
|
||||
def _get_delete_entities_from_search_output(self, search_output, data, filters):
|
||||
"""Get the entities to be deleted from the search output."""
|
||||
search_output_string = format_entities(search_output)
|
||||
system_prompt, user_prompt = get_delete_messages(search_output_string, data, filters["user_id"])
|
||||
|
||||
_tools = [DELETE_MEMORY_TOOL_GRAPH]
|
||||
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
|
||||
_tools = [
|
||||
DELETE_MEMORY_STRUCT_TOOL_GRAPH,
|
||||
]
|
||||
|
||||
memory_updates = self.llm.generate_response(
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
tools=_tools,
|
||||
)
|
||||
to_be_deleted = []
|
||||
for item in memory_updates["tool_calls"]:
|
||||
if item["name"] == "delete_graph_memory":
|
||||
to_be_deleted.append(item["arguments"])
|
||||
# in case if it is not in the correct format
|
||||
to_be_deleted = self._remove_spaces_from_entities(to_be_deleted)
|
||||
logger.debug(f"Deleted relationships: {to_be_deleted}")
|
||||
return to_be_deleted
|
||||
|
||||
def _delete_entities(self, to_be_deleted, filters):
|
||||
"""Delete the entities from the graph."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_deleted:
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# Build the agent filter for the query
|
||||
agent_filter = ""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"user_id": user_id,
|
||||
}
|
||||
|
||||
if agent_id:
|
||||
agent_filter = "AND n.agent_id = $agent_id AND m.agent_id = $agent_id"
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
# Delete the specific relationship between nodes
|
||||
cypher = f"""
|
||||
MATCH (n:Entity {{name: $source_name, user_id: $user_id}})
|
||||
-[r:{relationship}]->
|
||||
(m:Entity {{name: $dest_name, user_id: $user_id}})
|
||||
WHERE 1=1 {agent_filter}
|
||||
DELETE r
|
||||
RETURN
|
||||
n.name AS source,
|
||||
m.name AS target,
|
||||
type(r) AS relationship
|
||||
"""
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
# added Entity label to all nodes for vector search to work
|
||||
def _add_entities(self, to_be_added, filters, entity_type_map):
|
||||
"""Add the new entities to the graph. Merge the nodes if they already exist."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
results = []
|
||||
|
||||
for item in to_be_added:
|
||||
# entities
|
||||
source = item["source"]
|
||||
destination = item["destination"]
|
||||
relationship = item["relationship"]
|
||||
|
||||
# types
|
||||
source_type = entity_type_map.get(source, "__User__")
|
||||
destination_type = entity_type_map.get(destination, "__User__")
|
||||
|
||||
# embeddings
|
||||
source_embedding = self.embedding_model.embed(source)
|
||||
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=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 = ""
|
||||
if agent_id:
|
||||
agent_id_clause = ", agent_id: $agent_id"
|
||||
|
||||
# TODO: Create a cypher query and common params for all the cases
|
||||
if not destination_node_search_result and source_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source:Entity)
|
||||
WHERE id(source) = $source_id
|
||||
MERGE (destination:{destination_type}:Entity {{name: $destination_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET
|
||||
destination.created = timestamp(),
|
||||
destination.embedding = $destination_embedding,
|
||||
destination:Entity
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["id(source_candidate)"],
|
||||
"destination_name": destination,
|
||||
"destination_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
elif destination_node_search_result and not source_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (destination:Entity)
|
||||
WHERE id(destination) = $destination_id
|
||||
MERGE (source:{source_type}:Entity {{name: $source_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET
|
||||
source.created = timestamp(),
|
||||
source.embedding = $source_embedding,
|
||||
source:Entity
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
|
||||
params = {
|
||||
"destination_id": destination_node_search_result[0]["id(destination_candidate)"],
|
||||
"source_name": source,
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
elif source_node_search_result and destination_node_search_result:
|
||||
cypher = f"""
|
||||
MATCH (source:Entity)
|
||||
WHERE id(source) = $source_id
|
||||
MATCH (destination:Entity)
|
||||
WHERE id(destination) = $destination_id
|
||||
MERGE (source)-[r:{relationship}]->(destination)
|
||||
ON CREATE SET
|
||||
r.created_at = timestamp(),
|
||||
r.updated_at = timestamp()
|
||||
RETURN source.name AS source, type(r) AS relationship, destination.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_id": source_node_search_result[0]["id(source_candidate)"],
|
||||
"destination_id": destination_node_search_result[0]["id(destination_candidate)"],
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
else:
|
||||
cypher = f"""
|
||||
MERGE (n:{source_type}:Entity {{name: $source_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET n.created = timestamp(), n.embedding = $source_embedding, n:Entity
|
||||
ON MATCH SET n.embedding = $source_embedding
|
||||
MERGE (m:{destination_type}:Entity {{name: $dest_name, user_id: $user_id{agent_id_clause}}})
|
||||
ON CREATE SET m.created = timestamp(), m.embedding = $dest_embedding, m:Entity
|
||||
ON MATCH SET m.embedding = $dest_embedding
|
||||
MERGE (n)-[rel:{relationship}]->(m)
|
||||
ON CREATE SET rel.created = timestamp()
|
||||
RETURN n.name AS source, type(rel) AS relationship, m.name AS target
|
||||
"""
|
||||
params = {
|
||||
"source_name": source,
|
||||
"dest_name": destination,
|
||||
"source_embedding": source_embedding,
|
||||
"dest_embedding": dest_embedding,
|
||||
"user_id": user_id,
|
||||
}
|
||||
if agent_id:
|
||||
params["agent_id"] = agent_id
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _remove_spaces_from_entities(self, entity_list):
|
||||
return remove_spaces_from_entities(entity_list, sanitize_relationship=True)
|
||||
|
||||
def _search_source_node(self, source_embedding, filters, threshold=0.9):
|
||||
"""Search for source nodes with similar embeddings."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
|
||||
if agent_id:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $source_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS source_candidate, similarity
|
||||
WHERE source_candidate.user_id = $user_id
|
||||
AND source_candidate.agent_id = $agent_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(source_candidate);
|
||||
"""
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"agent_id": agent_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
else:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $source_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS source_candidate, similarity
|
||||
WHERE source_candidate.user_id = $user_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(source_candidate);
|
||||
"""
|
||||
params = {
|
||||
"source_embedding": source_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
def _search_destination_node(self, destination_embedding, filters, threshold=0.9):
|
||||
"""Search for destination nodes with similar embeddings."""
|
||||
user_id = filters["user_id"]
|
||||
agent_id = filters.get("agent_id", None)
|
||||
|
||||
if agent_id:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $destination_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS destination_candidate, similarity
|
||||
WHERE node.user_id = $user_id
|
||||
AND node.agent_id = $agent_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(destination_candidate);
|
||||
"""
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
"agent_id": agent_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
else:
|
||||
cypher = """
|
||||
CALL vector_search.search("memzero", 1, $destination_embedding)
|
||||
YIELD distance, node, similarity
|
||||
WITH node AS destination_candidate, similarity
|
||||
WHERE node.user_id = $user_id
|
||||
AND similarity >= $threshold
|
||||
RETURN id(destination_candidate);
|
||||
"""
|
||||
params = {
|
||||
"destination_embedding": destination_embedding,
|
||||
"user_id": user_id,
|
||||
"threshold": threshold,
|
||||
}
|
||||
|
||||
result = self.graph.query(cypher, params=params)
|
||||
return result
|
||||
|
||||
|
||||
def _vector_index_exists(self, index_info, index_name):
|
||||
"""
|
||||
Check if a vector index exists, compatible with both Memgraph versions.
|
||||
|
||||
Args:
|
||||
index_info (dict): Index information from _fetch_existing_indexes
|
||||
index_name (str): Name of the index to check
|
||||
|
||||
Returns:
|
||||
bool: True if index exists, False otherwise
|
||||
"""
|
||||
vector_indexes = index_info.get("vector_index_exists", [])
|
||||
|
||||
# Check for index by name regardless of version-specific format differences
|
||||
return any(
|
||||
idx.get("index_name") == index_name or
|
||||
idx.get("index name") == index_name or
|
||||
idx.get("name") == index_name
|
||||
for idx in vector_indexes
|
||||
)
|
||||
|
||||
def _label_property_index_exists(self, index_info, label, property_name):
|
||||
"""
|
||||
Check if a label+property index exists, compatible with both versions.
|
||||
|
||||
Args:
|
||||
index_info (dict): Index information from _fetch_existing_indexes
|
||||
label (str): Label name
|
||||
property_name (str): Property name
|
||||
|
||||
Returns:
|
||||
bool: True if index exists, False otherwise
|
||||
"""
|
||||
indexes = index_info.get("index_exists", [])
|
||||
|
||||
return any(
|
||||
(idx.get("index type") == "label+property" or idx.get("index_type") == "label+property") and
|
||||
(idx.get("label") == label) and
|
||||
(idx.get("property") == property_name or property_name in str(idx.get("properties", "")))
|
||||
for idx in indexes
|
||||
)
|
||||
|
||||
def _label_index_exists(self, index_info, label):
|
||||
"""
|
||||
Check if a label index exists, compatible with both versions.
|
||||
|
||||
Args:
|
||||
index_info (dict): Index information from _fetch_existing_indexes
|
||||
label (str): Label name
|
||||
|
||||
Returns:
|
||||
bool: True if index exists, False otherwise
|
||||
"""
|
||||
indexes = index_info.get("index_exists", [])
|
||||
|
||||
return any(
|
||||
(idx.get("index type") == "label" or idx.get("index_type") == "label") and
|
||||
(idx.get("label") == label)
|
||||
for idx in indexes
|
||||
)
|
||||
|
||||
def _fetch_existing_indexes(self):
|
||||
"""
|
||||
Retrieves information about existing indexes and vector indexes in the Memgraph database.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing lists of existing indexes and vector indexes.
|
||||
"""
|
||||
try:
|
||||
index_exists = list(self.graph.query("SHOW INDEX INFO;"))
|
||||
vector_index_exists = list(self.graph.query("SHOW VECTOR INDEX INFO;"))
|
||||
return {"index_exists": index_exists, "vector_index_exists": vector_index_exists}
|
||||
except Exception as e:
|
||||
logger.warning(f"Error fetching indexes: {e}. Returning empty index info.")
|
||||
return {"index_exists": [], "vector_index_exists": []}
|
||||
@@ -169,9 +169,6 @@ def capture_event(event_name, memory_instance, additional_data=None):
|
||||
"collection": memory_instance.collection_name,
|
||||
"vector_size": memory_instance.embedding_model.config.embedding_dims,
|
||||
"history_store": "sqlite",
|
||||
"graph_store": f"{memory_instance.graph.__class__.__module__}.{memory_instance.graph.__class__.__name__}"
|
||||
if memory_instance.config.graph_store.config
|
||||
else None,
|
||||
"vector_store": f"{memory_instance.vector_store.__class__.__module__}.{memory_instance.vector_store.__class__.__name__}",
|
||||
"llm": f"{memory_instance.llm.__class__.__module__}.{memory_instance.llm.__class__.__name__}",
|
||||
"embedding_model": f"{memory_instance.embedding_model.__class__.__module__}.{memory_instance.embedding_model.__class__.__name__}",
|
||||
|
||||
@@ -209,30 +209,6 @@ class VectorStoreFactory:
|
||||
return instance
|
||||
|
||||
|
||||
class GraphStoreFactory:
|
||||
"""
|
||||
Factory for creating MemoryGraph instances for different graph store providers.
|
||||
Usage: GraphStoreFactory.create(provider_name, config)
|
||||
"""
|
||||
|
||||
provider_to_class = {
|
||||
"memgraph": "mem0.memory.memgraph_memory.MemoryGraph",
|
||||
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
|
||||
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
|
||||
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
|
||||
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
|
||||
"default": "mem0.memory.graph_memory.MemoryGraph",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, provider_name, config):
|
||||
class_type = cls.provider_to_class.get(provider_name, cls.provider_to_class["default"])
|
||||
try:
|
||||
GraphClass = load_class(class_type)
|
||||
except (ImportError, AttributeError) as e:
|
||||
raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}")
|
||||
return GraphClass(config)
|
||||
|
||||
|
||||
class RerankerFactory:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user