klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[improvement]: Graph nodes extraction improved (#2035)

Browse Source
This commit is contained in:
Mayank
2024-11-21 12:27:39 +05:30
committed by GitHub
parent 62ca0ddbe2
commit 4b8e32830a
4 changed files with 171 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
mem0/graphs/tools.py
View File

@@ -82,11 +82,11 @@ NOOP_TOOL = {
} }
ADD_MESSAGE_TOOL = { RELATIONS_TOOL = {
"type": "function", "type": "function",
"function": { "function": {
"name": "add_query", "name": "establish_relations",
"description": "Add new entities and relationships to the graph based on the provided query.", "description": "Establish relationships among the entities based on the provided text.",
"parameters": { "parameters": {
"type": "object", "type": "object",
"properties": { "properties": {
@@ -95,18 +95,23 @@ ADD_MESSAGE_TOOL = {
"items": { "items": {
"type": "object", "type": "object",
"properties": { "properties": {
"source_node": {"type": "string"}, "source": {
"source_type": {"type": "string"}, "type": "string",
"relation": {"type": "string"}, "description": "The source entity of the relationship."
"destination_node": {"type": "string"}, },
"destination_type": {"type": "string"}, "relation": {
"type": "string",
"description": "The relationship between the source and destination entities."
},
"destination": {
"type": "string",
"description": "The destination entity of the relationship."
},
}, },
"required": [ "required": [
"source_node", "source_entity",
"source_type",
"relation", "relation",
"destination_node", "destination_entity",
"destination_type",
], ],
"additionalProperties": False, "additionalProperties": False,
}, },
@@ -119,29 +124,38 @@ ADD_MESSAGE_TOOL = {
} }
SEARCH_TOOL = { EXTRACT_ENTITIES_TOOL = {
"type": "function", "type": "function",
"function": { "function": {
"name": "search", "name": "extract_entities",
"description": "Search for nodes and relations in the graph.", "description": "Extract entities and their types from the text.",
"parameters": { "parameters": {
"type": "object", "type": "object",
"properties": { "properties": {
"nodes": { "entities": {
"type": "array", "type": "array",
"items": {"type": "string"}, "items": {
"description": "List of nodes to search for.", "type": "object",
}, "properties": {
"relations": { "entity": {
"type": "array", "type": "string",
"items": {"type": "string"}, "description": "The name or identifier of the entity."
"description": "List of relations to search for.", },
}, "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": ["nodes", "relations"], "required": ["entities"],
"additionalProperties": False, "additionalProperties": False
}, }
}, }
} }
UPDATE_MEMORY_STRUCT_TOOL_GRAPH = { UPDATE_MEMORY_STRUCT_TOOL_GRAPH = {
@@ -230,12 +244,11 @@ NOOP_STRUCT_TOOL = {
}, },
} }
RELATIONS_STRUCT_TOOL = {
ADD_MESSAGE_STRUCT_TOOL = {
"type": "function", "type": "function",
"function": { "function": {
"name": "add_query", "name": "establish_relations",
"description": "Add new entities and relationships to the graph based on the provided query.", "description": "Establish relationships among the entities based on the provided text.",
"strict": True, "strict": True,
"parameters": { "parameters": {
"type": "object", "type": "object",
@@ -245,18 +258,23 @@ ADD_MESSAGE_STRUCT_TOOL = {
"items": { "items": {
"type": "object", "type": "object",
"properties": { "properties": {
"source_node": {"type": "string"}, "source_entity": {
"source_type": {"type": "string"}, "type": "string",
"relation": {"type": "string"}, "description": "The source entity of the relationship."
"destination_node": {"type": "string"}, },
"destination_type": {"type": "string"}, "relation": {
"type": "string",
"description": "The relationship between the source and destination entities."
},
"destination_entity": {
"type": "string",
"description": "The destination entity of the relationship."
},
}, },
"required": [ "required": [
"source_node", "source_entity",
"source_type",
"relation", "relation",
"destination_node", "destination_entity",
"destination_type",
], ],
"additionalProperties": False, "additionalProperties": False,
}, },
@@ -269,28 +287,37 @@ ADD_MESSAGE_STRUCT_TOOL = {
} }
SEARCH_STRUCT_TOOL = { EXTRACT_ENTITIES_STRUCT_TOOL = {
"type": "function", "type": "function",
"function": { "function": {
"name": "search", "name": "extract_entities",
"description": "Search for nodes and relations in the graph.", "description": "Extract entities and their types from the text.",
"strict": True, "strict": True,
"parameters": { "parameters": {
"type": "object", "type": "object",
"properties": { "properties": {
"nodes": { "entities": {
"type": "array", "type": "array",
"items": {"type": "string"}, "items": {
"description": "List of nodes to search for.", "type": "object",
}, "properties": {
"relations": { "entity": {
"type": "array", "type": "string",
"items": {"type": "string"}, "description": "The name or identifier of the entity."
"description": "List of relations to search for.", },
}, "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": ["nodes", "relations"], "required": ["entities"],
"additionalProperties": False, "additionalProperties": False
}, }
}, }
} }

39
mem0/graphs/utils.py
View File

@@ -18,50 +18,51 @@ Guidelines:
7. Relationship Refinement: Look for opportunities to refine relationship descriptions for greater precision or clarity. 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. 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: Task Details:
- Existing Graph Memories: ======= Existing Graph Memories:=======
{existing_memories} {existing_memories}
- New Graph Memory: {memory} ======= New Graph Memory:=======
{new_memories}
Output: 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. 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 = """ 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 information while maintaining accuracy. Follow these key principles: 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. 1. Extract only explicitly stated information from the text.
2. Identify nodes (entities/concepts), their types, and relationships. 2. Establish relationships among the entities provided.
3. Use "USER_ID" as the source node for any self-references (I, me, my, etc.) in user messages. 3. Use "USER_ID" as the source entity for any self-references (e.g., "I," "me," "my," etc.) in user messages.
CUSTOM_PROMPT CUSTOM_PROMPT
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: Relationships:
- Use consistent, general, and timeless relationship types. - Use consistent, general, and timeless relationship types.
- Example: Prefer "PROFESSOR" over "BECAME_PROFESSOR". - Example: Prefer "PROFESSOR" over "BECAME_PROFESSOR."
- Relationships should only be established among the entities explicitly mentioned in the user message.
Entity Consistency: Entity Consistency:
- Use the most complete identifier for entities mentioned multiple times. - Ensure that relationships are coherent and logically align with the context of the message.
- Example: Always use "John Doe" instead of variations like "Joe" or pronouns. - Maintain consistent naming for entities across the extracted data.
Strive for a coherent, easily understandable knowledge graph by maintaining consistency in entity references and relationship types. Strive to construct a coherent and easily understandable knowledge graph by eshtablishing 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.""" Adhere strictly to these guidelines to ensure high-quality knowledge graph extraction."""
def get_update_memory_prompt(existing_memories, memory, template): def get_update_memory_prompt(existing_memories, new_memories, template):
return template.format(existing_memories=existing_memories, memory=memory) return template.format(existing_memories=existing_memories, new_memories=new_memories)
def get_update_memory_messages(existing_memories, memory): def get_update_memory_messages(existing_memories, new_memories):
return [ return [
{ {
"role": "user", "role": "user",
"content": get_update_memory_prompt(existing_memories, memory, UPDATE_GRAPH_PROMPT), "content": get_update_memory_prompt(existing_memories, new_memories, UPDATE_GRAPH_PROMPT),
}, },
] ]

119
mem0/memory/graph_memory.py
View File

@@ -15,16 +15,16 @@ except ImportError:
from mem0.graphs.tools import ( from mem0.graphs.tools import (
ADD_MEMORY_STRUCT_TOOL_GRAPH, ADD_MEMORY_STRUCT_TOOL_GRAPH,
ADD_MEMORY_TOOL_GRAPH, ADD_MEMORY_TOOL_GRAPH,
ADD_MESSAGE_STRUCT_TOOL, EXTRACT_ENTITIES_STRUCT_TOOL,
ADD_MESSAGE_TOOL, EXTRACT_ENTITIES_TOOL,
NOOP_STRUCT_TOOL, NOOP_STRUCT_TOOL,
NOOP_TOOL, NOOP_TOOL,
SEARCH_STRUCT_TOOL, RELATIONS_STRUCT_TOOL,
SEARCH_TOOL, RELATIONS_TOOL,
UPDATE_MEMORY_STRUCT_TOOL_GRAPH, UPDATE_MEMORY_STRUCT_TOOL_GRAPH,
UPDATE_MEMORY_TOOL_GRAPH, UPDATE_MEMORY_TOOL_GRAPH,
) )
from mem0.graphs.utils import EXTRACT_ENTITIES_PROMPT, get_update_memory_messages from mem0.graphs.utils import EXTRACT_RELATIONS_PROMPT, get_update_memory_messages
from mem0.utils.factory import EmbedderFactory, LlmFactory from mem0.utils.factory import EmbedderFactory, LlmFactory
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -60,44 +60,14 @@ class MemoryGraph:
""" """
# retrieve the search results # retrieve the search results
search_output = self._search(data, filters) search_output, entity_type_map = self._search(data, filters)
if self.config.graph_store.custom_prompt: # extract relations
messages = [ extracted_relations = self._extract_relations(data, filters, entity_type_map)
{
"role": "system",
"content": EXTRACT_ENTITIES_PROMPT.replace("USER_ID", self.user_id).replace(
"CUSTOM_PROMPT", f"4. {self.config.graph_store.custom_prompt}"
),
},
{"role": "user", "content": data},
]
else:
messages = [
{
"role": "system",
"content": EXTRACT_ENTITIES_PROMPT.replace("USER_ID", self.user_id),
},
{"role": "user", "content": data},
]
_tools = [ADD_MESSAGE_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [ADD_MESSAGE_STRUCT_TOOL]
extracted_entities = self.llm.generate_response(
messages=messages,
tools=_tools,
)
if extracted_entities["tool_calls"]:
extracted_entities = extracted_entities["tool_calls"][0]["arguments"]["entities"]
else:
extracted_entities = []
logger.debug(f"Extracted entities: {extracted_entities}")
search_output_string = format_entities(search_output) search_output_string = format_entities(search_output)
update_memory_prompt = get_update_memory_messages(search_output_string, extracted_entities) extracted_relations_string = format_entities(extracted_relations)
update_memory_prompt = get_update_memory_messages(search_output_string, extracted_relations_string)
_tools = [UPDATE_MEMORY_TOOL_GRAPH, ADD_MEMORY_TOOL_GRAPH, NOOP_TOOL] _tools = [UPDATE_MEMORY_TOOL_GRAPH, ADD_MEMORY_TOOL_GRAPH, NOOP_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]: if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
@@ -170,37 +140,33 @@ class MemoryGraph:
return returned_entities return returned_entities
def _search(self, query, filters, limit=100): def _search(self, query, filters, limit=100):
_tools = [SEARCH_TOOL] _tools = [EXTRACT_ENTITIES_TOOL]
if self.llm_provider in ["azure_openai_structured", "openai_structured"]: if self.llm_provider in ["azure_openai_structured", "openai_structured"]:
_tools = [SEARCH_STRUCT_TOOL] _tools = [EXTRACT_ENTITIES_STRUCT_TOOL]
search_results = self.llm.generate_response( search_results = self.llm.generate_response(
messages=[ messages=[
{ {
"role": "system", "role": "system",
"content": f"You are a smart assistant who understands the entities, their types, and relations in a given text. If user message contains self reference such as 'I', 'me', 'my' etc. then use {filters['user_id']} as the source node. Extract the entities. ***DO NOT*** answer the question itself if the given text is a question.", "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": query}, {"role": "user", "content": query},
], ],
tools=_tools, tools=_tools,
) )
node_list = [] entity_type_map = {}
for item in search_results["tool_calls"]: try:
if item["name"] == "search": for item in search_results["tool_calls"][0]["arguments"]["entities"]:
try: entity_type_map[item["entity"]] = item["entity_type"]
node_list.extend(item["arguments"]["nodes"]) except Exception as e:
except Exception as e: logger.error(f"Error in search tool: {e}")
logger.error(f"Error in search tool: {e}")
node_list = list(set(node_list)) logger.debug(f"Entity type map: {entity_type_map}")
node_list = [node.lower().replace(" ", "_") for node in node_list]
logger.debug(f"Node list for search query : {node_list}")
result_relations = [] result_relations = []
for node in node_list: for node in list(entity_type_map.keys()):
n_embedding = self.embedding_model.embed(node) n_embedding = self.embedding_model.embed(node)
cypher_query = """ cypher_query = """
@@ -235,7 +201,7 @@ class MemoryGraph:
ans = self.graph.query(cypher_query, params=params) ans = self.graph.query(cypher_query, params=params)
result_relations.extend(ans) result_relations.extend(ans)
return result_relations return result_relations, entity_type_map
def search(self, query, filters, limit=100): def search(self, query, filters, limit=100):
""" """
@@ -252,7 +218,7 @@ class MemoryGraph:
- "entities": List of related graph data based on the query. - "entities": List of related graph data based on the query.
""" """
search_output = self._search(query, filters, limit) search_output, entity_type_map = self._search(query, filters, limit)
if not search_output: if not search_output:
return [] return []
@@ -314,6 +280,45 @@ class MemoryGraph:
return final_results return final_results
def _extract_relations(self, data, filters, entity_type_map, limit=100):
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,
)
if extracted_entities["tool_calls"]:
extracted_entities = extracted_entities["tool_calls"][0]["arguments"]["entities"]
else:
extracted_entities = []
logger.debug(f"Extracted entities: {extracted_entities}")
return extracted_entities
def _update_relationship(self, source, target, relationship, filters): def _update_relationship(self, source, target, relationship, filters):
""" """
Update or create a relationship between two nodes in the graph. Update or create a relationship between two nodes in the graph.

8
mem0/memory/utils.py
View File

@@ -24,11 +24,7 @@ def format_entities(entities):
formatted_lines = [] formatted_lines = []
for entity in entities: for entity in entities:
simplified = { simplified = f"{entity['source']} -- {entity['relation'].upper()} -- {entity['destination']}"
"source": entity["source"], formatted_lines.append(simplified)
"relation": entity["relation"],
"destination": entity["destination"]
}
formatted_lines.append(json.dumps(simplified))
return "\n".join(formatted_lines) return "\n".join(formatted_lines)