mem0ai 2.0.2 → 2.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -29,8 +29,19 @@ var MemoryConfigSchema = z.object({
29
29
  }),
30
30
  historyDbPath: z.string().optional(),
31
31
  customPrompt: z.string().optional(),
32
+ enableGraph: z.boolean().optional(),
32
33
  graphStore: z.object({
33
- config: z.any().optional()
34
+ provider: z.string(),
35
+ config: z.object({
36
+ url: z.string(),
37
+ username: z.string(),
38
+ password: z.string()
39
+ }),
40
+ llm: z.object({
41
+ provider: z.string(),
42
+ config: z.record(z.string(), z.any())
43
+ }).optional(),
44
+ customPrompt: z.string().optional()
34
45
  }).optional()
35
46
  });
36
47
 
@@ -64,16 +75,28 @@ var OpenAILLM = class {
64
75
  this.openai = new OpenAI2({ apiKey: config.apiKey });
65
76
  this.model = config.model || "gpt-4-turbo-preview";
66
77
  }
67
- async generateResponse(messages, responseFormat) {
78
+ async generateResponse(messages, responseFormat, tools) {
68
79
  const completion = await this.openai.chat.completions.create({
69
80
  messages: messages.map((msg) => ({
70
81
  role: msg.role,
71
82
  content: msg.content
72
83
  })),
73
84
  model: this.model,
74
- response_format: responseFormat
85
+ response_format: responseFormat,
86
+ ...tools && { tools, tool_choice: "auto" }
75
87
  });
76
- return completion.choices[0].message.content || "";
88
+ const response = completion.choices[0].message;
89
+ if (response.tool_calls) {
90
+ return {
91
+ content: response.content || "",
92
+ role: response.role,
93
+ toolCalls: response.tool_calls.map((call) => ({
94
+ name: call.function.name,
95
+ arguments: call.function.arguments
96
+ }))
97
+ };
98
+ }
99
+ return response.content || "";
77
100
  }
78
101
  async generateChat(messages) {
79
102
  const completion = await this.openai.chat.completions.create({
@@ -95,37 +118,57 @@ var OpenAILLM = class {
95
118
  import OpenAI3 from "openai";
96
119
  var OpenAIStructuredLLM = class {
97
120
  constructor(config) {
98
- const apiKey = config.apiKey || process.env.OPENAI_API_KEY;
99
- if (!apiKey) {
100
- throw new Error("OpenAI API key is required");
101
- }
102
- const baseUrl = process.env.OPENAI_API_BASE || "https://api.openai.com/v1";
103
- this.client = new OpenAI3({ apiKey, baseURL: baseUrl });
104
- this.model = config.model || "gpt-4-0125-preview";
121
+ this.openai = new OpenAI3({ apiKey: config.apiKey });
122
+ this.model = config.model || "gpt-4-turbo-preview";
105
123
  }
106
- async generateResponse(messages, responseFormat) {
107
- const response = await this.client.chat.completions.create({
108
- model: this.model,
124
+ async generateResponse(messages, responseFormat, tools) {
125
+ const completion = await this.openai.chat.completions.create({
109
126
  messages: messages.map((msg) => ({
110
127
  role: msg.role,
111
128
  content: msg.content
112
129
  })),
113
- response_format: responseFormat
130
+ model: this.model,
131
+ ...tools ? {
132
+ tools: tools.map((tool) => ({
133
+ type: "function",
134
+ function: {
135
+ name: tool.function.name,
136
+ description: tool.function.description,
137
+ parameters: tool.function.parameters
138
+ }
139
+ })),
140
+ tool_choice: "auto"
141
+ } : responseFormat ? {
142
+ response_format: {
143
+ type: responseFormat.type
144
+ }
145
+ } : {}
114
146
  });
115
- return response.choices[0].message.content || "";
147
+ const response = completion.choices[0].message;
148
+ if (response.tool_calls) {
149
+ return {
150
+ content: response.content || "",
151
+ role: response.role,
152
+ toolCalls: response.tool_calls.map((call) => ({
153
+ name: call.function.name,
154
+ arguments: call.function.arguments
155
+ }))
156
+ };
157
+ }
158
+ return response.content || "";
116
159
  }
117
160
  async generateChat(messages) {
118
- const response = await this.client.chat.completions.create({
119
- model: this.model,
161
+ const completion = await this.openai.chat.completions.create({
120
162
  messages: messages.map((msg) => ({
121
163
  role: msg.role,
122
164
  content: msg.content
123
- }))
165
+ })),
166
+ model: this.model
124
167
  });
125
- const message = response.choices[0].message;
168
+ const response = completion.choices[0].message;
126
169
  return {
127
- content: message.content || "",
128
- role: message.role
170
+ content: response.content || "",
171
+ role: response.role
129
172
  };
130
173
  }
131
174
  };
@@ -966,7 +1009,7 @@ var EmbedderFactory = class {
966
1009
  };
967
1010
  var LLMFactory = class {
968
1011
  static create(provider, config) {
969
- switch (provider.toLowerCase()) {
1012
+ switch (provider) {
970
1013
  case "openai":
971
1014
  return new OpenAILLM(config);
972
1015
  case "openai_structured":
@@ -1032,7 +1075,7 @@ function getFactRetrievalMessages(parsedMessages) {
1032
1075
  Input: Me favourite movies are Inception and Interstellar.
1033
1076
  Output: {"facts" : ["Favourite movies are Inception and Interstellar"]}
1034
1077
 
1035
- Return the facts and preferences in a json format as shown above.
1078
+ 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.
1036
1079
 
1037
1080
  Remember the following:
1038
1081
  - Today's date is ${(/* @__PURE__ */ new Date()).toISOString().split("T")[0]}.
@@ -1041,16 +1084,16 @@ function getFactRetrievalMessages(parsedMessages) {
1041
1084
  - If the user asks where you fetched my information, answer that you found from publicly available sources on internet.
1042
1085
  - If you do not find anything relevant in the below conversation, you can return an empty list corresponding to the "facts" key.
1043
1086
  - Create the facts based on the user and assistant messages only. Do not pick anything from the system messages.
1044
- - 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.
1087
+ - 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.
1045
1088
  - DO NOT RETURN ANYTHING ELSE OTHER THAN THE JSON FORMAT.
1046
- - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALUD SUCH AS "\`\`\`json" OR "\`\`\`".
1089
+ - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALID SUCH AS "\`\`\`json" OR "\`\`\`".
1047
1090
  - You should detect the language of the user input and record the facts in the same language.
1048
1091
  - For basic factual statements, break them down into individual facts if they contain multiple pieces of information.
1049
1092
 
1050
- 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.
1093
+ 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.
1051
1094
  You should detect the language of the user input and record the facts in the same language.
1052
1095
  `;
1053
- 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.
1096
+ 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.
1054
1097
 
1055
1098
  Input:
1056
1099
  ${parsedMessages}`;
@@ -1213,14 +1256,14 @@ function getUpdateMemoryMessages(retrievedOldMemory, newRetrievedFacts) {
1213
1256
  ${JSON.stringify(newRetrievedFacts, null, 2)}
1214
1257
 
1215
1258
  Follow the instruction mentioned below:
1216
- - Do not return anything from the custom few shot prompts provided above.
1259
+ - Do not return anything from the custom few shot example prompts provided above.
1217
1260
  - If the current memory is empty, then you have to add the new retrieved facts to the memory.
1218
1261
  - 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.
1219
1262
  - If there is an addition, generate a new key and add the new memory corresponding to it.
1220
1263
  - If there is a deletion, the memory key-value pair should be removed from the memory.
1221
1264
  - If there is an update, the ID key should remain the same and only the value needs to be updated.
1222
1265
  - DO NOT RETURN ANYTHING ELSE OTHER THAN THE JSON FORMAT.
1223
- - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALUD SUCH AS "\`\`\`json" OR "\`\`\`".
1266
+ - DO NOT ADD ANY ADDITIONAL TEXT OR CODEBLOCK IN THE JSON FIELDS WHICH MAKE IT INVALID SUCH AS "\`\`\`json" OR "\`\`\`".
1224
1267
 
1225
1268
  Do not return anything except the JSON format.`;
1226
1269
  }
@@ -1320,6 +1363,21 @@ var DEFAULT_MEMORY_CONFIG = {
1320
1363
  model: "gpt-4-turbo-preview"
1321
1364
  }
1322
1365
  },
1366
+ enableGraph: false,
1367
+ graphStore: {
1368
+ provider: "neo4j",
1369
+ config: {
1370
+ url: process.env.NEO4J_URL || "neo4j://localhost:7687",
1371
+ username: process.env.NEO4J_USERNAME || "neo4j",
1372
+ password: process.env.NEO4J_PASSWORD || "password"
1373
+ },
1374
+ llm: {
1375
+ provider: "openai",
1376
+ config: {
1377
+ model: "gpt-4-turbo-preview"
1378
+ }
1379
+ }
1380
+ },
1323
1381
  historyDbPath: "memory.db"
1324
1382
  };
1325
1383
 
@@ -1353,12 +1411,742 @@ var ConfigManager = class {
1353
1411
  },
1354
1412
  historyDbPath: userConfig.historyDbPath || DEFAULT_MEMORY_CONFIG.historyDbPath,
1355
1413
  customPrompt: userConfig.customPrompt,
1356
- graphStore: userConfig.graphStore
1414
+ graphStore: {
1415
+ ...DEFAULT_MEMORY_CONFIG.graphStore,
1416
+ ...userConfig.graphStore
1417
+ },
1418
+ enableGraph: userConfig.enableGraph || DEFAULT_MEMORY_CONFIG.enableGraph
1357
1419
  };
1358
1420
  return MemoryConfigSchema.parse(mergedConfig);
1359
1421
  }
1360
1422
  };
1361
1423
 
1424
+ // src/oss/src/memory/graph_memory.ts
1425
+ import neo4j from "neo4j-driver";
1426
+
1427
+ // src/oss/src/utils/bm25.ts
1428
+ var BM25 = class {
1429
+ constructor(documents, k1 = 1.5, b = 0.75) {
1430
+ this.documents = documents;
1431
+ this.k1 = k1;
1432
+ this.b = b;
1433
+ this.docLengths = documents.map((doc) => doc.length);
1434
+ this.avgDocLength = this.docLengths.reduce((a, b2) => a + b2, 0) / documents.length;
1435
+ this.docFreq = /* @__PURE__ */ new Map();
1436
+ this.idf = /* @__PURE__ */ new Map();
1437
+ this.computeIdf();
1438
+ }
1439
+ computeIdf() {
1440
+ const N = this.documents.length;
1441
+ for (const doc of this.documents) {
1442
+ const terms = new Set(doc);
1443
+ for (const term of terms) {
1444
+ this.docFreq.set(term, (this.docFreq.get(term) || 0) + 1);
1445
+ }
1446
+ }
1447
+ for (const [term, freq] of this.docFreq) {
1448
+ this.idf.set(term, Math.log((N - freq + 0.5) / (freq + 0.5) + 1));
1449
+ }
1450
+ }
1451
+ score(query, doc, index) {
1452
+ let score = 0;
1453
+ const docLength = this.docLengths[index];
1454
+ for (const term of query) {
1455
+ const tf = doc.filter((t) => t === term).length;
1456
+ const idf = this.idf.get(term) || 0;
1457
+ score += idf * tf * (this.k1 + 1) / (tf + this.k1 * (1 - this.b + this.b * docLength / this.avgDocLength));
1458
+ }
1459
+ return score;
1460
+ }
1461
+ search(query) {
1462
+ const scores = this.documents.map((doc, idx) => ({
1463
+ doc,
1464
+ score: this.score(query, doc, idx)
1465
+ }));
1466
+ return scores.sort((a, b) => b.score - a.score).map((item) => item.doc);
1467
+ }
1468
+ };
1469
+
1470
+ // src/oss/src/graphs/tools.ts
1471
+ var RELATIONS_TOOL = {
1472
+ type: "function",
1473
+ function: {
1474
+ name: "establish_relationships",
1475
+ description: "Establish relationships among the entities based on the provided text.",
1476
+ parameters: {
1477
+ type: "object",
1478
+ properties: {
1479
+ entities: {
1480
+ type: "array",
1481
+ items: {
1482
+ type: "object",
1483
+ properties: {
1484
+ source: {
1485
+ type: "string",
1486
+ description: "The source entity of the relationship."
1487
+ },
1488
+ relationship: {
1489
+ type: "string",
1490
+ description: "The relationship between the source and destination entities."
1491
+ },
1492
+ destination: {
1493
+ type: "string",
1494
+ description: "The destination entity of the relationship."
1495
+ }
1496
+ },
1497
+ required: ["source", "relationship", "destination"],
1498
+ additionalProperties: false
1499
+ }
1500
+ }
1501
+ },
1502
+ required: ["entities"],
1503
+ additionalProperties: false
1504
+ }
1505
+ }
1506
+ };
1507
+ var EXTRACT_ENTITIES_TOOL = {
1508
+ type: "function",
1509
+ function: {
1510
+ name: "extract_entities",
1511
+ description: "Extract entities and their types from the text.",
1512
+ parameters: {
1513
+ type: "object",
1514
+ properties: {
1515
+ entities: {
1516
+ type: "array",
1517
+ items: {
1518
+ type: "object",
1519
+ properties: {
1520
+ entity: {
1521
+ type: "string",
1522
+ description: "The name or identifier of the entity."
1523
+ },
1524
+ entity_type: {
1525
+ type: "string",
1526
+ description: "The type or category of the entity."
1527
+ }
1528
+ },
1529
+ required: ["entity", "entity_type"],
1530
+ additionalProperties: false
1531
+ },
1532
+ description: "An array of entities with their types."
1533
+ }
1534
+ },
1535
+ required: ["entities"],
1536
+ additionalProperties: false
1537
+ }
1538
+ }
1539
+ };
1540
+ var DELETE_MEMORY_TOOL_GRAPH = {
1541
+ type: "function",
1542
+ function: {
1543
+ name: "delete_graph_memory",
1544
+ description: "Delete the relationship between two nodes.",
1545
+ parameters: {
1546
+ type: "object",
1547
+ properties: {
1548
+ source: {
1549
+ type: "string",
1550
+ description: "The identifier of the source node in the relationship."
1551
+ },
1552
+ relationship: {
1553
+ type: "string",
1554
+ description: "The existing relationship between the source and destination nodes that needs to be deleted."
1555
+ },
1556
+ destination: {
1557
+ type: "string",
1558
+ description: "The identifier of the destination node in the relationship."
1559
+ }
1560
+ },
1561
+ required: ["source", "relationship", "destination"],
1562
+ additionalProperties: false
1563
+ }
1564
+ }
1565
+ };
1566
+
1567
+ // src/oss/src/graphs/utils.ts
1568
+ var EXTRACT_RELATIONS_PROMPT = `
1569
+ 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:
1570
+
1571
+ 1. Extract only explicitly stated information from the text.
1572
+ 2. Establish relationships among the entities provided.
1573
+ 3. Use "USER_ID" as the source entity for any self-references (e.g., "I," "me," "my," etc.) in user messages.
1574
+ CUSTOM_PROMPT
1575
+
1576
+ Relationships:
1577
+ - Use consistent, general, and timeless relationship types.
1578
+ - Example: Prefer "professor" over "became_professor."
1579
+ - Relationships should only be established among the entities explicitly mentioned in the user message.
1580
+
1581
+ Entity Consistency:
1582
+ - Ensure that relationships are coherent and logically align with the context of the message.
1583
+ - Maintain consistent naming for entities across the extracted data.
1584
+
1585
+ 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.
1586
+
1587
+ Adhere strictly to these guidelines to ensure high-quality knowledge graph extraction.
1588
+ `;
1589
+ var DELETE_RELATIONS_SYSTEM_PROMPT = `
1590
+ 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.
1591
+ Input:
1592
+ 1. Existing Graph Memories: A list of current graph memories, each containing source, relationship, and destination information.
1593
+ 2. New Text: The new information to be integrated into the existing graph structure.
1594
+ 3. Use "USER_ID" as node for any self-references (e.g., "I," "me," "my," etc.) in user messages.
1595
+
1596
+ Guidelines:
1597
+ 1. Identification: Use the new information to evaluate existing relationships in the memory graph.
1598
+ 2. Deletion Criteria: Delete a relationship only if it meets at least one of these conditions:
1599
+ - Outdated or Inaccurate: The new information is more recent or accurate.
1600
+ - Contradictory: The new information conflicts with or negates the existing information.
1601
+ 3. DO NOT DELETE if their is a possibility of same type of relationship but different destination nodes.
1602
+ 4. Comprehensive Analysis:
1603
+ - Thoroughly examine each existing relationship against the new information and delete as necessary.
1604
+ - Multiple deletions may be required based on the new information.
1605
+ 5. Semantic Integrity:
1606
+ - Ensure that deletions maintain or improve the overall semantic structure of the graph.
1607
+ - Avoid deleting relationships that are NOT contradictory/outdated to the new information.
1608
+ 6. Temporal Awareness: Prioritize recency when timestamps are available.
1609
+ 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.
1610
+
1611
+ Note: DO NOT DELETE if their is a possibility of same type of relationship but different destination nodes.
1612
+
1613
+ For example:
1614
+ Existing Memory: alice -- loves_to_eat -- pizza
1615
+ New Information: Alice also loves to eat burger.
1616
+
1617
+ Do not delete in the above example because there is a possibility that Alice loves to eat both pizza and burger.
1618
+
1619
+ Memory Format:
1620
+ source -- relationship -- destination
1621
+
1622
+ Provide a list of deletion instructions, each specifying the relationship to be deleted.
1623
+ `;
1624
+ function getDeleteMessages(existingMemoriesString, data, userId) {
1625
+ return [
1626
+ DELETE_RELATIONS_SYSTEM_PROMPT.replace("USER_ID", userId),
1627
+ `Here are the existing memories: ${existingMemoriesString}
1628
+
1629
+ New Information: ${data}`
1630
+ ];
1631
+ }
1632
+
1633
+ // src/oss/src/utils/logger.ts
1634
+ var logger = {
1635
+ info: (message) => console.log(`[INFO] ${message}`),
1636
+ error: (message) => console.error(`[ERROR] ${message}`),
1637
+ debug: (message) => console.debug(`[DEBUG] ${message}`),
1638
+ warn: (message) => console.warn(`[WARN] ${message}`)
1639
+ };
1640
+
1641
+ // src/oss/src/memory/graph_memory.ts
1642
+ var MemoryGraph = class {
1643
+ constructor(config) {
1644
+ var _a, _b, _c, _d, _e, _f, _g, _h, _i;
1645
+ this.config = config;
1646
+ 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)) {
1647
+ throw new Error("Neo4j configuration is incomplete");
1648
+ }
1649
+ this.graph = neo4j.driver(
1650
+ config.graphStore.config.url,
1651
+ neo4j.auth.basic(
1652
+ config.graphStore.config.username,
1653
+ config.graphStore.config.password
1654
+ )
1655
+ );
1656
+ this.embeddingModel = EmbedderFactory.create(
1657
+ this.config.embedder.provider,
1658
+ this.config.embedder.config
1659
+ );
1660
+ this.llmProvider = "openai";
1661
+ if ((_g = this.config.llm) == null ? void 0 : _g.provider) {
1662
+ this.llmProvider = this.config.llm.provider;
1663
+ }
1664
+ if ((_i = (_h = this.config.graphStore) == null ? void 0 : _h.llm) == null ? void 0 : _i.provider) {
1665
+ this.llmProvider = this.config.graphStore.llm.provider;
1666
+ }
1667
+ this.llm = LLMFactory.create(this.llmProvider, this.config.llm.config);
1668
+ this.structuredLlm = LLMFactory.create(
1669
+ "openai_structured",
1670
+ this.config.llm.config
1671
+ );
1672
+ this.threshold = 0.7;
1673
+ }
1674
+ async add(data, filters) {
1675
+ const entityTypeMap = await this._retrieveNodesFromData(data, filters);
1676
+ const toBeAdded = await this._establishNodesRelationsFromData(
1677
+ data,
1678
+ filters,
1679
+ entityTypeMap
1680
+ );
1681
+ const searchOutput = await this._searchGraphDb(
1682
+ Object.keys(entityTypeMap),
1683
+ filters
1684
+ );
1685
+ const toBeDeleted = await this._getDeleteEntitiesFromSearchOutput(
1686
+ searchOutput,
1687
+ data,
1688
+ filters
1689
+ );
1690
+ const deletedEntities = await this._deleteEntities(
1691
+ toBeDeleted,
1692
+ filters["userId"]
1693
+ );
1694
+ const addedEntities = await this._addEntities(
1695
+ toBeAdded,
1696
+ filters["userId"],
1697
+ entityTypeMap
1698
+ );
1699
+ return {
1700
+ deleted_entities: deletedEntities,
1701
+ added_entities: addedEntities,
1702
+ relations: toBeAdded
1703
+ };
1704
+ }
1705
+ async search(query, filters, limit = 100) {
1706
+ const entityTypeMap = await this._retrieveNodesFromData(query, filters);
1707
+ const searchOutput = await this._searchGraphDb(
1708
+ Object.keys(entityTypeMap),
1709
+ filters
1710
+ );
1711
+ if (!searchOutput.length) {
1712
+ return [];
1713
+ }
1714
+ const searchOutputsSequence = searchOutput.map((item) => [
1715
+ item.source,
1716
+ item.relationship,
1717
+ item.destination
1718
+ ]);
1719
+ const bm25 = new BM25(searchOutputsSequence);
1720
+ const tokenizedQuery = query.split(" ");
1721
+ const rerankedResults = bm25.search(tokenizedQuery).slice(0, 5);
1722
+ const searchResults = rerankedResults.map((item) => ({
1723
+ source: item[0],
1724
+ relationship: item[1],
1725
+ destination: item[2]
1726
+ }));
1727
+ logger.info(`Returned ${searchResults.length} search results`);
1728
+ return searchResults;
1729
+ }
1730
+ async deleteAll(filters) {
1731
+ const session = this.graph.session();
1732
+ try {
1733
+ await session.run("MATCH (n {user_id: $user_id}) DETACH DELETE n", {
1734
+ user_id: filters["userId"]
1735
+ });
1736
+ } finally {
1737
+ await session.close();
1738
+ }
1739
+ }
1740
+ async getAll(filters, limit = 100) {
1741
+ const session = this.graph.session();
1742
+ try {
1743
+ const result = await session.run(
1744
+ `
1745
+ MATCH (n {user_id: $user_id})-[r]->(m {user_id: $user_id})
1746
+ RETURN n.name AS source, type(r) AS relationship, m.name AS target
1747
+ LIMIT toInteger($limit)
1748
+ `,
1749
+ { user_id: filters["userId"], limit: Math.floor(Number(limit)) }
1750
+ );
1751
+ const finalResults = result.records.map((record) => ({
1752
+ source: record.get("source"),
1753
+ relationship: record.get("relationship"),
1754
+ target: record.get("target")
1755
+ }));
1756
+ logger.info(`Retrieved ${finalResults.length} relationships`);
1757
+ return finalResults;
1758
+ } finally {
1759
+ await session.close();
1760
+ }
1761
+ }
1762
+ async _retrieveNodesFromData(data, filters) {
1763
+ const tools = [EXTRACT_ENTITIES_TOOL];
1764
+ const searchResults = await this.structuredLlm.generateResponse(
1765
+ [
1766
+ {
1767
+ role: "system",
1768
+ 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.`
1769
+ },
1770
+ { role: "user", content: data }
1771
+ ],
1772
+ { type: "json_object" },
1773
+ tools
1774
+ );
1775
+ let entityTypeMap = {};
1776
+ try {
1777
+ if (typeof searchResults !== "string" && searchResults.toolCalls) {
1778
+ for (const call of searchResults.toolCalls) {
1779
+ if (call.name === "extract_entities") {
1780
+ const args = JSON.parse(call.arguments);
1781
+ for (const item of args.entities) {
1782
+ entityTypeMap[item.entity] = item.entity_type;
1783
+ }
1784
+ }
1785
+ }
1786
+ }
1787
+ } catch (e) {
1788
+ logger.error(`Error in search tool: ${e}`);
1789
+ }
1790
+ entityTypeMap = Object.fromEntries(
1791
+ Object.entries(entityTypeMap).map(([k, v]) => [
1792
+ k.toLowerCase().replace(/ /g, "_"),
1793
+ v.toLowerCase().replace(/ /g, "_")
1794
+ ])
1795
+ );
1796
+ logger.debug(`Entity type map: ${JSON.stringify(entityTypeMap)}`);
1797
+ return entityTypeMap;
1798
+ }
1799
+ async _establishNodesRelationsFromData(data, filters, entityTypeMap) {
1800
+ var _a;
1801
+ let messages;
1802
+ if ((_a = this.config.graphStore) == null ? void 0 : _a.customPrompt) {
1803
+ messages = [
1804
+ {
1805
+ role: "system",
1806
+ content: EXTRACT_RELATIONS_PROMPT.replace(
1807
+ "USER_ID",
1808
+ filters["userId"]
1809
+ ).replace(
1810
+ "CUSTOM_PROMPT",
1811
+ `4. ${this.config.graphStore.customPrompt}`
1812
+ ) + "\nPlease provide your response in JSON format."
1813
+ },
1814
+ { role: "user", content: data }
1815
+ ];
1816
+ } else {
1817
+ messages = [
1818
+ {
1819
+ role: "system",
1820
+ content: EXTRACT_RELATIONS_PROMPT.replace("USER_ID", filters["userId"]) + "\nPlease provide your response in JSON format."
1821
+ },
1822
+ {
1823
+ role: "user",
1824
+ content: `List of entities: ${Object.keys(entityTypeMap)}.
1825
+
1826
+ Text: ${data}`
1827
+ }
1828
+ ];
1829
+ }
1830
+ const tools = [RELATIONS_TOOL];
1831
+ const extractedEntities = await this.structuredLlm.generateResponse(
1832
+ messages,
1833
+ { type: "json_object" },
1834
+ tools
1835
+ );
1836
+ let entities = [];
1837
+ if (typeof extractedEntities !== "string" && extractedEntities.toolCalls) {
1838
+ const toolCall = extractedEntities.toolCalls[0];
1839
+ if (toolCall && toolCall.arguments) {
1840
+ const args = JSON.parse(toolCall.arguments);
1841
+ entities = args.entities || [];
1842
+ }
1843
+ }
1844
+ entities = this._removeSpacesFromEntities(entities);
1845
+ logger.debug(`Extracted entities: ${JSON.stringify(entities)}`);
1846
+ return entities;
1847
+ }
1848
+ async _searchGraphDb(nodeList, filters, limit = 100) {
1849
+ const resultRelations = [];
1850
+ const session = this.graph.session();
1851
+ try {
1852
+ for (const node of nodeList) {
1853
+ const nEmbedding = await this.embeddingModel.embed(node);
1854
+ const cypher = `
1855
+ MATCH (n)
1856
+ WHERE n.embedding IS NOT NULL AND n.user_id = $user_id
1857
+ WITH n,
1858
+ round(reduce(dot = 0.0, i IN range(0, size(n.embedding)-1) | dot + n.embedding[i] * $n_embedding[i]) /
1859
+ (sqrt(reduce(l2 = 0.0, i IN range(0, size(n.embedding)-1) | l2 + n.embedding[i] * n.embedding[i])) *
1860
+ sqrt(reduce(l2 = 0.0, i IN range(0, size($n_embedding)-1) | l2 + $n_embedding[i] * $n_embedding[i]))), 4) AS similarity
1861
+ WHERE similarity >= $threshold
1862
+ MATCH (n)-[r]->(m)
1863
+ 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
1864
+ UNION
1865
+ MATCH (n)
1866
+ WHERE n.embedding IS NOT NULL AND n.user_id = $user_id
1867
+ WITH n,
1868
+ round(reduce(dot = 0.0, i IN range(0, size(n.embedding)-1) | dot + n.embedding[i] * $n_embedding[i]) /
1869
+ (sqrt(reduce(l2 = 0.0, i IN range(0, size(n.embedding)-1) | l2 + n.embedding[i] * n.embedding[i])) *
1870
+ sqrt(reduce(l2 = 0.0, i IN range(0, size($n_embedding)-1) | l2 + $n_embedding[i] * $n_embedding[i]))), 4) AS similarity
1871
+ WHERE similarity >= $threshold
1872
+ MATCH (m)-[r]->(n)
1873
+ 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
1874
+ ORDER BY similarity DESC
1875
+ LIMIT toInteger($limit)
1876
+ `;
1877
+ const result = await session.run(cypher, {
1878
+ n_embedding: nEmbedding,
1879
+ threshold: this.threshold,
1880
+ user_id: filters["userId"],
1881
+ limit: Math.floor(Number(limit))
1882
+ });
1883
+ resultRelations.push(
1884
+ ...result.records.map((record) => ({
1885
+ source: record.get("source"),
1886
+ source_id: record.get("source_id").toString(),
1887
+ relationship: record.get("relationship"),
1888
+ relation_id: record.get("relation_id").toString(),
1889
+ destination: record.get("destination"),
1890
+ destination_id: record.get("destination_id").toString(),
1891
+ similarity: record.get("similarity")
1892
+ }))
1893
+ );
1894
+ }
1895
+ } finally {
1896
+ await session.close();
1897
+ }
1898
+ return resultRelations;
1899
+ }
1900
+ async _getDeleteEntitiesFromSearchOutput(searchOutput, data, filters) {
1901
+ const searchOutputString = searchOutput.map(
1902
+ (item) => `${item.source} -- ${item.relationship} -- ${item.destination}`
1903
+ ).join("\n");
1904
+ const [systemPrompt, userPrompt] = getDeleteMessages(
1905
+ searchOutputString,
1906
+ data,
1907
+ filters["userId"]
1908
+ );
1909
+ const tools = [DELETE_MEMORY_TOOL_GRAPH];
1910
+ const memoryUpdates = await this.structuredLlm.generateResponse(
1911
+ [
1912
+ { role: "system", content: systemPrompt },
1913
+ { role: "user", content: userPrompt }
1914
+ ],
1915
+ { type: "json_object" },
1916
+ tools
1917
+ );
1918
+ const toBeDeleted = [];
1919
+ if (typeof memoryUpdates !== "string" && memoryUpdates.toolCalls) {
1920
+ for (const item of memoryUpdates.toolCalls) {
1921
+ if (item.name === "delete_graph_memory") {
1922
+ toBeDeleted.push(JSON.parse(item.arguments));
1923
+ }
1924
+ }
1925
+ }
1926
+ const cleanedToBeDeleted = this._removeSpacesFromEntities(toBeDeleted);
1927
+ logger.debug(
1928
+ `Deleted relationships: ${JSON.stringify(cleanedToBeDeleted)}`
1929
+ );
1930
+ return cleanedToBeDeleted;
1931
+ }
1932
+ async _deleteEntities(toBeDeleted, userId) {
1933
+ const results = [];
1934
+ const session = this.graph.session();
1935
+ try {
1936
+ for (const item of toBeDeleted) {
1937
+ const { source, destination, relationship } = item;
1938
+ const cypher = `
1939
+ MATCH (n {name: $source_name, user_id: $user_id})
1940
+ -[r:${relationship}]->
1941
+ (m {name: $dest_name, user_id: $user_id})
1942
+ DELETE r
1943
+ RETURN
1944
+ n.name AS source,
1945
+ m.name AS target,
1946
+ type(r) AS relationship
1947
+ `;
1948
+ const result = await session.run(cypher, {
1949
+ source_name: source,
1950
+ dest_name: destination,
1951
+ user_id: userId
1952
+ });
1953
+ results.push(result.records);
1954
+ }
1955
+ } finally {
1956
+ await session.close();
1957
+ }
1958
+ return results;
1959
+ }
1960
+ async _addEntities(toBeAdded, userId, entityTypeMap) {
1961
+ var _a, _b;
1962
+ const results = [];
1963
+ const session = this.graph.session();
1964
+ try {
1965
+ for (const item of toBeAdded) {
1966
+ const { source, destination, relationship } = item;
1967
+ const sourceType = entityTypeMap[source] || "unknown";
1968
+ const destinationType = entityTypeMap[destination] || "unknown";
1969
+ const sourceEmbedding = await this.embeddingModel.embed(source);
1970
+ const destEmbedding = await this.embeddingModel.embed(destination);
1971
+ const sourceNodeSearchResult = await this._searchSourceNode(
1972
+ sourceEmbedding,
1973
+ userId
1974
+ );
1975
+ const destinationNodeSearchResult = await this._searchDestinationNode(
1976
+ destEmbedding,
1977
+ userId
1978
+ );
1979
+ let cypher;
1980
+ let params;
1981
+ if (destinationNodeSearchResult.length === 0 && sourceNodeSearchResult.length > 0) {
1982
+ cypher = `
1983
+ MATCH (source)
1984
+ WHERE elementId(source) = $source_id
1985
+ MERGE (destination:${destinationType} {name: $destination_name, user_id: $user_id})
1986
+ ON CREATE SET
1987
+ destination.created = timestamp(),
1988
+ destination.embedding = $destination_embedding
1989
+ MERGE (source)-[r:${relationship}]->(destination)
1990
+ ON CREATE SET
1991
+ r.created = timestamp()
1992
+ RETURN source.name AS source, type(r) AS relationship, destination.name AS target
1993
+ `;
1994
+ params = {
1995
+ source_id: sourceNodeSearchResult[0].elementId,
1996
+ destination_name: destination,
1997
+ destination_embedding: destEmbedding,
1998
+ user_id: userId
1999
+ };
2000
+ } else if (destinationNodeSearchResult.length > 0 && sourceNodeSearchResult.length === 0) {
2001
+ cypher = `
2002
+ MATCH (destination)
2003
+ WHERE elementId(destination) = $destination_id
2004
+ MERGE (source:${sourceType} {name: $source_name, user_id: $user_id})
2005
+ ON CREATE SET
2006
+ source.created = timestamp(),
2007
+ source.embedding = $source_embedding
2008
+ MERGE (source)-[r:${relationship}]->(destination)
2009
+ ON CREATE SET
2010
+ r.created = timestamp()
2011
+ RETURN source.name AS source, type(r) AS relationship, destination.name AS target
2012
+ `;
2013
+ params = {
2014
+ destination_id: destinationNodeSearchResult[0].elementId,
2015
+ source_name: source,
2016
+ source_embedding: sourceEmbedding,
2017
+ user_id: userId
2018
+ };
2019
+ } else if (sourceNodeSearchResult.length > 0 && destinationNodeSearchResult.length > 0) {
2020
+ cypher = `
2021
+ MATCH (source)
2022
+ WHERE elementId(source) = $source_id
2023
+ MATCH (destination)
2024
+ WHERE elementId(destination) = $destination_id
2025
+ MERGE (source)-[r:${relationship}]->(destination)
2026
+ ON CREATE SET
2027
+ r.created_at = timestamp(),
2028
+ r.updated_at = timestamp()
2029
+ RETURN source.name AS source, type(r) AS relationship, destination.name AS target
2030
+ `;
2031
+ params = {
2032
+ source_id: (_a = sourceNodeSearchResult[0]) == null ? void 0 : _a.elementId,
2033
+ destination_id: (_b = destinationNodeSearchResult[0]) == null ? void 0 : _b.elementId,
2034
+ user_id: userId
2035
+ };
2036
+ } else {
2037
+ cypher = `
2038
+ MERGE (n:${sourceType} {name: $source_name, user_id: $user_id})
2039
+ ON CREATE SET n.created = timestamp(), n.embedding = $source_embedding
2040
+ ON MATCH SET n.embedding = $source_embedding
2041
+ MERGE (m:${destinationType} {name: $dest_name, user_id: $user_id})
2042
+ ON CREATE SET m.created = timestamp(), m.embedding = $dest_embedding
2043
+ ON MATCH SET m.embedding = $dest_embedding
2044
+ MERGE (n)-[rel:${relationship}]->(m)
2045
+ ON CREATE SET rel.created = timestamp()
2046
+ RETURN n.name AS source, type(rel) AS relationship, m.name AS target
2047
+ `;
2048
+ params = {
2049
+ source_name: source,
2050
+ dest_name: destination,
2051
+ source_embedding: sourceEmbedding,
2052
+ dest_embedding: destEmbedding,
2053
+ user_id: userId
2054
+ };
2055
+ }
2056
+ const result = await session.run(cypher, params);
2057
+ results.push(result.records);
2058
+ }
2059
+ } finally {
2060
+ await session.close();
2061
+ }
2062
+ return results;
2063
+ }
2064
+ _removeSpacesFromEntities(entityList) {
2065
+ return entityList.map((item) => ({
2066
+ ...item,
2067
+ source: item.source.toLowerCase().replace(/ /g, "_"),
2068
+ relationship: item.relationship.toLowerCase().replace(/ /g, "_"),
2069
+ destination: item.destination.toLowerCase().replace(/ /g, "_")
2070
+ }));
2071
+ }
2072
+ async _searchSourceNode(sourceEmbedding, userId, threshold = 0.9) {
2073
+ const session = this.graph.session();
2074
+ try {
2075
+ const cypher = `
2076
+ MATCH (source_candidate)
2077
+ WHERE source_candidate.embedding IS NOT NULL
2078
+ AND source_candidate.user_id = $user_id
2079
+
2080
+ WITH source_candidate,
2081
+ round(
2082
+ reduce(dot = 0.0, i IN range(0, size(source_candidate.embedding)-1) |
2083
+ dot + source_candidate.embedding[i] * $source_embedding[i]) /
2084
+ (sqrt(reduce(l2 = 0.0, i IN range(0, size(source_candidate.embedding)-1) |
2085
+ l2 + source_candidate.embedding[i] * source_candidate.embedding[i])) *
2086
+ sqrt(reduce(l2 = 0.0, i IN range(0, size($source_embedding)-1) |
2087
+ l2 + $source_embedding[i] * $source_embedding[i])))
2088
+ , 4) AS source_similarity
2089
+ WHERE source_similarity >= $threshold
2090
+
2091
+ WITH source_candidate, source_similarity
2092
+ ORDER BY source_similarity DESC
2093
+ LIMIT 1
2094
+
2095
+ RETURN elementId(source_candidate) as element_id
2096
+ `;
2097
+ const params = {
2098
+ source_embedding: sourceEmbedding,
2099
+ user_id: userId,
2100
+ threshold
2101
+ };
2102
+ const result = await session.run(cypher, params);
2103
+ return result.records.map((record) => ({
2104
+ elementId: record.get("element_id").toString()
2105
+ }));
2106
+ } finally {
2107
+ await session.close();
2108
+ }
2109
+ }
2110
+ async _searchDestinationNode(destinationEmbedding, userId, threshold = 0.9) {
2111
+ const session = this.graph.session();
2112
+ try {
2113
+ const cypher = `
2114
+ MATCH (destination_candidate)
2115
+ WHERE destination_candidate.embedding IS NOT NULL
2116
+ AND destination_candidate.user_id = $user_id
2117
+
2118
+ WITH destination_candidate,
2119
+ round(
2120
+ reduce(dot = 0.0, i IN range(0, size(destination_candidate.embedding)-1) |
2121
+ dot + destination_candidate.embedding[i] * $destination_embedding[i]) /
2122
+ (sqrt(reduce(l2 = 0.0, i IN range(0, size(destination_candidate.embedding)-1) |
2123
+ l2 + destination_candidate.embedding[i] * destination_candidate.embedding[i])) *
2124
+ sqrt(reduce(l2 = 0.0, i IN range(0, size($destination_embedding)-1) |
2125
+ l2 + $destination_embedding[i] * $destination_embedding[i])))
2126
+ , 4) AS destination_similarity
2127
+ WHERE destination_similarity >= $threshold
2128
+
2129
+ WITH destination_candidate, destination_similarity
2130
+ ORDER BY destination_similarity DESC
2131
+ LIMIT 1
2132
+
2133
+ RETURN elementId(destination_candidate) as element_id
2134
+ `;
2135
+ const params = {
2136
+ destination_embedding: destinationEmbedding,
2137
+ user_id: userId,
2138
+ threshold
2139
+ };
2140
+ const result = await session.run(cypher, params);
2141
+ return result.records.map((record) => ({
2142
+ elementId: record.get("element_id").toString()
2143
+ }));
2144
+ } finally {
2145
+ await session.close();
2146
+ }
2147
+ }
2148
+ };
2149
+
1362
2150
  // src/oss/src/memory/index.ts
1363
2151
  var Memory = class _Memory {
1364
2152
  constructor(config = {}) {
@@ -1379,6 +2167,10 @@ var Memory = class _Memory {
1379
2167
  this.db = new SQLiteManager(this.config.historyDbPath || ":memory:");
1380
2168
  this.collectionName = this.config.vectorStore.config.collectionName;
1381
2169
  this.apiVersion = this.config.version || "v1.0";
2170
+ this.enableGraph = this.config.enableGraph || false;
2171
+ if (this.enableGraph && this.config.graphStore) {
2172
+ this.graphMemory = new MemoryGraph(this.config);
2173
+ }
1382
2174
  }
1383
2175
  static fromConfig(configDict) {
1384
2176
  try {
@@ -1412,7 +2204,21 @@ var Memory = class _Memory {
1412
2204
  metadata,
1413
2205
  filters
1414
2206
  );
1415
- return { results: vectorStoreResult };
2207
+ let graphResult;
2208
+ if (this.graphMemory) {
2209
+ try {
2210
+ graphResult = await this.graphMemory.add(
2211
+ parsedMessages.map((m) => m.content).join("\n"),
2212
+ filters
2213
+ );
2214
+ } catch (error) {
2215
+ console.error("Error adding to graph memory:", error);
2216
+ }
2217
+ }
2218
+ return {
2219
+ results: vectorStoreResult,
2220
+ relations: graphResult == null ? void 0 : graphResult.relations
2221
+ };
1416
2222
  }
1417
2223
  async addToVectorStore(messages, metadata, filters) {
1418
2224
  const parsedMessages = messages.map((m) => m.content).join("\n");
@@ -1556,6 +2362,14 @@ ${parsedMessages}`] : getFactRetrievalMessages(parsedMessages);
1556
2362
  limit,
1557
2363
  filters
1558
2364
  );
2365
+ let graphResults;
2366
+ if (this.graphMemory) {
2367
+ try {
2368
+ graphResults = await this.graphMemory.search(query, filters);
2369
+ } catch (error) {
2370
+ console.error("Error searching graph memory:", error);
2371
+ }
2372
+ }
1559
2373
  const excludedKeys = /* @__PURE__ */ new Set([
1560
2374
  "userId",
1561
2375
  "agentId",
@@ -1577,7 +2391,10 @@ ${parsedMessages}`] : getFactRetrievalMessages(parsedMessages);
1577
2391
  ...mem.payload.agentId && { agentId: mem.payload.agentId },
1578
2392
  ...mem.payload.runId && { runId: mem.payload.runId }
1579
2393
  }));
1580
- return { results };
2394
+ return {
2395
+ results,
2396
+ relations: graphResults
2397
+ };
1581
2398
  }
1582
2399
  async update(memoryId, data) {
1583
2400
  const embedding = await this.embedder.embed(data);
@@ -1611,6 +2428,9 @@ ${parsedMessages}`] : getFactRetrievalMessages(parsedMessages);
1611
2428
  async reset() {
1612
2429
  await this.db.reset();
1613
2430
  await this.vectorStore.deleteCol();
2431
+ if (this.graphMemory) {
2432
+ await this.graphMemory.deleteAll({ userId: "default" });
2433
+ }
1614
2434
  this.vectorStore = VectorStoreFactory.create(
1615
2435
  this.config.vectorStore.provider,
1616
2436
  this.config.vectorStore.config