@mastra/pg 0.0.0-commonjs-20250227130920

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.
package/dist/index.cjs ADDED
@@ -0,0 +1,1043 @@
1
+ 'use strict';
2
+
3
+ var vector = require('@mastra/core/vector');
4
+ var pg = require('pg');
5
+ var filter = require('@mastra/core/filter');
6
+ var storage = require('@mastra/core/storage');
7
+ var pgPromise = require('pg-promise');
8
+
9
+ function _interopDefault (e) { return e && e.__esModule ? e : { default: e }; }
10
+
11
+ var pg__default = /*#__PURE__*/_interopDefault(pg);
12
+ var pgPromise__default = /*#__PURE__*/_interopDefault(pgPromise);
13
+
14
+ // src/vector/index.ts
15
+ var PGFilterTranslator = class extends filter.BaseFilterTranslator {
16
+ getSupportedOperators() {
17
+ return {
18
+ ...filter.BaseFilterTranslator.DEFAULT_OPERATORS,
19
+ custom: ["$contains", "$size"]
20
+ };
21
+ }
22
+ translate(filter) {
23
+ if (this.isEmpty(filter)) {
24
+ return filter;
25
+ }
26
+ this.validateFilter(filter);
27
+ return this.translateNode(filter);
28
+ }
29
+ translateNode(node, currentPath = "") {
30
+ const withPath = (result2) => currentPath ? { [currentPath]: result2 } : result2;
31
+ if (this.isPrimitive(node)) {
32
+ return withPath({ $eq: this.normalizeComparisonValue(node) });
33
+ }
34
+ if (Array.isArray(node)) {
35
+ return withPath({ $in: this.normalizeArrayValues(node) });
36
+ }
37
+ if (node instanceof RegExp) {
38
+ return withPath(this.translateRegexPattern(node.source, node.flags));
39
+ }
40
+ const entries = Object.entries(node);
41
+ const result = {};
42
+ if ("$options" in node && !("$regex" in node)) {
43
+ throw new Error("$options is not valid without $regex");
44
+ }
45
+ if ("$regex" in node) {
46
+ const options = node.$options || "";
47
+ return withPath(this.translateRegexPattern(node.$regex, options));
48
+ }
49
+ for (const [key, value] of entries) {
50
+ if (key === "$options") continue;
51
+ const newPath = currentPath ? `${currentPath}.${key}` : key;
52
+ if (this.isLogicalOperator(key)) {
53
+ result[key] = Array.isArray(value) ? value.map((filter) => this.translateNode(filter)) : this.translateNode(value);
54
+ } else if (this.isOperator(key)) {
55
+ if (this.isArrayOperator(key) && !Array.isArray(value) && key !== "$elemMatch") {
56
+ result[key] = [value];
57
+ } else if (this.isBasicOperator(key) && Array.isArray(value)) {
58
+ result[key] = JSON.stringify(value);
59
+ } else {
60
+ result[key] = value;
61
+ }
62
+ } else if (typeof value === "object" && value !== null) {
63
+ const hasOperators = Object.keys(value).some((k) => this.isOperator(k));
64
+ if (hasOperators) {
65
+ result[newPath] = this.translateNode(value);
66
+ } else {
67
+ Object.assign(result, this.translateNode(value, newPath));
68
+ }
69
+ } else {
70
+ result[newPath] = this.translateNode(value);
71
+ }
72
+ }
73
+ return result;
74
+ }
75
+ translateRegexPattern(pattern, options = "") {
76
+ if (!options) return { $regex: pattern };
77
+ const flags = options.split("").filter((f) => "imsux".includes(f)).join("");
78
+ return { $regex: flags ? `(?${flags})${pattern}` : pattern };
79
+ }
80
+ };
81
+
82
+ // src/vector/sql-builder.ts
83
+ var createBasicOperator = (symbol) => {
84
+ return (key, paramIndex) => ({
85
+ sql: `CASE
86
+ WHEN $${paramIndex}::text IS NULL THEN metadata#>>'{${handleKey(key)}}' IS ${symbol === "=" ? "" : "NOT"} NULL
87
+ ELSE metadata#>>'{${handleKey(key)}}' ${symbol} $${paramIndex}::text
88
+ END`,
89
+ needsValue: true
90
+ });
91
+ };
92
+ var createNumericOperator = (symbol) => {
93
+ return (key, paramIndex) => ({
94
+ sql: `(metadata#>>'{${handleKey(key)}}')::numeric ${symbol} $${paramIndex}`,
95
+ needsValue: true
96
+ });
97
+ };
98
+ function buildElemMatchConditions(value, paramIndex) {
99
+ if (typeof value !== "object" || Array.isArray(value)) {
100
+ throw new Error("$elemMatch requires an object with conditions");
101
+ }
102
+ const conditions = [];
103
+ const values = [];
104
+ Object.entries(value).forEach(([field, val]) => {
105
+ const nextParamIndex = paramIndex + values.length;
106
+ let paramOperator;
107
+ let paramKey;
108
+ let paramValue;
109
+ if (field.startsWith("$")) {
110
+ paramOperator = field;
111
+ paramKey = "";
112
+ paramValue = val;
113
+ } else if (typeof val === "object" && !Array.isArray(val)) {
114
+ const [op, opValue] = Object.entries(val || {})[0] || [];
115
+ paramOperator = op;
116
+ paramKey = field;
117
+ paramValue = opValue;
118
+ } else {
119
+ paramOperator = "$eq";
120
+ paramKey = field;
121
+ paramValue = val;
122
+ }
123
+ const operatorFn = FILTER_OPERATORS[paramOperator];
124
+ if (!operatorFn) {
125
+ throw new Error(`Invalid operator: ${paramOperator}`);
126
+ }
127
+ const result = operatorFn(paramKey, nextParamIndex, paramValue);
128
+ const sql = result.sql.replaceAll("metadata#>>", "elem#>>");
129
+ conditions.push(sql);
130
+ if (result.needsValue) {
131
+ values.push(paramValue);
132
+ }
133
+ });
134
+ return {
135
+ sql: conditions.join(" AND "),
136
+ values
137
+ };
138
+ }
139
+ var FILTER_OPERATORS = {
140
+ $eq: createBasicOperator("="),
141
+ $ne: createBasicOperator("!="),
142
+ $gt: createNumericOperator(">"),
143
+ $gte: createNumericOperator(">="),
144
+ $lt: createNumericOperator("<"),
145
+ $lte: createNumericOperator("<="),
146
+ // Array Operators
147
+ $in: (key, paramIndex) => ({
148
+ sql: `metadata#>>'{${handleKey(key)}}' = ANY($${paramIndex}::text[])`,
149
+ needsValue: true
150
+ }),
151
+ $nin: (key, paramIndex) => ({
152
+ sql: `metadata#>>'{${handleKey(key)}}' != ALL($${paramIndex}::text[])`,
153
+ needsValue: true
154
+ }),
155
+ $all: (key, paramIndex) => ({
156
+ sql: `CASE WHEN array_length($${paramIndex}::text[], 1) IS NULL THEN false
157
+ ELSE (metadata#>'{${handleKey(key)}}')::jsonb ?& $${paramIndex}::text[] END`,
158
+ needsValue: true
159
+ }),
160
+ $elemMatch: (key, paramIndex, value) => {
161
+ const { sql, values } = buildElemMatchConditions(value, paramIndex);
162
+ return {
163
+ sql: `(
164
+ CASE
165
+ WHEN jsonb_typeof(metadata->'${handleKey(key)}') = 'array' THEN
166
+ EXISTS (
167
+ SELECT 1
168
+ FROM jsonb_array_elements(metadata->'${handleKey(key)}') as elem
169
+ WHERE ${sql}
170
+ )
171
+ ELSE FALSE
172
+ END
173
+ )`,
174
+ needsValue: true,
175
+ transformValue: () => values
176
+ };
177
+ },
178
+ // Element Operators
179
+ $exists: (key) => ({
180
+ sql: `metadata ? '${key}'`,
181
+ needsValue: false
182
+ }),
183
+ // Logical Operators
184
+ $and: (key) => ({ sql: `(${key})`, needsValue: false }),
185
+ $or: (key) => ({ sql: `(${key})`, needsValue: false }),
186
+ $not: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
187
+ $nor: (key) => ({ sql: `NOT (${key})`, needsValue: false }),
188
+ // Regex Operators
189
+ $regex: (key, paramIndex) => ({
190
+ sql: `metadata#>>'{${handleKey(key)}}' ~ $${paramIndex}`,
191
+ needsValue: true
192
+ }),
193
+ $contains: (key, paramIndex) => ({
194
+ sql: `metadata @> $${paramIndex}::jsonb`,
195
+ needsValue: true,
196
+ transformValue: (value) => {
197
+ const parts = key.split(".");
198
+ return JSON.stringify(parts.reduceRight((value2, key2) => ({ [key2]: value2 }), value));
199
+ }
200
+ }),
201
+ $size: (key, paramIndex) => ({
202
+ sql: `(
203
+ CASE
204
+ WHEN jsonb_typeof(metadata#>'{${handleKey(key)}}') = 'array' THEN
205
+ jsonb_array_length(metadata#>'{${handleKey(key)}}') = $${paramIndex}
206
+ ELSE FALSE
207
+ END
208
+ )`,
209
+ needsValue: true
210
+ })
211
+ };
212
+ var handleKey = (key) => {
213
+ return key.replace(/\./g, ",");
214
+ };
215
+ function buildFilterQuery(filter, minScore) {
216
+ const values = [minScore];
217
+ function buildCondition(key, value, parentPath) {
218
+ if (["$and", "$or", "$not", "$nor"].includes(key)) {
219
+ return handleLogicalOperator(key, value);
220
+ }
221
+ if (!value || typeof value !== "object") {
222
+ values.push(value);
223
+ return `metadata#>>'{${handleKey(key)}}' = $${values.length}`;
224
+ }
225
+ const [[operator, operatorValue] = []] = Object.entries(value);
226
+ if (operator === "$not") {
227
+ const entries = Object.entries(operatorValue);
228
+ const conditions2 = entries.map(([nestedOp, nestedValue]) => {
229
+ if (!FILTER_OPERATORS[nestedOp]) {
230
+ throw new Error(`Invalid operator in $not condition: ${nestedOp}`);
231
+ }
232
+ const operatorFn2 = FILTER_OPERATORS[nestedOp];
233
+ const operatorResult2 = operatorFn2(key, values.length + 1);
234
+ if (operatorResult2.needsValue) {
235
+ values.push(nestedValue);
236
+ }
237
+ return operatorResult2.sql;
238
+ }).join(" AND ");
239
+ return `NOT (${conditions2})`;
240
+ }
241
+ const operatorFn = FILTER_OPERATORS[operator];
242
+ const operatorResult = operatorFn(key, values.length + 1, operatorValue);
243
+ if (operatorResult.needsValue) {
244
+ const transformedValue = operatorResult.transformValue ? operatorResult.transformValue(operatorValue) : operatorValue;
245
+ if (Array.isArray(transformedValue) && operator === "$elemMatch") {
246
+ values.push(...transformedValue);
247
+ } else {
248
+ values.push(transformedValue);
249
+ }
250
+ }
251
+ return operatorResult.sql;
252
+ }
253
+ function handleLogicalOperator(key, value, parentPath) {
254
+ if (key === "$not") {
255
+ const entries = Object.entries(value);
256
+ const conditions3 = entries.map(([fieldKey, fieldValue]) => buildCondition(fieldKey, fieldValue)).join(" AND ");
257
+ return `NOT (${conditions3})`;
258
+ }
259
+ if (!value || value.length === 0) {
260
+ switch (key) {
261
+ case "$and":
262
+ case "$nor":
263
+ return "true";
264
+ // Empty $and/$nor match everything
265
+ case "$or":
266
+ return "false";
267
+ // Empty $or matches nothing
268
+ default:
269
+ return "true";
270
+ }
271
+ }
272
+ const joinOperator = key === "$or" || key === "$nor" ? "OR" : "AND";
273
+ const conditions2 = value.map((f) => {
274
+ const entries = Object.entries(f);
275
+ if (entries.length === 0) return "";
276
+ const [firstKey, firstValue] = entries[0] || [];
277
+ if (["$and", "$or", "$not", "$nor"].includes(firstKey)) {
278
+ return buildCondition(firstKey, firstValue);
279
+ }
280
+ return entries.map(([k, v]) => buildCondition(k, v)).join(` ${joinOperator} `);
281
+ });
282
+ const joined = conditions2.join(` ${joinOperator} `);
283
+ const operatorFn = FILTER_OPERATORS[key];
284
+ return operatorFn(joined, 0, value).sql;
285
+ }
286
+ if (!filter) {
287
+ return { sql: "", values };
288
+ }
289
+ const conditions = Object.entries(filter).map(([key, value]) => buildCondition(key, value)).filter(Boolean).join(" AND ");
290
+ return { sql: conditions ? `WHERE ${conditions}` : "", values };
291
+ }
292
+
293
+ // src/vector/index.ts
294
+ var PgVector = class extends vector.MastraVector {
295
+ pool;
296
+ indexCache = /* @__PURE__ */ new Map();
297
+ constructor(connectionString) {
298
+ super();
299
+ const basePool = new pg__default.default.Pool({
300
+ connectionString,
301
+ max: 20,
302
+ // Maximum number of clients in the pool
303
+ idleTimeoutMillis: 3e4,
304
+ // Close idle connections after 30 seconds
305
+ connectionTimeoutMillis: 2e3
306
+ // Fail fast if can't connect
307
+ });
308
+ const telemetry = this.__getTelemetry();
309
+ this.pool = telemetry?.traceClass(basePool, {
310
+ spanNamePrefix: "pg-vector",
311
+ attributes: {
312
+ "vector.type": "postgres"
313
+ }
314
+ }) ?? basePool;
315
+ }
316
+ transformFilter(filter) {
317
+ const pgFilter = new PGFilterTranslator();
318
+ const translatedFilter = pgFilter.translate(filter ?? {});
319
+ return translatedFilter;
320
+ }
321
+ async getIndexInfo(indexName) {
322
+ if (!this.indexCache.has(indexName)) {
323
+ this.indexCache.set(indexName, await this.describeIndex(indexName));
324
+ }
325
+ return this.indexCache.get(indexName);
326
+ }
327
+ async query(indexName, queryVector, topK = 10, filter, includeVector = false, minScore = 0, options) {
328
+ const client = await this.pool.connect();
329
+ try {
330
+ const vectorStr = `[${queryVector.join(",")}]`;
331
+ const translatedFilter = this.transformFilter(filter);
332
+ const { sql: filterQuery, values: filterValues } = buildFilterQuery(translatedFilter, minScore);
333
+ const indexInfo = await this.getIndexInfo(indexName);
334
+ if (indexInfo.type === "hnsw") {
335
+ const calculatedEf = options?.ef ?? Math.max(topK, (indexInfo?.config?.m ?? 16) * topK);
336
+ const searchEf = Math.min(1e3, Math.max(1, calculatedEf));
337
+ await client.query(`SET LOCAL hnsw.ef_search = ${searchEf}`);
338
+ }
339
+ if (indexInfo.type === "ivfflat" && options?.probes) {
340
+ await client.query(`SET LOCAL ivfflat.probes = ${options.probes}`);
341
+ }
342
+ const query = `
343
+ WITH vector_scores AS (
344
+ SELECT
345
+ vector_id as id,
346
+ 1 - (embedding <=> '${vectorStr}'::vector) as score,
347
+ metadata
348
+ ${includeVector ? ", embedding" : ""}
349
+ FROM ${indexName}
350
+ ${filterQuery}
351
+ )
352
+ SELECT *
353
+ FROM vector_scores
354
+ WHERE score > $1
355
+ ORDER BY score DESC
356
+ LIMIT ${topK}`;
357
+ const result = await client.query(query, filterValues);
358
+ return result.rows.map(({ id, score, metadata, embedding }) => ({
359
+ id,
360
+ score,
361
+ metadata,
362
+ ...includeVector && embedding && { vector: JSON.parse(embedding) }
363
+ }));
364
+ } finally {
365
+ client.release();
366
+ }
367
+ }
368
+ async upsert(indexName, vectors, metadata, ids) {
369
+ const client = await this.pool.connect();
370
+ try {
371
+ await client.query("BEGIN");
372
+ const vectorIds = ids || vectors.map(() => crypto.randomUUID());
373
+ for (let i = 0; i < vectors.length; i++) {
374
+ const query = `
375
+ INSERT INTO ${indexName} (vector_id, embedding, metadata)
376
+ VALUES ($1, $2::vector, $3::jsonb)
377
+ ON CONFLICT (vector_id)
378
+ DO UPDATE SET
379
+ embedding = $2::vector,
380
+ metadata = $3::jsonb
381
+ RETURNING embedding::text
382
+ `;
383
+ await client.query(query, [vectorIds[i], `[${vectors[i]?.join(",")}]`, JSON.stringify(metadata?.[i] || {})]);
384
+ }
385
+ await client.query("COMMIT");
386
+ return vectorIds;
387
+ } catch (error) {
388
+ await client.query("ROLLBACK");
389
+ throw error;
390
+ } finally {
391
+ client.release();
392
+ }
393
+ }
394
+ async createIndex(indexName, dimension, metric = "cosine", indexConfig = {}, defineIndex = true) {
395
+ const client = await this.pool.connect();
396
+ try {
397
+ if (!indexName.match(/^[a-zA-Z_][a-zA-Z0-9_]*$/)) {
398
+ throw new Error("Invalid index name format");
399
+ }
400
+ if (!Number.isInteger(dimension) || dimension <= 0) {
401
+ throw new Error("Dimension must be a positive integer");
402
+ }
403
+ const extensionCheck = await client.query(`
404
+ SELECT EXISTS (
405
+ SELECT 1 FROM pg_available_extensions WHERE name = 'vector'
406
+ );
407
+ `);
408
+ if (!extensionCheck.rows[0].exists) {
409
+ throw new Error("PostgreSQL vector extension is not available. Please install it first.");
410
+ }
411
+ await client.query("CREATE EXTENSION IF NOT EXISTS vector");
412
+ await client.query(`
413
+ CREATE TABLE IF NOT EXISTS ${indexName} (
414
+ id SERIAL PRIMARY KEY,
415
+ vector_id TEXT UNIQUE NOT NULL,
416
+ embedding vector(${dimension}),
417
+ metadata JSONB DEFAULT '{}'::jsonb
418
+ );
419
+ `);
420
+ if (defineIndex) {
421
+ await this.defineIndex(indexName, metric, indexConfig);
422
+ }
423
+ } catch (error) {
424
+ console.error("Failed to create vector table:", error);
425
+ throw error;
426
+ } finally {
427
+ client.release();
428
+ }
429
+ }
430
+ /**
431
+ * @deprecated This function is deprecated. Use buildIndex instead
432
+ */
433
+ async defineIndex(indexName, metric = "cosine", indexConfig) {
434
+ return this.buildIndex(indexName, metric, indexConfig);
435
+ }
436
+ async buildIndex(indexName, metric = "cosine", indexConfig) {
437
+ const client = await this.pool.connect();
438
+ try {
439
+ await client.query(`DROP INDEX IF EXISTS ${indexName}_vector_idx`);
440
+ if (indexConfig.type === "flat") return;
441
+ const metricOp = metric === "cosine" ? "vector_cosine_ops" : metric === "euclidean" ? "vector_l2_ops" : "vector_ip_ops";
442
+ let indexSQL;
443
+ if (indexConfig.type === "hnsw") {
444
+ const m = indexConfig.hnsw?.m ?? 8;
445
+ const efConstruction = indexConfig.hnsw?.efConstruction ?? 32;
446
+ indexSQL = `
447
+ CREATE INDEX ${indexName}_vector_idx
448
+ ON ${indexName}
449
+ USING hnsw (embedding ${metricOp})
450
+ WITH (
451
+ m = ${m},
452
+ ef_construction = ${efConstruction}
453
+ )
454
+ `;
455
+ } else {
456
+ let lists;
457
+ if (indexConfig.ivf?.lists) {
458
+ lists = indexConfig.ivf.lists;
459
+ } else {
460
+ const size = (await client.query(`SELECT COUNT(*) FROM ${indexName}`)).rows[0].count;
461
+ lists = Math.max(100, Math.min(4e3, Math.floor(Math.sqrt(size) * 2)));
462
+ }
463
+ indexSQL = `
464
+ CREATE INDEX ${indexName}_vector_idx
465
+ ON ${indexName}
466
+ USING ivfflat (embedding ${metricOp})
467
+ WITH (lists = ${lists});
468
+ `;
469
+ }
470
+ await client.query(indexSQL);
471
+ this.indexCache.delete(indexName);
472
+ } finally {
473
+ client.release();
474
+ }
475
+ }
476
+ async listIndexes() {
477
+ const client = await this.pool.connect();
478
+ try {
479
+ const vectorTablesQuery = `
480
+ SELECT DISTINCT table_name
481
+ FROM information_schema.columns
482
+ WHERE table_schema = 'public'
483
+ AND udt_name = 'vector';
484
+ `;
485
+ const vectorTables = await client.query(vectorTablesQuery);
486
+ return vectorTables.rows.map((row) => row.table_name);
487
+ } finally {
488
+ client.release();
489
+ }
490
+ }
491
+ async describeIndex(indexName) {
492
+ const client = await this.pool.connect();
493
+ try {
494
+ const dimensionQuery = `
495
+ SELECT atttypmod as dimension
496
+ FROM pg_attribute
497
+ WHERE attrelid = $1::regclass
498
+ AND attname = 'embedding';
499
+ `;
500
+ const countQuery = `
501
+ SELECT COUNT(*) as count
502
+ FROM ${indexName};
503
+ `;
504
+ const indexQuery = `
505
+ SELECT
506
+ am.amname as index_method,
507
+ pg_get_indexdef(i.indexrelid) as index_def,
508
+ opclass.opcname as operator_class
509
+ FROM pg_index i
510
+ JOIN pg_class c ON i.indexrelid = c.oid
511
+ JOIN pg_am am ON c.relam = am.oid
512
+ JOIN pg_opclass opclass ON i.indclass[0] = opclass.oid
513
+ WHERE c.relname = '${indexName}_vector_idx';
514
+ `;
515
+ const [dimResult, countResult, indexResult] = await Promise.all([
516
+ client.query(dimensionQuery, [indexName]),
517
+ client.query(countQuery),
518
+ client.query(indexQuery)
519
+ ]);
520
+ const { index_method, index_def, operator_class } = indexResult.rows[0] || {
521
+ index_method: "flat",
522
+ index_def: "",
523
+ operator_class: "cosine"
524
+ };
525
+ const metric = operator_class.includes("l2") ? "euclidean" : operator_class.includes("ip") ? "dotproduct" : "cosine";
526
+ const config = {};
527
+ if (index_method === "hnsw") {
528
+ const m = index_def.match(/m\s*=\s*'?(\d+)'?/)?.[1];
529
+ const efConstruction = index_def.match(/ef_construction\s*=\s*'?(\d+)'?/)?.[1];
530
+ if (m) config.m = parseInt(m);
531
+ if (efConstruction) config.efConstruction = parseInt(efConstruction);
532
+ } else if (index_method === "ivfflat") {
533
+ const lists = index_def.match(/lists\s*=\s*'?(\d+)'?/)?.[1];
534
+ if (lists) config.lists = parseInt(lists);
535
+ }
536
+ return {
537
+ dimension: dimResult.rows[0].dimension,
538
+ count: parseInt(countResult.rows[0].count),
539
+ metric,
540
+ type: index_method,
541
+ config
542
+ };
543
+ } catch (e) {
544
+ await client.query("ROLLBACK");
545
+ throw new Error(`Failed to describe vector table: ${e.message}`);
546
+ } finally {
547
+ client.release();
548
+ }
549
+ }
550
+ async deleteIndex(indexName) {
551
+ const client = await this.pool.connect();
552
+ try {
553
+ await client.query(`DROP TABLE IF EXISTS ${indexName} CASCADE`);
554
+ } catch (error) {
555
+ await client.query("ROLLBACK");
556
+ throw new Error(`Failed to delete vector table: ${error.message}`);
557
+ } finally {
558
+ client.release();
559
+ }
560
+ }
561
+ async truncateIndex(indexName) {
562
+ const client = await this.pool.connect();
563
+ try {
564
+ await client.query(`TRUNCATE ${indexName}`);
565
+ } catch (e) {
566
+ await client.query("ROLLBACK");
567
+ throw new Error(`Failed to truncate vector table: ${e.message}`);
568
+ } finally {
569
+ client.release();
570
+ }
571
+ }
572
+ async disconnect() {
573
+ await this.pool.end();
574
+ }
575
+ };
576
+ var PostgresStore = class extends storage.MastraStorage {
577
+ db;
578
+ pgp;
579
+ constructor(config) {
580
+ super({ name: "PostgresStore" });
581
+ this.pgp = pgPromise__default.default();
582
+ this.db = this.pgp(
583
+ `connectionString` in config ? { connectionString: config.connectionString } : {
584
+ host: config.host,
585
+ port: config.port,
586
+ database: config.database,
587
+ user: config.user,
588
+ password: config.password
589
+ }
590
+ );
591
+ }
592
+ getEvalsByAgentName(_agentName, _type) {
593
+ throw new Error("Method not implemented.");
594
+ }
595
+ async batchInsert({ tableName, records }) {
596
+ try {
597
+ await this.db.query("BEGIN");
598
+ for (const record of records) {
599
+ await this.insert({ tableName, record });
600
+ }
601
+ await this.db.query("COMMIT");
602
+ } catch (error) {
603
+ console.error(`Error inserting into ${tableName}:`, error);
604
+ await this.db.query("ROLLBACK");
605
+ throw error;
606
+ }
607
+ }
608
+ async getTraces({
609
+ name,
610
+ scope,
611
+ page,
612
+ perPage,
613
+ attributes
614
+ }) {
615
+ let idx = 1;
616
+ const limit = perPage;
617
+ const offset = page * perPage;
618
+ const args = [];
619
+ const conditions = [];
620
+ if (name) {
621
+ conditions.push(`name LIKE CONCAT($${idx++}, '%')`);
622
+ }
623
+ if (scope) {
624
+ conditions.push(`scope = $${idx++}`);
625
+ }
626
+ if (attributes) {
627
+ Object.keys(attributes).forEach((key) => {
628
+ conditions.push(`attributes->>'${key}' = $${idx++}`);
629
+ });
630
+ }
631
+ const whereClause = conditions.length > 0 ? `WHERE ${conditions.join(" AND ")}` : "";
632
+ if (name) {
633
+ args.push(name);
634
+ }
635
+ if (scope) {
636
+ args.push(scope);
637
+ }
638
+ if (attributes) {
639
+ for (const [_key, value] of Object.entries(attributes)) {
640
+ args.push(value);
641
+ }
642
+ }
643
+ console.log(
644
+ "QUERY",
645
+ `SELECT * FROM ${storage.MastraStorage.TABLE_TRACES} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`,
646
+ args
647
+ );
648
+ const result = await this.db.manyOrNone(
649
+ `SELECT * FROM ${storage.MastraStorage.TABLE_TRACES} ${whereClause} ORDER BY "createdAt" DESC LIMIT ${limit} OFFSET ${offset}`,
650
+ args
651
+ );
652
+ if (!result) {
653
+ return [];
654
+ }
655
+ return result.map((row) => ({
656
+ id: row.id,
657
+ parentSpanId: row.parentSpanId,
658
+ traceId: row.traceId,
659
+ name: row.name,
660
+ scope: row.scope,
661
+ kind: row.kind,
662
+ status: row.status,
663
+ events: row.events,
664
+ links: row.links,
665
+ attributes: row.attributes,
666
+ startTime: row.startTime,
667
+ endTime: row.endTime,
668
+ other: row.other,
669
+ createdAt: row.createdAt
670
+ }));
671
+ }
672
+ async createTable({
673
+ tableName,
674
+ schema
675
+ }) {
676
+ try {
677
+ const columns = Object.entries(schema).map(([name, def]) => {
678
+ const constraints = [];
679
+ if (def.primaryKey) constraints.push("PRIMARY KEY");
680
+ if (!def.nullable) constraints.push("NOT NULL");
681
+ return `"${name}" ${def.type.toUpperCase()} ${constraints.join(" ")}`;
682
+ }).join(",\n");
683
+ const sql = `
684
+ CREATE TABLE IF NOT EXISTS ${tableName} (
685
+ ${columns}
686
+ );
687
+ ${tableName === storage.MastraStorage.TABLE_WORKFLOW_SNAPSHOT ? `
688
+ DO $$ BEGIN
689
+ IF NOT EXISTS (
690
+ SELECT 1 FROM pg_constraint WHERE conname = 'mastra_workflow_snapshot_workflow_name_run_id_key'
691
+ ) THEN
692
+ ALTER TABLE ${tableName}
693
+ ADD CONSTRAINT mastra_workflow_snapshot_workflow_name_run_id_key
694
+ UNIQUE (workflow_name, run_id);
695
+ END IF;
696
+ END $$;
697
+ ` : ""}
698
+ `;
699
+ await this.db.none(sql);
700
+ } catch (error) {
701
+ console.error(`Error creating table ${tableName}:`, error);
702
+ throw error;
703
+ }
704
+ }
705
+ async clearTable({ tableName }) {
706
+ try {
707
+ await this.db.none(`TRUNCATE TABLE ${tableName} CASCADE`);
708
+ } catch (error) {
709
+ console.error(`Error clearing table ${tableName}:`, error);
710
+ throw error;
711
+ }
712
+ }
713
+ async insert({ tableName, record }) {
714
+ try {
715
+ const columns = Object.keys(record);
716
+ const values = Object.values(record);
717
+ const placeholders = values.map((_, i) => `$${i + 1}`).join(", ");
718
+ await this.db.none(
719
+ `INSERT INTO ${tableName} (${columns.map((c) => `"${c}"`).join(", ")}) VALUES (${placeholders})`,
720
+ values
721
+ );
722
+ } catch (error) {
723
+ console.error(`Error inserting into ${tableName}:`, error);
724
+ throw error;
725
+ }
726
+ }
727
+ async load({ tableName, keys }) {
728
+ try {
729
+ const keyEntries = Object.entries(keys);
730
+ const conditions = keyEntries.map(([key], index) => `"${key}" = $${index + 1}`).join(" AND ");
731
+ const values = keyEntries.map(([_, value]) => value);
732
+ const result = await this.db.oneOrNone(`SELECT * FROM ${tableName} WHERE ${conditions}`, values);
733
+ if (!result) {
734
+ return null;
735
+ }
736
+ if (tableName === storage.MastraStorage.TABLE_WORKFLOW_SNAPSHOT) {
737
+ const snapshot = result;
738
+ if (typeof snapshot.snapshot === "string") {
739
+ snapshot.snapshot = JSON.parse(snapshot.snapshot);
740
+ }
741
+ return snapshot;
742
+ }
743
+ return result;
744
+ } catch (error) {
745
+ console.error(`Error loading from ${tableName}:`, error);
746
+ throw error;
747
+ }
748
+ }
749
+ async getThreadById({ threadId }) {
750
+ try {
751
+ const thread = await this.db.oneOrNone(
752
+ `SELECT
753
+ id,
754
+ "resourceId",
755
+ title,
756
+ metadata,
757
+ "createdAt",
758
+ "updatedAt"
759
+ FROM "${storage.MastraStorage.TABLE_THREADS}"
760
+ WHERE id = $1`,
761
+ [threadId]
762
+ );
763
+ if (!thread) {
764
+ return null;
765
+ }
766
+ return {
767
+ ...thread,
768
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
769
+ createdAt: thread.createdAt,
770
+ updatedAt: thread.updatedAt
771
+ };
772
+ } catch (error) {
773
+ console.error(`Error getting thread ${threadId}:`, error);
774
+ throw error;
775
+ }
776
+ }
777
+ async getThreadsByResourceId({ resourceId }) {
778
+ try {
779
+ const threads = await this.db.manyOrNone(
780
+ `SELECT
781
+ id,
782
+ "resourceId",
783
+ title,
784
+ metadata,
785
+ "createdAt",
786
+ "updatedAt"
787
+ FROM "${storage.MastraStorage.TABLE_THREADS}"
788
+ WHERE "resourceId" = $1`,
789
+ [resourceId]
790
+ );
791
+ return threads.map((thread) => ({
792
+ ...thread,
793
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
794
+ createdAt: thread.createdAt,
795
+ updatedAt: thread.updatedAt
796
+ }));
797
+ } catch (error) {
798
+ console.error(`Error getting threads for resource ${resourceId}:`, error);
799
+ throw error;
800
+ }
801
+ }
802
+ async saveThread({ thread }) {
803
+ try {
804
+ await this.db.none(
805
+ `INSERT INTO "${storage.MastraStorage.TABLE_THREADS}" (
806
+ id,
807
+ "resourceId",
808
+ title,
809
+ metadata,
810
+ "createdAt",
811
+ "updatedAt"
812
+ ) VALUES ($1, $2, $3, $4, $5, $6)
813
+ ON CONFLICT (id) DO UPDATE SET
814
+ "resourceId" = EXCLUDED."resourceId",
815
+ title = EXCLUDED.title,
816
+ metadata = EXCLUDED.metadata,
817
+ "createdAt" = EXCLUDED."createdAt",
818
+ "updatedAt" = EXCLUDED."updatedAt"`,
819
+ [
820
+ thread.id,
821
+ thread.resourceId,
822
+ thread.title,
823
+ thread.metadata ? JSON.stringify(thread.metadata) : null,
824
+ thread.createdAt,
825
+ thread.updatedAt
826
+ ]
827
+ );
828
+ return thread;
829
+ } catch (error) {
830
+ console.error("Error saving thread:", error);
831
+ throw error;
832
+ }
833
+ }
834
+ async updateThread({
835
+ id,
836
+ title,
837
+ metadata
838
+ }) {
839
+ try {
840
+ const existingThread = await this.getThreadById({ threadId: id });
841
+ if (!existingThread) {
842
+ throw new Error(`Thread ${id} not found`);
843
+ }
844
+ const mergedMetadata = {
845
+ ...existingThread.metadata,
846
+ ...metadata
847
+ };
848
+ const thread = await this.db.one(
849
+ `UPDATE "${storage.MastraStorage.TABLE_THREADS}"
850
+ SET title = $1,
851
+ metadata = $2,
852
+ "updatedAt" = $3
853
+ WHERE id = $4
854
+ RETURNING *`,
855
+ [title, mergedMetadata, (/* @__PURE__ */ new Date()).toISOString(), id]
856
+ );
857
+ return {
858
+ ...thread,
859
+ metadata: typeof thread.metadata === "string" ? JSON.parse(thread.metadata) : thread.metadata,
860
+ createdAt: thread.createdAt,
861
+ updatedAt: thread.updatedAt
862
+ };
863
+ } catch (error) {
864
+ console.error("Error updating thread:", error);
865
+ throw error;
866
+ }
867
+ }
868
+ async deleteThread({ threadId }) {
869
+ try {
870
+ await this.db.tx(async (t) => {
871
+ await t.none(`DELETE FROM "${storage.MastraStorage.TABLE_MESSAGES}" WHERE thread_id = $1`, [threadId]);
872
+ await t.none(`DELETE FROM "${storage.MastraStorage.TABLE_THREADS}" WHERE id = $1`, [threadId]);
873
+ });
874
+ } catch (error) {
875
+ console.error("Error deleting thread:", error);
876
+ throw error;
877
+ }
878
+ }
879
+ async getMessages({ threadId, selectBy }) {
880
+ try {
881
+ const messages = [];
882
+ const limit = typeof selectBy?.last === `number` ? selectBy.last : 40;
883
+ const include = selectBy?.include || [];
884
+ if (include.length) {
885
+ const includeResult = await this.db.manyOrNone(
886
+ `
887
+ WITH ordered_messages AS (
888
+ SELECT
889
+ *,
890
+ ROW_NUMBER() OVER (ORDER BY "createdAt") as row_num
891
+ FROM "${storage.MastraStorage.TABLE_MESSAGES}"
892
+ WHERE thread_id = $1
893
+ )
894
+ SELECT DISTINCT ON (m.id)
895
+ m.id,
896
+ m.content,
897
+ m.role,
898
+ m.type,
899
+ m."createdAt",
900
+ m.thread_id AS "threadId"
901
+ FROM ordered_messages m
902
+ WHERE m.id = ANY($2)
903
+ OR EXISTS (
904
+ SELECT 1 FROM ordered_messages target
905
+ WHERE target.id = ANY($2)
906
+ AND (
907
+ -- Get previous messages based on the max withPreviousMessages
908
+ (m.row_num >= target.row_num - $3 AND m.row_num < target.row_num)
909
+ OR
910
+ -- Get next messages based on the max withNextMessages
911
+ (m.row_num <= target.row_num + $4 AND m.row_num > target.row_num)
912
+ )
913
+ )
914
+ ORDER BY m.id, m."createdAt"
915
+ `,
916
+ [
917
+ threadId,
918
+ include.map((i) => i.id),
919
+ Math.max(...include.map((i) => i.withPreviousMessages || 0)),
920
+ Math.max(...include.map((i) => i.withNextMessages || 0))
921
+ ]
922
+ );
923
+ messages.push(...includeResult);
924
+ }
925
+ const result = await this.db.manyOrNone(
926
+ `
927
+ SELECT
928
+ id,
929
+ content,
930
+ role,
931
+ type,
932
+ "createdAt",
933
+ thread_id AS "threadId"
934
+ FROM "${storage.MastraStorage.TABLE_MESSAGES}"
935
+ WHERE thread_id = $1
936
+ AND id != ALL($2)
937
+ ORDER BY "createdAt" DESC
938
+ LIMIT $3
939
+ `,
940
+ [threadId, messages.map((m) => m.id), limit]
941
+ );
942
+ messages.push(...result);
943
+ messages.sort((a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime());
944
+ messages.forEach((message) => {
945
+ if (typeof message.content === "string") {
946
+ try {
947
+ message.content = JSON.parse(message.content);
948
+ } catch {
949
+ }
950
+ }
951
+ });
952
+ return messages;
953
+ } catch (error) {
954
+ console.error("Error getting messages:", error);
955
+ throw error;
956
+ }
957
+ }
958
+ async saveMessages({ messages }) {
959
+ if (messages.length === 0) return messages;
960
+ try {
961
+ const threadId = messages[0]?.threadId;
962
+ if (!threadId) {
963
+ throw new Error("Thread ID is required");
964
+ }
965
+ const thread = await this.getThreadById({ threadId });
966
+ if (!thread) {
967
+ throw new Error(`Thread ${threadId} not found`);
968
+ }
969
+ await this.db.tx(async (t) => {
970
+ for (const message of messages) {
971
+ await t.none(
972
+ `INSERT INTO "${storage.MastraStorage.TABLE_MESSAGES}" (id, thread_id, content, "createdAt", role, type)
973
+ VALUES ($1, $2, $3, $4, $5, $6)`,
974
+ [
975
+ message.id,
976
+ threadId,
977
+ typeof message.content === "string" ? message.content : JSON.stringify(message.content),
978
+ message.createdAt || (/* @__PURE__ */ new Date()).toISOString(),
979
+ message.role,
980
+ message.type
981
+ ]
982
+ );
983
+ }
984
+ });
985
+ return messages;
986
+ } catch (error) {
987
+ console.error("Error saving messages:", error);
988
+ throw error;
989
+ }
990
+ }
991
+ async persistWorkflowSnapshot({
992
+ workflowName,
993
+ runId,
994
+ snapshot
995
+ }) {
996
+ try {
997
+ const now = (/* @__PURE__ */ new Date()).toISOString();
998
+ await this.db.none(
999
+ `INSERT INTO "${storage.MastraStorage.TABLE_WORKFLOW_SNAPSHOT}" (
1000
+ workflow_name,
1001
+ run_id,
1002
+ snapshot,
1003
+ "createdAt",
1004
+ "updatedAt"
1005
+ ) VALUES ($1, $2, $3, $4, $5)
1006
+ ON CONFLICT (workflow_name, run_id) DO UPDATE
1007
+ SET snapshot = EXCLUDED.snapshot,
1008
+ "updatedAt" = EXCLUDED."updatedAt"`,
1009
+ [workflowName, runId, JSON.stringify(snapshot), now, now]
1010
+ );
1011
+ } catch (error) {
1012
+ console.error("Error persisting workflow snapshot:", error);
1013
+ throw error;
1014
+ }
1015
+ }
1016
+ async loadWorkflowSnapshot({
1017
+ workflowName,
1018
+ runId
1019
+ }) {
1020
+ try {
1021
+ const result = await this.load({
1022
+ tableName: storage.MastraStorage.TABLE_WORKFLOW_SNAPSHOT,
1023
+ keys: {
1024
+ workflow_name: workflowName,
1025
+ run_id: runId
1026
+ }
1027
+ });
1028
+ if (!result) {
1029
+ return null;
1030
+ }
1031
+ return result.snapshot;
1032
+ } catch (error) {
1033
+ console.error("Error loading workflow snapshot:", error);
1034
+ throw error;
1035
+ }
1036
+ }
1037
+ async close() {
1038
+ this.pgp.end();
1039
+ }
1040
+ };
1041
+
1042
+ exports.PgVector = PgVector;
1043
+ exports.PostgresStore = PostgresStore;