Merge pull request #64 from FormulaMonks/fix/vertex-ai-schema-constra…

…ined-tokens-error

fix: make KurtVertexAI properly throw KurtCapabilityError for the schema-constrained tokens case

packages/kurt-vertex-ai/spec/generateStructuredData.spec.ts

@@ -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 () => {
@@ -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) =>

packages/kurt-vertex-ai/src/KurtVertexAI.ts

@@ -19,6 +19,7 @@ import {
   KurtResultValidateError,
   KurtResultLimitError,
   KurtResultBlockedError,
+  KurtCapabilityError,
 } from "@formula-monks/kurt"
 import type {
   VertexAI,
@@ -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")