diff --git a/src/tools/mongodb/mongodbSchemas.ts b/src/tools/mongodb/mongodbSchemas.ts index 392273ba..d9c16a09 100644 --- a/src/tools/mongodb/mongodbSchemas.ts +++ b/src/tools/mongodb/mongodbSchemas.ts @@ -5,29 +5,28 @@ export const zVoyageModels = z .enum(["voyage-3-large", "voyage-3.5", "voyage-3.5-lite", "voyage-code-3"]) .default("voyage-3-large"); -// Zod does not undestand JS boxed numbers (like Int32) as integer literals, -// so we preprocess them to unwrap them so Zod understands them. -function unboxNumber(v: unknown): number { - if (v && typeof v === "object" && typeof v.valueOf === "function") { - const n = Number(v.valueOf()); - if (!Number.isNaN(n)) return n; - } - return v as number; -} - export const zVoyageEmbeddingParameters = z.object({ + // OpenAPI JSON Schema supports enum only as string so the public facing + // parameters that are fed to LLM providers should expect the dimensions as + // stringified numbers which are then transformed to actual numbers. outputDimension: z - .preprocess( - unboxNumber, - z.union([z.literal(256), z.literal(512), z.literal(1024), z.literal(2048), z.literal(4096)]) - ) - .optional() - .default(1024), + .union([z.literal("256"), z.literal("512"), z.literal("1024"), z.literal("2048"), z.literal("4096")]) + .default("1024") + .transform((value): number => Number.parseInt(value)) + .optional(), outputDtype: z.enum(["float", "int8", "uint8", "binary", "ubinary"]).optional().default("float"), }); export const zVoyageAPIParameters = zVoyageEmbeddingParameters .extend({ + // Unlike public facing parameters, `zVoyageEmbeddingParameters`, the + // api parameters need to be correct number and because we do an + // additional parsing before calling the API, we override the + // outputDimension schema to expect a union of numbers. + outputDimension: z + .union([z.literal(256), z.literal(512), z.literal(1024), z.literal(2048), z.literal(4096)]) + .default(1024) + .optional(), inputType: z.enum(["query", "document"]), }) .strip(); diff --git a/tests/integration/tools/mongodb/create/insertMany.test.ts b/tests/integration/tools/mongodb/create/insertMany.test.ts index dda4ae71..2170efb9 100644 --- a/tests/integration/tools/mongodb/create/insertMany.test.ts +++ b/tests/integration/tools/mongodb/create/insertMany.test.ts @@ -203,7 +203,7 @@ describeWithMongoDB( // embeddingParameters so that we can simulate the idea // of unknown or mismatched quantization. - // embeddingParameters: { outputDimension: 256, + // embeddingParameters: { outputDimension: "256", // outputDtype: "float", model: "voyage-3-large", input: // [ // { @@ -558,7 +558,7 @@ describeWithMongoDB( documents: [{ title: "The Matrix" }], embeddingParameters: { model: "voyage-3.5-lite", - outputDimension: 256, + outputDimension: "256", input: [{ titleEmbeddings: "The Matrix" }], }, }, diff --git a/tests/integration/tools/mongodb/read/aggregate.test.ts b/tests/integration/tools/mongodb/read/aggregate.test.ts index ff45433d..f35ada51 100644 --- a/tests/integration/tools/mongodb/read/aggregate.test.ts +++ b/tests/integration/tools/mongodb/read/aggregate.test.ts @@ -444,7 +444,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", }, }, }, @@ -508,7 +508,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: dataType, }, }, @@ -575,7 +575,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: dataType, }, }, @@ -642,7 +642,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: dataType, }, }, @@ -709,7 +709,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: dataType, }, }, @@ -774,7 +774,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: "float", }, filter: { name: 10 }, @@ -837,7 +837,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: "float", }, filter: { name: 10 }, @@ -900,7 +900,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi limit: 10, embeddingParameters: { model: "voyage-3-large", - outputDimension: 256, + outputDimension: "256", outputDType: "float", }, filter: { name: 10 },