@mastra/pg 1.23.0 → 1.24.0-alpha.1

package/dist/index.js

@@ -7084,6 +7084,12 @@ var ExperimentsPG = class ExperimentsPG extends ExperimentsStorage {
 				name: "idx_experiment_results_org_project",
 				table: TABLE_EXPERIMENT_RESULTS,
 				columns: ["organizationId", "projectId"]
+			},
+			{
+				name: "idx_experiment_results_tags_gin",
+				table: TABLE_EXPERIMENT_RESULTS,
+				columns: ["tags"],
+				method: "gin"
 			}
 		];
 	}
@@ -7692,6 +7698,10 @@ var ExperimentsPG = class ExperimentsPG extends ExperimentsStorage {
 				conditions.push(`"status" = $${paramIndex++}`);
 				queryParams.push(args.status);
 			}
+			for (const tag of args.tags ?? []) {
+				conditions.push(`"tags" @> $${paramIndex++}::jsonb`);
+				queryParams.push(JSON.stringify([tag]));
+			}
 			if (args.filters) {
 				const { organizationId, projectId } = args.filters;
 				if (organizationId !== void 0) {
@@ -11040,13 +11050,36 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 	* asked for a page after the last row, because there is then no row to carry the
 	* count on.
 	*/
-	async #fetchMessagePage({ selectStatement, tableName, whereClause, orderByStatement, queryParams, perPageInput, perPage, offset }) {
+	async #fetchMessagePage({ selectStatement, tableName, whereClause, orderByStatement, queryParams, perPageInput, perPage, offset, includeTotal = true }) {
+		if (includeTotal === false && perPageInput !== false) {
+			const peekLimit = perPage + 1;
+			const rows = await this.#db.readClient.manyOrNone(`${selectStatement} FROM ${tableName} ${whereClause} ${orderByStatement} LIMIT $${queryParams.length + 1} OFFSET $${queryParams.length + 2}`, [
+				...queryParams,
+				peekLimit,
+				offset
+			]) || [];
+			const hasMore = rows.length > perPage;
+			const messages = hasMore ? rows.slice(0, perPage) : rows;
+			return {
+				total: offset + messages.length,
+				messages,
+				hasMore
+			};
+		}
 		const limitClause = perPageInput === false ? "" : ` LIMIT $${queryParams.length + 1} OFFSET $${queryParams.length + 2}`;
 		const dataParams = perPageInput === false ? queryParams : [
 			...queryParams,
 			perPage,
 			offset
 		];
