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

[Misc] Lint code and fix code smells (#1871)

Browse Source
This commit is contained in:
Deshraj Yadav
2024-09-16 17:39:54 -07:00
committed by GitHub
parent 0a78cb9f7a
commit 55c54beeab
57 changed files with 1178 additions and 1357 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
mem0/graphs/configs.py
View File

@@ -18,28 +18,16 @@ class Neo4jConfig(BaseModel):
values.get("password"),
)
if not url or not username or not password:
raise ValueError(
"Please provide 'url', 'username' and '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
)
llm: Optional[LlmConfig] = Field(
description="LLM configuration for querying the graph store",
default=None
)
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)
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
description="Custom prompt to fetch entities from the given text", default=None
)
@field_validator("config")
@@ -49,4 +37,3 @@ class GraphStoreConfig(BaseModel):
return Neo4jConfig(**v.model_dump())
else:
raise ValueError(f"Unsupported graph store provider: {provider}")

173
mem0/graphs/tools.py
View File

@@ -1,4 +1,3 @@
UPDATE_MEMORY_TOOL_GRAPH = {
"type": "function",
"function": {
@@ -9,21 +8,21 @@ UPDATE_MEMORY_TOOL_GRAPH = {
"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."
"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."
"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."
}
"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
}
}
"additionalProperties": False,
},
},
}
ADD_MEMORY_TOOL_GRAPH = {
@@ -36,29 +35,35 @@ ADD_MEMORY_TOOL_GRAPH = {
"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."
"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."
"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."
"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."
"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."
}
"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
}
}
"required": [
"source",
"destination",
"relationship",
"source_type",
"destination_type",
],
"additionalProperties": False,
},
},
}
@@ -71,9 +76,9 @@ NOOP_TOOL = {
"type": "object",
"properties": {},
"required": [],
"additionalProperties": False
}
}
"additionalProperties": False,
},
},
}
@@ -94,17 +99,23 @@ ADD_MESSAGE_TOOL = {
"source_type": {"type": "string"},
"relation": {"type": "string"},
"destination_node": {"type": "string"},
"destination_type": {"type": "string"}
"destination_type": {"type": "string"},
},
"required": ["source_node", "source_type", "relation", "destination_node", "destination_type"],
"additionalProperties": False
}
"required": [
"source_node",
"source_type",
"relation",
"destination_node",
"destination_type",
],
"additionalProperties": False,
},
}
},
"required": ["entities"],
"additionalProperties": False
}
}
"additionalProperties": False,
},
},
}
@@ -118,23 +129,19 @@ SEARCH_TOOL = {
"properties": {
"nodes": {
"type": "array",
"items": {
"type": "string"
},
"description": "List of nodes to search for."
"items": {"type": "string"},
"description": "List of nodes to search for.",
},
"relations": {
"type": "array",
"items": {
"type": "string"
},
"description": "List of relations to search for."
}
"items": {"type": "string"},
"description": "List of relations to search for.",
},
},
"required": ["nodes", "relations"],
"additionalProperties": False
}
}
"additionalProperties": False,
},
},
}
UPDATE_MEMORY_STRUCT_TOOL_GRAPH = {
@@ -148,21 +155,21 @@ UPDATE_MEMORY_STRUCT_TOOL_GRAPH = {
"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."
"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."
"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."
}
"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
}
}
"additionalProperties": False,
},
},
}
ADD_MEMORY_STRUCT_TOOL_GRAPH = {
@@ -176,29 +183,35 @@ ADD_MEMORY_STRUCT_TOOL_GRAPH = {
"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."
"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."
"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."
"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."
"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."
}
"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
}
}
"required": [
"source",
"destination",
"relationship",
"source_type",
"destination_type",
],
"additionalProperties": False,
},
},
}
@@ -212,9 +225,9 @@ NOOP_STRUCT_TOOL = {
"type": "object",
"properties": {},
"required": [],
"additionalProperties": False
}
}
"additionalProperties": False,
},
},
}
@@ -236,17 +249,23 @@ ADD_MESSAGE_STRUCT_TOOL = {
"source_type": {"type": "string"},
"relation": {"type": "string"},
"destination_node": {"type": "string"},
"destination_type": {"type": "string"}
"destination_type": {"type": "string"},
},
"required": ["source_node", "source_type", "relation", "destination_node", "destination_type"],
"additionalProperties": False
}
"required": [
"source_node",
"source_type",
"relation",
"destination_node",
"destination_type",
],
"additionalProperties": False,
},
}
},
"required": ["entities"],
"additionalProperties": False
}
}
"additionalProperties": False,
},
},
}
@@ -261,21 +280,17 @@ SEARCH_STRUCT_TOOL = {
"properties": {
"nodes": {
"type": "array",
"items": {
"type": "string"
},
"description": "List of nodes to search for."
"items": {"type": "string"},
"description": "List of nodes to search for.",
},
"relations": {
"type": "array",
"items": {
"type": "string"
},
"description": "List of relations to search for."
}
"items": {"type": "string"},
"description": "List of relations to search for.",
},
},
"required": ["nodes", "relations"],
"additionalProperties": False
}
}
"additionalProperties": False,
},
},
}

3
mem0/graphs/utils.py
View File

@@ -1,4 +1,3 @@
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.
@@ -55,10 +54,10 @@ Strive for a coherent, easily understandable knowledge graph by maintaining cons
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 [
{