[Mem0] Integrate Graph Memory (#1718)
Co-authored-by: Deshraj Yadav <deshrajdry@gmail.com>
This commit is contained in:
Prateek Chhikara
40
mem0/graphs/configs.py
Normal file
40
mem0/graphs/configs.py
Normal file
@@ -0,0 +1,40 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||
|
||||
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")
|
||||
|
||||
@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 and not username and not password:
|
||||
raise ValueError(
|
||||
"Please provide 'url', 'username' and 'password'."
|
||||
)
|
||||
return values
|
||||
|
||||
|
||||
class GraphStoreConfig(BaseModel):
|
||||
provider: str = Field(
|
||||
description="Provider of the data store (e.g., 'neo4j')",
|
||||
default="neo4j"
|
||||
)
|
||||
config: Neo4jConfig = Field(
|
||||
description="Configuration for the specific data store",
|
||||
default=None
|
||||
)
|
||||
|
||||
@field_validator("config")
|
||||
def validate_config(cls, v, values):
|
||||
provider = values.data.get("provider")
|
||||
if provider == "neo4j":
|
||||
return Neo4jConfig(**v.model_dump())
|
||||
else:
|
||||
raise ValueError(f"Unsupported graph store provider: {provider}")
|
||||
80
mem0/graphs/tools.py
Normal file
80
mem0/graphs/tools.py
Normal file
@@ -0,0 +1,80 @@
|
||||
|
||||
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.",
|
||||
"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_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_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
|
||||
}
|
||||
}
|
||||
}
|
||||
107
mem0/graphs/utils.py
Normal file
107
mem0/graphs/utils.py
Normal file
@@ -0,0 +1,107 @@
|
||||
from mem0.llms.openai import OpenAILLM
|
||||
|
||||
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.
|
||||
|
||||
Task Details:
|
||||
- Existing Graph Memories:
|
||||
{existing_memories}
|
||||
|
||||
- New Graph Memory: {memory}
|
||||
|
||||
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_ENTITIES_PROMPT = """
|
||||
|
||||
You are an advanced algorithm designed to extract structured information from text to construct knowledge graphs. Your goal is to capture comprehensive information while maintaining accuracy. Follow these key principles:
|
||||
|
||||
1. Extract only explicitly stated information from the text.
|
||||
2. Identify nodes (entities/concepts), their types, and relationships.
|
||||
3. Use "USER_ID" as the source node for any self-references (I, me, my, etc.) in user messages.
|
||||
|
||||
Nodes and Types:
|
||||
- Aim for simplicity and clarity in node representation.
|
||||
- Use basic, general types for node labels (e.g. "person" instead of "mathematician").
|
||||
|
||||
Relationships:
|
||||
- Use consistent, general, and timeless relationship types.
|
||||
- Example: Prefer "PROFESSOR" over "BECAME_PROFESSOR".
|
||||
|
||||
Entity Consistency:
|
||||
- Use the most complete identifier for entities mentioned multiple times.
|
||||
- Example: Always use "John Doe" instead of variations like "Joe" or pronouns.
|
||||
|
||||
Strive for a coherent, easily understandable knowledge graph by maintaining consistency in entity references and relationship types.
|
||||
|
||||
Adhere strictly to these guidelines to ensure high-quality knowledge graph extraction."""
|
||||
|
||||
|
||||
|
||||
def get_update_memory_prompt(existing_memories, memory, template):
|
||||
return template.format(existing_memories=existing_memories, memory=memory)
|
||||
|
||||
def get_update_memory_messages(existing_memories, memory):
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": get_update_memory_prompt(existing_memories, memory, UPDATE_GRAPH_PROMPT),
|
||||
},
|
||||
]
|
||||
|
||||
def get_search_results(entities, query):
|
||||
|
||||
search_graph_prompt = f"""
|
||||
You are an expert at searching through graph entity memories.
|
||||
When provided with existing graph entities and a query, your task is to search through the provided graph entities to find the most relevant information from the graph entities related to the query.
|
||||
The output should be from the graph entities only.
|
||||
|
||||
Here are the details of the task:
|
||||
- Existing Graph Entities (source -> relationship -> target):
|
||||
{entities}
|
||||
|
||||
- Query: {query}
|
||||
|
||||
The output should be from the graph entities only.
|
||||
The output should be in the following JSON format:
|
||||
{{
|
||||
"search_results": [
|
||||
{{
|
||||
"source_node": "source_node",
|
||||
"relationship": "relationship",
|
||||
"target_node": "target_node"
|
||||
}}
|
||||
]
|
||||
}}
|
||||
"""
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": search_graph_prompt,
|
||||
}
|
||||
]
|
||||
|
||||
llm = OpenAILLM()
|
||||
|
||||
results = llm.generate_response(messages=messages, response_format={"type": "json_object"})
|
||||
|
||||
return results
|
||||
Reference in New Issue
Block a user