Skip to content

Commit

Permalink
fix: make KurtVertexAI properly throw KurtCapabilityError for the sch…
Browse files Browse the repository at this point in the history
…ema-constrained tokens case

Vertex AI is not yet capable of this feature, and we want to proactively warn the user that they
cannot expect a schema conformance guarantee.
  • Loading branch information
jemc committed Dec 9, 2024
1 parent 9ba8bdb commit 623cfe7
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 1 deletion.
32 changes: 31 additions & 1 deletion packages/kurt-vertex-ai/spec/generateStructuredData.spec.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
import { describe, test, expect } from "@jest/globals"
import { z } from "zod"
import { snapshotAndMock, snapshotAndMockWithError } from "./snapshots"
import { KurtResultValidateError } from "@formula-monks/kurt"
import {
KurtCapabilityError,
KurtResultValidateError,
} from "@formula-monks/kurt"

describe("KurtVertexAI generateStructuredData", () => {
test("says hello (response format 1)", async () => {
Expand Down Expand Up @@ -46,6 +49,33 @@ describe("KurtVertexAI generateStructuredData", () => {
expect(result.data).toEqual({ say: "hello" })
})

test("throws a capability error for schema constrained tokens", async () => {
await snapshotAndMockWithError(
(kurt) =>
kurt.generateStructuredData({
prompt: "Say hello!",
schema: z
.object({
say: z.string().describe("A single word to say"),
})
.describe("Say a word"),
sampling: {
// This is not available as a capability of Vertex AI.
forceSchemaConstrainedTokens: true,
},
}),

(errorAny) => {
expect(errorAny).toBeInstanceOf(KurtCapabilityError)
const error = errorAny as KurtCapabilityError
expect(error.missingCapability).toEqual(
"forceSchemaConstrainedTokens is not available for Vertex AI"
)
expect(error.message).toContain(error.missingCapability)
}
)
})

test("throws a validate error from an impossible schema", async () => {
await snapshotAndMockWithError(
(kurt) =>
Expand Down
8 changes: 8 additions & 0 deletions packages/kurt-vertex-ai/src/KurtVertexAI.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import {
KurtResultValidateError,
KurtResultLimitError,
KurtResultBlockedError,
KurtCapabilityError,
} from "@formula-monks/kurt"
import type {
VertexAI,
Expand Down Expand Up @@ -91,6 +92,13 @@ export class KurtVertexAI
model: this.options.model,
}) as VertexAIGenerativeModel

if (options.sampling.forceSchemaConstrainedTokens) {
throw new KurtCapabilityError(
this,
"forceSchemaConstrainedTokens is not available for Vertex AI"
)
}

// VertexAI requires that system messages be sent as a single message,
// so we filter them out from the main messages array to send separately.
const normalMessages = options.messages.filter((m) => m.role !== "system")
Expand Down

0 comments on commit 623cfe7

Please sign in to comment.