mem0ai 2.0.1 → 2.1.0

package/dist/oss/index.js

@@ -77,8 +77,19 @@ var MemoryConfigSchema = import_zod.z.object({
   }),
   historyDbPath: import_zod.z.string().optional(),
   customPrompt: import_zod.z.string().optional(),
+  enableGraph: import_zod.z.boolean().optional(),
   graphStore: import_zod.z.object({
-    config: import_zod.z.any().optional()
+    provider: import_zod.z.string(),
+    config: import_zod.z.object({
+      url: import_zod.z.string(),
+      username: import_zod.z.string(),
+      password: import_zod.z.string()
+    }),
+    llm: import_zod.z.object({
+      provider: import_zod.z.string(),
+      config: import_zod.z.record(import_zod.z.string(), import_zod.z.any())
+    }).optional(),
+    customPrompt: import_zod.z.string().optional()
   }).optional()
 });
 
@@ -112,16 +123,28 @@ var OpenAILLM = class {
     this.openai = new import_openai2.default({ apiKey: config.apiKey });
     this.model = config.model || "gpt-4-turbo-preview";
   }
-  async generateResponse(messages, responseFormat) {
+  async generateResponse(messages, responseFormat, tools) {
     const completion = await this.openai.chat.completions.create({
       messages: messages.map((msg) => ({
         role: msg.role,
         content: msg.content
       })),
       model: this.model,
-      response_format: responseFormat
+      response_format: responseFormat,
+      ...tools && { tools, tool_choice: "auto" }
     });
-    return completion.choices[0].message.content || "";
+    const response = completion.choices[0].message;
+    if (response.tool_calls) {
+      return {
+        content: response.content || "",
+        role: response.role,
+        toolCalls: response.tool_calls.map((call) => ({
+          name: call.function.name,
+          arguments: call.function.arguments
+        }))
+      };
+    }
+    return response.content || "";
   }
   async generateChat(messages) {
     const completion = await this.openai.chat.completions.create({
@@ -143,37 +166,57 @@ var OpenAILLM = class {
 var import_openai3 = __toESM(require("openai"));
 var OpenAIStructuredLLM = class {
   constructor(config) {
-    const apiKey = config.apiKey || process.env.OPENAI_API_KEY;
-    if (!apiKey) {
-      throw new Error("OpenAI API key is required");
-    }
-    const baseUrl = process.env.OPENAI_API_BASE || "https://api.openai.com/v1";
-    this.client = new import_openai3.default({ apiKey, baseURL: baseUrl });
-    this.model = config.model || "gpt-4-0125-preview";
+    this.openai = new import_openai3.default({ apiKey: config.apiKey });
+    this.model = config.model || "gpt-4-turbo-preview";
   }
-  async generateResponse(messages, responseFormat) {
-    const response = await this.client.chat.completions.create({
-      model: this.model,
+  async generateResponse(messages, responseFormat, tools) {
+    const completion = await this.openai.chat.completions.create({
       messages: messages.map((msg) => ({
         role: msg.role,
         content: msg.content
       })),
-      response_format: responseFormat
+      model: this.model,
+      ...tools ? {
+        tools: tools.map((tool) => ({
+          type: "function",
+          function: {
+            name: tool.function.name,
+            description: tool.function.description,
+            parameters: tool.function.parameters
+          }
+        })),
+        tool_choice: "auto"
+      } : responseFormat ? {
+        response_format: {
+          type: responseFormat.type
+        }
+      } : {}
     });
-    return response.choices[0].message.content || "";
+    const response = completion.choices[0].message;
+    if (response.tool_calls) {
+      return {
+        content: response.content || "",
+        role: response.role,
+        toolCalls: response.tool_calls.map((call) => ({
+          name: call.function.name,
+          arguments: call.function.arguments
+        }))
+      };
+    }
+    return response.content || "";
   }
   async generateChat(messages) {
-    const response = await this.client.chat.completions.create({
-      model: this.model,
+    const completion = await this.openai.chat.completions.create({
       messages: messages.map((msg) => ({
         role: msg.role,
         content: msg.content
-      }))
+      })),
+      model: this.model
     });
-    const message = response.choices[0].message;
+    const response = completion.choices[0].message;
     return {
-      content: message.content || "",
-      role: message.role
+      content: response.content || "",
+      role: response.role
     };
   }
 };
@@ -251,10 +294,50 @@ var GroqLLM = class {
 };
 
 // src/oss/src/vector_stores/memory.ts
+var import_sqlite3 = __toESM(require("sqlite3"));
+var import_path = __toESM(require("path"));
 var MemoryVectorStore = class {
   constructor(config) {
-    this.vectors = /* @__PURE__ */ new Map();
     this.dimension = config.dimension || 1536;
+    this.dbPath = import_path.default.join(process.cwd(), "vector_store.db");
+    if (config.dbPath) {
+      this.dbPath = config.dbPath;
+    }
+    this.db = new import_sqlite3.default.Database(this.dbPath);
+    this.init().catch(console.error);
+  }
+  async init() {
+    await this.run(`
+      CREATE TABLE IF NOT EXISTS vectors (
+        id TEXT PRIMARY KEY,
+        vector BLOB NOT NULL,
+        payload TEXT NOT NULL
+      )
+    `);
+  }
+  async run(sql, params = []) {
+    return new Promise((resolve, reject) => {
+      this.db.run(sql, params, (err) => {
+        if (err) reject(err);
+        else resolve();
+      });
+    });
+  }
+  async all(sql, params = []) {
+    return new Promise((resolve, reject) => {
+      this.db.all(sql, params, (err, rows) => {
+        if (err) reject(err);
+        else resolve(rows);
+      });
+    });
+  }
+  async getOne(sql, params = []) {
+    return new Promise((resolve, reject) => {
+      this.db.get(sql, params, (err, row) => {
+        if (err) reject(err);
+        else resolve(row);
+      });
+    });
   }
   cosineSimilarity(a, b) {
     let dotProduct = 0;
@@ -280,11 +363,11 @@ var MemoryVectorStore = class {
           `Vector dimension mismatch. Expected ${this.dimension}, got ${vectors[i].length}`
         );
       }
-      this.vectors.set(ids[i], {
-        id: ids[i],
-        vector: vectors[i],
-        payload: payloads[i]
-      });
+      const vectorBuffer = Buffer.from(new Float32Array(vectors[i]).buffer);
+      await this.run(
+        `INSERT OR REPLACE INTO vectors (id, vector, payload) VALUES (?, ?, ?)`,
+        [ids[i], vectorBuffer, JSON.stringify(payloads[i])]
+      );
     }
   }
   async search(query, limit = 10, filters) {
@@ -293,13 +376,21 @@ var MemoryVectorStore = class {
         `Query dimension mismatch. Expected ${this.dimension}, got ${query.length}`
       );
     }
+    const rows = await this.all(`SELECT * FROM vectors`);
     const results = [];
-    for (const vector of this.vectors.values()) {
-      if (this.filterVector(vector, filters)) {
-        const score = this.cosineSimilarity(query, vector.vector);
+    for (const row of rows) {
+      const vector = new Float32Array(row.vector.buffer);
+      const payload = JSON.parse(row.payload);
+      const memoryVector = {
+        id: row.id,
+        vector: Array.from(vector),
+        payload
+      };
+      if (this.filterVector(memoryVector, filters)) {
+        const score = this.cosineSimilarity(query, Array.from(vector));
         results.push({
-          id: vector.id,
-          payload: vector.payload,
+          id: memoryVector.id,
+          payload: memoryVector.payload,
           score
         });
       }
@@ -308,11 +399,14 @@ var MemoryVectorStore = class {
     return results.slice(0, limit);
   }
   async get(vectorId) {
-    const vector = this.vectors.get(vectorId);
-    if (!vector) return null;
+    const row = await this.getOne(`SELECT * FROM vectors WHERE id = ?`, [
+      vectorId
+    ]);
+    if (!row) return null;
+    const payload = JSON.parse(row.payload);
     return {
-      id: vector.id,
-      payload: vector.payload
+      id: row.id,
+      payload
     };
   }
   async update(vectorId, vector, payload) {
@@ -321,27 +415,34 @@ var MemoryVectorStore = class {
         `Vector dimension mismatch. Expected ${this.dimension}, got ${vector.length}`
       );
     }
-    const existing = this.vectors.get(vectorId);
-    if (!existing) throw new Error(`Vector with ID ${vectorId} not found`);
-    this.vectors.set(vectorId, {
-      id: vectorId,
-      vector,
-      payload
-    });
+    const vectorBuffer = Buffer.from(new Float32Array(vector).buffer);
+    await this.run(`UPDATE vectors SET vector = ?, payload = ? WHERE id = ?`, [
+      vectorBuffer,
+      JSON.stringify(payload),
+      vectorId
+    ]);
   }
   async delete(vectorId) {
-    this.vectors.delete(vectorId);
+    await this.run(`DELETE FROM vectors WHERE id = ?`, [vectorId]);
   }
   async deleteCol() {
-    this.vectors.clear();
+    await this.run(`DROP TABLE IF EXISTS vectors`);
+    await this.init();
   }
   async list(filters, limit = 100) {
+    const rows = await this.all(`SELECT * FROM vectors`);
     const results = [];
-    for (const vector of this.vectors.values()) {
-      if (this.filterVector(vector, filters)) {
+    for (const row of rows) {
+      const payload = JSON.parse(row.payload);
+      const memoryVector = {
+        id: row.id,
+        vector: Array.from(new Float32Array(row.vector.buffer)),
+        payload
+      };
+      if (this.filterVector(memoryVector, filters)) {
         results.push({
-          id: vector.id,
-          payload: vector.payload
+          id: memoryVector.id,
+          payload: memoryVector.payload
         });
       }
     }
@@ -956,7 +1057,7 @@ var EmbedderFactory = class {
 };
 var LLMFactory = class {
   static create(provider, config) {
-    switch (provider.toLowerCase()) {
+    switch (provider) {
       case "openai":
         return new OpenAILLM(config);
       case "openai_structured":
@@ -1022,7 +1123,7 @@ function getFactRetrievalMessages(parsedMessages) {
   Input: Me favourite movies are Inception and Interstellar.
   Output: {"facts" : ["Favourite movies are Inception and Interstellar"]}
   
-  Return the facts and preferences in a json format as shown above.
+  Return the facts and preferences in a JSON format as shown above. You MUST return a valid JSON object with a 'facts' key containing an array of strings.
   
   Remember the following:
   - Today's date is ${(/* @__PURE__ */ new Date()).toISOString().split("T")[0]}.
@@ -1031,16 +1132,16 @@ function getFactRetrievalMessages(parsedMessages) {
   - If the user asks where you fetched my information, answer that you found from publicly available sources on internet.
   - If you do not find anything relevant in the below conversation, you can return an empty list corresponding to the "facts" key.
   - Create the facts based on the user and assistant messages only. Do not pick anything from the system messages.
-  - Make sure to return the response in the format mentioned in the examples. The response should be in json with a key as "facts" and corresponding value will be a list of strings.
+  - Make sure to return the response in the JSON format mentioned in the examples. The response should be in JSON with a key as "facts" and corresponding value will be a list of strings.
   - DO NOT RETURN ANYTHING ELSE OTHER THAN THE JSON FORMAT.
-  - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALUD SUCH AS "\`\`\`json" OR "\`\`\`".
+  - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALID SUCH AS "\`\`\`json" OR "\`\`\`".
   - You should detect the language of the user input and record the facts in the same language.
   - For basic factual statements, break them down into individual facts if they contain multiple pieces of information.
   
-  Following is a conversation between the user and the assistant. You have to extract the relevant facts and preferences about the user, if any, from the conversation and return them in the json format as shown above.
+  Following is a conversation between the user and the assistant. You have to extract the relevant facts and preferences about the user, if any, from the conversation and return them in the JSON format as shown above.
   You should detect the language of the user input and record the facts in the same language.
   `;
-  const userPrompt = `Following is a conversation between the user and the assistant. You have to extract the relevant facts and preferences about the user, if any, from the conversation and return them in the json format as shown above.
+  const userPrompt = `Following is a conversation between the user and the assistant. You have to extract the relevant facts and preferences about the user, if any, from the conversation and return them in the JSON format as shown above.
 
 Input:
 ${parsedMessages}`;
@@ -1203,14 +1304,14 @@ function getUpdateMemoryMessages(retrievedOldMemory, newRetrievedFacts) {
   ${JSON.stringify(newRetrievedFacts, null, 2)}
   
   Follow the instruction mentioned below:
-  - Do not return anything from the custom few shot prompts provided above.
+  - Do not return anything from the custom few shot example prompts provided above.
   - If the current memory is empty, then you have to add the new retrieved facts to the memory.
   - You should return the updated memory in only JSON format as shown below. The memory key should be the same if no changes are made.
   - If there is an addition, generate a new key and add the new memory corresponding to it.
   - If there is a deletion, the memory key-value pair should be removed from the memory.
   - If there is an update, the ID key should remain the same and only the value needs to be updated.
   - DO NOT RETURN ANYTHING ELSE OTHER THAN THE JSON FORMAT.
-  - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALUD SUCH AS "\`\`\`json" OR "\`\`\`".
+  - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALID SUCH AS "\`\`\`json" OR "\`\`\`".
   
   Do not return anything except the JSON format.`;
 }
@@ -1219,10 +1320,10 @@ function removeCodeBlocks(text) {
 }
 
 // src/oss/src/storage/SQLiteManager.ts
-var import_sqlite3 = __toESM(require("sqlite3"));
+var import_sqlite32 = __toESM(require("sqlite3"));
 var SQLiteManager = class {
   constructor(dbPath) {
-    this.db = new import_sqlite3.default.Database(dbPath);
+    this.db = new import_sqlite32.default.Database(dbPath);
     this.init().catch(console.error);
   }
   async init() {
@@ -1310,6 +1411,21 @@ var DEFAULT_MEMORY_CONFIG = {
       model: "gpt-4-turbo-preview"
     }
   },
+  enableGraph: false,
+  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"
 };
 
@@ -1343,12 +1459,742 @@ var ConfigManager = class {
       },
       historyDbPath: userConfig.historyDbPath || DEFAULT_MEMORY_CONFIG.historyDbPath,
       customPrompt: userConfig.customPrompt,
-      graphStore: userConfig.graphStore
+      graphStore: {
+        ...DEFAULT_MEMORY_CONFIG.graphStore,
+        ...userConfig.graphStore
+      },
+      enableGraph: userConfig.enableGraph || DEFAULT_MEMORY_CONFIG.enableGraph
     };
     return MemoryConfigSchema.parse(mergedConfig);
   }
 };
 
+// src/oss/src/memory/graph_memory.ts
+var import_neo4j_driver = __toESM(require("neo4j-driver"));
+
+// src/oss/src/utils/bm25.ts
+var BM25 = class {
+  constructor(documents, k1 = 1.5, b = 0.75) {
+    this.documents = documents;
+    this.k1 = k1;
+    this.b = b;
+    this.docLengths = documents.map((doc) => doc.length);
+    this.avgDocLength = this.docLengths.reduce((a, b2) => a + b2, 0) / documents.length;
+    this.docFreq = /* @__PURE__ */ new Map();
+    this.idf = /* @__PURE__ */ new Map();
+    this.computeIdf();
+  }
+  computeIdf() {
+    const N = this.documents.length;
+    for (const doc of this.documents) {
+      const terms = new Set(doc);
+      for (const term of terms) {
+        this.docFreq.set(term, (this.docFreq.get(term) || 0) + 1);
+      }
+    }
+    for (const [term, freq] of this.docFreq) {
+      this.idf.set(term, Math.log((N - freq + 0.5) / (freq + 0.5) + 1));
+    }
+  }
+  score(query, doc, index) {
+    let score = 0;
+    const docLength = this.docLengths[index];
+    for (const term of query) {
+      const tf = doc.filter((t) => t === term).length;
+      const idf = this.idf.get(term) || 0;
+      score += idf * tf * (this.k1 + 1) / (tf + this.k1 * (1 - this.b + this.b * docLength / this.avgDocLength));
+    }
+    return score;
+  }
+  search(query) {
+    const scores = this.documents.map((doc, idx) => ({
+      doc,
+      score: this.score(query, doc, idx)
+    }));
+    return scores.sort((a, b) => b.score - a.score).map((item) => item.doc);
+  }
+};
+
+// src/oss/src/graphs/tools.ts
+var 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
+    }
+  }
+};
+var 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
+    }
+  }
+};
+var 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
+    }
+  }
+};
+
+// src/oss/src/graphs/utils.ts
+var 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.
+`;
+var 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.
+`;
+function getDeleteMessages(existingMemoriesString, data, userId) {
+  return [
+    DELETE_RELATIONS_SYSTEM_PROMPT.replace("USER_ID", userId),
+    `Here are the existing memories: ${existingMemoriesString} 
+
+ New Information: ${data}`
+  ];
+}
+
+// src/oss/src/utils/logger.ts
+var logger = {
+  info: (message) => console.log(`[INFO] ${message}`),
+  error: (message) => console.error(`[ERROR] ${message}`),
+  debug: (message) => console.debug(`[DEBUG] ${message}`),
+  warn: (message) => console.warn(`[WARN] ${message}`)
+};
+
+// src/oss/src/memory/graph_memory.ts
+var MemoryGraph = class {
+  constructor(config) {
+    var _a, _b, _c, _d, _e, _f, _g, _h, _i;
+    this.config = config;
+    if (!((_b = (_a = config.graphStore) == null ? void 0 : _a.config) == null ? void 0 : _b.url) || !((_d = (_c = config.graphStore) == null ? void 0 : _c.config) == null ? void 0 : _d.username) || !((_f = (_e = config.graphStore) == null ? void 0 : _e.config) == null ? void 0 : _f.password)) {
+      throw new Error("Neo4j configuration is incomplete");
+    }
+    this.graph = import_neo4j_driver.default.driver(
+      config.graphStore.config.url,
+      import_neo4j_driver.default.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";
+    if ((_g = this.config.llm) == null ? void 0 : _g.provider) {
+      this.llmProvider = this.config.llm.provider;
+    }
+    if ((_i = (_h = this.config.graphStore) == null ? void 0 : _h.llm) == null ? void 0 : _i.provider) {
+      this.llmProvider = this.config.graphStore.llm.provider;
+    }
+    this.llm = LLMFactory.create(this.llmProvider, this.config.llm.config);
+    this.structuredLlm = LLMFactory.create(
+      "openai_structured",
+      this.config.llm.config
+    );
+    this.threshold = 0.7;
+  }
+  async add(data, filters) {
+    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, filters, limit = 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) {
+    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, limit = 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(limit)) }
+      );
+      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();
+    }
+  }
+  async _retrieveNodesFromData(data, filters) {
+    const tools = [EXTRACT_ENTITIES_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.`
+        },
+        { role: "user", content: data }
+      ],
+      { type: "json_object" },
+      tools
+    );
+    let entityTypeMap = {};
+    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;
+  }
+  async _establishNodesRelationsFromData(data, filters, entityTypeMap) {
+    var _a;
+    let messages;
+    if ((_a = this.config.graphStore) == null ? void 0 : _a.customPrompt) {
+      messages = [
+        {
+          role: "system",
+          content: EXTRACT_RELATIONS_PROMPT.replace(
+            "USER_ID",
+            filters["userId"]
+          ).replace(
+            "CUSTOM_PROMPT",
+            `4. ${this.config.graphStore.customPrompt}`
+          ) + "\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)}. 
+
+Text: ${data}`
+        }
+      ];
+    }
+    const tools = [RELATIONS_TOOL];
+    const extractedEntities = await this.structuredLlm.generateResponse(
+      messages,
+      { type: "json_object" },
+      tools
+    );
+    let entities = [];
+    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;
+  }
+  async _searchGraphDb(nodeList, filters, limit = 100) {
+    const resultRelations = [];
+    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(limit))
+        });
+        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;
+  }
+  async _getDeleteEntitiesFromSearchOutput(searchOutput, data, filters) {
+    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];
+    const memoryUpdates = await this.structuredLlm.generateResponse(
+      [
+        { role: "system", content: systemPrompt },
+        { role: "user", content: userPrompt }
+      ],
+      { type: "json_object" },
+      tools
+    );
+    const toBeDeleted = [];
+    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;
+  }
+  async _deleteEntities(toBeDeleted, userId) {
+    const results = [];
+    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;
+  }
+  async _addEntities(toBeAdded, userId, entityTypeMap) {
+    var _a, _b;
+    const results = [];
+    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;
+        let params;
+        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: (_a = sourceNodeSearchResult[0]) == null ? void 0 : _a.elementId,
+            destination_id: (_b = destinationNodeSearchResult[0]) == null ? void 0 : _b.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;
+  }
+  _removeSpacesFromEntities(entityList) {
+    return entityList.map((item) => ({
+      ...item,
+      source: item.source.toLowerCase().replace(/ /g, "_"),
+      relationship: item.relationship.toLowerCase().replace(/ /g, "_"),
+      destination: item.destination.toLowerCase().replace(/ /g, "_")
+    }));
+  }
+  async _searchSourceNode(sourceEmbedding, userId, 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();
+    }
+  }
+  async _searchDestinationNode(destinationEmbedding, userId, 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();
+    }
+  }
+};
+
 // src/oss/src/memory/index.ts
 var Memory = class _Memory {
   constructor(config = {}) {
@@ -1369,6 +2215,10 @@ var Memory = class _Memory {
     this.db = new SQLiteManager(this.config.historyDbPath || ":memory:");
     this.collectionName = this.config.vectorStore.config.collectionName;
     this.apiVersion = this.config.version || "v1.0";
+    this.enableGraph = this.config.enableGraph || false;
+    if (this.enableGraph && this.config.graphStore) {
+      this.graphMemory = new MemoryGraph(this.config);
+    }
   }
   static fromConfig(configDict) {
     try {
@@ -1402,7 +2252,21 @@ var Memory = class _Memory {
       metadata,
       filters
     );
-    return { results: vectorStoreResult };
+    let graphResult;
+    if (this.graphMemory) {
+      try {
+        graphResult = await this.graphMemory.add(
+          parsedMessages.map((m) => m.content).join("\n"),
+          filters
+        );
+      } catch (error) {
+        console.error("Error adding to graph memory:", error);
+      }
+    }
+    return {
+      results: vectorStoreResult,
+      relations: graphResult == null ? void 0 : graphResult.relations
+    };
   }
   async addToVectorStore(messages, metadata, filters) {
     const parsedMessages = messages.map((m) => m.content).join("\n");
@@ -1546,6 +2410,14 @@ ${parsedMessages}`] : getFactRetrievalMessages(parsedMessages);
       limit,
       filters
     );
+    let graphResults;
+    if (this.graphMemory) {
+      try {
+        graphResults = await this.graphMemory.search(query, filters);
+      } catch (error) {
+        console.error("Error searching graph memory:", error);
+      }
+    }
     const excludedKeys = /* @__PURE__ */ new Set([
       "userId",
       "agentId",
@@ -1567,7 +2439,10 @@ ${parsedMessages}`] : getFactRetrievalMessages(parsedMessages);
       ...mem.payload.agentId && { agentId: mem.payload.agentId },
       ...mem.payload.runId && { runId: mem.payload.runId }
     }));
-    return { results };
+    return {
+      results,
+      relations: graphResults
+    };
   }
   async update(memoryId, data) {
     const embedding = await this.embedder.embed(data);
@@ -1601,6 +2476,9 @@ ${parsedMessages}`] : getFactRetrievalMessages(parsedMessages);
   async reset() {
     await this.db.reset();
     await this.vectorStore.deleteCol();
+    if (this.graphMemory) {
+      await this.graphMemory.deleteAll({ userId: "default" });
+    }
     this.vectorStore = VectorStoreFactory.create(
       this.config.vectorStore.provider,
       this.config.vectorStore.config