+		if (includeTotal === false) {
+			const rows = await this.#db.readClient.manyOrNone(`${selectStatement} FROM ${tableName} ${whereClause} ${orderByStatement}`, dataParams) || [];
+			return {
+				total: rows.length,
+				messages: rows,
+				hasMore: false
+			};
+		}
 		const rows = await this.#db.readClient.manyOrNone(`${selectStatement}, (SELECT COUNT(*) FROM ${tableName} ${whereClause}) AS "__total" FROM ${tableName} ${whereClause} ${orderByStatement}${limitClause}`, dataParams) || [];
 		if (rows.length > 0) return {
 			total: Number(rows[0].__total),
@@ -11063,7 +11096,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 		};
 	}
 	async listMessages(args) {
-		const { threadId, resourceId, include, filter, perPage: perPageInput, page = 0, orderBy } = args;
+		const { threadId, resourceId, include, filter, perPage: perPageInput, page = 0, orderBy, includeTotal = true } = args;
 		const threadIds = (Array.isArray(threadId) ? threadId : [threadId]).filter((id) => typeof id === "string");
 		if (threadIds.length === 0 || threadIds.some((id) => !id.trim())) throw new MastraError({
 			id: createStorageErrorId("PG", "LIST_MESSAGES", "INVALID_THREAD_ID"),
@@ -11149,20 +11182,27 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			}) : null;
 			let total;
 			let messages;
+			let peekedHasMore;
 			if (metadataFilter) {
 				const filteredRows = (await this.#db.readClient.manyOrNone(`${selectStatement} FROM ${tableName} ${whereClause} ${orderByStatement}`, queryParams) || []).filter((row) => storageMessageMatchesMetadataFilter(row.content, metadataFilter));
 				total = filteredRows.length;
 				messages = perPageInput === false ? filteredRows : filteredRows.slice(offset, offset + perPage);
-			} else ({total, messages} = await this.#fetchMessagePage({
-				selectStatement,
-				tableName,
-				whereClause,
-				orderByStatement,
-				queryParams,
-				perPageInput,
-				perPage,
-				offset
-			}));
+			} else {
+				const pageResult = await this.#fetchMessagePage({
+					selectStatement,
+					tableName,
+					whereClause,
+					orderByStatement,
+					queryParams,
+					perPageInput,
+					perPage,
+					offset,
+					includeTotal
+				});
+				total = pageResult.total;
+				messages = pageResult.messages;
+				peekedHasMore = pageResult.hasMore;
+			}
 			const primaryPageCount = messages.length;
 			if (total === 0 && messages.length === 0 && (!include || include.length === 0)) return {
 				messages: [],
@@ -11187,7 +11227,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			const finalMessages = this._sortMessages(list.get.all.db(), field, direction);
 			const threadIdSet = new Set(threadIds);
 			const allThreadMessagesReturned = new Set(finalMessages.filter((m) => m.threadId && threadIdSet.has(m.threadId)).map((m) => m.id)).size >= total;
-			const hasMore = metadataFilter ? perPageInput !== false && offset + primaryPageCount < total : perPageInput !== false && !allThreadMessagesReturned && offset + perPage < total;
+			const hasMore = peekedHasMore !== void 0 ? peekedHasMore : metadataFilter ? perPageInput !== false && offset + primaryPageCount < total : perPageInput !== false && !allThreadMessagesReturned && offset + perPage < total;
 			return {
 				messages: finalMessages,
 				total,
@@ -11212,7 +11252,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 		}
 	}
 	async listMessagesByResourceId(args) {
-		const { resourceId, include, filter, perPage: perPageInput, page = 0, orderBy } = args;
+		const { resourceId, include, filter, perPage: perPageInput, page = 0, orderBy, includeTotal = true } = args;
 		if (!(resourceId !== void 0 && resourceId !== null && resourceId.trim() !== "")) throw new MastraError({
 			id: createStorageErrorId("PG", "LIST_MESSAGES_BY_RESOURCE_ID", "INVALID_QUERY"),
 			domain: ErrorDomain.STORAGE,
@@ -11295,20 +11335,27 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			}) : null;
 			let total;
 			let messages;
+			let peekedHasMore;
 			if (metadataFilter) {
 				const filteredRows = (await this.#db.readClient.manyOrNone(`${selectStatement} FROM ${tableName} ${whereClause} ${orderByStatement}`, queryParams) || []).filter((row) => storageMessageMatchesMetadataFilter(row.content, metadataFilter));
 				total = filteredRows.length;
 				messages = perPageInput === false ? filteredRows : filteredRows.slice(offset, offset + perPage);
-			} else ({total, messages} = await this.#fetchMessagePage({
-				selectStatement,
-				tableName,
-				whereClause,
-				orderByStatement,
-				queryParams,
-				perPageInput,
-				perPage,
-				offset
-			}));
+			} else {
+				const pageResult = await this.#fetchMessagePage({
+					selectStatement,
+					tableName,
+					whereClause,
+					orderByStatement,
+					queryParams,
+					perPageInput,
+					perPage,
+					offset,
+					includeTotal
+				});
+				total = pageResult.total;
+				messages = pageResult.messages;
+				peekedHasMore = pageResult.hasMore;
+			}
 			if (total === 0 && messages.length === 0 && (!include || include.length === 0)) return {
 				messages: [],
 				total: 0,
@@ -11330,7 +11377,7 @@ var MemoryPG = class MemoryPG extends MemoryStorage {
 			const messagesWithParsedContent = messages.map((row) => this.parseRow(row));
 			const list = new MessageList().add(messagesWithParsedContent, "memory");
 			const finalMessages = this._sortMessages(list.get.all.db(), field, direction);
-			const hasMore = perPageInput !== false && offset + perPage < total;
+			const hasMore = peekedHasMore !== void 0 ? peekedHasMore : perPageInput !== false && offset + perPage < total;
 			return {
 				messages: finalMessages,
 				total,
@@ -17020,12 +17067,32 @@ const TRACE_FIELDS = {
 	status: TRACE_STATUS_SQL
 };
 const SPAN_FIELDS = {
+	name: "s.\"name\"",
 	spanType: "s.\"spanType\"",
-	error: "s.\"error\""
+	model: "s.\"model\"",
+	provider: "s.\"provider\"",
+	startedAt: "s.\"startedAt\"",
+	endedAt: "s.\"endedAt\"",
+	durationMs: "s.\"durationMs\"",
+	status: "s.\"status\"",
+	error: "s.\"error\"",
+	entityType: "s.\"entityType\"",
+	entityId: "s.\"entityId\"",
+	entityName: "s.\"entityName\"",
+	entityVersionId: "s.\"entityVersionId\"",
+	parentEntityVersionId: "s.\"parentEntityVersionId\"",
+	rootEntityVersionId: "s.\"rootEntityVersionId\""
 };
 const SCORE_FIELDS = {
 	scorerId: "s.\"scorerId\"",
-	score: "s.\"score\""
+	scorerVersion: "s.\"scorerVersion\"",
+	scoreSource: "s.\"scoreSource\"",
+	score: "s.\"score\"",
+	timestamp: "s.\"timestamp\"",
+	spanId: "s.\"spanId\"",
+	entityVersionId: "s.\"entityVersionId\"",
+	parentEntityVersionId: "s.\"parentEntityVersionId\"",
+	rootEntityVersionId: "s.\"rootEntityVersionId\""
 };
 const TRACE_SELECT = `
   r."traceId" AS "traceId",
@@ -17043,15 +17110,18 @@ function fieldSql(registry, field) {
 	if (sql === void 0) throw new Error(`Unsupported trusted trace-query field: ${field}`);
 	return sql;
 }
+function isMetadataField(field) {
+	return field.startsWith("metadata.");
+}
 function placeholders(values, offset) {
 	return values.map((_, index) => `$${offset + index}`).join(", ");
 }
-function compileScalarPredicate(predicate, registry, parameterOffset) {
+function compileScalarPredicate(predicate, registry, parameterOffset, allowMetadata = false) {
 	if (predicate.type === "boolean") {
 		const values = [];
 		return {
 			sql: predicate.args.map((arg) => {
-				const compiled = compileScalarPredicate(arg, registry, parameterOffset + values.length);
+				const compiled = compileScalarPredicate(arg, registry, parameterOffset + values.length, allowMetadata);
 				values.push(...compiled.values);
 				return `(${compiled.sql})`;
 			}).join(predicate.operator === "and" ? " AND " : " OR "),
@@ -17059,26 +17129,36 @@ function compileScalarPredicate(predicate, registry, parameterOffset) {
 		};
 	}
 	if (predicate.type === "not") {
-		const compiled = compileScalarPredicate(predicate.arg, registry, parameterOffset);
+		const compiled = compileScalarPredicate(predicate.arg, registry, parameterOffset, allowMetadata);
 		return {
 			sql: `NOT (${compiled.sql})`,
 			values: compiled.values
 		};
 	}
-	const field = fieldSql(registry, predicate.field);
+	let field;
+	let fieldValues = [];
+	if (isMetadataField(predicate.field)) {
+		if (!allowMetadata) throw new Error(`Unsupported trusted trace-query field: ${predicate.field}`);
+		const keyParameter = `$${parameterOffset++}`;
+		field = `COALESCE(
+      CASE WHEN jsonb_typeof(r."metadataSearch" -> ${keyParameter}) = 'string' THEN r."metadataSearch" ->> ${keyParameter} END,
+      CASE WHEN jsonb_typeof(r."metadataRaw" -> ${keyParameter}) = 'string' THEN NULLIF(btrim(r."metadataRaw" ->> ${keyParameter}), '') END
+    )`;
+		fieldValues = [predicate.field.slice(9)];
+	} else field = fieldSql(registry, predicate.field);
 	if (predicate.type === "presence") return {
 		sql: `${field} IS ${predicate.operator === "exists" ? "NOT " : ""}NULL`,
-		values: []
+		values: fieldValues
 	};
 	if (predicate.type === "membership") {
 		const list = placeholders(predicate.values, parameterOffset);
 		if (predicate.operator === "in") return {
 			sql: `${field} IS NOT NULL AND ${field} IN (${list})`,
-			values: predicate.values
+			values: [...fieldValues, ...predicate.values]
 		};
 		return {
 			sql: `${field} IS NULL OR ${field} NOT IN (${list})`,
-			values: predicate.values
+			values: [...fieldValues, ...predicate.values]
 		};
 	}
 	const parameter = `$${parameterOffset}`;
@@ -17090,17 +17170,17 @@ function compileScalarPredicate(predicate, registry, parameterOffset) {
 	};
 	if (predicate.operator === "eq") return {
 		sql: `${field} IS NOT DISTINCT FROM ${parameter}`,
-		values: [predicate.value]
+		values: [...fieldValues, predicate.value]
 	};
 	if (predicate.operator === "ne") return {
 		sql: `${field} IS DISTINCT FROM ${parameter}`,
-		values: [predicate.value]
+		values: [...fieldValues, predicate.value]
 	};
 	const operator = operators[predicate.operator];
 	if (operator === void 0) throw new Error(`Unsupported trusted trace-query operator: ${predicate.operator}`);
 	return {
 		sql: `${field} IS NOT NULL AND ${field} ${operator} ${parameter}`,
-		values: [predicate.value]
+		values: [...fieldValues, predicate.value]
 	};
 }
 function latestRootPredicate$1(spanTable) {
@@ -17165,7 +17245,7 @@ function compilePredicate(predicate, parameterOffset) {
 			values: compiled.values
 		};
 	}
-	return compileScalarPredicate(predicate, TRACE_FIELDS, parameterOffset);
+	return compileScalarPredicate(predicate, TRACE_FIELDS, parameterOffset, true);
 }
 function compilePostgresTraceQuery(schema, plan) {
 	const spanTable = qualifiedTable(schema, TABLE_SPAN_EVENTS);
@@ -17185,14 +17265,43 @@ function compilePostgresTraceQuery(schema, plan) {
 	].join("\n      AND ")}
   )`];
 	if (relationCollections.has("spans")) ctes.push(`current_spans AS MATERIALIZED (
-    SELECT s."traceId", s."spanType", s."error"
+    SELECT
+      s."traceId",
+      s."name",
+      s."spanType",
+      CASE WHEN jsonb_typeof(s."attributes" -> 'model') = 'string' THEN s."attributes" ->> 'model' END AS "model",
+      CASE WHEN jsonb_typeof(s."attributes" -> 'provider') = 'string' THEN s."attributes" ->> 'provider' END AS "provider",
+      s."startedAt",
+      CASE WHEN s."isPending" THEN NULL ELSE s."endedAt" END AS "endedAt",
+      CASE
+        WHEN s."isPending" THEN NULL
+        ELSE EXTRACT(EPOCH FROM (s."endedAt" - s."startedAt")) * 1000
+      END AS "durationMs",
+      CASE WHEN s."error" IS NOT NULL THEN 'error' ELSE 'success' END AS "status",
+      s."error",
+      s."entityType",
+      s."entityId",
+      s."entityName",
+      s."entityVersionId",
+      s."parentEntityVersionId",
+      s."rootEntityVersionId"
     FROM ${spanTable} s
     WHERE s."traceId" IS NOT NULL
       AND s."traceId" IN (SELECT "traceId" FROM root_scope)
       AND ${latestSpanPredicate$1(spanTable)}
   )`);
 	if (relationCollections.has("scores")) ctes.push(`current_scores AS MATERIALIZED (
-    SELECT s."traceId", s."scorerId", s."score"
+    SELECT
+      s."traceId",
+      s."spanId",
+      s."timestamp",
+      s."scorerId",
+      s."scorerVersion",
+      s."scoreSource",
+      s."score",
+      s."entityVersionId",
+      s."parentEntityVersionId",
+      s."rootEntityVersionId"
     FROM ${scoreTable} s
     WHERE s."traceId" IS NOT NULL
       AND s."traceId" IN (SELECT "traceId" FROM root_scope)