Skip to content

Commit 313ee7f

Browse files
himanshusinghshimanshusinghs
authored
fix: fix the schema for embeddings parameter MCP-281 (#763)
1 parent 3a21cd8 commit 313ee7f

File tree

3 files changed

+25
-26
lines changed

3 files changed

+25
-26
lines changed

‎src/tools/mongodb/mongodbSchemas.ts‎

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,29 +5,28 @@ export const zVoyageModels = z
55
.enum(["voyage-3-large", "voyage-3.5", "voyage-3.5-lite", "voyage-code-3"])
66
.default("voyage-3-large");
77

8-
// Zod does not undestand JS boxed numbers (like Int32) as integer literals,
9-
// so we preprocess them to unwrap them so Zod understands them.
10-
function unboxNumber(v: unknown): number {
11-
if (v && typeof v === "object" && typeof v.valueOf === "function") {
12-
const n = Number(v.valueOf());
13-
if (!Number.isNaN(n)) return n;
14-
}
15-
return v as number;
16-
}
17-
188
export const zVoyageEmbeddingParameters = z.object({
9+
// OpenAPI JSON Schema supports enum only as string so the public facing
10+
// parameters that are fed to LLM providers should expect the dimensions as
11+
// stringified numbers which are then transformed to actual numbers.
1912
outputDimension: z
20-
.preprocess(
21-
unboxNumber,
22-
z.union([z.literal(256), z.literal(512), z.literal(1024), z.literal(2048), z.literal(4096)])
23-
)
24-
.optional()
25-
.default(1024),
13+
.union([z.literal("256"), z.literal("512"), z.literal("1024"), z.literal("2048"), z.literal("4096")])
14+
.default("1024")
15+
.transform((value): number => Number.parseInt(value))
16+
.optional(),
2617
outputDtype: z.enum(["float", "int8", "uint8", "binary", "ubinary"]).optional().default("float"),
2718
});
2819

2920
export const zVoyageAPIParameters = zVoyageEmbeddingParameters
3021
.extend({
22+
// Unlike public facing parameters, `zVoyageEmbeddingParameters`, the
23+
// api parameters need to be correct number and because we do an
24+
// additional parsing before calling the API, we override the
25+
// outputDimension schema to expect a union of numbers.
26+
outputDimension: z
27+
.union([z.literal(256), z.literal(512), z.literal(1024), z.literal(2048), z.literal(4096)])
28+
.default(1024)
29+
.optional(),
3130
inputType: z.enum(["query", "document"]),
3231
})
3332
.strip();

‎tests/integration/tools/mongodb/create/insertMany.test.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ describeWithMongoDB(
203203
// embeddingParameters so that we can simulate the idea
204204
// of unknown or mismatched quantization.
205205

206-
// embeddingParameters: { outputDimension: 256,
206+
// embeddingParameters: { outputDimension: "256",
207207
// outputDtype: "float", model: "voyage-3-large", input:
208208
// [
209209
// {
@@ -558,7 +558,7 @@ describeWithMongoDB(
558558
documents: [{ title: "The Matrix" }],
559559
embeddingParameters: {
560560
model: "voyage-3.5-lite",
561-
outputDimension: 256,
561+
outputDimension: "256",
562562
input: [{ titleEmbeddings: "The Matrix" }],
563563
},
564564
},

‎tests/integration/tools/mongodb/read/aggregate.test.ts‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
444444
limit: 10,
445445
embeddingParameters: {
446446
model: "voyage-3-large",
447-
outputDimension: 256,
447+
outputDimension: "256",
448448
},
449449
},
450450
},
@@ -508,7 +508,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
508508
limit: 10,
509509
embeddingParameters: {
510510
model: "voyage-3-large",
511-
outputDimension: 256,
511+
outputDimension: "256",
512512
outputDType: dataType,
513513
},
514514
},
@@ -575,7 +575,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
575575
limit: 10,
576576
embeddingParameters: {
577577
model: "voyage-3-large",
578-
outputDimension: 256,
578+
outputDimension: "256",
579579
outputDType: dataType,
580580
},
581581
},
@@ -642,7 +642,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
642642
limit: 10,
643643
embeddingParameters: {
644644
model: "voyage-3-large",
645-
outputDimension: 256,
645+
outputDimension: "256",
646646
outputDType: dataType,
647647
},
648648
},
@@ -709,7 +709,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
709709
limit: 10,
710710
embeddingParameters: {
711711
model: "voyage-3-large",
712-
outputDimension: 256,
712+
outputDimension: "256",
713713
outputDType: dataType,
714714
},
715715
},
@@ -774,7 +774,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
774774
limit: 10,
775775
embeddingParameters: {
776776
model: "voyage-3-large",
777-
outputDimension: 256,
777+
outputDimension: "256",
778778
outputDType: "float",
779779
},
780780
filter: { name: 10 },
@@ -837,7 +837,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
837837
limit: 10,
838838
embeddingParameters: {
839839
model: "voyage-3-large",
840-
outputDimension: 256,
840+
outputDimension: "256",
841841
outputDType: "float",
842842
},
843843
filter: { name: 10 },
@@ -900,7 +900,7 @@ If the user requests additional filtering, include filters in \`$vectorSearch.fi
900900
limit: 10,
901901
embeddingParameters: {
902902
model: "voyage-3-large",
903-
outputDimension: 256,
903+
outputDimension: "256",
904904
outputDType: "float",
905905
},
906906
filter: { name: 10 },

0 commit comments

Comments
 (0)