From 68abe3ce79b739ca3ba907c6aa074d0f207623fe Mon Sep 17 00:00:00 2001 From: Morten Olsen Date: Thu, 11 Dec 2025 08:36:40 +0100 Subject: [PATCH] feat: improved zod handling --- packages/query-dsl/src/exports.ts | 3 +- packages/query-dsl/src/query-parser.codec.ts | 20 +++++++++++++ packages/query-dsl/src/query-parser.ts | 4 ++- .../document-chunks.mappings.ts | 8 ++--- .../document-chunks.schemas.ts | 29 ++++++++++++------- .../document-chunks/document-chunks.ts | 10 ++----- .../services/documents/documents.mapping.ts | 9 ++---- .../services/documents/documents.schemas.ts | 8 ++--- packages/runtime/src/utils/utils.schema.ts | 12 +++++++- packages/server/src/api.ts | 2 +- 10 files changed, 66 insertions(+), 39 deletions(-) create mode 100644 packages/query-dsl/src/query-parser.codec.ts diff --git a/packages/query-dsl/src/exports.ts b/packages/query-dsl/src/exports.ts index dd7a312..2cf0af0 100644 --- a/packages/query-dsl/src/exports.ts +++ b/packages/query-dsl/src/exports.ts @@ -1,3 +1,4 @@ export * from './query-parser.schemas.js'; -export { QueryParser } from './query-parser.js'; +export { QueryParser, queryParser } from './query-parser.js'; export * from './utils.filter.js'; +export * from './query-parser.codec.js'; diff --git a/packages/query-dsl/src/query-parser.codec.ts b/packages/query-dsl/src/query-parser.codec.ts new file mode 100644 index 0000000..59113ff --- /dev/null +++ b/packages/query-dsl/src/query-parser.codec.ts @@ -0,0 +1,20 @@ +import { z } from 'zod'; + +import { queryFilterSchema } from './query-parser.schemas.js'; +import { queryParser } from './query-parser.js'; + +const queryStringSchema: typeof queryFilterSchema = z + .codec(z.string(), queryFilterSchema, { + encode: (filter) => { + return queryParser.stringify(filter); + }, + decode: (input) => { + return queryParser.parse(input); + }, + // eslint-disable-next-line + }).meta({ id: 'QueryString', examples: ["metadata.foo = 'bar'"] }) as any; + +// eslint-disable-next-line +const querySchema: typeof queryFilterSchema = z.union([queryStringSchema, queryFilterSchema]) as any + +export { querySchema }; diff --git a/packages/query-dsl/src/query-parser.ts b/packages/query-dsl/src/query-parser.ts index 755662b..6d95c1b 100644 --- a/packages/query-dsl/src/query-parser.ts +++ b/packages/query-dsl/src/query-parser.ts @@ -19,4 +19,6 @@ class QueryParser { }; } -export { QueryParser }; +const queryParser = new QueryParser(); + +export { QueryParser, queryParser }; diff --git a/packages/runtime/src/services/document-chunks/document-chunks.mappings.ts b/packages/runtime/src/services/document-chunks/document-chunks.mappings.ts index 6a38e46..ed1520a 100644 --- a/packages/runtime/src/services/document-chunks/document-chunks.mappings.ts +++ b/packages/runtime/src/services/document-chunks/document-chunks.mappings.ts @@ -1,13 +1,11 @@ import type { TableRows } from '../database/database.js'; -import type { DocumentChunk } from './document-chunks.schemas.js'; +import { documentChunkWithDistanceSchema } from './document-chunks.schemas.js'; const mapFromDocumentChunkRow = ( row: TableRows['documentChunks'] & { - metadata: unknown; + metadata?: unknown; }, -): DocumentChunk => ({ - ...row, -}); +) => documentChunkWithDistanceSchema.parse(row); export { mapFromDocumentChunkRow }; diff --git a/packages/runtime/src/services/document-chunks/document-chunks.schemas.ts b/packages/runtime/src/services/document-chunks/document-chunks.schemas.ts index 1fe5932..bcdc38a 100644 --- a/packages/runtime/src/services/document-chunks/document-chunks.schemas.ts +++ b/packages/runtime/src/services/document-chunks/document-chunks.schemas.ts @@ -1,7 +1,7 @@ import { z } from 'zod'; -import { queryFilterSchema } from '@morten-olsen/stash-query-dsl'; +import { querySchema } from '@morten-olsen/stash-query-dsl'; -import { createListResultSchema, queryDSLSchema } from '../../utils/utils.schema.js'; +import { createListResultSchema } from '../../utils/utils.schema.js'; const documentChunkSchema = z .object({ @@ -19,21 +19,28 @@ const documentChunkFilterSchema = z limit: z.number().default(20), offset: z.number().default(0), semanticText: z.string().optional(), - conditions: z.union([queryDSLSchema, queryFilterSchema]).optional(), + conditions: querySchema.optional(), }) .meta({ id: 'DocumentChunkFilter' }); type DocumentChunkFilter = z.infer; -const documentChunksFindResultSchema = createListResultSchema( - documentChunkSchema - .extend({ - distance: z.number().optional(), - }) - .meta({ id: 'DocumentChunkWithDistance' }), -).meta({ id: 'DocumentChunkFindResult' }); +const documentChunkWithDistanceSchema = documentChunkSchema + .extend({ + distance: z.number().optional(), + }) + .meta({ id: 'DocumentChunkWithDistance' }); + +const documentChunksFindResultSchema = createListResultSchema(documentChunkWithDistanceSchema).meta({ + id: 'DocumentChunkFindResult', +}); type DocumentChunksFindResult = z.infer; export type { DocumentChunk, DocumentChunkFilter, DocumentChunksFindResult }; -export { documentChunkSchema, documentChunkFilterSchema, documentChunksFindResultSchema }; +export { + documentChunkSchema, + documentChunkFilterSchema, + documentChunksFindResultSchema, + documentChunkWithDistanceSchema, +}; diff --git a/packages/runtime/src/services/document-chunks/document-chunks.ts b/packages/runtime/src/services/document-chunks/document-chunks.ts index 3715e06..b687e12 100644 --- a/packages/runtime/src/services/document-chunks/document-chunks.ts +++ b/packages/runtime/src/services/document-chunks/document-chunks.ts @@ -1,5 +1,3 @@ -import { QueryParser } from '@morten-olsen/stash-query-dsl'; - import { DatabaseService, tableNames, type TableRows } from '../database/database.js'; import { EmbeddingsService } from '../embeddings/embeddings.js'; import type { Services } from '../../utils/utils.services.js'; @@ -44,11 +42,7 @@ class DocumentChunksService { query = query.orderBy('createdAt', 'desc'); } if (filter.conditions) { - const parser = this.#services.get(QueryParser); - query = applyQueryFilter( - query, - typeof filter.conditions === 'string' ? parser.parse(filter.conditions) : filter.conditions, - ); + query = applyQueryFilter(query, filter.conditions); } query = query.limit(filter.limit).offset(filter.offset); @@ -56,7 +50,7 @@ class DocumentChunksService { const items = await query; return { - items: items.map(mapFromDocumentChunkRow as ExplicitAny), + items: items.map(mapFromDocumentChunkRow), }; }; } diff --git a/packages/runtime/src/services/documents/documents.mapping.ts b/packages/runtime/src/services/documents/documents.mapping.ts index af5a036..4997c66 100644 --- a/packages/runtime/src/services/documents/documents.mapping.ts +++ b/packages/runtime/src/services/documents/documents.mapping.ts @@ -1,12 +1,7 @@ import type { TableRows } from '../database/database.js'; -import type { Document } from './documents.schemas.js'; +import { documentSchema, type Document } from './documents.schemas.js'; -const mapFromDocumentRow = (row: TableRows['documents']): Document => ({ - ...row, - createdAt: row.createdAt.toISOString?.(), - updatedAt: row.updatedAt.toISOString?.(), - deletedAt: row.deletedAt?.toISOString?.() || null, -}); +const mapFromDocumentRow = (row: TableRows['documents']): Document => documentSchema.parse(row); export { mapFromDocumentRow }; diff --git a/packages/runtime/src/services/documents/documents.schemas.ts b/packages/runtime/src/services/documents/documents.schemas.ts index 4dc6409..747ae3c 100644 --- a/packages/runtime/src/services/documents/documents.schemas.ts +++ b/packages/runtime/src/services/documents/documents.schemas.ts @@ -6,10 +6,10 @@ import { createListResultSchema, queryDSLSchema } from '../../utils/utils.schema const documentSchema = z .object({ id: z.guid(), - owner: z.string().nullable(), - createdAt: z.iso.datetime(), - updatedAt: z.iso.datetime(), - deletedAt: z.iso.datetime().nullable(), + owner: z.guid().nullable(), + createdAt: z.coerce.date(), + updatedAt: z.coerce.date(), + deletedAt: z.coerce.date().nullable(), contentType: z.string().nullable(), text: z.string().nullable(), source: z.string().nullable(), diff --git a/packages/runtime/src/utils/utils.schema.ts b/packages/runtime/src/utils/utils.schema.ts index 0b29aa9..989509e 100644 --- a/packages/runtime/src/utils/utils.schema.ts +++ b/packages/runtime/src/utils/utils.schema.ts @@ -1,5 +1,8 @@ +import { QueryParser } from '@morten-olsen/stash-query-dsl'; import { z, type ZodType } from 'zod'; +const parser = new QueryParser(); + const createListResultSchema = (schema: T) => z.object({ items: z.array(schema), @@ -8,8 +11,15 @@ const createListResultSchema = (schema: T) => const queryDSLSchema = z .string() .describe('Query DSL based filter') + .superRefine((value, context) => { + try { + parser.parse(value); + } catch (err) { + context.addIssue(String(err)); + } + }) .meta({ - id: 'QueryDSLString', + id: 'QueryDQLString', examples: ["metadata.foo = 'bar'"], }); diff --git a/packages/server/src/api.ts b/packages/server/src/api.ts index 360c161..e84d56e 100644 --- a/packages/server/src/api.ts +++ b/packages/server/src/api.ts @@ -31,7 +31,7 @@ const createApi = async (runtime: StashRuntime = new StashRuntime()) => { runtime.documents.subscribe({ filter: "metadata.foo = 'bar'", fn: (document) => { - console.log(document); + // console.log(document); }, }); const app = fastify().withTypeProvider();