Improvements to Graph Memory (#1779)

mem0/graphs/configs.py

@@ -1,6 +1,6 @@
 from typing import Optional
-
 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")
@@ -30,6 +30,14 @@ class GraphStoreConfig(BaseModel):
         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
+    )
 
     @field_validator("config")
     def validate_config(cls, v, values):
@@ -38,3 +46,4 @@ class GraphStoreConfig(BaseModel):
             return Neo4jConfig(**v.model_dump())
         else:
             raise ValueError(f"Unsupported graph store provider: {provider}")
+        

mem0/graphs/tools.py

@@ -78,3 +78,66 @@ NOOP_TOOL = {
         }
     }
 }
+
+
+ADD_MESSAGE_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "add_query",
+        "description": "Add new entities and relationships to the graph based on the provided query.",
+        "strict": True,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "entities": {
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "source_node": {"type": "string"},
+                            "source_type": {"type": "string"},
+                            "relation": {"type": "string"},
+                            "destination_node": {"type": "string"},
+                            "destination_type": {"type": "string"}
+                        },
+                        "required": ["source_node", "source_type", "relation", "destination_node", "destination_type"],
+                        "additionalProperties": False
+                    }
+                }
+            },
+            "required": ["entities"],
+            "additionalProperties": False
+        }
+    }
+}
+
+
+SEARCH_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "search",
+        "description": "Search for nodes and relations in the graph.",
+        "strict": True,
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "nodes": {
+                    "type": "array",
+                    "items": {
+                        "type": "string"
+                    },
+                    "description": "List of nodes to search for."
+                },
+                "relations": {
+                    "type": "array",
+                    "items": {
+                        "type": "string"
+                    },
+                    "description": "List of relations to search for."
+                }
+            },
+            "required": ["nodes", "relations"],
+            "additionalProperties": False
+        }
+    }
+}

mem0/graphs/utils.py

@@ -37,6 +37,7 @@ You are an advanced algorithm designed to extract structured information from te
 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.
+CUSTOM_PROMPT
 
 Nodes and Types:
 - Aim for simplicity and clarity in node representation.