JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: TS SDK v3 port, graph store removal, client v3 API migration

Browse Source 
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
2026-04-13 10:45:40 -07:00
co-authored by Claude Opus 4.6
parent 3402c67805
commit 59880a6f8f
59 changed files with 5237 additions and 11748 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mem0-ts/package.json
+4 -2
View File
@@ -120,10 +120,12 @@
"better-sqlite3": "^12.6.2",
"cloudflare": "^4.2.0",
"groq-sdk": "0.3.0",
"neo4j-driver": "^5.28.1",
"ollama": "^0.5.14",
"pg": "8.11.3",
"redis": "^4.6.13"
"redis": "^4.6.13",
"compromise": "^14.0.0",
"natural": "^8.0.1"
},
"engines": {
"node": ">=18"
mem0-ts/pnpm-lock.yaml
Generated
+2360 -4094
View File
File diff suppressed because it is too large Load Diff
mem0-ts/src/client/mem0.ts
+3 -5
View File
@@ -190,7 +190,7 @@ export default class MemoryClient {
this._captureEvent("add", [payloadKeys]);
const response = await this._fetchWithErrorHandling(
`${this.host}/v1/memories/`,
`${this.host}/v3/memories/`,
{
method: "POST",
headers: this.headers,
@@ -286,20 +286,18 @@ export default class MemoryClient {
this._captureEvent("search", [payloadKeys]);
const payload: Record<string, any> = {
query,
output_format: "v1.1",
...camelToSnakeKeys(options ?? {}),
};
const response = await this._fetchWithErrorHandling(
`${this.host}/v2/memories/search/`,
`${this.host}/v3/memories/search/`,
{
method: "POST",
headers: this.headers,
body: JSON.stringify(payload),
},
);
// Unwrap v1.1 format: { results: [...] } → [...]
return Array.isArray(response) ? response : (response?.results ?? response);
return response;
}
async delete(memoryId: string): Promise<{ message: string }> {
mem0-ts/src/oss/examples/basic.ts
-93
View File
@@ -292,86 +292,6 @@ async function demoRedis() {
await runTests(memory);
}
async function demoGraphMemory() {
console.log("\n=== Testing Graph Memory Store ===\n");
const memory = new Memory({
version: "v1.1",
embedder: {
provider: "openai",
config: {
apiKey: process.env.OPENAI_API_KEY || "",
model: "text-embedding-3-small",
},
},
vectorStore: {
provider: "memory",
config: {
collectionName: "memories",
dimension: 1536,
},
},
llm: {
provider: "openai",
config: {
apiKey: process.env.OPENAI_API_KEY || "",
model: "gpt-4-turbo-preview",
},
},
graphStore: {
provider: "neo4j",
config: {
url: process.env.NEO4J_URL || "neo4j://localhost:7687",
username: process.env.NEO4J_USERNAME || "neo4j",
password: process.env.NEO4J_PASSWORD || "password",
},
llm: {
provider: "openai",
config: {
model: "gpt-4-turbo-preview",
},
},
},
historyDbPath: "memory.db",
});
try {
// Reset all memories
await memory.reset();
// Add memories with relationships
const result = await memory.add(
[
{
role: "user",
content: "Alice is Bob's sister and works as a doctor.",
},
{
role: "assistant",
content:
"I understand that Alice and Bob are siblings and Alice is a medical professional.",
},
{ role: "user", content: "Bob is married to Carol who is a teacher." },
],
{
userId: "john",
},
);
console.log("Added memories with relationships:", result);
// Search for connected information
const searchResult = await memory.search(
"Tell me about Bob's family connections",
{
userId: "john",
},
);
console.log("Search results with graph relationships:", searchResult);
} catch (error) {
console.error("Error in graph memory demo:", error);
}
}
async function main() {
// Test in-memory store
await demoMemoryStore();
@@ -379,19 +299,6 @@ async function main() {
// Test in-memory store with Ollama
await demoLocalMemory();
// Test graph memory if Neo4j environment variables are set
if (
process.env.NEO4J_URL &&
process.env.NEO4J_USERNAME &&
process.env.NEO4J_PASSWORD
) {
await demoGraphMemory();
} else {
console.log(
"\nSkipping Graph Memory test - Neo4j environment variables not set",
);
}
// Test PGVector store if environment variables are set
if (process.env.PGVECTOR_DB) {
await demoPGVector();
mem0-ts/src/oss/src/config/defaults.ts
+1 -1
View File
@@ -22,7 +22,7 @@ export const DEFAULT_MEMORY_CONFIG: MemoryConfig = {
config: {
baseURL: "https://api.openai.com/v1",
apiKey: process.env.OPENAI_API_KEY || "",
model: "gpt-4-turbo-preview",
model: "gpt-4.1-nano-2025-04-14",
modelProperties: undefined,
},
},
mem0-ts/src/oss/src/config/manager.ts
-3
View File
@@ -132,9 +132,6 @@ export class ConfigManager {
userConfig.historyStore?.config?.historyDbPath ||
DEFAULT_MEMORY_CONFIG.historyStore?.config?.historyDbPath,
customInstructions: userConfig.customInstructions,
graphStore: userConfig.graphStore
? { ...userConfig.graphStore }
: undefined,
historyStore: (() => {
const defaultHistoryStore = DEFAULT_MEMORY_CONFIG.historyStore!;
const historyProvider =
mem0-ts/src/oss/src/embeddings/azure.ts
+18 -8
View File
@@ -35,13 +35,23 @@ export class AzureOpenAIEmbedder implements Embedder {
}
async embedBatch(texts: string[]): Promise<number[][]> {
const response = await this.client.embeddings.create({
model: this.model,
input: texts,
...(this.embeddingDims !== undefined && {
dimensions: this.embeddingDims,
}),
});
return response.data.map((item) => item.embedding);
const MAX_BATCH = 100;
const allEmbeddings: number[][] = [];
for (let i = 0; i < texts.length; i += MAX_BATCH) {
const chunk = texts.slice(i, i + MAX_BATCH);
const response = await this.client.embeddings.create({
model: this.model,
input: chunk,
...(this.embeddingDims !== undefined && {
dimensions: this.embeddingDims,
}),
});
allEmbeddings.push(
...response.data
.sort((a, b) => a.index - b.index)
.map((item) => item.embedding),
);
}
return allEmbeddings;
}
}
mem0-ts/src/oss/src/embeddings/openai.ts
+18 -8
View File
@@ -28,13 +28,23 @@ export class OpenAIEmbedder implements Embedder {
}
async embedBatch(texts: string[]): Promise<number[][]> {
const response = await this.openai.embeddings.create({
model: this.model,
input: texts,
...(this.embeddingDims !== undefined && {
dimensions: this.embeddingDims,
}),
});
return response.data.map((item) => item.embedding);
const MAX_BATCH = 100;
const allEmbeddings: number[][] = [];
for (let i = 0; i < texts.length; i += MAX_BATCH) {
const chunk = texts.slice(i, i + MAX_BATCH);
const response = await this.openai.embeddings.create({
model: this.model,
input: chunk,
...(this.embeddingDims !== undefined && {
dimensions: this.embeddingDims,
}),
});
allEmbeddings.push(
...response.data
.sort((a, b) => a.index - b.index)
.map((item) => item.embedding),
);
}
return allEmbeddings;
}
}
mem0-ts/src/oss/src/graphs/configs.ts
-30
View File
@@ -1,30 +0,0 @@
import { LLMConfig } from "../types";
export interface Neo4jConfig {
url: string | null;
username: string | null;
password: string | null;
}
export interface GraphStoreConfig {
provider: string;
config: Neo4jConfig;
llm?: LLMConfig;
customInstructions?: string;
}
export function validateNeo4jConfig(config: Neo4jConfig): void {
const { url, username, password } = config;
if (!url || !username || !password) {
throw new Error("Please provide 'url', 'username' and 'password'.");
}
}
export function validateGraphStoreConfig(config: GraphStoreConfig): void {
const { provider } = config;
if (provider === "neo4j") {
validateNeo4jConfig(config.config);
} else {
throw new Error(`Unsupported graph store provider: ${provider}`);
}
}
mem0-ts/src/oss/src/graphs/tools.ts
-267
View File
@@ -1,267 +0,0 @@
import { z } from "zod";
export interface GraphToolParameters {
source: string;
destination: string;
relationship: string;
source_type?: string;
destination_type?: string;
}
export interface GraphEntitiesParameters {
entities: Array<{
entity: string;
entity_type: string;
}>;
}
export interface GraphRelationsParameters {
entities: Array<{
source: string;
relationship: string;
destination: string;
}>;
}
// --- Zod Schemas for Tool Arguments ---
// Schema for simple relationship arguments (Update, Delete)
export const GraphSimpleRelationshipArgsSchema = z.object({
source: z
.string()
.describe("The identifier of the source node in the relationship."),
relationship: z
.string()
.describe("The relationship between the source and destination nodes."),
destination: z
.string()
.describe("The identifier of the destination node in the relationship."),
});
// Schema for adding a relationship (includes types)
export const GraphAddRelationshipArgsSchema =
GraphSimpleRelationshipArgsSchema.extend({
source_type: z
.string()
.describe("The type or category of the source node."),
destination_type: z
.string()
.describe("The type or category of the destination node."),
});
// Schema for extracting entities
export const GraphExtractEntitiesArgsSchema = z.object({
entities: z
.array(
z.object({
entity: z.string().describe("The name or identifier of the entity."),
entity_type: z.string().describe("The type or category of the entity."),
}),
)
.describe("An array of entities with their types."),
});
// Schema for establishing relationships
export const GraphRelationsArgsSchema = z.object({
entities: z
.array(GraphSimpleRelationshipArgsSchema)
.describe("An array of relationships (source, relationship, destination)."),
});
// --- Tool Definitions (using JSON schema, keep as is) ---
// Note: The tool definitions themselves still use JSON schema format
// as expected by the LLM APIs. The Zod schemas above are for internal
// validation and potentially for use with Langchain's .withStructuredOutput
// if we adapt it to handle tool calls via schema.
export const 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.",
parameters: {
type: "object",
properties: {
source: {
type: "string",
description:
"The identifier of the source node in the relationship to be updated.",
},
destination: {
type: "string",
description:
"The identifier of the destination node in the relationship to be updated.",
},
relationship: {
type: "string",
description:
"The new or updated relationship between the source and destination nodes.",
},
},
required: ["source", "destination", "relationship"],
additionalProperties: false,
},
},
};
export const ADD_MEMORY_TOOL_GRAPH = {
type: "function",
function: {
name: "add_graph_memory",
description: "Add a new graph memory to the knowledge graph.",
parameters: {
type: "object",
properties: {
source: {
type: "string",
description:
"The identifier of the source node in the new relationship.",
},
destination: {
type: "string",
description:
"The identifier of the destination node in the new relationship.",
},
relationship: {
type: "string",
description:
"The type of relationship between the source and destination nodes.",
},
source_type: {
type: "string",
description: "The type or category of the source node.",
},
destination_type: {
type: "string",
description: "The type or category of the destination node.",
},
},
required: [
"source",
"destination",
"relationship",
"source_type",
"destination_type",
],
additionalProperties: false,
},
},
};
export const NOOP_TOOL = {
type: "function",
function: {
name: "noop",
description: "No operation should be performed to the graph entities.",
parameters: {
type: "object",
properties: {},
required: [],
additionalProperties: false,
},
},
};
export const 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,
},
},
};
export const 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,
},
},
};
export const DELETE_MEMORY_TOOL_GRAPH = {
type: "function",
function: {
name: "delete_graph_memory",
description: "Delete the relationship between two nodes.",
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,
},
},
};
mem0-ts/src/oss/src/graphs/utils.ts
-116
View File
@@ -1,116 +0,0 @@
export const 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.
`;
export const 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 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.
`;
export const 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.
Respond in JSON format.
`;
export function getDeleteMessages(
existingMemoriesString: string,
data: string,
userId: string,
): [string, string] {
return [
DELETE_RELATIONS_SYSTEM_PROMPT.replace("USER_ID", userId),
`Here are the existing memories: ${existingMemoriesString} \n\n New Information: ${data}`,
];
}
export function formatEntities(
entities: Array<{
source: string;
relationship: string;
destination: string;
}>,
): string {
return entities
.map((e) => `${e.source} -- ${e.relationship} -- ${e.destination}`)
.join("\n");
}
mem0-ts/src/oss/src/llms/langchain.ts
+14 -70
View File
@@ -10,12 +10,6 @@ import { LLM, LLMResponse } from "./base";
import { LLMConfig, Message } from "../types/index";
// Import the schemas directly into LangchainLLM
import { FactRetrievalSchema, MemoryUpdateSchema } from "../prompts";
// Import graph tool argument schemas
import {
GraphExtractEntitiesArgsSchema,
GraphRelationsArgsSchema,
GraphSimpleRelationshipArgsSchema, // Used for delete tool
} from "../graphs/tools";
const convertToLangchainMessages = (messages: Message[]): BaseMessage[] => {
return messages.map((msg) => {
@@ -73,28 +67,14 @@ export class LangchainLLM implements LLM {
const invokeOptions: Record<string, any> = {};
let isStructuredOutput = false;
let selectedSchema: z.ZodSchema<any> | null = null;
let isToolCallResponse = false;
// --- Internal Schema Selection Logic (runs regardless of response_format) ---
const systemPromptContent =
(messages.find((m) => m.role === "system")?.content as string) || "";
const userPromptContent =
(messages.find((m) => m.role === "user")?.content as string) || "";
const toolNames = tools?.map((t) => t.function.name) || [];
// Prioritize tool call argument schemas
if (toolNames.includes("extract_entities")) {
selectedSchema = GraphExtractEntitiesArgsSchema;
isToolCallResponse = true;
} else if (toolNames.includes("establish_relationships")) {
selectedSchema = GraphRelationsArgsSchema;
isToolCallResponse = true;
} else if (toolNames.includes("delete_graph_memory")) {
selectedSchema = GraphSimpleRelationshipArgsSchema;
isToolCallResponse = true;
}
// Check for memory prompts if no tool schema matched
else if (
// Check for memory prompts
if (
systemPromptContent.includes("Personal Information Organizer") &&
systemPromptContent.includes("extract relevant pieces of information")
) {
@@ -111,26 +91,18 @@ export class LangchainLLM implements LLM {
selectedSchema &&
typeof (this.llmInstance as any).withStructuredOutput === "function"
) {
// Apply if a schema was selected (for memory or single tool calls)
if (
!isToolCallResponse ||
(isToolCallResponse && tools && tools.length === 1)
) {
try {
runnable = (this.llmInstance as any).withStructuredOutput(
selectedSchema,
{ name: tools?.[0]?.function.name },
);
isStructuredOutput = true;
} catch (e) {
isStructuredOutput = false; // Ensure flag is false on error
// No fallback to response_format here unless explicitly passed
if (response_format?.type === "json_object") {
invokeOptions.response_format = { type: "json_object" };
}
try {
runnable = (this.llmInstance as any).withStructuredOutput(
selectedSchema,
{ name: tools?.[0]?.function.name },
);
isStructuredOutput = true;
} catch (e) {
isStructuredOutput = false; // Ensure flag is false on error
// No fallback to response_format here unless explicitly passed
if (response_format?.type === "json_object") {
invokeOptions.response_format = { type: "json_object" };
}
} else if (isToolCallResponse) {
// If multiple tools, don't apply structured output, handle via tool binding below
}
} else if (selectedSchema && response_format?.type === "json_object") {
// Schema selected, but no .withStructuredOutput. Try basic response_format only if explicitly requested.
@@ -164,37 +136,9 @@ export class LangchainLLM implements LLM {
try {
const response = await runnable.invoke(langchainMessages, invokeOptions);
if (isStructuredOutput && !isToolCallResponse) {
if (isStructuredOutput) {
// Memory prompt with structured output
return JSON.stringify(response);
} else if (isStructuredOutput && isToolCallResponse) {
// Tool call with structured arguments
if (response?.tool_calls && Array.isArray(response.tool_calls)) {
const mappedToolCalls = response.tool_calls.map((call: any) => ({
name: call.name || tools?.[0]?.function.name || "unknown_tool",
arguments:
typeof call.args === "string"
? call.args
: JSON.stringify(call.args),
}));
return {
content: response.content || "",
role: "assistant",
toolCalls: mappedToolCalls,
};
} else {
// Direct object response for tool args
return {
content: "",
role: "assistant",
toolCalls: [
{
name: tools?.[0]?.function.name || "unknown_tool",
arguments: JSON.stringify(response),
},
],
};
}
} else if (
response &&
response.tool_calls &&
mem0-ts/src/oss/src/memory/graph_memory.ts
-675
View File
@@ -1,675 +0,0 @@
import neo4j, { Driver } from "neo4j-driver";
import { BM25 } from "../utils/bm25";
import { GraphStoreConfig } from "../graphs/configs";
import { MemoryConfig } from "../types";
import { EmbedderFactory, LLMFactory } from "../utils/factory";
import { Embedder } from "../embeddings/base";
import { LLM } from "../llms/base";
import {
DELETE_MEMORY_TOOL_GRAPH,
EXTRACT_ENTITIES_TOOL,
RELATIONS_TOOL,
} from "../graphs/tools";
import { EXTRACT_RELATIONS_PROMPT, getDeleteMessages } from "../graphs/utils";
import { logger } from "../utils/logger";
interface SearchOutput {
source: string;
source_id: string;
relationship: string;
relation_id: string;
destination: string;
destination_id: string;
similarity: number;
}
interface ToolCall {
name: string;
arguments: string;
}
interface LLMResponse {
toolCalls?: ToolCall[];
}
interface Tool {
type: string;
function: {
name: string;
description: string;
parameters: Record<string, any>;
};
}
interface GraphMemoryResult {
deleted_entities: any[];
added_entities: any[];
relations?: any[];
}
export class MemoryGraph {
private config: MemoryConfig;
private graph: Driver;
private embeddingModel: Embedder;
private llm: LLM;
private structuredLlm: LLM;
private llmProvider: string;
private threshold: number;
constructor(config: MemoryConfig) {
this.config = config;
if (
!config.graphStore?.config?.url ||
!config.graphStore?.config?.username ||
!config.graphStore?.config?.password
) {
throw new Error("Neo4j configuration is incomplete");
}
this.graph = neo4j.driver(
config.graphStore.config.url,
neo4j.auth.basic(
config.graphStore.config.username,
config.graphStore.config.password,
),
);
this.embeddingModel = EmbedderFactory.create(
this.config.embedder.provider,
this.config.embedder.config,
);
this.llmProvider = "openai";
let llmConfig = this.config.llm.config;
if (this.config.llm?.provider) {
this.llmProvider = this.config.llm.provider;
}
if (this.config.graphStore?.llm?.provider) {
this.llmProvider = this.config.graphStore.llm.provider;
llmConfig = this.config.graphStore.llm.config ?? llmConfig;
}
this.llm = LLMFactory.create(this.llmProvider, llmConfig);
this.structuredLlm = LLMFactory.create(this.llmProvider, llmConfig);
this.threshold = 0.7;
}
async add(
data: string,
filters: Record<string, any>,
): Promise<GraphMemoryResult> {
const entityTypeMap = await this._retrieveNodesFromData(data, filters);
const toBeAdded = await this._establishNodesRelationsFromData(
data,
filters,
entityTypeMap,
);
const searchOutput = await this._searchGraphDb(
Object.keys(entityTypeMap),
filters,
);
const toBeDeleted = await this._getDeleteEntitiesFromSearchOutput(
searchOutput,
data,
filters,
);
const deletedEntities = await this._deleteEntities(
toBeDeleted,
filters["userId"],
);
const addedEntities = await this._addEntities(
toBeAdded,
filters["userId"],
entityTypeMap,
);
return {
deleted_entities: deletedEntities,
added_entities: addedEntities,
relations: toBeAdded,
};
}
async search(query: string, filters: Record<string, any>, topK = 100) {
const entityTypeMap = await this._retrieveNodesFromData(query, filters);
const searchOutput = await this._searchGraphDb(
Object.keys(entityTypeMap),
filters,
);
if (!searchOutput.length) {
return [];
}
const searchOutputsSequence = searchOutput.map((item) => [
item.source,
item.relationship,
item.destination,
]);
const bm25 = new BM25(searchOutputsSequence);
const tokenizedQuery = query.split(" ");
const rerankedResults = bm25.search(tokenizedQuery).slice(0, 5);
const searchResults = rerankedResults.map((item) => ({
source: item[0],
relationship: item[1],
destination: item[2],
}));
logger.info(`Returned ${searchResults.length} search results`);
return searchResults;
}
async deleteAll(filters: Record<string, any>) {
const session = this.graph.session();
try {
await session.run("MATCH (n {user_id: $user_id}) DETACH DELETE n", {
user_id: filters["userId"],
});
} finally {
await session.close();
}
}
async getAll(filters: Record<string, any>, topK = 100) {
const session = this.graph.session();
try {
const result = await session.run(
`
MATCH (n {user_id: $user_id})-[r]->(m {user_id: $user_id})
RETURN n.name AS source, type(r) AS relationship, m.name AS target
LIMIT toInteger($limit)
`,
{ user_id: filters["userId"], limit: Math.floor(Number(topK)) },
);
const finalResults = result.records.map((record) => ({
source: record.get("source"),
relationship: record.get("relationship"),
target: record.get("target"),
}));
logger.info(`Retrieved ${finalResults.length} relationships`);
return finalResults;
} finally {
await session.close();
}
}
private async _retrieveNodesFromData(
data: string,
filters: Record<string, any>,
) {
const tools = [EXTRACT_ENTITIES_TOOL] as Tool[];
const searchResults = await this.structuredLlm.generateResponse(
[
{
role: "system",
content: `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["userId"]} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question. Respond in JSON format.`,
},
{ role: "user", content: data },
],
{ type: "json_object" },
tools,
);
let entityTypeMap: Record<string, string> = {};
try {
if (typeof searchResults !== "string" && searchResults.toolCalls) {
for (const call of searchResults.toolCalls) {
if (call.name === "extract_entities") {
const args = JSON.parse(call.arguments);
for (const item of args.entities) {
entityTypeMap[item.entity] = item.entity_type;
}
}
}
}
} catch (e) {
logger.error(`Error in search tool: ${e}`);
}
entityTypeMap = Object.fromEntries(
Object.entries(entityTypeMap).map(([k, v]) => [
k.toLowerCase().replace(/ /g, "_"),
v.toLowerCase().replace(/ /g, "_"),
]),
);
logger.debug(`Entity type map: ${JSON.stringify(entityTypeMap)}`);
return entityTypeMap;
}
private async _establishNodesRelationsFromData(
data: string,
filters: Record<string, any>,
entityTypeMap: Record<string, string>,
) {
let messages;
if (this.config.graphStore?.customInstructions) {
messages = [
{
role: "system",
content:
EXTRACT_RELATIONS_PROMPT.replace(
"USER_ID",
filters["userId"],
).replace(
"CUSTOM_PROMPT",
`4. ${this.config.graphStore.customInstructions}`,
) + "\nPlease provide your response in JSON format.",
},
{ role: "user", content: data },
];
} else {
messages = [
{
role: "system",
content:
EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["userId"]) +
"\nPlease provide your response in JSON format.",
},
{
role: "user",
content: `List of entities: ${Object.keys(entityTypeMap)}. \n\nText: ${data}`,
},
];
}
const tools = [RELATIONS_TOOL] as Tool[];
const extractedEntities = await this.structuredLlm.generateResponse(
messages,
{ type: "json_object" },
tools,
);
let entities: any[] = [];
if (typeof extractedEntities !== "string" && extractedEntities.toolCalls) {
const toolCall = extractedEntities.toolCalls[0];
if (toolCall && toolCall.arguments) {
const args = JSON.parse(toolCall.arguments);
entities = args.entities || [];
}
}
entities = this._removeSpacesFromEntities(entities);
logger.debug(`Extracted entities: ${JSON.stringify(entities)}`);
return entities;
}
private async _searchGraphDb(
nodeList: string[],
filters: Record<string, any>,
topK = 100,
): Promise<SearchOutput[]> {
const resultRelations: SearchOutput[] = [];
const session = this.graph.session();
try {
for (const node of nodeList) {
const nEmbedding = await this.embeddingModel.embed(node);
const cypher = `
MATCH (n)
WHERE n.embedding IS NOT NULL AND n.user_id = $user_id
WITH n,
round(reduce(dot = 0.0, i IN range(0, size(n.embedding)-1) | dot + n.embedding[i] * $n_embedding[i]) /
(sqrt(reduce(l2 = 0.0, i IN range(0, size(n.embedding)-1) | l2 + n.embedding[i] * n.embedding[i])) *
sqrt(reduce(l2 = 0.0, i IN range(0, size($n_embedding)-1) | l2 + $n_embedding[i] * $n_embedding[i]))), 4) AS similarity
WHERE similarity >= $threshold
MATCH (n)-[r]->(m)
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, similarity
UNION
MATCH (n)
WHERE n.embedding IS NOT NULL AND n.user_id = $user_id
WITH n,
round(reduce(dot = 0.0, i IN range(0, size(n.embedding)-1) | dot + n.embedding[i] * $n_embedding[i]) /
(sqrt(reduce(l2 = 0.0, i IN range(0, size(n.embedding)-1) | l2 + n.embedding[i] * n.embedding[i])) *
sqrt(reduce(l2 = 0.0, i IN range(0, size($n_embedding)-1) | l2 + $n_embedding[i] * $n_embedding[i]))), 4) AS similarity
WHERE similarity >= $threshold
MATCH (m)-[r]->(n)
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, similarity
ORDER BY similarity DESC
LIMIT toInteger($limit)
`;
const result = await session.run(cypher, {
n_embedding: nEmbedding,
threshold: this.threshold,
user_id: filters["userId"],
limit: Math.floor(Number(topK)),
});
resultRelations.push(
...result.records.map((record) => ({
source: record.get("source"),
source_id: record.get("source_id").toString(),
relationship: record.get("relationship"),
relation_id: record.get("relation_id").toString(),
destination: record.get("destination"),
destination_id: record.get("destination_id").toString(),
similarity: record.get("similarity"),
})),
);
}
} finally {
await session.close();
}
return resultRelations;
}
private async _getDeleteEntitiesFromSearchOutput(
searchOutput: SearchOutput[],
data: string,
filters: Record<string, any>,
) {
const searchOutputString = searchOutput
.map(
(item) =>
`${item.source} -- ${item.relationship} -- ${item.destination}`,
)
.join("\n");
const [systemPrompt, userPrompt] = getDeleteMessages(
searchOutputString,
data,
filters["userId"],
);
const tools = [DELETE_MEMORY_TOOL_GRAPH] as Tool[];
const memoryUpdates = await this.structuredLlm.generateResponse(
[
{ role: "system", content: systemPrompt },
{ role: "user", content: userPrompt },
],
{ type: "json_object" },
tools,
);
const toBeDeleted: any[] = [];
if (typeof memoryUpdates !== "string" && memoryUpdates.toolCalls) {
for (const item of memoryUpdates.toolCalls) {
if (item.name === "delete_graph_memory") {
toBeDeleted.push(JSON.parse(item.arguments));
}
}
}
const cleanedToBeDeleted = this._removeSpacesFromEntities(toBeDeleted);
logger.debug(
`Deleted relationships: ${JSON.stringify(cleanedToBeDeleted)}`,
);
return cleanedToBeDeleted;
}
private async _deleteEntities(toBeDeleted: any[], userId: string) {
const results: any[] = [];
const session = this.graph.session();
try {
for (const item of toBeDeleted) {
const { source, destination, relationship } = item;
const cypher = `
MATCH (n {name: $source_name, user_id: $user_id})
-[r:${relationship}]->
(m {name: $dest_name, user_id: $user_id})
DELETE r
RETURN
n.name AS source,
m.name AS target,
type(r) AS relationship
`;
const result = await session.run(cypher, {
source_name: source,
dest_name: destination,
user_id: userId,
});
results.push(result.records);
}
} finally {
await session.close();
}
return results;
}
private async _addEntities(
toBeAdded: any[],
userId: string,
entityTypeMap: Record<string, string>,
) {
const results: any[] = [];
const session = this.graph.session();
try {
for (const item of toBeAdded) {
const { source, destination, relationship } = item;
const sourceType = entityTypeMap[source] || "unknown";
const destinationType = entityTypeMap[destination] || "unknown";
const sourceEmbedding = await this.embeddingModel.embed(source);
const destEmbedding = await this.embeddingModel.embed(destination);
const sourceNodeSearchResult = await this._searchSourceNode(
sourceEmbedding,
userId,
);
const destinationNodeSearchResult = await this._searchDestinationNode(
destEmbedding,
userId,
);
let cypher: string;
let params: Record<string, any>;
if (
destinationNodeSearchResult.length === 0 &&
sourceNodeSearchResult.length > 0
) {
cypher = `
MATCH (source)
WHERE elementId(source) = $source_id
MERGE (destination:${destinationType} {name: $destination_name, user_id: $user_id})
ON CREATE SET
destination.created = timestamp(),
destination.embedding = $destination_embedding
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: sourceNodeSearchResult[0].elementId,
destination_name: destination,
destination_embedding: destEmbedding,
user_id: userId,
};
} else if (
destinationNodeSearchResult.length > 0 &&
sourceNodeSearchResult.length === 0
) {
cypher = `
MATCH (destination)
WHERE elementId(destination) = $destination_id
MERGE (source:${sourceType} {name: $source_name, user_id: $user_id})
ON CREATE SET
source.created = timestamp(),
source.embedding = $source_embedding
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: destinationNodeSearchResult[0].elementId,
source_name: source,
source_embedding: sourceEmbedding,
user_id: userId,
};
} else if (
sourceNodeSearchResult.length > 0 &&
destinationNodeSearchResult.length > 0
) {
cypher = `
MATCH (source)
WHERE elementId(source) = $source_id
MATCH (destination)
WHERE elementId(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: sourceNodeSearchResult[0]?.elementId,
destination_id: destinationNodeSearchResult[0]?.elementId,
user_id: userId,
};
} else {
cypher = `
MERGE (n:${sourceType} {name: $source_name, user_id: $user_id})
ON CREATE SET n.created = timestamp(), n.embedding = $source_embedding
ON MATCH SET n.embedding = $source_embedding
MERGE (m:${destinationType} {name: $dest_name, user_id: $user_id})
ON CREATE SET m.created = timestamp(), m.embedding = $dest_embedding
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: sourceEmbedding,
dest_embedding: destEmbedding,
user_id: userId,
};
}
const result = await session.run(cypher, params);
results.push(result.records);
}
} finally {
await session.close();
}
return results;
}
private _removeSpacesFromEntities(entityList: any[]) {
return entityList.map((item) => ({
...item,
source: item.source.toLowerCase().replace(/ /g, "_"),
relationship: item.relationship.toLowerCase().replace(/ /g, "_"),
destination: item.destination.toLowerCase().replace(/ /g, "_"),
}));
}
private async _searchSourceNode(
sourceEmbedding: number[],
userId: string,
threshold = 0.9,
) {
const session = this.graph.session();
try {
const cypher = `
MATCH (source_candidate)
WHERE source_candidate.embedding IS NOT NULL
AND source_candidate.user_id = $user_id
WITH source_candidate,
round(
reduce(dot = 0.0, i IN range(0, size(source_candidate.embedding)-1) |
dot + source_candidate.embedding[i] * $source_embedding[i]) /
(sqrt(reduce(l2 = 0.0, i IN range(0, size(source_candidate.embedding)-1) |
l2 + source_candidate.embedding[i] * source_candidate.embedding[i])) *
sqrt(reduce(l2 = 0.0, i IN range(0, size($source_embedding)-1) |
l2 + $source_embedding[i] * $source_embedding[i])))
, 4) AS source_similarity
WHERE source_similarity >= $threshold
WITH source_candidate, source_similarity
ORDER BY source_similarity DESC
LIMIT 1
RETURN elementId(source_candidate) as element_id
`;
const params = {
source_embedding: sourceEmbedding,
user_id: userId,
threshold,
};
const result = await session.run(cypher, params);
return result.records.map((record) => ({
elementId: record.get("element_id").toString(),
}));
} finally {
await session.close();
}
}
private async _searchDestinationNode(
destinationEmbedding: number[],
userId: string,
threshold = 0.9,
) {
const session = this.graph.session();
try {
const cypher = `
MATCH (destination_candidate)
WHERE destination_candidate.embedding IS NOT NULL
AND destination_candidate.user_id = $user_id
WITH destination_candidate,
round(
reduce(dot = 0.0, i IN range(0, size(destination_candidate.embedding)-1) |
dot + destination_candidate.embedding[i] * $destination_embedding[i]) /
(sqrt(reduce(l2 = 0.0, i IN range(0, size(destination_candidate.embedding)-1) |
l2 + destination_candidate.embedding[i] * destination_candidate.embedding[i])) *
sqrt(reduce(l2 = 0.0, i IN range(0, size($destination_embedding)-1) |
l2 + $destination_embedding[i] * $destination_embedding[i])))
, 4) AS destination_similarity
WHERE destination_similarity >= $threshold
WITH destination_candidate, destination_similarity
ORDER BY destination_similarity DESC
LIMIT 1
RETURN elementId(destination_candidate) as element_id
`;
const params = {
destination_embedding: destinationEmbedding,
user_id: userId,
threshold,
};
const result = await session.run(cypher, params);
return result.records.map((record) => ({
elementId: record.get("element_id").toString(),
}));
} finally {
await session.close();
}
}
}
mem0-ts/src/oss/src/memory/index.ts
+642 -168
View File
@@ -20,13 +20,17 @@ import {
getUpdateMemoryMessages,
parseMessages,
extractJson,
ADDITIVE_EXTRACTION_PROMPT,
AGENT_CONTEXT_SUFFIX,
AdditiveExtractionSchema,
generateAdditiveExtractionPrompt,
} from "../prompts";
import { DummyHistoryManager } from "../storage/DummyHistoryManager";
import { Embedder } from "../embeddings/base";
import { LLM } from "../llms/base";
import { VectorStore } from "../vector_stores/base";
import { ConfigManager } from "../config/manager";
import { MemoryGraph } from "./graph_memory";
import {
AddMemoryOptions,
SearchMemoryOptions,
@@ -36,6 +40,18 @@ import {
import { parse_vision_messages } from "../utils/memory";
import { HistoryManager } from "../storage/base";
import { captureClientEvent } from "../utils/telemetry";
import { lemmatizeForBm25 } from "../utils/lemmatization";
import {
extractEntities,
extractEntitiesBatch,
} from "../utils/entity_extraction";
import {
scoreAndRank,
getBm25Params,
normalizeBm25,
ENTITY_BOOST_WEIGHT,
ScoredResult,
} from "../utils/scoring";
export class Memory {
private config: MemoryConfig;
@@ -46,10 +62,10 @@ export class Memory {
private db: HistoryManager;
private collectionName: string | undefined;
private apiVersion: string;
private graphMemory?: MemoryGraph;
telemetryId: string;
private _initPromise: Promise<void>;
private _initError?: Error;
private _entityStore?: VectorStore;
constructor(config: Partial<MemoryConfig> = {}) {
// Merge and validate config
@@ -80,11 +96,6 @@ export class Memory {
this.apiVersion = this.config.version || "v1.0";
this.telemetryId = "anonymous";
// Initialize graph memory if graphStore is configured
if (this.config.graphStore) {
this.graphMemory = new MemoryGraph(this.config);
}
// Auto-detect embedding dimension (if needed), create vector store,
// and initialize it. All public methods await this before proceeding.
this._initPromise = this._autoInitialize().catch((error) => {
@@ -148,6 +159,35 @@ export class Memory {
}
}
private async getEntityStore(): Promise<VectorStore> {
if (!this._entityStore) {
const entityCollectionName = `${this.collectionName}_entities`;
const entityConfig = {
...this.config.vectorStore.config,
collectionName: entityCollectionName,
};
// For file-based stores (memory/SQLite), use a separate DB path for entities
if (entityConfig.dbPath) {
entityConfig.dbPath = entityConfig.dbPath.replace(/\.db$/, "_entities.db");
}
this._entityStore = VectorStoreFactory.create(
this.config.vectorStore.provider,
entityConfig,
);
await this._entityStore.initialize();
}
return this._entityStore;
}
private buildSessionScope(filters: SearchFilters): string {
const parts: string[] = [];
for (const key of ["agentId", "runId", "userId"].sort()) {
const val = (filters as any)[key];
if (val) parts.push(`${key}=${val}`);
}
return parts.join("&");
}
private async _initializeTelemetry() {
try {
await this._getTelemetryId();
@@ -244,22 +284,8 @@ export class Memory {
infer,
);
// Add to graph store if available
let graphResult;
if (this.graphMemory) {
try {
graphResult = await this.graphMemory.add(
final_parsedMessages.map((m) => m.content).join("\n"),
filters,
);
} catch (error) {
console.error("Error adding to graph memory:", error);
}
}
return {
results: vectorStoreResult,
relations: graphResult?.relations,
};
}
@@ -288,145 +314,440 @@ export class Memory {
}
return returnedMemories;
}
// === V3 PHASED BATCH PIPELINE ===
// Phase 0: Context gathering
const sessionScope = this.buildSessionScope(filters);
let lastMessages: Array<{
role: string;
content: string;
name?: string;
}> = [];
if (typeof this.db.getLastMessages === "function") {
try {
lastMessages = await this.db.getLastMessages(sessionScope, 10);
} catch {
// getLastMessages not supported — proceed without context
}
}
const parsedMessages = messages.map((m) => m.content).join("\n");
const [systemPrompt, userPrompt] = this.customInstructions
? [
this.customInstructions.toLowerCase().includes("json")
? this.customInstructions
: `${this.customInstructions}\n\nYou MUST return a valid JSON object with a 'facts' key containing an array of strings.`,
`Input:\n${parsedMessages}`,
]
: getFactRetrievalMessages(parsedMessages);
// Phase 1: Existing memory retrieval
const searchFilters: SearchFilters = {};
if (filters.userId) searchFilters.userId = filters.userId;
if (filters.agentId) searchFilters.agentId = filters.agentId;
if (filters.runId) searchFilters.runId = filters.runId;
const response = await this.llm.generateResponse(
[
{ role: "system", content: systemPrompt },
{ role: "user", content: userPrompt },
],
{ type: "json_object" },
const queryEmbedding = await this.embedder.embed(parsedMessages);
const existingResults = await this.vectorStore.search(
queryEmbedding,
10,
searchFilters,
);
const cleanResponse = extractJson(response as string);
let facts: string[] = [];
try {
const parsed = FactRetrievalSchema.parse(JSON.parse(cleanResponse));
facts = parsed.facts;
} catch (e) {
console.error(
"Failed to parse facts from LLM response:",
cleanResponse,
e,
);
facts = [];
// Map UUIDs to integers (anti-hallucination)
const existingMemories: Array<{ id: string; text: string }> = [];
const uuidMapping: Record<string, string> = {};
for (let idx = 0; idx < existingResults.length; idx++) {
const mem = existingResults[idx];
uuidMapping[String(idx)] = mem.id;
existingMemories.push({
id: String(idx),
text: mem.payload?.data ?? "",
});
}
// Get embeddings for new facts
const newMessageEmbeddings: Record<string, number[]> = {};
const retrievedOldMemory: Array<{ id: string; text: string }> = [];
// Create embeddings and search for similar memories
for (const fact of facts) {
const embedding = await this.embedder.embed(fact);
newMessageEmbeddings[fact] = embedding;
const existingMemories = await this.vectorStore.search(
embedding,
5,
filters,
);
for (const mem of existingMemories) {
retrievedOldMemory.push({ id: mem.id, text: mem.payload.data });
}
// Phase 2: LLM extraction (single call)
const isAgentScoped =
!!filters.agentId && !filters.userId;
let systemPrompt = ADDITIVE_EXTRACTION_PROMPT;
if (isAgentScoped) {
systemPrompt += AGENT_CONTEXT_SUFFIX;
}
// Remove duplicates from old memories
const uniqueOldMemories = retrievedOldMemory.filter(
(mem, index) =>
retrievedOldMemory.findIndex((m) => m.id === mem.id) === index,
);
// Create UUID mapping for handling UUID hallucinations
const tempUuidMapping: Record<string, string> = {};
uniqueOldMemories.forEach((item, idx) => {
tempUuidMapping[String(idx)] = item.id;
uniqueOldMemories[idx].id = String(idx);
const userPrompt = generateAdditiveExtractionPrompt({
existingMemories,
newMessages: parsedMessages,
lastKMessages: lastMessages,
customInstructions: this.customInstructions,
});
// Get memory update decisions
const updatePrompt = getUpdateMemoryMessages(uniqueOldMemories, facts);
const updateResponse = await this.llm.generateResponse(
[{ role: "user", content: updatePrompt }],
{ type: "json_object" },
);
const cleanUpdateResponse = extractJson(updateResponse as string);
let memoryActions: any[] = [];
let response: string;
try {
memoryActions = JSON.parse(cleanUpdateResponse).memory || [];
response = (await this.llm.generateResponse(
[
{ role: "system", content: systemPrompt },
{ role: "user", content: userPrompt },
],
{ type: "json_object" },
)) as string;
} catch (e) {
console.error(
"Failed to parse memory actions from LLM response:",
cleanUpdateResponse,
e,
);
memoryActions = [];
console.error("LLM extraction failed:", e);
return [];
}
// Process memory actions
const results: MemoryItem[] = [];
for (const action of memoryActions) {
try {
switch (action.event) {
case "ADD": {
const memoryId = await this.createMemory(
action.text,
newMessageEmbeddings,
metadata,
);
results.push({
id: memoryId,
memory: action.text,
metadata: { event: action.event },
});
break;
}
case "UPDATE": {
const realMemoryId = tempUuidMapping[action.id];
await this.updateMemory(
realMemoryId,
action.text,
newMessageEmbeddings,
metadata,
);
results.push({
id: realMemoryId,
memory: action.text,
metadata: {
event: action.event,
previousMemory: action.old_memory,
},
});
break;
}
case "DELETE": {
const realMemoryId = tempUuidMapping[action.id];
await this.deleteMemory(realMemoryId);
results.push({
id: realMemoryId,
memory: action.text,
metadata: { event: action.event },
});
break;
}
// Parse response
let extractedMemories: Array<{
id?: string;
text?: string;
attributed_to?: string;
linked_memory_ids?: string[];
}> = [];
try {
const cleanResponse = extractJson(response);
if (cleanResponse && cleanResponse.trim()) {
try {
const parsed = AdditiveExtractionSchema.parse(
JSON.parse(cleanResponse),
);
extractedMemories = parsed.memory;
} catch {
const fallbackJson = extractJson(cleanResponse);
extractedMemories =
JSON.parse(fallbackJson)?.memory ?? [];
}
}
} catch (e) {
console.error("Error parsing extraction response:", e);
extractedMemories = [];
}
if (extractedMemories.length === 0) {
// Save messages even if nothing extracted
if (typeof this.db.saveMessages === "function") {
try {
await this.db.saveMessages(
messages.map((m) => ({
role: m.role,
content: m.content as string,
})),
sessionScope,
);
} catch {}
}
return [];
}
// Phase 3: Batch embed all extracted memory texts
const memTexts = extractedMemories
.map((m) => m.text ?? "")
.filter((t) => t.length > 0);
let embedMap: Record<string, number[]> = {};
try {
const memEmbeddingsList = await this.embedder.embedBatch(memTexts);
for (let i = 0; i < memTexts.length; i++) {
embedMap[memTexts[i]] = memEmbeddingsList[i];
}
} catch {
// Fallback: embed individually
for (const text of memTexts) {
try {
embedMap[text] = await this.embedder.embed(text);
} catch (e) {
console.warn(`Failed to embed memory text: ${e}`);
}
} catch (error) {
console.error(`Error processing memory action: ${error}`);
}
}
return results;
// Phase 4-5: CPU processing + hash dedup
const existingHashes = new Set<string>();
for (const mem of existingResults) {
const h = mem.payload?.hash;
if (h) existingHashes.add(h);
}
const records: Array<{
memoryId: string;
text: string;
embedding: number[];
payload: Record<string, any>;
}> = [];
const seenHashes = new Set<string>();
for (const mem of extractedMemories) {
const text = mem.text;
if (!text || !(text in embedMap)) continue;
const memHash = createHash("md5").update(text).digest("hex");
if (existingHashes.has(memHash) || seenHashes.has(memHash)) {
continue;
}
seenHashes.add(memHash);
const textLemmatized = lemmatizeForBm25(text);
const memoryId = uuidv4();
const now = new Date().toISOString();
const memPayload: Record<string, any> = {
...metadata,
data: text,
textLemmatized,
hash: memHash,
createdAt: now,
updatedAt: now,
};
if (mem.attributed_to) {
memPayload.attributedTo = mem.attributed_to;
}
if (filters.userId) memPayload.userId = filters.userId;
if (filters.agentId) memPayload.agentId = filters.agentId;
if (filters.runId) memPayload.runId = filters.runId;
records.push({
memoryId,
text,
embedding: embedMap[text],
payload: memPayload,
});
}
if (records.length === 0) {
if (typeof this.db.saveMessages === "function") {
try {
await this.db.saveMessages(
messages.map((m) => ({
role: m.role,
content: m.content as string,
})),
sessionScope,
);
} catch {}
}
return [];
}
// Phase 6: Batch persist
const allVectors = records.map((r) => r.embedding);
const allIds = records.map((r) => r.memoryId);
const allPayloads = records.map((r) => r.payload);
try {
await this.vectorStore.insert(allVectors, allIds, allPayloads);
} catch {
// Fallback: insert one by one
for (let i = 0; i < allIds.length; i++) {
try {
await this.vectorStore.insert(
[allVectors[i]],
[allIds[i]],
[allPayloads[i]],
);
} catch (e) {
console.error(`Failed to insert memory ${allIds[i]}: ${e}`);
}
}
}
// Batch history
const historyRecords = records.map((r) => ({
memoryId: r.memoryId,
previousValue: null as string | null,
newValue: r.text as string | null,
action: "ADD",
createdAt: r.payload.createdAt as string | undefined,
updatedAt: undefined as string | undefined,
isDeleted: 0,
}));
if (typeof this.db.batchAddHistory === "function") {
try {
await this.db.batchAddHistory(historyRecords);
} catch {
// Fallback: add one by one
for (const hr of historyRecords) {
try {
await this.db.addHistory(
hr.memoryId,
null,
hr.newValue,
"ADD",
hr.createdAt,
);
} catch (e) {
console.error(
`Failed to add history for ${hr.memoryId}: ${e}`,
);
}
}
}
} else {
for (const hr of historyRecords) {
try {
await this.db.addHistory(
hr.memoryId,
null,
hr.newValue,
"ADD",
hr.createdAt,
);
} catch (e) {
console.error(
`Failed to add history for ${hr.memoryId}: ${e}`,
);
}
}
}
// Phase 7: Batch entity linking
try {
const allTexts = records.map((r) => r.text);
const allEntities = extractEntitiesBatch(allTexts);
// 7a: Global dedup — collect unique entities across all memories
const globalEntities: Record<
string,
{ entityType: string; entityText: string; memoryIds: Set<string> }
> = {};
for (let idx = 0; idx < records.length; idx++) {
const memoryId = records[idx].memoryId;
const entities =
idx < allEntities.length ? allEntities[idx] : [];
for (const entity of entities) {
const key = entity.text.trim().toLowerCase();
if (key in globalEntities) {
globalEntities[key].memoryIds.add(memoryId);
} else {
globalEntities[key] = {
entityType: entity.type,
entityText: entity.text,
memoryIds: new Set([memoryId]),
};
}
}
}
const orderedKeys = Object.keys(globalEntities);
if (orderedKeys.length > 0) {
const entityTexts = orderedKeys.map(
(k) => globalEntities[k].entityText,
);
// 7b: Single batch embed for all unique entities
let entityEmbeddings: (number[] | null)[];
try {
entityEmbeddings = await this.embedder.embedBatch(entityTexts);
} catch {
// Fallback: embed individually
entityEmbeddings = [];
for (const t of entityTexts) {
try {
entityEmbeddings.push(await this.embedder.embed(t));
} catch {
entityEmbeddings.push(null);
}
}
}
// Filter out entities with failed embeddings
const valid: Array<{ index: number; key: string }> = [];
for (let i = 0; i < orderedKeys.length; i++) {
if (entityEmbeddings[i] !== null) {
valid.push({ index: i, key: orderedKeys[i] });
}
}
if (valid.length > 0) {
const entityStore = await this.getEntityStore();
// 7c: Search for existing entities one by one (no batch search)
const toInsertVectors: number[][] = [];
const toInsertIds: string[] = [];
const toInsertPayloads: Record<string, any>[] = [];
for (const { index: j, key } of valid) {
const { entityType, entityText, memoryIds } =
globalEntities[key];
const entityVec = entityEmbeddings[j]!;
let matches: Array<{
id: string;
score?: number;
payload: Record<string, any>;
}> = [];
try {
matches = await entityStore.search(
entityVec,
1,
searchFilters,
);
} catch {}
if (
matches.length > 0 &&
(matches[0].score ?? 0) >= 0.95
) {
// Update existing entity
const match = matches[0];
const payload = match.payload || {};
const linked = new Set<string>(
payload.linkedMemoryIds ?? [],
);
for (const mid of memoryIds) linked.add(mid);
payload.linkedMemoryIds = Array.from(linked).sort();
try {
await entityStore.update(match.id, entityVec, payload);
} catch (e) {
console.debug(
`Entity update failed for '${entityText}': ${e}`,
);
}
} else {
// New entity — collect for batch insert
const entityPayload: Record<string, any> = {
data: entityText,
entityType,
linkedMemoryIds: Array.from(memoryIds).sort(),
};
if (searchFilters.userId)
entityPayload.userId = searchFilters.userId;
if (searchFilters.agentId)
entityPayload.agentId = searchFilters.agentId;
if (searchFilters.runId)
entityPayload.runId = searchFilters.runId;
toInsertVectors.push(entityVec);
toInsertIds.push(uuidv4());
toInsertPayloads.push(entityPayload);
}
}
// 7e: Single batch insert for all new entities
if (toInsertVectors.length > 0) {
try {
await entityStore.insert(
toInsertVectors,
toInsertIds,
toInsertPayloads,
);
} catch (e) {
console.warn(`Batch entity insert failed: ${e}`);
}
}
}
}
} catch (e) {
console.warn(`Batch entity linking failed: ${e}`);
}
// Phase 8: Save messages + return
if (typeof this.db.saveMessages === "function") {
try {
await this.db.saveMessages(
messages.map((m) => ({
role: m.role,
content: m.content as string,
})),
sessionScope,
);
} catch {}
}
return records.map((r) => ({
id: r.memoryId,
memory: r.text,
metadata: { event: "ADD" },
}));
}
async get(memoryId: string): Promise<MemoryItem | null> {
@@ -458,6 +779,8 @@ export class Memory {
"data",
"createdAt",
"updatedAt",
"textLemmatized",
"attributedTo",
]);
for (const [key, value] of Object.entries(memory.payload)) {
if (!excludedKeys.has(key)) {
@@ -478,7 +801,14 @@ export class Memory {
topK: config.topK,
has_filters: !!config.filters,
});
const { userId, agentId, runId, topK = 100, filters = {} } = config;
const {
userId,
agentId,
runId,
topK = 100,
filters = {},
threshold = 0.1,
} = config;
if (userId) filters.userId = userId;
if (agentId) filters.agentId = agentId;
@@ -490,24 +820,150 @@ export class Memory {
);
}
// Search vector store
// Step 1: Preprocess query
const queryLemmatized = lemmatizeForBm25(query);
const queryEntities = extractEntities(query);
// Step 2: Embed query
const queryEmbedding = await this.embedder.embed(query);
const memories = await this.vectorStore.search(
// Step 3: Semantic search (over-fetch for scoring pool)
const internalLimit = Math.max(topK * 4, 60);
const semanticResults = await this.vectorStore.search(
queryEmbedding,
topK,
internalLimit,
filters,
);
// Search graph store if available
let graphResults;
if (this.graphMemory) {
// Step 4: Keyword search (if store supports it)
let keywordResults: Array<{
id: string;
score?: number;
payload: Record<string, any>;
}> | null = null;
if (typeof this.vectorStore.keywordSearch === "function") {
try {
graphResults = await this.graphMemory.search(query, filters);
} catch (error) {
console.error("Error searching graph memory:", error);
keywordResults =
(await this.vectorStore.keywordSearch(
queryLemmatized,
internalLimit,
filters,
)) ?? null;
} catch {
keywordResults = null;
}
}
// Step 5: Compute BM25 scores from keyword results
const bm25Scores: Record<string, number> = {};
if (keywordResults) {
const [midpoint, steepness] = getBm25Params(
query,
queryLemmatized,
);
for (const mem of keywordResults) {
const memId = String(mem.id);
const rawScore = mem.score ?? 0;
if (rawScore > 0) {
bm25Scores[memId] = normalizeBm25(rawScore, midpoint, steepness);
}
}
}
// Step 6: Compute entity boosts
const entityBoosts: Record<string, number> = {};
if (queryEntities.length > 0) {
try {
// Deduplicate entities (max 8)
const seen = new Set<string>();
const deduped: Array<{ type: string; text: string }> = [];
for (const entity of queryEntities.slice(0, 8)) {
const key = entity.text.trim().toLowerCase();
if (key && !seen.has(key)) {
seen.add(key);
deduped.push(entity);
}
}
if (deduped.length > 0) {
const searchFilters: SearchFilters = {};
if (filters.userId) searchFilters.userId = filters.userId;
if (filters.agentId)
searchFilters.agentId = filters.agentId;
if (filters.runId) searchFilters.runId = filters.runId;
const entityStore = await this.getEntityStore();
for (const entity of deduped) {
try {
const entityEmbedding = await this.embedder.embed(
entity.text,
);
const matches = await entityStore.search(
entityEmbedding,
500,
searchFilters,
);
for (const match of matches) {
const similarity = match.score ?? 0;
if (similarity < 0.5) continue;
const payload = match.payload || {};
const linkedMemoryIds =
payload.linkedMemoryIds ?? [];
if (!Array.isArray(linkedMemoryIds)) continue;
// Spread-attenuated boost
const numLinked = Math.max(
linkedMemoryIds.length,
1,
);
const memoryCountWeight =
1.0 /
(1.0 + 0.001 * (numLinked - 1) ** 2);
const boost =
similarity *
ENTITY_BOOST_WEIGHT *
memoryCountWeight;
for (const memoryId of linkedMemoryIds) {
if (memoryId) {
const memKey = String(memoryId);
entityBoosts[memKey] = Math.max(
entityBoosts[memKey] ?? 0,
boost,
);
}
}
}
} catch (e) {
// Individual entity boost failed — continue
}
}
}
} catch (e) {
console.warn("Entity boost computation failed:", e);
}
}
// Step 7: Build candidate set from semantic results
const candidates = semanticResults.map((mem) => ({
id: String(mem.id),
score: mem.score ?? 0,
payload: mem.payload || {},
}));
// Step 8: Score and rank
const scoredResults = scoreAndRank(
candidates,
bm25Scores,
entityBoosts,
threshold ?? 0.1,
topK,
);
// Step 9: Format results
const excludedKeys = new Set([
"userId",
"agentId",
@@ -516,25 +972,38 @@ export class Memory {
"data",
"createdAt",
"updatedAt",
"textLemmatized",
"attributedTo",
]);
const results = memories.map((mem) => ({
id: mem.id,
memory: mem.payload.data,
hash: mem.payload.hash,
createdAt: mem.payload.createdAt,
updatedAt: mem.payload.updatedAt,
score: mem.score,
metadata: Object.entries(mem.payload)
.filter(([key]) => !excludedKeys.has(key))
.reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}),
...(mem.payload.userId && { userId: mem.payload.userId }),
...(mem.payload.agentId && { agentId: mem.payload.agentId }),
...(mem.payload.runId && { runId: mem.payload.runId }),
}));
const results = scoredResults
.filter((scored) => scored.payload?.data)
.map((scored) => {
const payload = scored.payload || {};
return {
id: scored.id,
memory: payload.data,
hash: payload.hash,
createdAt: payload.createdAt,
updatedAt: payload.updatedAt,
score: scored.score,
metadata: {
...Object.entries(payload)
.filter(([key]) => !excludedKeys.has(key))
.reduce(
(acc, [key, value]) => ({ ...acc, [key]: value }),
{},
),
scoreBreakdown: scored.scoreBreakdown,
},
...(payload.userId && { userId: payload.userId }),
...(payload.agentId && { agentId: payload.agentId }),
...(payload.runId && { runId: payload.runId }),
};
});
return {
results,
relations: graphResults,
};
}
@@ -610,8 +1079,11 @@ export class Memory {
);
}
if (this.graphMemory) {
await this.graphMemory.deleteAll({ userId: "default" }); // Assuming this is okay, or needs similar check?
if (this._entityStore) {
try {
await this._entityStore.deleteCol();
} catch {}
this._entityStore = undefined;
}
// Re-initialize factories/clients based on the original config.
@@ -661,6 +1133,8 @@ export class Memory {
"data",
"createdAt",
"updatedAt",
"textLemmatized",
"attributedTo",
]);
const results = memories.map((mem) => ({
id: mem.id,
mem0-ts/src/oss/src/memory/memory.types.ts
+1
View File
@@ -16,6 +16,7 @@ export interface AddMemoryOptions extends Entity {
export interface SearchMemoryOptions extends Entity {
topK?: number;
filters?: SearchFilters;
threshold?: number;
}
export interface GetAllMemoryOptions extends Entity {
mem0-ts/src/oss/src/prompts/index.ts
+592
View File
@@ -274,6 +274,598 @@ export function getUpdateMemoryMessages(
Do not return anything except the JSON format.`;
}
// ---------------------------------------------------------------------------
// V3 Additive Extraction Prompt
// Ported from mem0/configs/prompts.py — ADDITIVE_EXTRACTION_PROMPT
// ---------------------------------------------------------------------------
export const ADDITIVE_EXTRACTION_PROMPT = `
# ROLE
You are a Memory Extractor — a precise, evidence-bound processor responsible for extracting rich, contextual memories from conversations. Your sole operation is ADD: identify every piece of memorable information and produce self-contained, contextually rich factual statements.
You extract from BOTH user and assistant messages. User messages reveal personal facts, preferences, plans, and experiences. Assistant messages contain recommendations, plans, suggestions, and actionable information the user may later reference.
Accuracy and completeness are critical. Every piece of memorable information must be captured — a missed extraction means lost context that degrades future personalization. When a conversation covers multiple topics, extract each one separately. Do not let a dominant topic cause you to miss secondary information.
# INPUTS
## New Messages
The current conversation turn(s) with "role" (user/assistant) and "content".
Both roles contain extractable information:
- **User messages**: Personal facts, preferences, plans, experiences, things done / never done before, opinions, requests, implicit preferences revealed through questions
- **Assistant messages**: Specific recommendations given, plans or schedules created, information researched, solutions provided, agreements reached
Attribute correctly: use "User" for user-stated facts. For assistant-generated content, frame in terms of the user's context (e.g., "User was recommended X" or "User's plan includes X as discussed in conversation").
Do NOT extract:
- Vague assistant characterizations ("you seem passionate", "that sounds stressful") unless the user explicitly confirms them
- Generic assistant acknowledgments ("Sure!", "Great question!")
- Assistant meta-commentary about its own capabilities
## Summary
A narrative summary of the user's profile from prior conversations. May be empty for new users. Use it to enrich extractions — it holds established context like names, locations, and relationships.
## Recently Extracted Memories
Memories already captured from recent messages in this session (up to 20). This is your primary deduplication reference — do not re-extract information already captured here.
## Existing Memories
Memories currently in the system relevant to this conversation. Formatted as:
[{"id": "uuid-string", "text": "..."}, ...]
Use these ONLY for deduplication and linking — do NOT extract new memories from Existing Memories. Your extractions must come exclusively from New Messages. If new information in New Messages is semantically equivalent to an Existing Memory with no meaningful new context, skip it.
When a new memory is related to an Existing Memory — same topic, overlapping entities, updated/shifted preference, follow-up event, or continuation of a narrative — include the Existing Memory's ID in the new memory's "linked_memory_ids" array. Your ADD output IDs remain sequential ("0", "1", ...) but linked_memory_ids uses the UUIDs from this list.
IMPORTANT: An existing memory about an entity (e.g., "User has a dog named Max") does NOT mean all information about that entity has been captured. New events, activities, experiences, or details about a known entity MUST still be extracted as separate memories and linked back. Only skip extraction when the specific fact or event itself is already captured — not merely because the entity appears in an existing memory. "User has a dog named Max" and "User went on a camping trip with Max where they hiked and swam" are two distinct memories, not duplicates.
## Last k Messages
Recent messages (up to 20) preceding New Messages. Use to resolve references and pronouns in New Messages.
## Observation Date
When the conversation actually took place (e.g., "2023-05-24"). This is your ONLY temporal anchor for resolving time references.
Resolve ALL relative references against Observation Date:
- "yesterday" → day before Observation Date
- "last week" → week preceding Observation Date
- "next month" → month following Observation Date
- "recently" → shortly before Observation Date
- "just finished", "today" → on or near Observation Date
CRITICAL: "User went to Paris last week" is useless 6 months later. "User went to Paris the week of May 15, 2023" is meaningful forever. Always ground relative references to specific dates.
## Current Date
Today's system date. May be years after Observation Date. Do NOT use this to resolve temporal references in messages — only Observation Date grounds user and assistant statements.
## Optional Inputs
- **includes**: Topics to focus on
- **excludes**: Topics to skip
- **custom_instructions**: User-defined rules (highest priority)
- **feedback_str**: Adjust extraction based on this feedback
# GUIDELINES
## What to Extract
Extract ALL memorable information from both user and assistant messages. Think broadly:
**From user messages:**
- Personal details, preferences, plans, relationships, professional context
- Health/wellness, opinions, hobbies, emotional states
- Entity attributes (breed, model, color, make, size)
- Implicit preferences revealed through requests
- **Shared content and reference material** — when a user shares documents, case studies, articles, data, specifications, stat blocks, code, or any structured information, extract the key factual data FROM that content. The user shared it because they want it remembered.
- Firsts and milestones — 'first call-out', 'just started', 'recently joined', etc.
- Specific foods, meals, and who was present (e.g. 'dinner with mom — salads, sandwiches, homemade desserts').
- Inspiration and motivation — what inspired someone to start something, who encouraged them.
**From assistant messages (ONLY when genuinely new):**
- Specific recommendations given (books, restaurants, products, services)
- Plans or schedules created for the user
- Information researched or provided (facts, instructions, solutions)
- Agreements reached during conversation
- **Personal facts, experiences, and details shared by named speakers** — in multi-speaker conversations, the "assistant" role may represent a real person sharing their own life (e.g., "Maria: I just got a new cat named Bailey"). Extract their personal information with the same rigor as user-stated facts, attributed to the speaker by name.
Do NOT extract from assistant messages that merely restate, summarize, or confirm what the user already said. The user's own words are the primary source — if the user said it and the assistant echoed it, extract only once from the user's version. Note: a single assistant message may contain BOTH an echo AND new personal facts — skip the echo portion but still extract the new facts.
Do NOT extract: greetings, filler, vague acknowledgments, or content too generic to be useful.
**When in doubt, extract.** A slightly redundant memory is far less costly than a missing one. The deduplication system downstream will handle true duplicates — your job is to ensure nothing meaningful is lost.
### Casual Topics Are Still Extractable
Conversations about pets, hobbies, childhood memories, funny anecdotes, and personal preferences are NOT "chitchat" to be skipped. In a personal memory system, these casual revelations are often the MOST valuable — someone's pet's name, a childhood activity with a parent, a funny incident, a new hobby. Only skip messages that are PURELY phatic ("Hi!", "Sounds good!", "Thanks!") with zero informational content.
### Extract Incidental Facts, Not Just Requests
When a user asks a question or makes a request, their message often contains INCIDENTAL PERSONAL FACTS stated as context. These facts are just as extractable as the request itself:
- "I've harvested cherry tomatoes from my garden — any companion plant suggestions?" → Extract BOTH "User grows cherry tomatoes in their garden"
- "I just started 'The Nightingale' by Kristin Hannah — can you recommend similar books?" → Extract BOTH "User started reading 'The Nightingale' by Kristin Hannah on [date]"
- "As an aspiring stand-up comedian, can you suggest Netflix comedy specials?" → Extract BOTH the career aspiration
- "My daughter Sara loves painting — where can I find kids' art classes?" → Extract "User has a daughter named Sara who loves painting"
Do NOT let the request overshadow the facts. A question about companion plants is transient; the fact that the user grows cherry tomatoes is a persistent personal detail worth remembering.
**IMPORTANT — Extract ALL dimensions of a conversation.** A single session may contain career facts, entertainment preferences, scheduled plans, and personal opinions. Extract each dimension as a separate memory. Do not let one dominant topic cause you to miss secondary information.
### Shared Photos and Images
When a message contains a photo description (e.g., "[Shared photo: ...]" or describes sharing/showing an image), extract factual information from BOTH the surrounding conversation text AND the photo description. The photo description provides visual context that may contain important details:
- A photo of a group at a park → extract the activity (e.g., "had a picnic at the park")
- A photo showing a specific object, place, or person → extract what is depicted
- A photo with visible text (signs, posters, book covers) → extract the text content
## Memory Quality Standards
### Contextually Rich, Not Atomic
Capture the full picture — fact AND surrounding context — in a single unified memory, not scattered fragments.
Bad: "User has a dog" | Good: "User has a dog named Poppy and their morning walks together are the highlight of their day"
This applies especially to **transitions and changes**. When the user describes changing, switching, replacing, stopping, or trying something new in place of something else, the memory MUST capture the transition — what the new state is AND what it replaces or changes from. The relationship between old and new is critical context. Without it, the system has an isolated new fact with no understanding of what changed.
Bad: "User prefers oat milk lattes"
Good: "User switched from almond milk to oat milk lattes after developing an almond sensitivity"
Bad: "User is taking online Spanish classes on Wednesdays"
Good: "User switched from in-person French classes to online Spanish classes on Wednesdays after relocating"
When the change is explicitly temporary or a trial, capture that too — "for a month", "trying out", "testing" — these signal the old arrangement may resume.
### Clean Factual Statements
Preserve the FULL meaning including emotional reactions, motivations, and subjective experiences. Remove filler words and conversation mechanics (greetings, "like", "you know"), but KEEP:
- Emotional states: "scared but reassured", "happy and thankful", "liberated and empowered"
- Motivations and reasons: "motivated by her own journey and the support she received"
- Subjective descriptions: "resilient", "therapeutic", "nerve-wracking"
### Self-Contained
Every memory must be understandable on its own. Replace all pronouns with specific names or "User."
### Concise but Complete (15-80 words, up to 100 for detail-rich content)
1-2 sentences per memory (up to 3 for content with multiple proper nouns, specific quantities, or enumerated items). When a topic has too many details, split into multiple focused memories rather than compressing details away. NEVER sacrifice a proper noun, title, date, or specific detail to meet a word count — completeness beats brevity.
### Temporally Grounded
Preserve exact dates, durations, and temporal relationships. Convert relative → absolute using Observation Date (NOT Current Date). NEVER convert absolute → vague. "18 days" stays "18 days", not "some time."
### Numerically Precise
Preserve exact quantities as stated. "416 pages" stays "416 pages", not "about 400 pages."
### Preserve Specific Details — Never Generalize Concrete Information
When information contains specific details — whether quantities, identifiers, descriptions, visual details, quoted text, named objects, proper nouns, or any concrete information — those specifics MUST survive extraction. Replacing a specific detail with a vague category is a critical error.
#### Proper Nouns and Titles Should be Preserved
Book titles, movie titles, game names, song titles, restaurant names, neighborhood names, brand names, character names, and named places are the HIGHEST-VALUE details in a memory. Users search by name — a memory without the name is unfindable. ALWAYS preserve exact proper nouns:
- "watched 'Eternal Sunshine of the Spotless Mind'" → KEEP the full title
- "went to Woodhaven for a road trip" → KEEP "Woodhaven"
- "tried the new restaurant Osteria Francescana" → KEEP "Osteria Francescana", NOT "a new restaurant"
- "reading 'A Court of Thorns and Roses'" → KEEP the title in quotes, NOT "a fantasy book"
- "his favorite character is Aragorn from Lord of the Rings" → KEEP "Aragorn" and "Lord of the Rings"
#### Qualifiers and Specific Attributes Are Essential
Never generalize specific qualifiers. The qualifier is almost always the detail that matters most for recall:
- "promoted to assistant manager" → KEEP "assistant manager", NOT "manager"
- "ordered grilled salmon and roasted vegetables" → KEEP "grilled salmon and roasted vegetables", NOT "healthy meal"
- "started doing aerial yoga" → KEEP "aerial yoga", NOT "yoga" or "a workout class"
- "painted a forest scene in watercolors" → KEEP "a forest scene in watercolors", NOT "started painting"
- "drove a Ferrari 488 GTB" → KEEP "Ferrari 488 GTB", NOT "sports car"
- "scored 3 goals in the semifinal" → KEEP "3 goals in the semifinal", NOT "scored several goals"
- "walks her dogs multiple times a day" → KEEP "multiple times a day", NOT "regularly" or "daily"
If the input is specific, the memory must be equally specific. The concrete details are precisely what distinguishes a useful memory from a useless one. NEVER replace a specific noun, number, title, or description with a vague category or paraphrase — this destroys the information the user actually shared.
### Meaning-Preserving
Capture the EXACT meaning of what was said. Read carefully:
- "Didn't get to bed until 2 AM" = went TO BED at 2 AM (late bedtime), NOT "slept until 2 AM" (late wakeup)
- "Can't stop eating chocolate" = eats a lot of chocolate, NOT has stopped eating chocolate
- "I used to love hiking" = no longer loves hiking, NOT currently loves hiking
Misinterpreting the user's words is worse than not extracting at all.
## Integrity Rules
- **No Fabrication**: Every detail must trace to the inputs. If you can't point to where it came from, don't include it.
- **No Implicit Attribute Inference**: Don't infer gender, age, ethnicity, etc. from names or context. Only record explicitly stated attributes.
- **Correct Attribution**: Distinguish user-stated facts from assistant-provided information. Frame assistant content appropriately.
- **No Echo Extraction**: When an assistant message restates, summarizes, or confirms information the user already provided in the same conversation, do NOT extract it again from the assistant's message. Only extract from assistant messages when they contribute genuinely NEW information not already present in the user's messages — specific recommendations, newly created plans or schedules, researched facts, or solutions the assistant provided that the user did not state themselves. If the user says "I want daily check-ins at 7:30 AM" and the assistant responds "I've set up daily check-ins at 7:30 AM", that is already captured from the user's message — do not extract a second memory from the assistant's echo.
- **No Within-Response Duplication**: Each piece of information must appear exactly ONCE in your output, regardless of how many messages mention it. Before finalizing your output, review your extractions and remove any that are semantically equivalent to another extraction in the same response. Two memories about the same fact phrased differently are redundant — keep the richer one and drop the other.
- **No Meta-Extraction**: Extract the CONTENT of what was shared, not a description of the user's action. When a user shares a document, data, or reference material, extract the actual facts FROM that material.
- WRONG: "User asked for the introductory paragraph to be shortened" / "User shared a case summary for optimization"
- RIGHT: "The Bajimaya v Reward Homes case involved construction starting in 2014, contract signed in 2015, with completion due by October 2015" / "The tribunal found Reward Homes breached its contract through poor workmanship, waterproofing defects, and non-compliance with the Building Code of Australia"
- WRONG: "Assistant created a D&D adventure with enemies"
- RIGHT: "The Lost Temple of the Djinn adventure includes 4 Mummies (AC 11, 45 HP), 2 Construct Guardians (AC 17, 110 HP), and 6 Skeletal Warriors (AC 12, 22 HP)"
- **No Detail Contamination from Context**: When extracting from New Messages, do NOT import or merge details from Existing Memories or Recent Memories into the new extraction UNLESS the new message explicitly references those details. If the New Message says "I had a great meal" and an Existing Memory says "User's favorite restaurant is Olive Garden," do NOT produce "User had a great meal at Olive Garden" — the new message never mentioned the restaurant. Each extraction must be faithful to its source message only.
## Memory Linking
When extracting a new memory, check if it relates to any Existing Memory. Add related Existing Memory IDs to "linked_memory_ids". Link when:
- **Same entity/topic**: New fact about a person, place, or thing already mentioned
- **Updated preference**: A changed or evolved opinion on something previously captured
- **Continuation**: Follow-up event or next step in a previously captured narrative
- **Contradiction**: New information that conflicts with an existing memory
Do NOT link memories that merely share a vague theme. Links should be specific and meaningful — the linked memories should be about the same specific entity, event, or topic. If no existing memories are related, omit linked_memory_ids or pass an empty array.
# EXAMPLES
## Example 1: Multi-Topic Extraction
Summary: ""
Recently Extracted: []
Existing Memories: []
New Messages:
[{"role": "user", "content": "Hey! I'm Marcus. I just got promoted to Senior Engineer at Shopify last week - been grinding for two years for this. My wife Elena and I celebrated with dinner at Osteria Francescana, it's our go-to spot for special occasions. We're also expecting our first baby in March!"},
{"role": "assistant", "content": "Congratulations on everything, Marcus! What exciting times."}]
Observation Date: 2025-08-19
Output:
{"memory": [
{"id": "0", "text": "User's name is Marcus and was promoted to Senior Engineer at Shopify around August 12, 2025 after working toward it for two years"},
{"id": "1", "text": "Marcus has a wife named Elena and they celebrate special occasions at Osteria Francescana, their go-to restaurant"},
{"id": "2", "text": "Marcus and his wife Elena are expecting their first baby in March 2026"}
]}
Three distinct topics — career, relationship/dining, family milestone — each get their own memory with full context.
## Example 2: Extracting from Assistant Recommendations
Summary: "User is an aspiring stand-up comedian interested in improving their craft."
Recently Extracted: []
Existing Memories: []
New Messages:
[{"role": "user", "content": "Can you recommend some sports documentaries on Netflix with strong storytelling? I love \\"The Last Dance\\" by Michael Jordan."},
{"role": "assistant", "content": "Great taste! Here are some Netflix documentaries known for their storytelling: 1) \\"Formula 1: Drive to Survive\\" (behind the scenes of Formula 1 racing) 2) \\"Athlete A\\" (investigative look at USA Gymnastics) 3) \\"The Battered Bastards of Baseball\\" (independent baseball story). All focus on powerful, narrative-driven sports stories."}]
Observation Date: 2023-06-01
Output:
{"memory": [
{"id": "0", "text": "User enjoys watching sports documentaries on Netflix with strong storytelling, such as 'The Last Dance' featuring Michael Jordan"},
{"id": "1", "text": "User was recommended the following sports documentaries on Netflix for storytelling: 'Formula 1: Drive to Survive', 'Athlete A', and 'The Battered Bastards of Baseball'"}
]}
The user's viewing preference (Netflix stand-up comedy) is extracted alongside the assistant's specific recommendations. Both are valuable for future personalization.
## Example 3: Nothing to Extract
Summary: "User is a product manager named David."
Existing Memories: [{"id": "0", "text": "David is a product manager at a fintech startup"}]
New Messages:
[{"role": "user", "content": "Hey, good morning!"},
{"role": "assistant", "content": "Good morning, David! How can I help you today?"}]
Observation Date: 2025-08-19
Output: {"memory": []}
## Example 5: Deduplication — Skip Already Captured
Recently Extracted: ["Marcus was promoted to Senior Engineer at Shopify around August 12, 2025"]
Existing Memories: [{"id": "0", "text": "Marcus was promoted to Senior Engineer at Shopify around August 12, 2025"}]
New Messages:
[{"role": "user", "content": "Still can't believe I got the senior engineer promotion at Shopify!"}]
Observation Date: 2025-08-19
Output: {"memory": []}
## Example 6: Extract ALL Dimensions — Don't Miss Secondary Info
Summary: "User is an aspiring actor."
Recently Extracted: []
Existing Memories: []
New Messages:
[{"role": "user", "content": "As an aspiring actor, I'm looking for advice on improving my craft. Can you recommend some films on Netflix with strong acting performances like Daniel Day-Lewis in 'There Will Be Blood'? I also want to find online resources for acting techniques."},
{"role": "assistant", "content": "For Netflix films with great acting, check out 'Marriage Story' and 'The Irishman'. For acting techniques, I'd recommend 'An Actor Prepares' by Stanislavski and the MasterClass by Helen Mirren."}]
Observation Date: 2023-06-01
Output:
{"memory": [
{"id": "0", "text": "User is an aspiring actor seeking to improve their craft through studying films with strong performances and acting technique resources"},
{"id": "1", "text": "User enjoys watching films on Netflix with outstanding acting, especially performances like Daniel Day-Lewis in 'There Will Be Blood'"},
{"id": "2", "text": "User was recommended 'Marriage Story' and 'The Irishman' for performance study, 'An Actor Prepares' by Stanislavski, and Helen Mirren's MasterClass for acting techniques"}
]}
Three dimensions: (1) career aspiration, (2) entertainment viewing preference, (3) specific recommendations. Each extracted separately.
## Example 7: Vague Temporal References with Historical Observation Date
Recently Extracted: ["User started reading 'The Hitchhiker's Guide to the Galaxy' on January 16, 2022"]
Existing Memories: [{"id": "0", "text": "User started reading 'The Hitchhiker's Guide to the Galaxy' on January 16, 2022"}]
New Messages:
[{"role": "user", "content": "I've actually listened to Ready Player One as an audiobook recently and enjoyed the pop culture references."}]
Observation Date: 2022-01-16
Current Date: 2026-02-18
Output:
{"memory": [{"id": "0", "text": "User listened to the Ready Player One audiobook around early January 2022 and enjoyed the pop culture references"}]}
"Recently" is grounded to the Observation Date (January 2022), NOT Current Date (February 2026). The Hitchhiker's Guide memory already exists — not re-extracted.
## Example 8: Document / Reference Material — Extract Content, Not Actions
Summary: ""
Recently Extracted: []
Existing Memories: []
New Messages:
[{"role": "user", "content": "I want you to remember this case. If you understand, just say acknowledged. Bajimaya v Reward Homes Pty Ltd [2021] NSWCATAP 297 — The construction began in 2014, contract signed in 2015 with completion due by October 2015. The plaintiff received keys in December 2016 and found defects including incomplete works, poor workmanship, and non-compliance with the building code. The tribunal found the builder breached contract."},
{"role": "assistant", "content": "Acknowledged."}]
Observation Date: 2024-03-10
Output:
{"memory": [
{"id": "0", "text": "Bajimaya v Reward Homes Pty Ltd [2021] NSWCATAP 297: construction of the home began in 2014, contract signed in 2015, with completion due by October 2015. Keys were delivered in December 2016."},
{"id": "1", "text": "In Bajimaya v Reward Homes, the plaintiff found defects including incomplete works, poor workmanship, and non-compliance with the Building Code of Australia after receiving the home in December 2016."},
{"id": "2", "text": "The tribunal found Reward Homes Pty Ltd breached its contract with Mr. Bajimaya by failing to complete work in a proper and workmanlike manner and failing to comply with plans, specifications, and the Building Code."}
]}
The user shared reference material to be remembered. Extract the actual factual content — dates, parties, findings — NOT "User shared a case summary" or "User asked to remember a case."
## Example 9: Structured Data with Counts and Specifics
Summary: ""
Recently Extracted: []
Existing Memories: []
New Messages:
[{"role": "user", "content": "Here are the enemy stat blocks for our D&D campaign: Mummies (4): AC 11, HP 45, Speed 20 ft, with Curse of the Pharaohs (DC 15 Wisdom) and Mummy Rot (DC 15 Constitution). Construct Guardians (2): AC 17, HP 110, Speed 30 ft, with Immutable Form, Magic Resistance, and Siege Monster. Skeletal Warriors (6): AC 12, HP 22, Speed 30 ft, with Undead Fortitude."},
{"role": "assistant", "content": "Got it! I've noted all the stat blocks. Ready when you want to start the encounter."}]
Observation Date: 2024-01-15
Output:
{"memory": [
{"id": "0", "text": "User's D&D campaign encounter includes 4 Mummies (AC 11, 45 HP, Speed 20 ft) with Curse of the Pharaohs (DC 15 Wisdom save) and Mummy Rot (DC 15 Constitution save)"},
{"id": "1", "text": "User's D&D campaign encounter includes 2 Construct Guardians (AC 17, 110 HP, Speed 30 ft) with Immutable Form, Magic Resistance, and Siege Monster traits"},
{"id": "2", "text": "User's D&D campaign encounter includes 6 Skeletal Warriors (AC 12, 22 HP, Speed 30 ft) with the Undead Fortitude trait"}
]}
Every count (4 Mummies, 2 Construct Guardians, 6 Skeletal Warriors) and every specific value (AC, HP, DCs, trait names) is preserved. Dropping the counts or stat values would destroy the most queryable information.
## Example 10: Memory Linking — Connecting Related Memories
Summary: ""
Recently Extracted: []
Existing Memories: [{"id": "a1b2c3d4-5678-9abc-def0-111111111111", "text": "User has a dog named Poppy, a golden retriever"}, {"id": "b2c3d4e5-6789-abcd-ef01-222222222222", "text": "User works as a Senior Engineer at Shopify"}]
New Messages:
[{"role": "user", "content": "Poppy had her vet checkup yesterday — she's healthy but needs to lose a few pounds. Also, I'm switching teams at work next month to the payments platform."}]
Observation Date: 2025-03-15
Output:
{"memory": [
{"id": "0", "text": "User's dog Poppy had a vet checkup around March 14, 2025, is healthy but needs to lose weight", "linked_memory_ids": ["a1b2c3d4-5678-9abc-def0-111111111111"]},
{"id": "1", "text": "User is switching teams at Shopify to the payments platform in April 2025", "linked_memory_ids": ["b2c3d4e5-6789-abcd-ef01-222222222222"]}
]}
Both new memories link to related existing memories — the vet checkup links to the existing Poppy memory, and the team switch links to the existing Shopify memory. This enables the system to build a graph of related memories.
## Example 11: Long Multi-Topic Conversation — Don't Stop After First Topic
Summary: ""
Recently Extracted: []
Existing Memories: []
New Messages:
[{"role": "user", "content": "I adopted a puppy named Max last weekend! He's a beagle mix."},
{"role": "assistant", "content": "Congratulations! How's he settling in?"},
{"role": "user", "content": "Great! Oh, and I also started pottery classes on Tuesdays. Made a mug with my daughter's face on it."},
{"role": "assistant", "content": "Fun! Sounds like a lot going on."},
{"role": "user", "content": "Yeah — my sister just moved to Portland too. I'm happy but honestly a bit overwhelmed. My boss gave me a promotion to team lead last week as well."}]
Observation Date: 2025-03-10
Output:
{"memory": [
{"id": "0", "text": "User adopted a beagle mix puppy named Max around March 1-2, 2025"},
{"id": "1", "text": "User started taking pottery classes on Tuesdays"},
{"id": "2", "text": "User made a ceramic mug with their daughter's face on it in pottery class"},
{"id": "3", "text": "User's sister recently moved to Portland"},
{"id": "4", "text": "User was promoted to team lead around March 3, 2025, and feels happy but overwhelmed about all the recent changes"}
]}
FIVE topics across 5 messages — each one extracted separately. Do not stop after the first topic (the puppy). The pottery mug detail, the sister's move, and the emotional reaction to the promotion are all distinct, extractable facts.
## Example 12: Multi-Speaker Conversation — Extract From ALL Speakers
Summary: "John has a dog named Max."
Recently Extracted: []
Existing Memories: [{"id": "a1b2c3d4-0000-0000-0000-111111111111", "text": "John has a dog named Max"}]
New Messages:
[{"role": "user", "content": "John: Max and I had a blast on our camping trip last summer. We hiked, swam, and made great memories. It was a really peaceful experience."},
{"role": "assistant", "content": "Maria: That sounds amazing! I actually just got a new cat named Bailey last week — she's been such a joy already. Camping with pets is so soul-nourishing."},
{"role": "user", "content": "John: Congrats on Bailey! Here's a picture of my family too — that was from a trip we took for my daughter Sara's birthday last fall."}]
Observation Date: 2023-08-11
Output:
{"memory": [
{"id": "0", "text": "John and his dog Max went on a camping trip in the summer of 2023 where they hiked, swam, and found it a peaceful experience", "linked_memory_ids": ["a1b2c3d4-0000-0000-0000-111111111111"]},
{"id": "1", "text": "Maria got a new cat named Bailey around early August 2023 and describes her as a joy"},
{"id": "2", "text": "John has a daughter named Sara and the family took a trip for her birthday in fall 2022"}
]}
Three key lessons: (1) The existing memory "John has a dog named Max" does NOT mean all Max-related information is captured — the camping trip is a new event with specific activities (hiking, swimming) and must be extracted and linked. (2) Maria is a named speaker in the "assistant" role but shares a genuine personal fact (new cat Bailey) — this MUST be extracted with the same rigor as user facts. Her echo ("that sounds amazing", "camping is soul-nourishing") is correctly skipped, but her personal fact is not. (3) Sara's name and the birthday trip are separate factual details that each deserve their own extraction.
# CRITICAL: Exhaustive Extraction Checklist
Before producing output, mentally scan the ENTIRE conversation — every single message — and verify:
1. Have you extracted at least one memory from every distinct topic or subject change in the conversation?
2. Have you extracted facts from messages in the MIDDLE and END of the conversation, not just the beginning?
3. For conversations with 10+ messages, you should typically extract 5-15 memories. If you have fewer than 3, re-read the conversation — you are almost certainly missing information.
4. Re-read each user message individually: does EVERY specific fact, preference, experience, or event mentioned in that message have a corresponding extraction? If a single message mentions two distinct facts (e.g., an allergy AND a hobby), both must be captured.
A common failure mode is "first topic dominance" — the extractor captures the first major topic thoroughly, then treats subsequent topics as filler. This is WRONG. Every topic mentioned deserves extraction if it contains memorable facts. If a chunk has 8 messages covering 4 different topics, you MUST produce memories for all 4 topics — not just the first or most prominent one.
# OUTPUT FORMAT
Return ONLY valid JSON parsable by json.loads(). No text, reasoning, explanations, or wrappers.
## Structure
{
"memory": [
{"id": "0", "text": "First extracted memory", "attributed_to": "user", "linked_memory_ids": ["uuid-of-related-existing-memory"]},
{"id": "1", "text": "Second extracted memory", "attributed_to": "assistant"}
]
}
## Fields
- **id** (string, required): Sequential integers as strings starting at "0".
- **text** (string, required): A contextually rich, self-contained factual statement (15-80 words).
- **attributed_to** (string, required): Who this memory is about. Use "user" for facts stated by or about the user (preferences, plans, personal facts). Use "assistant" for information provided by the assistant (recommendations, confirmations, plans created, information researched).
- **linked_memory_ids** (array of strings, optional): IDs of Existing Memories that this new memory relates to. Use the exact IDs from the Existing Memories list. Omit or pass [] if no existing memories are related.
## Rules
- Extract every piece of memorable information as a separate memory object.
- If nothing is worth extracting, return: {"memory": []}
- No duplicate IDs. Use double quotes. No trailing commas.
`;
export const AGENT_CONTEXT_SUFFIX = `
## Entity Context
The primary entity is an AI agent. Frame memories from the agent's perspective:
- For user-stated facts, frame as agent knowledge: "Agent was informed that [fact]" or "Agent learned that [fact]"
- For agent actions, use direct statements: "Agent recommended [X]" or "Agent specializes in [domain]"
- For agent configuration or instructions, capture directly: "Agent is configured to [behavior]"
The attributed_to field should still reflect the original source: "user" for facts the user stated, "assistant" for things the agent said or did.
`;
// ---------------------------------------------------------------------------
// V3 Additive Extraction Schema
// ---------------------------------------------------------------------------
export const AdditiveExtractionSchema = z.object({
memory: z.array(
z.object({
id: z.string(),
text: z.string(),
attributed_to: z.enum(["user", "assistant"]).optional(),
linked_memory_ids: z.array(z.string()).optional(),
}),
),
});
// ---------------------------------------------------------------------------
// V3 Prompt Builder — generates the user-side prompt for additive extraction
// Ported from mem0/configs/prompts.py generate_additive_extraction_prompt()
// ---------------------------------------------------------------------------
const PAST_MESSAGE_TRUNCATION_LIMIT = 300;
function truncateContent(
text: string,
limit = PAST_MESSAGE_TRUNCATION_LIMIT,
): string {
if (text.length <= limit) return text;
return text.slice(0, limit) + "...";
}
function formatConversationHistory(
messages?: Array<{ role: string; content: string }>,
): string {
if (!messages || messages.length === 0) return "";
let result = "";
for (const msg of messages) {
const role = msg.role ?? "";
const content = msg.content ?? "";
if (role && content) {
result += `${role}: ${truncateContent(content)}\n`;
}
}
return result;
}
function serializeMemories(
memories?: Array<{ id: string; text: string }>,
): string {
return JSON.stringify(memories ?? []);
}
export function generateAdditiveExtractionPrompt(options: {
existingMemories?: Array<{ id: string; text: string }>;
newMessages?: string;
lastKMessages?: Array<{ role: string; content: string }>;
customInstructions?: string;
currentDate?: string;
observationDate?: string;
}): string {
const now = new Date().toISOString().split("T")[0];
const currentDate = options.currentDate ?? now;
const observationDate = options.observationDate ?? currentDate;
const sections: string[] = [];
// Summary — empty for now; callers can extend later
sections.push("## Summary\n");
sections.push(
`## Last k Messages\n${formatConversationHistory(options.lastKMessages)}`,
);
// Recently Extracted Memories — empty for now
sections.push("## Recently Extracted Memories\n[]");
sections.push(
`## Existing Memories\n${serializeMemories(options.existingMemories)}`,
);
sections.push(`## New Messages\n${options.newMessages ?? "[]"}`);
sections.push(`## Observation Date\n${observationDate}`);
sections.push(`## Current Date\n${currentDate}`);
if (options.customInstructions) {
sections.push(`## Custom Instructions\n${options.customInstructions}`);
}
sections.push("# Output:");
return sections.join("\n\n");
}
// ---------------------------------------------------------------------------
// Legacy helpers (kept for backward compatibility)
// ---------------------------------------------------------------------------
export function parseMessages(messages: string[]): string {
return messages.join("\n");
}
mem0-ts/src/oss/src/storage/SQLiteManager.ts
+106
View File
@@ -1,4 +1,5 @@
import Database from "better-sqlite3";
import { randomUUID } from "crypto";
import { HistoryManager } from "./base";
import { ensureSQLiteDirectory } from "../utils/sqlite";
@@ -26,6 +27,16 @@ export class SQLiteManager implements HistoryManager {
is_deleted INTEGER DEFAULT 0
)
`);
this.db.exec(`
CREATE TABLE IF NOT EXISTS messages (
id TEXT PRIMARY KEY,
session_scope TEXT,
role TEXT,
content TEXT,
name TEXT,
created_at TEXT
)
`);
this.stmtInsert = this.db.prepare(
`INSERT INTO memory_history
(memory_id, previous_value, new_value, action, created_at, updated_at, is_deleted)
@@ -60,8 +71,103 @@ export class SQLiteManager implements HistoryManager {
return this.stmtSelect.all(memoryId) as any[];
}
async saveMessages(
messages: Array<{ role: string; content: string; name?: string }>,
sessionScope: string,
): Promise<void> {
if (!messages.length) return;
const insertMsg = this.db.prepare(
`INSERT INTO messages (id, session_scope, role, content, name, created_at)
VALUES (?, ?, ?, ?, ?, ?)`,
);
const evict = this.db.prepare(
`DELETE FROM messages WHERE session_scope = ? AND id NOT IN (
SELECT id FROM (
SELECT id FROM messages WHERE session_scope = ? ORDER BY created_at DESC LIMIT 10
)
)`,
);
const txn = this.db.transaction(() => {
const now = new Date().toISOString();
for (const msg of messages) {
insertMsg.run(
randomUUID(),
sessionScope,
msg.role,
msg.content,
msg.name ?? null,
now,
);
}
evict.run(sessionScope, sessionScope);
});
txn();
}
async getLastMessages(
sessionScope: string,
limit = 10,
): Promise<
Array<{ role: string; content: string; name?: string; createdAt: string }>
> {
const rows = this.db
.prepare(
`SELECT role, content, name, created_at FROM (
SELECT role, content, name, created_at
FROM messages
WHERE session_scope = ?
ORDER BY created_at DESC
LIMIT ?
) ORDER BY created_at ASC`,
)
.all(sessionScope, limit) as Array<{
role: string;
content: string;
name: string | null;
created_at: string;
}>;
return rows.map((r) => ({
role: r.role,
content: r.content,
...(r.name != null ? { name: r.name } : {}),
createdAt: r.created_at,
}));
}
async batchAddHistory(
records: Array<{
memoryId: string;
previousValue: string | null;
newValue: string | null;
action: string;
createdAt?: string;
updatedAt?: string;
isDeleted?: number;
}>,
): Promise<void> {
const txn = this.db.transaction(() => {
for (const record of records) {
this.stmtInsert.run(
record.memoryId,
record.previousValue,
record.newValue,
record.action,
record.createdAt ?? null,
record.updatedAt ?? null,
record.isDeleted ?? 0,
);
}
});
txn();
}
async reset(): Promise<void> {
this.db.exec("DROP TABLE IF EXISTS memory_history");
this.db.exec("DROP TABLE IF EXISTS messages");
this.init();
}
mem0-ts/src/oss/src/storage/base.ts
+23
View File
@@ -11,4 +11,27 @@ export interface HistoryManager {
getHistory(memoryId: string): Promise<any[]>;
reset(): Promise<void>;
close(): void;
// V3 optional methods — implementations that don't need them can omit these.
saveMessages?(
messages: Array<{ role: string; content: string; name?: string }>,
sessionScope: string,
): Promise<void>;
getLastMessages?(
sessionScope: string,
limit?: number,
): Promise<
Array<{ role: string; content: string; name?: string; createdAt: string }>
>;
batchAddHistory?(
records: Array<{
memoryId: string;
previousValue: string | null;
newValue: string | null;
action: string;
createdAt?: string;
updatedAt?: string;
isDeleted?: number;
}>,
): Promise<void>;
}
mem0-ts/src/oss/src/tests/sqlite-backward-compat.test.ts
-15
View File
@@ -38,7 +38,6 @@ describe("backward compat: ConfigManager.mergeConfig", () => {
expect(cfg.historyStore!.provider).toBe("sqlite");
expect(cfg.historyStore!.config.historyDbPath).toBe("memory.db");
expect(cfg.disableHistory).toBe(false);
expect(cfg.graphStore).toBeUndefined();
});
it("workaround: explicit historyStore still works (existing user pattern)", () => {
@@ -104,20 +103,6 @@ describe("backward compat: ConfigManager.mergeConfig", () => {
expect(cfg.vectorStore.config.dimension).toBe(768);
});
it("graphStore config passes through unchanged", () => {
const cfg = ConfigManager.mergeConfig({
graphStore: {
provider: "neo4j",
config: {
url: "neo4j://custom:7687",
username: "admin",
password: "pass",
},
},
});
expect(cfg.graphStore!.config.url).toBe("neo4j://custom:7687");
});
it("customInstructions passes through unchanged", () => {
const cfg = ConfigManager.mergeConfig({
customInstructions: "You are a helpful assistant",
mem0-ts/src/oss/src/types/index.ts
-32
View File
@@ -50,19 +50,6 @@ export interface LLMConfig {
modelProperties?: Record<string, any>;
}
export interface Neo4jConfig {
url: string;
username: string;
password: string;
}
export interface GraphStoreConfig {
provider: string;
config: Neo4jConfig;
llm?: LLMConfig;
customInstructions?: string;
}
export interface MemoryConfig {
version?: string;
embedder: {
@@ -81,7 +68,6 @@ export interface MemoryConfig {
disableHistory?: boolean;
historyDbPath?: string;
customInstructions?: string;
graphStore?: GraphStoreConfig;
}
export interface MemoryItem {
@@ -103,7 +89,6 @@ export interface SearchFilters {
export interface SearchResult {
results: MemoryItem[];
relations?: any[];
}
export interface VectorStoreResult {
@@ -148,23 +133,6 @@ export const MemoryConfigSchema = z.object({
}),
historyDbPath: z.string().optional(),
customInstructions: z.string().optional(),
graphStore: z
.object({
provider: z.string(),
config: z.object({
url: z.string(),
username: z.string(),
password: z.string(),
}),
llm: z
.object({
provider: z.string(),
config: z.record(z.string(), z.any()),
})
.optional(),
customInstructions: z.string().optional(),
})
.optional(),
historyStore: z
.object({
provider: z.string(),
mem0-ts/src/oss/src/utils/entity_extraction.ts
+718
View File
@@ -0,0 +1,718 @@
/**
* Entity extraction from text using NLP and regex heuristics.
*
* Extracts four types of entities from text:
* - PROPER: Capitalized multi-word sequences (person names, places, brands)
* - QUOTED: Text in single or double quotes (titles, specific terms)
* - COMPOUND: Multi-word noun phrases with specific modifiers (e.g., "machine learning")
* - NOUN: Single nouns from circumstantial compound patterns
*
* Uses the `compromise` npm package for NLP-based extraction when available.
* Falls back to regex-only extraction if `compromise` is not installed.
*/
// ---------------------------------------------------------------------------
// Filter lists (ported from Python)
// ---------------------------------------------------------------------------
/** Words that are too generic to be useful as entity heads. */
const GENERIC_HEADS: Set<string> = new Set([
"thing",
"stuff",
"way",
"time",
"experience",
"situation",
"case",
"fact",
"matter",
"issue",
"idea",
"thought",
"feeling",
"place",
"area",
"part",
"kind",
"type",
"sort",
"lot",
"bit",
"day",
"year",
"week",
"month",
"moment",
"instance",
"example",
"technique",
"method",
"approach",
"process",
"step",
"tool",
"result",
"outcome",
"goal",
"task",
"item",
"topic",
"scale",
"size",
"level",
"degree",
"amount",
"number",
"style",
"look",
"color",
"colour",
"shape",
"form",
"piece",
"section",
"side",
"end",
"edge",
"surface",
"point",
]);
/** Adjectives too vague to make a compound entity specific. */
const NON_SPECIFIC_ADJ: Set<string> = new Set([
"many",
"few",
"several",
"some",
"any",
"all",
"most",
"more",
"less",
"much",
"little",
"enough",
"various",
"numerous",
"multiple",
"countless",
"great",
"good",
"bad",
"nice",
"terrible",
"awful",
"awesome",
"amazing",
"wonderful",
"horrible",
"excellent",
"poor",
"best",
"worst",
"fine",
"okay",
"new",
"old",
"recent",
"past",
"future",
"current",
"previous",
"next",
"last",
"first",
"latest",
"early",
"late",
"former",
"modern",
"ancient",
"big",
"small",
"large",
"tiny",
"huge",
"enormous",
"long",
"short",
"tall",
"high",
"low",
"wide",
"narrow",
"thick",
"thin",
"deep",
"shallow",
"similar",
"different",
"same",
"other",
"another",
"such",
"certain",
"important",
"main",
"major",
"minor",
"key",
"primary",
"real",
"actual",
"true",
"whole",
"entire",
"full",
"complete",
"total",
"basic",
"simple",
"interesting",
"boring",
"exciting",
"special",
"particular",
"general",
"common",
"unique",
"rare",
"typical",
"usual",
"normal",
"regular",
"possible",
"likely",
"potential",
"available",
"necessary",
"only",
"solo",
"individual",
"team",
"group",
"joint",
"collaborative",
"final",
"initial",
"side",
]);
/** Generic tail words to strip from compound entities. */
const GENERIC_ENDINGS: Set<string> = new Set([
"work",
"works",
"job",
"jobs",
"task",
"tasks",
"stuff",
"things",
"thing",
"info",
"information",
"details",
"data",
"content",
"material",
"materials",
"activities",
"activity",
"efforts",
"effort",
"options",
"option",
"choices",
"choice",
"results",
"result",
"output",
"outputs",
"products",
"product",
"items",
"item",
]);
/** Capitalized single words that are too generic to be proper nouns. */
const GENERIC_CAPS: Set<string> = new Set([
"works",
"items",
"things",
"stuff",
"resources",
"options",
"tips",
"ideas",
"steps",
"ways",
"methods",
"tools",
"features",
"benefits",
"examples",
"details",
"notes",
"instructions",
"guidelines",
"recommendations",
"suggestions",
"overview",
"summary",
"conclusion",
"introduction",
"pros",
"cons",
"advantages",
"disadvantages",
]);
/** Markdown/formatting markers to skip during extraction. */
const FORMATTING_MARKERS: Set<string> = new Set([
"*",
"-",
"+",
"\u2022",
"\u2013",
"\u2014",
"#",
"##",
"###",
"**",
"__",
]);
// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------
export interface ExtractedEntity {
type: "PROPER" | "QUOTED" | "COMPOUND" | "NOUN";
text: string;
}
// ---------------------------------------------------------------------------
// compromise dynamic import
// ---------------------------------------------------------------------------
let nlp: any;
try {
nlp = require("compromise");
} catch {
// compromise not installed -- use regex-only fallback
}
// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------
/** Check for formatting artifacts that indicate non-entity text. */
function hasArtifacts(txt: string): boolean {
if (txt.includes("**") || txt.includes("__") || txt.includes(":*")) {
return true;
}
if (/\s\*\s|\s\*$|^\*\s/.test(txt)) {
return true;
}
if (txt.includes(" ") || txt.includes("\n") || txt.includes("\t")) {
return true;
}
if (txt.length > 100) {
return true;
}
if (/^[\u2022\-+\u2013\u2014]/.test(txt)) {
return true;
}
return false;
}
/** Strip generic trailing words from a phrase's word list. */
function stripGenericEnding(words: string[]): string[] {
if (words.length <= 1) {
return words;
}
const last = words[words.length - 1].toLowerCase();
if (GENERIC_ENDINGS.has(last) && words.length > 2) {
return words.slice(0, -1);
}
return words;
}
/**
* Determine if a token position is at the start of a sentence.
* Simple heuristic: index 0, or preceded by sentence-ending punctuation
* or formatting markers.
*/
function isSentenceStart(
tokens: string[],
idx: number,
rawText: string,
): boolean {
if (idx === 0) {
return true;
}
const prev = tokens[idx - 1];
if (/[.!?:]$/.test(prev)) {
return true;
}
if (FORMATTING_MARKERS.has(prev)) {
return true;
}
// Check for newline before this token in the raw text
const tokenStart = rawText.indexOf(tokens[idx]);
if (tokenStart > 0 && rawText.charAt(tokenStart - 1) === "\n") {
return true;
}
return false;
}
// ---------------------------------------------------------------------------
// Extraction strategies
// ---------------------------------------------------------------------------
/** Extract quoted entities via regex. */
function extractQuoted(text: string): ExtractedEntity[] {
const entities: ExtractedEntity[] = [];
// Double-quoted
const doubleQuoteRe = /"([^"]+)"/g;
let match: RegExpExecArray | null;
while ((match = doubleQuoteRe.exec(text)) !== null) {
const inner = match[1].trim();
if (inner.length > 2) {
entities.push({ type: "QUOTED", text: inner });
}
}
// Single-quoted (with boundary constraints to avoid apostrophes)
const singleQuoteRe = /(?:^|[\s([{,;])'([^']+)'(?=[\s.,;:!?)\]]|$)/g;
while ((match = singleQuoteRe.exec(text)) !== null) {
const inner = match[1].trim();
if (inner.length > 2) {
entities.push({ type: "QUOTED", text: inner });
}
}
return entities;
}
/**
* Extract proper noun sequences using capitalization heuristics.
* Finds sequences of capitalized words that are not at sentence starts.
*/
function extractProper(text: string): ExtractedEntity[] {
const entities: ExtractedEntity[] = [];
// Tokenize on whitespace, preserving order
const tokens = text.split(/\s+/).filter(Boolean);
const functionWords = new Set([
"'s",
"of",
"the",
"in",
"and",
"for",
"at",
"is",
]);
let i = 0;
while (i < tokens.length) {
const tok = tokens[i];
// Skip formatting markers
if (FORMATTING_MARKERS.has(tok)) {
i++;
continue;
}
const isLabel = i + 1 < tokens.length && tokens[i + 1] === ":";
const isCap = tok.length > 0 && tok.charAt(0) === tok.charAt(0).toUpperCase() && /[A-Z]/.test(tok.charAt(0));
if (isCap && !isLabel) {
const seq: Array<{ token: string; idx: number }> = [
{ token: tok, idx: i },
];
let j = i + 1;
while (j < tokens.length) {
const t = tokens[j];
const tIsCap =
t.length > 0 &&
t.charAt(0) === t.charAt(0).toUpperCase() &&
/[A-Z]/.test(t.charAt(0));
if (tIsCap || functionWords.has(t.toLowerCase())) {
seq.push({ token: t, idx: j });
j++;
} else {
break;
}
}
// Strip trailing function words
while (
seq.length > 0 &&
functionWords.has(seq[seq.length - 1].token.toLowerCase())
) {
seq.pop();
}
if (seq.length > 0) {
// Check for at least one mid-sentence capitalized word
const hasMidCap = seq.some(({ token, idx: tokenIdx }) => {
const isCapWord =
/[A-Z]/.test(token.charAt(0)) &&
!functionWords.has(token.toLowerCase());
return isCapWord && !isSentenceStart(tokens, tokenIdx, text);
});
if (hasMidCap) {
const phrase = seq.map((s) => s.token).join(" ");
if (phrase.length > 2) {
entities.push({ type: "PROPER", text: phrase });
}
}
}
i = j;
} else {
i++;
}
}
return entities;
}
/**
* Extract compound noun phrases using the `compromise` NLP library.
* Returns COMPOUND and NOUN entities derived from noun chunks.
*/
function extractCompoundsWithNlp(text: string): ExtractedEntity[] {
if (!nlp) {
return [];
}
const entities: ExtractedEntity[] = [];
const doc = nlp(text);
const nouns = doc.nouns().out("array") as string[];
for (const nounPhrase of nouns) {
const trimmed = nounPhrase.trim();
if (!trimmed || trimmed.length <= 3) {
continue;
}
const words = trimmed.split(/\s+/);
if (words.length < 2) {
continue;
}
// Filter out phrases where the head is generic
const head = words[words.length - 1].toLowerCase();
if (GENERIC_HEADS.has(head)) {
// Check if there's a specific modifier
const hasSpecificMod = words.some(
(w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()) && w !== words[words.length - 1],
);
if (!hasSpecificMod) {
continue;
}
}
// Filter non-specific adjectives from the beginning
const filtered = words.filter((w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()));
const cleaned = stripGenericEnding(filtered);
if (cleaned.length >= 2) {
const phrase = cleaned.join(" ");
if (phrase.length > 3) {
entities.push({ type: "COMPOUND", text: phrase });
}
}
}
return entities;
}
/**
* Regex-only fallback for compound extraction when compromise is not available.
* Finds multi-word capitalized sequences and common compound patterns.
*/
function extractCompoundsRegex(text: string): ExtractedEntity[] {
const entities: ExtractedEntity[] = [];
// Multi-word sequences with at least one non-trivial word
// Match sequences like "machine learning", "New York", "data science"
const compoundRe =
/\b([A-Z][a-z]+(?:\s+(?:of|and|the|for|in)\s+)?[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b/g;
let match: RegExpExecArray | null;
while ((match = compoundRe.exec(text)) !== null) {
const phrase = match[1].trim();
if (phrase.length > 3 && phrase.includes(" ")) {
const words = phrase.split(/\s+/);
const head = words[words.length - 1].toLowerCase();
if (!GENERIC_HEADS.has(head)) {
const filtered = words.filter(
(w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()),
);
const cleaned = stripGenericEnding(filtered);
if (cleaned.length >= 2) {
entities.push({ type: "COMPOUND", text: cleaned.join(" ") });
}
}
}
}
// Also try lowercase compound patterns (e.g., "machine learning", "deep learning")
const lowerCompoundRe =
/\b([a-z]+(?:\s+[a-z]+){1,3})\b/g;
while ((match = lowerCompoundRe.exec(text)) !== null) {
const phrase = match[1].trim();
const words = phrase.split(/\s+/);
if (
words.length >= 2 &&
words.length <= 4 &&
phrase.length > 5
) {
const head = words[words.length - 1].toLowerCase();
const allGeneric = words.every(
(w) =>
NON_SPECIFIC_ADJ.has(w.toLowerCase()) ||
GENERIC_HEADS.has(w.toLowerCase()),
);
if (!allGeneric && !GENERIC_HEADS.has(head)) {
// Only include if it looks like a meaningful compound
const hasContentWord = words.some(
(w) =>
!NON_SPECIFIC_ADJ.has(w.toLowerCase()) &&
!GENERIC_HEADS.has(w.toLowerCase()) &&
w.length > 2,
);
if (hasContentWord) {
const filtered = words.filter(
(w) => !NON_SPECIFIC_ADJ.has(w.toLowerCase()),
);
const cleaned = stripGenericEnding(filtered);
if (cleaned.length >= 2) {
entities.push({ type: "COMPOUND", text: cleaned.join(" ") });
}
}
}
}
}
return entities;
}
// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
/**
* Extract named entities, quoted text, and noun compounds from text.
*
* Uses `compromise` for NLP-based noun phrase extraction when available,
* falling back to regex-only heuristics otherwise.
*
* Entity types (in priority order for deduplication):
* PROPER - Capitalized multi-word sequences not at sentence start
* COMPOUND - Multi-word noun phrases with specific modifiers
* QUOTED - Text in single or double quotes (min 3 chars)
* NOUN - Single nouns from circumstantial patterns
*
* @param text - Input text to extract entities from.
* @returns Deduplicated list of extracted entities.
*/
export function extractEntities(text: string): ExtractedEntity[] {
const raw: ExtractedEntity[] = [];
// 1. QUOTED entities (always regex)
raw.push(...extractQuoted(text));
// 2. PROPER entities (capitalization heuristics)
raw.push(...extractProper(text));
// 3. COMPOUND entities (NLP or regex fallback)
if (nlp) {
raw.push(...extractCompoundsWithNlp(text));
} else {
raw.push(...extractCompoundsRegex(text));
}
// === DEDUPLICATION & CLEANUP ===
// First pass: deduplicate by lowercase text
const seen = new Set<string>();
const deduped: ExtractedEntity[] = [];
for (const entity of raw) {
const key = entity.text.toLowerCase().trim();
if (key.length > 2 && !seen.has(key)) {
seen.add(key);
deduped.push(entity);
}
}
// Clean up formatting artifacts
const cleaned: ExtractedEntity[] = [];
for (const entity of deduped) {
let txt = entity.text.trim();
// Strip leading/trailing asterisks
txt = txt.replace(/^\*+\s*|\s*\*+$/g, "");
// Strip trailing colons
txt = txt.replace(/\s*:+$/, "");
// Strip leading numbered list markers
txt = txt.replace(/^\d+\s*\.\s*/, "");
if (!txt || txt.length <= 2 || hasArtifacts(txt)) {
continue;
}
// Filter generic single-word PROPER nouns
if (
entity.type === "PROPER" &&
!txt.includes(" ") &&
GENERIC_CAPS.has(txt.toLowerCase())
) {
continue;
}
cleaned.push({ type: entity.type, text: txt });
}
// Keep best type per entity (PROPER > COMPOUND > QUOTED > NOUN)
const typePriority: Record<string, number> = {
PROPER: 0,
COMPOUND: 1,
QUOTED: 2,
NOUN: 3,
};
const best = new Map<string, ExtractedEntity>();
for (const entity of cleaned) {
const key = entity.text.toLowerCase();
const existing = best.get(key);
if (
!existing ||
(typePriority[entity.type] ?? 99) < (typePriority[existing.type] ?? 99)
) {
best.set(key, entity);
}
}
const bestEntities = Array.from(best.values());
// Remove entities that are substrings of longer entities
const allLower = bestEntities.map((e) => e.text.toLowerCase());
return bestEntities.filter(
(entity) =>
!allLower.some(
(other) =>
entity.text.toLowerCase() !== other &&
other.includes(entity.text.toLowerCase()),
),
);
}
/**
* Extract entities from multiple texts.
*
* @param texts - List of input texts to extract entities from.
* @returns List of entity lists, one per input text.
*/
export function extractEntitiesBatch(texts: string[]): ExtractedEntity[][] {
return texts.map(extractEntities);
}
mem0-ts/src/oss/src/utils/lemmatization.ts
+282
View File
@@ -0,0 +1,282 @@
/**
* BM25 lemmatization for consistent keyword matching.
*
* Uses the `natural` npm package for Porter stemming when available.
* Falls back to simple lowercasing + stop word removal if `natural`
* is not installed.
*
* Also includes original -ing forms alongside stems to handle cases
* where stemming produces inconsistent results (e.g., "meeting" as
* noun vs verb -> different stems).
*/
/** Standard English stop words (based on NLTK stop word list). */
const STOP_WORDS: Set<string> = new Set([
"a",
"about",
"above",
"after",
"again",
"against",
"all",
"am",
"an",
"and",
"any",
"are",
"aren't",
"as",
"at",
"be",
"because",
"been",
"before",
"being",
"below",
"between",
"both",
"but",
"by",
"can",
"can't",
"cannot",
"could",
"couldn't",
"did",
"didn't",
"do",
"does",
"doesn't",
"doing",
"don't",
"down",
"during",
"each",
"few",
"for",
"from",
"further",
"get",
"got",
"had",
"hadn't",
"has",
"hasn't",
"have",
"haven't",
"having",
"he",
"her",
"here",
"hers",
"herself",
"him",
"himself",
"his",
"how",
"i",
"if",
"in",
"into",
"is",
"isn't",
"it",
"it's",
"its",
"itself",
"just",
"let's",
"me",
"might",
"more",
"most",
"mustn't",
"must",
"my",
"myself",
"no",
"nor",
"not",
"of",
"off",
"on",
"once",
"only",
"or",
"other",
"ought",
"our",
"ours",
"ourselves",
"out",
"over",
"own",
"same",
"shall",
"shan't",
"she",
"should",
"shouldn't",
"so",
"some",
"such",
"than",
"that",
"the",
"their",
"theirs",
"them",
"themselves",
"then",
"there",
"these",
"they",
"this",
"those",
"through",
"to",
"too",
"under",
"until",
"up",
"very",
"was",
"wasn't",
"we",
"were",
"weren't",
"what",
"when",
"where",
"which",
"while",
"who",
"whom",
"why",
"will",
"with",
"won't",
"would",
"wouldn't",
"you",
"your",
"yours",
"yourself",
"yourselves",
]);
/**
* Attempt to load the Porter stemmer from the `natural` package.
* Returns null if the package is not installed.
*/
let _porterStemmer: { stem: (word: string) => string } | null | undefined;
function getPorterStemmer(): { stem: (word: string) => string } | null {
if (_porterStemmer !== undefined) {
return _porterStemmer;
}
try {
// eslint-disable-next-line @typescript-eslint/no-var-requires
const natural = require("natural");
_porterStemmer = natural.PorterStemmer;
return _porterStemmer!;
} catch {
_porterStemmer = null;
return null;
}
}
/**
* Simple built-in Porter-like stemmer for common English suffixes.
* Used only when the `natural` package is not available.
*/
function simpleStem(word: string): string {
if (word.length <= 3) {
return word;
}
// Step-like suffix stripping (simplified Porter rules)
let w = word;
if (w.endsWith("ies") && w.length > 4) {
w = w.slice(0, -3) + "i";
} else if (w.endsWith("sses")) {
w = w.slice(0, -2);
} else if (w.endsWith("ness")) {
w = w.slice(0, -4);
} else if (w.endsWith("ment") && w.length > 5) {
w = w.slice(0, -4);
} else if (w.endsWith("ation") && w.length > 6) {
w = w.slice(0, -5) + "e";
} else if (w.endsWith("ting") && w.length > 5) {
w = w.slice(0, -3);
} else if (w.endsWith("ing") && w.length > 5) {
w = w.slice(0, -3);
} else if (w.endsWith("ed") && w.length > 4) {
w = w.slice(0, -2);
} else if (w.endsWith("ly") && w.length > 4) {
w = w.slice(0, -2);
} else if (w.endsWith("er") && w.length > 4) {
w = w.slice(0, -2);
} else if (w.endsWith("est") && w.length > 4) {
w = w.slice(0, -3);
} else if (w.endsWith("s") && !w.endsWith("ss") && w.length > 3) {
w = w.slice(0, -1);
}
return w;
}
/**
* Lemmatize (stem) text for BM25 matching.
*
* Processing steps:
* 1. Lowercase the text.
* 2. Tokenize into words (alphanumeric sequences).
* 3. Remove stop words.
* 4. Apply Porter stemming to each word.
* 5. For words ending in -ing, keep both the stemmed and original form.
* 6. Return space-joined result.
*
* Falls back to simple suffix stripping if `natural` is not installed.
*
* @param text - Input text to lemmatize.
* @returns Space-joined lemmatized/stemmed tokens.
*/
export function lemmatizeForBm25(text: string): string {
const lower = text.toLowerCase();
const words = lower.match(/[a-z0-9]+/g);
if (!words) {
return text.toLowerCase();
}
const stemmer = getPorterStemmer();
const stemFn = stemmer
? (w: string) => stemmer.stem(w).toLowerCase()
: simpleStem;
const tokens: string[] = [];
for (const word of words) {
if (STOP_WORDS.has(word)) {
continue;
}
const stemmed = stemFn(word);
if (stemmed && /^[a-z0-9]+$/.test(stemmed)) {
tokens.push(stemmed);
}
// Also add original if it ends in -ing and differs from stem.
// This handles noun/verb ambiguity (meeting/meet, attending/attend).
if (
word.endsWith("ing") &&
word !== stemmed &&
/^[a-z0-9]+$/.test(word)
) {
tokens.push(word);
}
}
return tokens.join(" ");
}
mem0-ts/src/oss/src/utils/scoring.ts
+148
View File
@@ -0,0 +1,148 @@
/**
* Scoring utilities for hybrid retrieval.
*
* Provides:
* - BM25 normalization: Sigmoid normalization of raw BM25 scores to [0, 1].
* - BM25 parameter selection: Query-length-adaptive sigmoid parameters.
* - Additive scoring: Combined scoring with semantic + BM25 + entity boost.
*/
export const ENTITY_BOOST_WEIGHT = 0.5;
/**
* Get BM25 sigmoid parameters based on query length.
*
* Longer queries tend to have higher raw BM25 scores, so we adjust
* the sigmoid midpoint and steepness accordingly.
*
* @param query - The original query string.
* @param lemmatized - Optional pre-lemmatized query string. If not provided,
* the term count is estimated from the raw query.
* @returns A tuple of [midpoint, steepness] for sigmoid normalization.
*/
export function getBm25Params(
query: string,
lemmatized?: string,
): [number, number] {
const text = lemmatized ?? query;
const numTerms = text.trim().split(/\s+/).filter(Boolean).length || 1;
if (numTerms <= 3) {
return [5.0, 0.7];
} else if (numTerms <= 6) {
return [7.0, 0.6];
} else if (numTerms <= 9) {
return [9.0, 0.5];
} else if (numTerms <= 15) {
return [10.0, 0.5];
} else {
return [12.0, 0.5];
}
}
/**
* Normalize a raw BM25 score to [0, 1] using logistic sigmoid.
*
* @param rawScore - Raw BM25 score (unbounded, typically 0-20+).
* @param midpoint - Score at which sigmoid outputs 0.5.
* @param steepness - Controls how quickly sigmoid transitions.
* @returns Normalized score in range [0, 1].
*/
export function normalizeBm25(
rawScore: number,
midpoint: number,
steepness: number,
): number {
return 1.0 / (1.0 + Math.exp(-steepness * (rawScore - midpoint)));
}
export interface ScoredResult {
id: string;
score: number;
scoreBreakdown: {
semantic: number;
bm25: number;
entityBoost: number;
};
payload: Record<string, any>;
}
/**
* Score candidates additively and return top-k results.
*
* For each candidate:
* combined = (semantic + bm25 + entity_boost) / max_possible
*
* Threshold gates the semantic score BEFORE combining -- candidates
* below the threshold are excluded even if BM25/entity would boost them.
*
* The divisor adapts based on which signals are active:
* - Semantic only: max_possible = 1.0
* - Semantic + BM25: max_possible = 2.0
* - Semantic + BM25 + entity: max_possible = 2.5
* - Semantic + entity (no BM25): max_possible = 1.5
*
* @param semanticResults - Candidate results with id, score, and payload.
* @param bm25Scores - Map of memory ID to normalized BM25 score.
* @param entityBoosts - Map of memory ID to entity boost score.
* @param threshold - Minimum semantic score to include a candidate.
* @param topK - Maximum number of results to return.
* @returns Sorted list of scored results, highest score first.
*/
export function scoreAndRank(
semanticResults: Array<{
id: string;
score: number;
payload: Record<string, any>;
}>,
bm25Scores: Record<string, number>,
entityBoosts: Record<string, number>,
threshold: number,
topK: number,
): ScoredResult[] {
const hasBm25 = Object.keys(bm25Scores).length > 0;
const hasEntity = Object.keys(entityBoosts).length > 0;
let maxPossible = 1.0;
if (hasBm25) {
maxPossible += 1.0;
}
if (hasEntity) {
maxPossible += ENTITY_BOOST_WEIGHT;
}
const scored: ScoredResult[] = [];
for (const result of semanticResults) {
const memId = result.id;
if (memId == null) {
continue;
}
const semanticScore = result.score ?? 0.0;
if (semanticScore < threshold) {
continue;
}
const memIdStr = String(memId);
const bm25Score = bm25Scores[memIdStr] ?? 0.0;
const entityBoost = entityBoosts[memIdStr] ?? 0.0;
const rawCombined = semanticScore + bm25Score + entityBoost;
const combined = Math.min(rawCombined / maxPossible, 1.0);
scored.push({
id: memIdStr,
score: combined,
scoreBreakdown: {
semantic: semanticScore,
bm25: bm25Score,
entityBoost: entityBoost,
},
payload: result.payload,
});
}
scored.sort((a, b) => b.score - a.score);
return scored.slice(0, topK);
}
mem0-ts/src/oss/src/vector_stores/azure_ai_search.ts
+39
View File
@@ -326,6 +326,45 @@ export class AzureAISearch implements VectorStore {
}
}
/**
* Keyword search using Azure AI Search native full-text (BM25) capabilities
*/
async keywordSearch(
query: string,
topK: number = 5,
filters?: SearchFilters,
): Promise<VectorStoreResult[] | null> {
try {
const filterExpression = filters
? this.buildFilterExpression(filters)
: undefined;
const searchResults = await this.searchClient.search(query, {
filter: filterExpression,
top: topK,
searchFields: ["payload"],
});
const results: VectorStoreResult[] = [];
for await (const result of searchResults.results) {
const payloadStr = result.document.payload as string;
const payload = JSON.parse(this.extractJson(payloadStr));
results.push({
id: result.document.id as string,
score: result.score,
payload,
});
}
return results;
} catch (error) {
console.error("Error during keyword search:", error);
return null;
}
}
/**
* Search for similar vectors
*/
mem0-ts/src/oss/src/vector_stores/base.ts
+5
View File
@@ -11,6 +11,11 @@ export interface VectorStore {
topK?: number,
filters?: SearchFilters,
): Promise<VectorStoreResult[]>;
keywordSearch?(
query: string,
topK?: number,
filters?: SearchFilters,
): Promise<VectorStoreResult[] | null>;
get(vectorId: string): Promise<VectorStoreResult | null>;
update(
vectorId: string,
mem0-ts/src/oss/src/vector_stores/langchain.ts
+4
View File
@@ -99,6 +99,10 @@ export class LangchainVectorStore implements VectorStore {
}
}
async keywordSearch(): Promise<null> {
return null;
}
async search(
query: number[],
topK: number = 5,
mem0-ts/src/oss/src/vector_stores/memory.ts
+105
View File
@@ -98,6 +98,111 @@ export class MemoryVectorStore implements VectorStore {
insertMany(vectors, ids, payloads);
}
private tokenize(text: string): string[] {
return text.toLowerCase().split(/\s+/).filter(Boolean);
}
async keywordSearch(
query: string,
topK: number = 10,
filters?: SearchFilters,
): Promise<VectorStoreResult[] | null> {
try {
const rows = this.db.prepare(`SELECT * FROM vectors`).all() as any[];
// Collect documents that pass the filter
const candidates: {
id: string;
payload: Record<string, any>;
tokens: string[];
}[] = [];
for (const row of rows) {
const payload = JSON.parse(row.payload);
const memoryVector: MemoryVector = {
id: row.id,
vector: Array.from(
new Float32Array(
row.vector.buffer,
row.vector.byteOffset,
row.vector.byteLength / 4,
),
),
payload,
};
if (this.filterVector(memoryVector, filters)) {
const text = payload.text_lemmatized || payload.data || "";
candidates.push({ id: row.id, payload, tokens: this.tokenize(text) });
}
}
if (candidates.length === 0) {
return [];
}
const tokenizedQuery = this.tokenize(query);
if (tokenizedQuery.length === 0) {
return [];
}
// Compute BM25 scores inline
const k1 = 1.5;
const b = 0.75;
const N = candidates.length;
const avgDocLength =
candidates.reduce((sum, c) => sum + c.tokens.length, 0) / N;
// Compute document frequency for query terms
const docFreq = new Map<string, number>();
for (const term of tokenizedQuery) {
if (!docFreq.has(term)) {
let count = 0;
for (const c of candidates) {
if (c.tokens.includes(term)) count++;
}
docFreq.set(term, count);
}
}
// Compute IDF for query terms
const idf = new Map<string, number>();
for (const [term, freq] of docFreq) {
idf.set(term, Math.log((N - freq + 0.5) / (freq + 0.5) + 1));
}
// Score each candidate
const scored = candidates.map((candidate) => {
let score = 0;
const docLength = candidate.tokens.length;
for (const term of tokenizedQuery) {
const tf = candidate.tokens.filter((t) => t === term).length;
const termIdf = idf.get(term) || 0;
score +=
(termIdf * tf * (k1 + 1)) /
(tf + k1 * (1 - b + (b * docLength) / avgDocLength));
}
return { ...candidate, score };
});
// Filter out zero-score documents and sort descending
const results = scored
.filter((s) => s.score > 0)
.sort((a, b) => b.score - a.score)
.slice(0, topK)
.map((s) => ({
id: s.id,
payload: s.payload,
score: s.score,
}));
return results;
} catch (error) {
console.error("Error during keyword search:", error);
return null;
}
}
async search(
query: number[],
topK: number = 10,
mem0-ts/src/oss/src/vector_stores/pgvector.ts
+45
View File
@@ -162,6 +162,51 @@ export class PGVector implements VectorStore {
);
}
async keywordSearch(
query: string,
topK: number = 5,
filters?: SearchFilters,
): Promise<VectorStoreResult[] | null> {
try {
const filterConditions: string[] = [];
const filterValues: any[] = [query, topK];
let filterIndex = 3;
if (filters) {
for (const [key, value] of Object.entries(filters)) {
filterConditions.push(`payload->>'${key}' = $${filterIndex}`);
filterValues.push(value);
filterIndex++;
}
}
const filterClause =
filterConditions.length > 0
? "AND " + filterConditions.join(" AND ")
: "";
const searchQuery = `
SELECT id, ts_rank_cd(to_tsvector('simple', payload->>'text_lemmatized'), plainto_tsquery('simple', $1)) AS score, payload
FROM ${this.collectionName}
WHERE to_tsvector('simple', payload->>'text_lemmatized') @@ plainto_tsquery('simple', $1)
${filterClause}
ORDER BY score DESC
LIMIT $2
`;
const result = await this.client.query(searchQuery, filterValues);
return result.rows.map((row) => ({
id: row.id,
payload: row.payload,
score: row.score,
}));
} catch (error) {
console.error("Error during keyword search:", error);
return null;
}
}
async search(
query: number[],
topK: number = 5,
mem0-ts/src/oss/src/vector_stores/qdrant.ts
+4
View File
@@ -137,6 +137,10 @@ export class Qdrant implements VectorStore {
});
}
async keywordSearch(): Promise<null> {
return null;
}
async search(
query: number[],
topK: number = 5,
mem0-ts/src/oss/src/vector_stores/redis.ts
+4
View File
@@ -357,6 +357,10 @@ export class RedisDB implements VectorStore {
}
}
async keywordSearch(): Promise<null> {
return null;
}
async search(
query: number[],
topK: number = 5,
mem0-ts/src/oss/src/vector_stores/supabase.ts
+4
View File
@@ -229,6 +229,10 @@ See the SQL migration instructions in the code comments.`,
}
}
async keywordSearch(): Promise<null> {
return null;
}
async search(
query: number[],
topK: number = 5,
mem0-ts/src/oss/src/vector_stores/vectorize.ts
+4
View File
@@ -74,6 +74,10 @@ export class VectorizeDB implements VectorStore {
}
}
async keywordSearch(): Promise<null> {
return null;
}
async search(
query: number[],
topK: number = 5,
mem0-ts/src/oss/tests/config-manager.test.ts
-72
View File
@@ -349,78 +349,6 @@ describe("ConfigManager", () => {
});
});
// ─────────────────────────────────────────────────────────────────────────
// Graph store LLM config propagation (issue #3425)
// ─────────────────────────────────────────────────────────────────────────
describe("mergeConfig - graph store LLM config (issue #3425)", () => {
const baseEmbedder = {
provider: "openai",
config: { apiKey: "test-key" },
};
const baseVectorStore = {
provider: "memory",
config: { collectionName: "test" },
};
const graphStoreNeo4j = {
provider: "neo4j",
config: {
url: "neo4j://localhost:7687",
username: "neo4j",
password: "password",
},
};
it("should NOT have a default graphStore.llm — root llm should be the fallback", () => {
const config = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514" },
},
graphStore: graphStoreNeo4j,
});
// graphStore should NOT have its own llm after merge
expect(config.graphStore?.llm).toBeUndefined();
// root llm should be anthropic
expect(config.llm.provider).toBe("anthropic");
expect(config.llm.config.model).toBe("claude-sonnet-4-20250514");
});
it("should preserve explicit graphStore.llm when user provides it", () => {
const config = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514" },
},
graphStore: {
...graphStoreNeo4j,
llm: { provider: "openai", config: { model: "gpt-4o" } },
},
});
// graphStore should have its own llm
expect(config.graphStore?.llm?.provider).toBe("openai");
expect(config.graphStore?.llm?.config).toEqual({ model: "gpt-4o" });
// root llm should still be anthropic
expect(config.llm.provider).toBe("anthropic");
});
it("should not have graphStore.llm when user does not provide one", () => {
const config = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: { provider: "openai", config: { model: "gpt-4o" } },
});
// Default graphStore should not have llm
expect(config.graphStore?.llm).toBeUndefined();
});
});
// ─────────────────────────────────────────────────────────────────────────
// Memory class – LM Studio end-to-end flow (mocked factories)
// ─────────────────────────────────────────────────────────────────────────
mem0-ts/src/oss/tests/graph-memory-parsing.test.ts
-681
View File
@@ -1,681 +0,0 @@
/**
* Regression tests for graph_memory.ts response parsing (issue #4248).
*
* Exercises the three json_object call sites in MemoryGraph with a mocked LLM:
* 1. _retrieveNodesFromData → entity extraction
* 2. _establishNodesRelationsFromData → relation extraction
* 3. _getDeleteEntitiesFromSearchOutput → deletion identification
*
* Covers: malformed LLM responses, missing fields, bad JSON in toolCalls,
* string-only responses, empty tool calls, and prompt construction.
*
* See: https://github.com/mem0ai/mem0/issues/4248
*/
import { MemoryGraph } from "../src/memory/graph_memory";
import {
EXTRACT_RELATIONS_PROMPT,
getDeleteMessages,
} from "../src/graphs/utils";
// ---------------------------------------------------------------------------
// Mocks – we replace heavy dependencies so tests run without Neo4j / OpenAI
// ---------------------------------------------------------------------------
// Mock neo4j-driver: provides a fake Driver with a no-op session
jest.mock("neo4j-driver", () => ({
__esModule: true,
default: {
driver: jest.fn(() => ({
session: () => ({
run: jest.fn().mockResolvedValue({ records: [] }),
close: jest.fn(),
}),
})),
auth: { basic: jest.fn() },
},
}));
// Mock factory so constructor doesn't try to instantiate real LLMs / embedders
const mockGenerateResponse = jest.fn();
const mockGenerateChat = jest.fn();
const mockEmbed = jest.fn().mockResolvedValue([0.1, 0.2, 0.3]);
jest.mock("../src/utils/factory", () => ({
LLMFactory: {
create: jest.fn(() => ({
generateResponse: mockGenerateResponse,
generateChat: mockGenerateChat,
})),
},
EmbedderFactory: {
create: jest.fn(() => ({
embed: mockEmbed,
})),
},
}));
// Minimal config that satisfies the MemoryGraph constructor
function makeConfig(overrides: Record<string, any> = {}) {
return {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
...overrides,
},
embedder: { provider: "openai", config: {} },
llm: { provider: "openai", config: {} },
} as any;
}
// Helper to access private methods via `any` cast
function graph(overrides: Record<string, any> = {}): any {
return new MemoryGraph(makeConfig(overrides));
}
const FILTERS = { userId: "test-user" };
beforeEach(() => {
jest.clearAllMocks();
});
// ═══════════════════════════════════════════════════════════════════════════
// 1. _retrieveNodesFromData – entity extraction
// ═══════════════════════════════════════════════════════════════════════════
describe("_retrieveNodesFromData", () => {
it("parses a well-formed extract_entities tool call", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "extract_entities",
arguments: JSON.stringify({
entities: [
{ entity: "Alice", entity_type: "person" },
{ entity: "Pizza", entity_type: "food" },
],
}),
},
],
});
const mg = graph();
const result = await mg._retrieveNodesFromData(
"Alice likes pizza",
FILTERS,
);
expect(result).toEqual({ alice: "person", pizza: "food" });
});
it("returns empty map when LLM returns a plain string", async () => {
mockGenerateResponse.mockResolvedValueOnce("I am a string, not an object");
const mg = graph();
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({});
});
it("returns empty map when toolCalls is undefined", async () => {
mockGenerateResponse.mockResolvedValueOnce({});
const mg = graph();
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({});
});
it("returns empty map when toolCalls is an empty array", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({});
});
it("handles malformed JSON in tool call arguments gracefully", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{ name: "extract_entities", arguments: "NOT VALID JSON {{{" },
],
});
const mg = graph();
// Should not throw — the catch block in the source logs the error
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({});
});
it("handles missing entities array in arguments", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "extract_entities",
arguments: JSON.stringify({ wrong_key: [] }),
},
],
});
const mg = graph();
// args.entities is undefined → for..of on undefined throws → caught
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({});
});
it("skips tool calls with unrelated names", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "some_other_tool",
arguments: JSON.stringify({
entities: [{ entity: "X", entity_type: "Y" }],
}),
},
],
});
const mg = graph();
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({});
});
it("normalises entity names to lowercase with underscores", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "extract_entities",
arguments: JSON.stringify({
entities: [{ entity: "New York City", entity_type: "City Name" }],
}),
},
],
});
const mg = graph();
const result = await mg._retrieveNodesFromData("anything", FILTERS);
expect(result).toEqual({ new_york_city: "city_name" });
});
it("passes json_object response format and the correct system prompt", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
await mg._retrieveNodesFromData("test data", FILTERS);
const [messages, responseFormat] = mockGenerateResponse.mock.calls[0];
expect(responseFormat).toEqual({ type: "json_object" });
const systemMsg = messages[0].content as string;
expect(systemMsg.toLowerCase()).toContain("json");
expect(systemMsg).toContain("test-user");
});
});
// ═══════════════════════════════════════════════════════════════════════════
// 2. _establishNodesRelationsFromData – relation extraction
// ═══════════════════════════════════════════════════════════════════════════
describe("_establishNodesRelationsFromData", () => {
it("parses a well-formed establish_relationships tool call", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "establish_relationships",
arguments: JSON.stringify({
entities: [
{ source: "Alice", relationship: "likes", destination: "Pizza" },
],
}),
},
],
});
const mg = graph();
const result = await mg._establishNodesRelationsFromData(
"Alice likes pizza",
FILTERS,
{ alice: "person", pizza: "food" },
);
expect(result).toEqual([
{ source: "alice", relationship: "likes", destination: "pizza" },
]);
});
it("returns empty array when LLM returns a string", async () => {
mockGenerateResponse.mockResolvedValueOnce("just a string");
const mg = graph();
const result = await mg._establishNodesRelationsFromData("x", FILTERS, {});
expect(result).toEqual([]);
});
it("returns empty array when toolCalls is empty", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
const result = await mg._establishNodesRelationsFromData("x", FILTERS, {});
expect(result).toEqual([]);
});
it("returns empty array when entities key is missing from arguments", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "establish_relationships",
arguments: JSON.stringify({ not_entities: [] }),
},
],
});
const mg = graph();
const result = await mg._establishNodesRelationsFromData("x", FILTERS, {});
// args.entities is undefined → falls back to []
expect(result).toEqual([]);
});
it("throws on malformed JSON in tool call arguments (no try/catch in source)", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [{ name: "establish_relationships", arguments: "<<BROKEN>>" }],
});
const mg = graph();
// _establishNodesRelationsFromData does JSON.parse without try/catch
await expect(
mg._establishNodesRelationsFromData("x", FILTERS, {}),
).rejects.toThrow();
});
it("appends JSON format suffix to system prompt (no custom prompt)", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
await mg._establishNodesRelationsFromData("data", FILTERS, { a: "b" });
const [messages, responseFormat] = mockGenerateResponse.mock.calls[0];
expect(responseFormat).toEqual({ type: "json_object" });
const systemContent = messages[0].content as string;
expect(systemContent.toLowerCase()).toContain("json");
expect(systemContent).toContain("test-user");
expect(systemContent).not.toContain("USER_ID");
// CUSTOM_PROMPT placeholder stays when no custom prompt is configured
// (only replaced when config.graphStore.customInstructions is set)
});
it("appends JSON format suffix and custom prompt when configured", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph({
customInstructions: "Focus on food relationships only.",
});
await mg._establishNodesRelationsFromData("data", FILTERS, {});
const [messages] = mockGenerateResponse.mock.calls[0];
const systemContent = messages[0].content as string;
expect(systemContent.toLowerCase()).toContain("json");
expect(systemContent).toContain("Focus on food relationships only.");
});
});
// ═══════════════════════════════════════════════════════════════════════════
// 3. _getDeleteEntitiesFromSearchOutput – deletion identification
// ═══════════════════════════════════════════════════════════════════════════
describe("_getDeleteEntitiesFromSearchOutput", () => {
const SEARCH_OUTPUT = [
{
source: "alice",
source_id: "1",
relationship: "likes",
relation_id: "r1",
destination: "pizza",
destination_id: "2",
similarity: 0.95,
},
];
it("parses a well-formed delete_graph_memory tool call", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "delete_graph_memory",
arguments: JSON.stringify({
source: "Alice",
relationship: "likes",
destination: "Pizza",
}),
},
],
});
const mg = graph();
const result = await mg._getDeleteEntitiesFromSearchOutput(
SEARCH_OUTPUT,
"Alice hates pizza",
FILTERS,
);
expect(result).toEqual([
{ source: "alice", relationship: "likes", destination: "pizza" },
]);
});
it("returns empty array when LLM returns a string", async () => {
mockGenerateResponse.mockResolvedValueOnce("string response");
const mg = graph();
const result = await mg._getDeleteEntitiesFromSearchOutput(
SEARCH_OUTPUT,
"x",
FILTERS,
);
expect(result).toEqual([]);
});
it("returns empty array when no tool calls are present", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
const result = await mg._getDeleteEntitiesFromSearchOutput(
SEARCH_OUTPUT,
"x",
FILTERS,
);
expect(result).toEqual([]);
});
it("skips non-delete_graph_memory tool calls", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "noop",
arguments: JSON.stringify({}),
},
],
});
const mg = graph();
const result = await mg._getDeleteEntitiesFromSearchOutput(
SEARCH_OUTPUT,
"x",
FILTERS,
);
expect(result).toEqual([]);
});
it("collects multiple delete tool calls", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "delete_graph_memory",
arguments: JSON.stringify({
source: "A",
relationship: "r1",
destination: "B",
}),
},
{
name: "delete_graph_memory",
arguments: JSON.stringify({
source: "C",
relationship: "r2",
destination: "D",
}),
},
],
});
const mg = graph();
const result = await mg._getDeleteEntitiesFromSearchOutput(
SEARCH_OUTPUT,
"x",
FILTERS,
);
expect(result).toHaveLength(2);
expect(result[0].source).toBe("a");
expect(result[1].source).toBe("c");
});
it("passes json_object format and includes 'json' in system prompt", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
await mg._getDeleteEntitiesFromSearchOutput(SEARCH_OUTPUT, "data", FILTERS);
const [messages, responseFormat] = mockGenerateResponse.mock.calls[0];
expect(responseFormat).toEqual({ type: "json_object" });
const systemContent = messages[0].content as string;
expect(systemContent.toLowerCase()).toContain("json");
expect(systemContent).toContain("test-user");
expect(systemContent).not.toContain("USER_ID");
});
it("handles empty searchOutput array", async () => {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
const result = await mg._getDeleteEntitiesFromSearchOutput(
[],
"data",
FILTERS,
);
expect(result).toEqual([]);
});
});
// ═══════════════════════════════════════════════════════════════════════════
// 4. Prompt construction — JSON keyword present in every json_object site
// ═══════════════════════════════════════════════════════════════════════════
describe("Prompt construction — all json_object sites include 'json'", () => {
it("_retrieveNodesFromData system message includes 'json' for any userId", async () => {
for (const userId of ["", "user-1", "special<>chars", "ユーザー"]) {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
await mg._retrieveNodesFromData("test", { userId });
const systemMsg = mockGenerateResponse.mock.calls.at(-1)![0][0].content;
expect(systemMsg.toLowerCase()).toContain("json");
}
});
it("_establishNodesRelationsFromData system message includes 'json' for any userId", async () => {
for (const userId of ["", "user-1", "special<>chars"]) {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
await mg._establishNodesRelationsFromData("test", { userId }, {});
const systemMsg = mockGenerateResponse.mock.calls.at(-1)![0][0].content;
expect(systemMsg.toLowerCase()).toContain("json");
}
});
it("_getDeleteEntitiesFromSearchOutput system message includes 'json' for any userId", async () => {
for (const userId of ["", "user-1", "special<>chars"]) {
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
const mg = graph();
await mg._getDeleteEntitiesFromSearchOutput([], "test", { userId });
const systemMsg = mockGenerateResponse.mock.calls.at(-1)![0][0].content;
expect(systemMsg.toLowerCase()).toContain("json");
}
});
});
// ═══════════════════════════════════════════════════════════════════════════
// 5. Edge cases – malformed entity fields in _removeSpacesFromEntities
// ═══════════════════════════════════════════════════════════════════════════
// ═══════════════════════════════════════════════════════════════════════════
// 5a. LLM config propagation — graph store uses correct provider & config
// Regression test for https://github.com/mem0ai/mem0/issues/3425
// ═══════════════════════════════════════════════════════════════════════════
describe("LLM config propagation to graph store (issue #3425)", () => {
const { LLMFactory } = require("../src/utils/factory");
beforeEach(() => {
(LLMFactory.create as jest.Mock).mockClear();
});
it("uses root llm config when no graphStore.llm is provided", () => {
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
},
embedder: { provider: "openai", config: {} },
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514", apiKey: "sk-ant-test" },
},
} as any;
new MemoryGraph(config);
expect(LLMFactory.create).toHaveBeenCalledWith("anthropic", {
model: "claude-sonnet-4-20250514",
apiKey: "sk-ant-test",
});
// Both llm and structuredLlm should use the same config
expect(LLMFactory.create).toHaveBeenCalledTimes(2);
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "anthropic", {
model: "claude-sonnet-4-20250514",
apiKey: "sk-ant-test",
});
expect(LLMFactory.create).toHaveBeenNthCalledWith(2, "anthropic", {
model: "claude-sonnet-4-20250514",
apiKey: "sk-ant-test",
});
});
it("uses graphStore.llm config when provided, overriding root llm", () => {
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
llm: {
provider: "openai",
config: { model: "gpt-4o", apiKey: "sk-openai-test" },
},
},
embedder: { provider: "openai", config: {} },
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514", apiKey: "sk-ant-test" },
},
} as any;
new MemoryGraph(config);
// Should use graphStore.llm, NOT root llm
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "openai", {
model: "gpt-4o",
apiKey: "sk-openai-test",
});
expect(LLMFactory.create).toHaveBeenNthCalledWith(2, "openai", {
model: "gpt-4o",
apiKey: "sk-openai-test",
});
});
it("falls back to root llm config when graphStore.llm.config is undefined", () => {
// Note: in practice, Zod schema requires config when graphStore.llm is
// present. This tests the defensive fallback in MemoryGraph itself.
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
llm: {
provider: "openai",
// config explicitly undefined
config: undefined,
},
},
embedder: { provider: "openai", config: {} },
llm: {
provider: "anthropic",
config: { model: "claude-sonnet-4-20250514" },
},
} as any;
new MemoryGraph(config);
// Provider from graphStore.llm, but config falls back to root llm.config
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "openai", {
model: "claude-sonnet-4-20250514",
});
});
it("defaults to openai when neither root nor graphStore llm provider is set", () => {
const config = {
graphStore: {
config: {
url: "bolt://localhost:7687",
username: "neo4j",
password: "test",
},
},
embedder: { provider: "openai", config: {} },
llm: { config: { model: "gpt-4" } },
} as any;
new MemoryGraph(config);
expect(LLMFactory.create).toHaveBeenNthCalledWith(1, "openai", {
model: "gpt-4",
});
});
});
describe("_removeSpacesFromEntities (via _establishNodesRelationsFromData)", () => {
it("normalises spaces and case in entity source/relationship/destination", async () => {
mockGenerateResponse.mockResolvedValueOnce({
toolCalls: [
{
name: "establish_relationships",
arguments: JSON.stringify({
entities: [
{
source: "New York",
relationship: "Capital Of",
destination: "United States",
},
],
}),
},
],
});
const mg = graph();
const result = await mg._establishNodesRelationsFromData(
"test",
FILTERS,
{},
);
expect(result).toEqual([
{
source: "new_york",
relationship: "capital_of",
destination: "united_states",
},
]);
});
});
mem0-ts/src/oss/tests/graph-prompts.test.ts
-177
View File
@@ -1,177 +0,0 @@
import {
DELETE_RELATIONS_SYSTEM_PROMPT,
EXTRACT_RELATIONS_PROMPT,
UPDATE_GRAPH_PROMPT,
getDeleteMessages,
formatEntities,
} from "../src/graphs/utils";
/**
* Regression tests for graph prompts (issue #4248).
*
* When response_format: { type: "json_object" } is used, OpenAI requires
* the word "json" (case-insensitive) to appear in at least one message.
* Missing it produces a 400 error.
*
* Three call sites use json_object today:
* 1. _getDeleteEntitiesFromSearchOutput → DELETE_RELATIONS_SYSTEM_PROMPT
* 2. _retrieveNodesFromData → inline prompt (graph_memory.ts)
* 3. _getRelatedEntities → EXTRACT_RELATIONS_PROMPT + suffix
*
* See: https://github.com/mem0ai/mem0/issues/4248
*/
// ─── JSON keyword presence ────────────────────────────────────────────────────
describe("Graph prompts — JSON keyword requirement", () => {
it("DELETE_RELATIONS_SYSTEM_PROMPT contains 'json'", () => {
expect(DELETE_RELATIONS_SYSTEM_PROMPT.toLowerCase()).toContain("json");
});
it("EXTRACT_RELATIONS_PROMPT produces a message containing 'json' once the suffix is appended", () => {
// graph_memory.ts appends "\nPlease provide your response in JSON format."
const withSuffix =
EXTRACT_RELATIONS_PROMPT +
"\nPlease provide your response in JSON format.";
expect(withSuffix.toLowerCase()).toContain("json");
});
it("getDeleteMessages system message contains 'json' after USER_ID substitution", () => {
const [systemContent] = getDeleteMessages(
"alice -- loves -- pizza",
"Alice now hates pizza",
"user-42",
);
expect(systemContent.toLowerCase()).toContain("json");
});
it("entity extraction inline prompt contains 'json' (simulated from graph_memory.ts)", () => {
// Mirrors the template in _retrieveNodesFromData()
const userId = "user-1";
const prompt = `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 ${userId} as the source entity. Extract all the entities from the text. ***DO NOT*** answer the question itself if the given text is a question. Respond in JSON format.`;
expect(prompt.toLowerCase()).toContain("json");
});
});
// ─── getDeleteMessages ────────────────────────────────────────────────────────
describe("getDeleteMessages", () => {
it("replaces USER_ID with the provided userId in the system prompt", () => {
const [system] = getDeleteMessages("mem", "data", "alice-123");
expect(system).toContain("alice-123");
expect(system).not.toContain("USER_ID");
});
it("includes existing memories and new data in the user prompt", () => {
const existing = "bob -- knows -- carol";
const newData = "Bob no longer knows Carol";
const [, user] = getDeleteMessages(existing, newData, "u1");
expect(user).toContain(existing);
expect(user).toContain(newData);
});
it("returns a 2-tuple [system, user]", () => {
const result = getDeleteMessages("a", "b", "c");
expect(result).toHaveLength(2);
expect(typeof result[0]).toBe("string");
expect(typeof result[1]).toBe("string");
});
// — Malformed / edge-case inputs —
it("handles empty strings without throwing", () => {
expect(() => getDeleteMessages("", "", "")).not.toThrow();
const [system, user] = getDeleteMessages("", "", "");
expect(system.toLowerCase()).toContain("json");
expect(typeof user).toBe("string");
});
it("handles special characters in userId (e.g. angle brackets, quotes)", () => {
const [system] = getDeleteMessages(
"mem",
"data",
'<script>alert("xss")</script>',
);
expect(system).toContain('<script>alert("xss")</script>');
expect(system).not.toContain("USER_ID");
});
it("handles unicode input", () => {
const [system, user] = getDeleteMessages(
"日本語メモリ",
"新しい情報",
"ユーザー1",
);
expect(system).toContain("ユーザー1");
expect(user).toContain("日本語メモリ");
expect(user).toContain("新しい情報");
});
it("handles very long input strings", () => {
const longStr = "x".repeat(100_000);
expect(() => getDeleteMessages(longStr, longStr, "u")).not.toThrow();
const [system] = getDeleteMessages(longStr, longStr, "u");
expect(system.toLowerCase()).toContain("json");
});
});
// ─── formatEntities ───────────────────────────────────────────────────────────
describe("formatEntities", () => {
it("formats a single entity triplet", () => {
const result = formatEntities([
{ source: "Alice", relationship: "knows", destination: "Bob" },
]);
expect(result).toBe("Alice -- knows -- Bob");
});
it("joins multiple entities with newlines", () => {
const result = formatEntities([
{ source: "A", relationship: "r1", destination: "B" },
{ source: "C", relationship: "r2", destination: "D" },
]);
expect(result).toBe("A -- r1 -- B\nC -- r2 -- D");
});
it("returns empty string for empty array", () => {
expect(formatEntities([])).toBe("");
});
it("preserves special characters in entity fields", () => {
const result = formatEntities([
{ source: "O'Brien", relationship: 'said "hello"', destination: "café" },
]);
expect(result).toContain("O'Brien");
expect(result).toContain('said "hello"');
expect(result).toContain("café");
});
});
// ─── Prompt structural invariants ─────────────────────────────────────────────
describe("Prompt structural invariants", () => {
it("DELETE_RELATIONS_SYSTEM_PROMPT contains USER_ID placeholder", () => {
expect(DELETE_RELATIONS_SYSTEM_PROMPT).toContain("USER_ID");
});
it("EXTRACT_RELATIONS_PROMPT contains USER_ID placeholder", () => {
expect(EXTRACT_RELATIONS_PROMPT).toContain("USER_ID");
});
it("EXTRACT_RELATIONS_PROMPT contains CUSTOM_PROMPT placeholder", () => {
expect(EXTRACT_RELATIONS_PROMPT).toContain("CUSTOM_PROMPT");
});
it("UPDATE_GRAPH_PROMPT contains memory template placeholders", () => {
expect(UPDATE_GRAPH_PROMPT).toContain("{existing_memories}");
expect(UPDATE_GRAPH_PROMPT).toContain("{new_memories}");
});
it("DELETE_RELATIONS_SYSTEM_PROMPT is non-empty and reasonably sized", () => {
expect(DELETE_RELATIONS_SYSTEM_PROMPT.length).toBeGreaterThan(100);
});
it("EXTRACT_RELATIONS_PROMPT is non-empty and reasonably sized", () => {
expect(EXTRACT_RELATIONS_PROMPT.length).toBeGreaterThan(100);
});
});
mem0-ts/src/oss/tests/memory.add.test.ts
+27 -16
View File
@@ -22,18 +22,21 @@ jest.mock("../src/llms/openai", () => ({
.fn()
.mockImplementation(
(messages: Array<{ role: string; content: string }>) => {
const hasSystemRole = messages.some((m) => m.role === "system");
if (hasSystemRole) {
return JSON.stringify({ facts: ["extracted fact from input"] });
}
// V3 pipeline: single LLM call with additive extraction prompt.
const userMsg = messages.find((m) => m.role === "user");
const content = userMsg?.content ?? "";
const newMsgMatch = content.match(
/## New Messages\n([\s\S]*?)(?=\n##|$)/,
);
const extracted = newMsgMatch
? newMsgMatch[1].trim()
: "extracted fact from input";
return JSON.stringify({
memory: [
{
id: "new",
event: "ADD",
text: "extracted fact from input",
old_memory: "",
new_memory: "extracted fact from input",
id: "0",
text: extracted,
attributed_to: "user",
},
],
});
@@ -42,9 +45,15 @@ jest.mock("../src/llms/openai", () => ({
})),
}));
const mockEmbedding = new Array(1536).fill(0.1);
jest.mock("../src/embeddings/openai", () => ({
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
embed: jest.fn().mockResolvedValue(new Array(1536).fill(0.1)),
embed: jest.fn().mockResolvedValue(mockEmbedding),
embedBatch: jest
.fn()
.mockImplementation((texts: string[]) =>
Promise.resolve(texts.map(() => mockEmbedding)),
),
embeddingDims: 1536,
})),
}));
@@ -93,17 +102,19 @@ describe("Memory - add()", () => {
});
test("returns at least one result with an id", async () => {
const result: SearchResult = await memory.add("I am a software engineer", {
userId,
});
const result: SearchResult = await memory.add(
"I enjoy hiking in the mountains",
{ userId },
);
expect(result.results.length).toBeGreaterThan(0);
expect(result.results[0].id).toBeDefined();
});
test("result item has a memory string field", async () => {
const result: SearchResult = await memory.add("I am a software engineer", {
userId,
});
const result: SearchResult = await memory.add(
"My favorite color is blue",
{ userId },
);
expect(typeof result.results[0].memory).toBe("string");
});
mem0-ts/src/oss/tests/memory.crud.test.ts
+21 -10
View File
@@ -22,18 +22,23 @@ jest.mock("../src/llms/openai", () => ({
.fn()
.mockImplementation(
(messages: Array<{ role: string; content: string }>) => {
const hasSystemRole = messages.some((m) => m.role === "system");
if (hasSystemRole) {
return JSON.stringify({ facts: ["stored fact"] });
}
// V3 pipeline: single LLM call with additive extraction prompt.
// Extract the user input from the prompt to produce unique memories.
const userMsg = messages.find((m) => m.role === "user");
const content = userMsg?.content ?? "";
// Pull the text between "## New Messages" and the next "##"
const newMsgMatch = content.match(
/## New Messages\n([\s\S]*?)(?=\n##|$)/,
);
const extracted = newMsgMatch
? newMsgMatch[1].trim()
: "stored fact";
return JSON.stringify({
memory: [
{
id: "new",
event: "ADD",
text: "stored fact",
old_memory: "",
new_memory: "stored fact",
id: "0",
text: extracted,
attributed_to: "user",
},
],
});
@@ -42,9 +47,15 @@ jest.mock("../src/llms/openai", () => ({
})),
}));
const mockEmbedding = new Array(1536).fill(0.1);
jest.mock("../src/embeddings/openai", () => ({
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
embed: jest.fn().mockResolvedValue(new Array(1536).fill(0.1)),
embed: jest.fn().mockResolvedValue(mockEmbedding),
embedBatch: jest
.fn()
.mockImplementation((texts: string[]) =>
Promise.resolve(texts.map(() => mockEmbedding)),
),
embeddingDims: 1536,
})),
}));
mem0-ts/src/oss/tests/memory.init.test.ts
+19 -12
View File
@@ -16,26 +16,27 @@ jest.mock("../src/llms/google", () => ({
GoogleLLM: jest.fn(),
}));
// ─── Content-based LLM mock (reviewer #9) ────────────────
// Returns facts for system-prompt calls, memory actions for user-only calls.
// ─── Content-based LLM mock (V3 additive extraction pipeline) ─────────
jest.mock("../src/llms/openai", () => ({
OpenAILLM: jest.fn().mockImplementation(() => ({
generateResponse: jest
.fn()
.mockImplementation(
(messages: Array<{ role: string; content: string }>) => {
const hasSystemRole = messages.some((m) => m.role === "system");
if (hasSystemRole) {
return JSON.stringify({ facts: ["test fact"] });
}
const userMsg = messages.find((m) => m.role === "user");
const content = userMsg?.content ?? "";
const newMsgMatch = content.match(
/## New Messages\n([\s\S]*?)(?=\n##|$)/,
);
const extracted = newMsgMatch
? newMsgMatch[1].trim()
: "test fact";
return JSON.stringify({
memory: [
{
id: "new",
event: "ADD",
text: "test fact",
old_memory: "",
new_memory: "test fact",
id: "0",
text: extracted,
attributed_to: "user",
},
],
});
@@ -44,9 +45,15 @@ jest.mock("../src/llms/openai", () => ({
})),
}));
const mockEmbedding = new Array(1536).fill(0.1);
jest.mock("../src/embeddings/openai", () => ({
OpenAIEmbedder: jest.fn().mockImplementation(() => ({
embed: jest.fn().mockResolvedValue(new Array(1536).fill(0.1)),
embed: jest.fn().mockResolvedValue(mockEmbedding),
embedBatch: jest
.fn()
.mockImplementation((texts: string[]) =>
Promise.resolve(texts.map(() => mockEmbedding)),
),
embeddingDims: 1536,
})),
}));
mem0-ts/tsup.config.ts
+2 -1
View File
@@ -13,13 +13,14 @@ const external = [
"ollama",
"@google/genai",
"@mistralai/mistralai",
"neo4j-driver",
"@supabase/supabase-js",
"@azure/search-documents",
"@azure/identity",
"cloudflare",
"@cloudflare/workers-types",
"@langchain/core",
"compromise",
"natural",
];
export default defineConfig([
mem0/client/main.py
+6 -22
View File
@@ -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(
mem0/configs/base.py
-5
View File
@@ -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,
mem0/exceptions.py
-18
View File
@@ -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):
mem0/graphs/__init__.py
View File
mem0/graphs/configs.py
-136
View File
@@ -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}")
mem0/graphs/neptune/__init__.py
View File
mem0/graphs/neptune/base.py
-515
View File
@@ -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})
mem0/graphs/neptune/neptunedb.py
-511
View File
@@ -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
mem0/graphs/neptune/neptunegraph.py
-475
View File
@@ -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
mem0/graphs/tools.py
-371
View File
@@ -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,
},
},
}
mem0/graphs/utils.py
-97
View File
@@ -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}"
mem0/memory/apache_age_memory.py
-595
View File
@@ -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()
mem0/memory/graph_memory.py
-744
View File
@@ -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)
mem0/memory/kuzu_memory.py
-732
View File
@@ -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)
mem0/memory/main.py
+18 -211
View File
@@ -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):
mem0/memory/memgraph_memory.py
-708
View File
@@ -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": []}
mem0/memory/telemetry.py
-3
View File
@@ -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__}",
mem0/utils/factory.py
-24
View File
@@ -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:
"""
pyproject.toml
-13
View File
@@ -27,15 +27,6 @@ dependencies = [
nlp = [
"spacy>=3.7.0",
]
graph = [
"langchain-neo4j>=0.4.0",
"langchain-aws>=0.2.23",
"langchain-memgraph>=0.1.0",
"neo4j>=5.23.1",
"rank-bm25>=0.2.2",
"kuzu>=0.11.0",
"apache-age-python>=0.0.6",
]
vector_stores = [
"vecs>=0.4.0",
"chromadb>=0.4.24",
@@ -110,7 +101,6 @@ only-include = ["mem0"]
python = "3.9"
features = [
"test",
"graph",
"vector_stores",
"llms",
"extras",
@@ -120,7 +110,6 @@ features = [
python = "3.10"
features = [
"test",
"graph",
"vector_stores",
"llms",
"extras",
@@ -130,7 +119,6 @@ features = [
python = "3.11"
features = [
"test",
"graph",
"vector_stores",
"llms",
"extras",
@@ -140,7 +128,6 @@ features = [
python = "3.12"
features = [
"test",
"graph",
"vector_stores",
"llms",
"extras",
server/main.py
-12
View File
@@ -40,14 +40,6 @@ POSTGRES_USER = os.environ.get("POSTGRES_USER", "postgres")
POSTGRES_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "postgres")
POSTGRES_COLLECTION_NAME = os.environ.get("POSTGRES_COLLECTION_NAME", "memories")
NEO4J_URI = os.environ.get("NEO4J_URI", "bolt://neo4j:7687")
NEO4J_USERNAME = os.environ.get("NEO4J_USERNAME", "neo4j")
NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD", "mem0graph")
MEMGRAPH_URI = os.environ.get("MEMGRAPH_URI", "bolt://localhost:7687")
MEMGRAPH_USERNAME = os.environ.get("MEMGRAPH_USERNAME", "memgraph")
MEMGRAPH_PASSWORD = os.environ.get("MEMGRAPH_PASSWORD", "mem0graph")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
HISTORY_DB_PATH = os.environ.get("HISTORY_DB_PATH", "/app/history/history.db")
@@ -64,10 +56,6 @@ DEFAULT_CONFIG = {
"collection_name": POSTGRES_COLLECTION_NAME,
},
},
"graph_store": {
"provider": "neo4j",
"config": {"url": NEO4J_URI, "username": NEO4J_USERNAME, "password": NEO4J_PASSWORD},
},
"llm": {"provider": "openai", "config": {"api_key": OPENAI_API_KEY, "temperature": 0.2, "model": "gpt-4.1-nano-2025-04-14"}},
"embedder": {"provider": "openai", "config": {"api_key": OPENAI_API_KEY, "model": "text-embedding-3-small"}},
"history_db_path": HISTORY_DB_PATH,