feat: add validated llm json helper

2026-06-21 23:43:02 +08:00
parent 44e9e04533
commit b9cfa3faf6
2 changed files with 99 additions and 0 deletions
@@ -0,0 +1,69 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { z } from "zod";
+
+import * as client from "../client";
+
+describe("generateValidatedJson", () => {
+  const originalProvider = process.env.LLM_PROVIDER;
+  const originalDeepSeekKey = process.env.DEEPSEEK_API_KEY;
+
+    afterEach(() => {
+      process.env.LLM_PROVIDER = originalProvider;
+      process.env.DEEPSEEK_API_KEY = originalDeepSeekKey;
+      client.setGenerateJsonForValidation(client.generateJson);
+      vi.restoreAllMocks();
+    });
+
+  it("returns null when no provider key is configured", async () => {
+    process.env.LLM_PROVIDER = "deepseek";
+    delete process.env.DEEPSEEK_API_KEY;
+
+    const result = await client.generateValidatedJson({
+      schema: z.object({ value: z.string() }),
+      prompt: "Return JSON.",
+    });
+
+    expect(result).toBeNull();
+  });
+
+    it("returns parsed data when the model response matches the schema", async () => {
+      process.env.LLM_PROVIDER = "deepseek";
+      process.env.DEEPSEEK_API_KEY = "test-key";
+      client.setGenerateJsonForValidation(async () => ({ value: "from-llm" }));
+
+    const result = await client.generateValidatedJson({
+      schema: z.object({ value: z.string() }),
+      prompt: "Return JSON.",
+    });
+
+    expect(result).toEqual({ value: "from-llm" });
+  });
+
+    it("returns null when the model response fails schema validation", async () => {
+      process.env.LLM_PROVIDER = "deepseek";
+      process.env.DEEPSEEK_API_KEY = "test-key";
+      client.setGenerateJsonForValidation(async () => ({ value: 42 }));
+
+    const result = await client.generateValidatedJson({
+      schema: z.object({ value: z.string() }),
+      prompt: "Return JSON.",
+    });
+
+    expect(result).toBeNull();
+  });
+
+    it("returns null when the provider call rejects", async () => {
+      process.env.LLM_PROVIDER = "deepseek";
+      process.env.DEEPSEEK_API_KEY = "test-key";
+      client.setGenerateJsonForValidation(async () => {
+        throw new Error("provider down");
+      });
+
+    const result = await client.generateValidatedJson({
+      schema: z.object({ value: z.string() }),
+      prompt: "Return JSON.",
+    });
+
+    expect(result).toBeNull();
+  });
+});