From d3909796ec01df78a84822220c2131e0fb63c749 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Fri, 19 Jun 2026 12:54:00 +0100 Subject: [PATCH 1/6] feat: structured output for ask mode via Zod schemas or JSON Schema Adds an `outputFormat` option to `QueryAgent.ask` / `askStream`, mirroring the Python client's `output_format`: - a Zod schema (the Pydantic-model equivalent): the final answer is parsed and validated into `z.infer`, returned as `ParsedAskModeResponse` - a raw Draft 2020-12 JSON Schema object: the final answer is parsed as JSON, returned as `ParsedAskModeResponse>` - nothing (default): plain text on `finalAnswer`, returned as `AskModeResponse` Overloads select the precise return type per input. Zod schemas are serialised to JSON Schema with `z.toJSONSchema` and sent as `output_format` in the request body. Adds zod as a dependency and tests for all three modes. Co-Authored-By: Claude Opus 4.8 --- package.json | 6 +- pnpm-lock.yaml | 8 ++ src/query/agent.test.ts | 149 +++++++++++++++++++++++++ src/query/agent.ts | 45 +++++++- src/query/response/response-mapping.ts | 62 +++++++++- src/query/response/response.ts | 23 ++++ 6 files changed, 280 insertions(+), 13 deletions(-) diff --git a/package.json b/package.json index 14f7ed7..bd931b1 100644 --- a/package.json +++ b/package.json @@ -27,6 +27,9 @@ "files": [ "dist" ], + "dependencies": { + "zod": "^4.0.0" + }, "peerDependencies": { "weaviate-client": "^3.5.4" }, @@ -41,6 +44,7 @@ "typedoc": "^0.28.15", "typedoc-plugin-extras": "^4.0.1", "typescript": "^5.8.3", - "typescript-eslint": "^8.32.1" + "typescript-eslint": "^8.32.1", + "zod": "^4.4.3" } } diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index fedbffe..eb40d40 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -11,6 +11,9 @@ importers: weaviate-client: specifier: ^3.5.4 version: 3.5.4 + zod: + specifier: ^4.0.0 + version: 4.4.3 devDependencies: '@eslint/js': specifier: ^9.26.0 @@ -2179,6 +2182,9 @@ packages: zod@3.24.4: resolution: {integrity: sha512-OdqJE9UDRPwWsrHjLN2F8bPxvwJBK22EHLWtanu0LSYr5YqzsaaW3RMgmjwr8Rypg5k+meEJdSPXJZXE/yqOMg==} + zod@4.4.3: + resolution: {integrity: sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ==} + snapshots: '@ampproject/remapping@2.3.0': @@ -4654,3 +4660,5 @@ snapshots: zod: 3.24.4 zod@3.24.4: {} + + zod@4.4.3: {} diff --git a/src/query/agent.test.ts b/src/query/agent.test.ts index fbca8e8..88a799b 100644 --- a/src/query/agent.test.ts +++ b/src/query/agent.test.ts @@ -1,3 +1,4 @@ +import { z } from "zod"; import { WeaviateClient } from "weaviate-client"; import { QueryAgent } from "./agent.js"; import { ApiQueryAgentResponse } from "./response/api-response.js"; @@ -6,6 +7,7 @@ import { ComparisonOperator, AskModeResponse, SuggestQueryResponse, + ParsedAskModeResponse, } from "./response/response.js"; import { ApiSearchModeResponse, @@ -887,3 +889,150 @@ it("suggest queries mode failure propagates QueryAgentError", async () => { }); } }); + +const mockClient = () => + ({ + getConnectionDetails: jest.fn().mockResolvedValue({ + host: "test-cluster", + bearerToken: "test-token", + headers: { "X-Provider": "test-key" }, + }), + }) as unknown as WeaviateClient; + +const askApiResponse = (finalAnswer: string): ApiAskModeResponse => ({ + searches: [], + aggregations: [], + usage: { + model_units: 1, + usage_in_plan: true, + remaining_plan_requests: 2, + }, + total_time: 1.5, + is_partial_answer: false, + missing_information: [], + final_answer: finalAnswer, + sources: [], +}); + +it("ask with a Zod output format parses and validates the final answer", async () => { + const Answer = z.object({ + answer: z.string(), + score: z.number(), + }); + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const capturedBodies: any[] = []; + + const finalAnswer = JSON.stringify({ answer: "Paris", score: 0.9 }); + global.fetch = jest.fn((url, init?: RequestInit) => { + if (init && init.body) { + capturedBodies.push(JSON.parse(init.body as string)); + } + return Promise.resolve({ + ok: true, + json: () => Promise.resolve(askApiResponse(finalAnswer)), + } as Response); + }) as jest.Mock; + + const agent = new QueryAgent(mockClient()); + + const response = await agent.ask("What is the capital of France?", { + collections: ["test-collection"], + outputFormat: Answer, + }); + + // The request carries the schema serialised to a Draft 2020-12 JSON Schema. + expect(capturedBodies[0].output_format).toEqual( + z.toJSONSchema(Answer, { target: "draft-2020-12" }), + ); + + // The raw string is still available, plus the parsed (and validated) object. + expect(response.finalAnswer).toBe(finalAnswer); + expect(response.finalAnswerParsed).toEqual({ answer: "Paris", score: 0.9 }); + + // Compile-time: the parsed type is inferred from the schema. + const parsed: ParsedAskModeResponse> = response; + const score: number = parsed.finalAnswerParsed.score; + expect(score).toBe(0.9); +}); + +it("ask with a Zod output format throws when the answer violates the schema", async () => { + const Answer = z.object({ answer: z.string(), score: z.number() }); + + global.fetch = jest.fn(() => + Promise.resolve({ + ok: true, + // `score` is missing -> Zod validation must reject it. + json: () => Promise.resolve(askApiResponse('{"answer":"Paris"}')), + } as Response), + ) as jest.Mock; + + const agent = new QueryAgent(mockClient()); + + await expect( + agent.ask("What is the capital of France?", { + collections: ["test-collection"], + outputFormat: Answer, + }), + ).rejects.toThrow(); +}); + +it("ask with a raw JSON Schema output format parses the final answer as JSON", async () => { + const jsonSchema: Record = { + type: "object", + properties: { answer: { type: "string" } }, + required: ["answer"], + additionalProperties: false, + }; + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const capturedBodies: any[] = []; + + const finalAnswer = JSON.stringify({ answer: "Paris" }); + global.fetch = jest.fn((url, init?: RequestInit) => { + if (init && init.body) { + capturedBodies.push(JSON.parse(init.body as string)); + } + return Promise.resolve({ + ok: true, + json: () => Promise.resolve(askApiResponse(finalAnswer)), + } as Response); + }) as jest.Mock; + + const agent = new QueryAgent(mockClient()); + + const response = await agent.ask("What is the capital of France?", { + collections: ["test-collection"], + outputFormat: jsonSchema, + }); + + // A raw JSON Schema is forwarded verbatim. + expect(capturedBodies[0].output_format).toEqual(jsonSchema); + expect(response.finalAnswer).toBe(finalAnswer); + expect(response.finalAnswerParsed).toEqual({ answer: "Paris" }); +}); + +it("ask without an output format omits output_format and returns plain text", async () => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const capturedBodies: any[] = []; + + global.fetch = jest.fn((url, init?: RequestInit) => { + if (init && init.body) { + capturedBodies.push(JSON.parse(init.body as string)); + } + return Promise.resolve({ + ok: true, + json: () => Promise.resolve(askApiResponse("Plain text answer")), + } as Response); + }) as jest.Mock; + + const agent = new QueryAgent(mockClient()); + + const response = await agent.ask("What is the capital of France?", { + collections: ["test-collection"], + }); + + expect(capturedBodies[0].output_format).toBeUndefined(); + expect(response.finalAnswer).toBe("Plain text answer"); + expect("finalAnswerParsed" in response).toBe(false); +}); diff --git a/src/query/agent.ts b/src/query/agent.ts index 81ab68a..ef450c2 100644 --- a/src/query/agent.ts +++ b/src/query/agent.ts @@ -1,10 +1,13 @@ import { WeaviateClient } from "weaviate-client"; +import { z } from "zod"; import { QueryAgentResponse, ProgressMessage, StreamedTokens, AskModeResponse, SuggestQueryResponse, + OutputFormat, + ParsedAskModeResponse, } from "./response/response.js"; import { mapResponse, @@ -12,6 +15,7 @@ import { mapStreamedTokensFromSSE, mapAskModeResponse, mapSuggestQueryResponse, + isZodSchema, } from "./response/response-mapping.js"; import { mapApiResponse } from "./response/api-response-mapping.js"; import { fetchServerSentEvents } from "./response/server-sent-events.js"; @@ -136,8 +140,12 @@ export class QueryAgent { */ async ask( query: QueryAgentQuery, - { collections, resultEvaluation }: QueryAgentAskOptions = {}, - ): Promise { + { + collections, + resultEvaluation, + outputFormat, + }: QueryAgentAskOptions & { outputFormat?: OutputFormat } = {}, + ): Promise> { const targetCollections = this.validateCollections(collections); const { requestHeaders, connectionHeaders } = await getHeaders(this.client); @@ -150,6 +158,7 @@ export class QueryAgent { collections: mapCollections(targetCollections), system_prompt: this.systemPrompt, result_evaluation: resultEvaluation ?? "none", + output_format: mapOutputFormat(outputFormat), }), }); @@ -271,6 +280,7 @@ export class QueryAgent { askStream( query: QueryAgentQuery, options: QueryAgentAskStreamOptions & { + outputFormat?: undefined; includeProgress: false; includeFinalState: false; }, @@ -279,6 +289,7 @@ export class QueryAgent { askStream( query: QueryAgentQuery, options: QueryAgentAskStreamOptions & { + outputFormat?: undefined; includeProgress: false; includeFinalState?: true; }, @@ -287,6 +298,7 @@ export class QueryAgent { askStream( query: QueryAgentQuery, options: QueryAgentAskStreamOptions & { + outputFormat?: undefined; includeProgress?: true; includeFinalState: false; }, @@ -338,6 +350,7 @@ export class QueryAgent { askStream( query: QueryAgentQuery, options?: QueryAgentAskStreamOptions & { + outputFormat?: undefined; includeProgress?: true; includeFinalState?: true; }, @@ -349,8 +362,14 @@ export class QueryAgent { includeProgress, includeFinalState, resultEvaluation, - }: QueryAgentAskStreamOptions = {}, - ): AsyncGenerator { + outputFormat, + }: QueryAgentAskStreamOptions & { outputFormat?: OutputFormat } = {}, + ): AsyncGenerator< + | ProgressMessage + | StreamedTokens + | AskModeResponse + | ParsedAskModeResponse + > { const targetCollections = collections ?? this.collections; if (!targetCollections) { @@ -378,6 +397,7 @@ export class QueryAgent { include_progress: includeProgress ?? true, include_final_state: includeFinalState ?? true, result_evaluation: resultEvaluation ?? "none", + output_format: mapOutputFormat(outputFormat), }), }, ); @@ -387,13 +407,26 @@ export class QueryAgent { await handleError(event.data); } - let output: ProgressMessage | StreamedTokens | AskModeResponse; + let output: + | ProgressMessage + | StreamedTokens + | AskModeResponse + | ParsedAskModeResponse; if (event.event === "progress_message") { output = mapProgressMessageFromSSE(event); } else if (event.event === "streamed_tokens") { output = mapStreamedTokensFromSSE(event); } else if (event.event === "final_state") { - output = mapAskModeResponse(JSON.parse(event.data)); + const finalState = JSON.parse(event.data); + if (outputFormat === undefined) { + output = mapAskModeResponse(finalState); + } else { + // Both arms are the same call; the branch only narrows the type + // (Zod schema vs raw JSON Schema) to select the right overload. + output = isZodSchema(outputFormat) + ? mapAskModeResponse(finalState, outputFormat) + : mapAskModeResponse(finalState, outputFormat); + } } else { throw new Error(`Unexpected event type: ${event.event}: ${event.data}`); } diff --git a/src/query/response/response-mapping.ts b/src/query/response/response-mapping.ts index 0a0e937..487f668 100644 --- a/src/query/response/response-mapping.ts +++ b/src/query/response/response-mapping.ts @@ -18,8 +18,12 @@ import { ModelUnitUsage, QuerySort, SuggestQueryResponse, + OutputFormat, + ParsedAskModeResponse, } from "./response.js"; +import { z } from "zod"; + import { ApiQueryAgentResponse, ApiPropertyFilter, @@ -40,9 +44,29 @@ import { import { ServerSentEvent } from "./server-sent-events.js"; -export const mapAskModeResponse = ( +/** + * Duck-types a value as a Zod schema (as opposed to a raw JSON Schema object): + * only a Zod schema carries a `parse` method, and a function value can never + * appear in a JSON-serialisable schema, so this reliably tells the two apart. + */ +export const isZodSchema = (value: OutputFormat): value is z.ZodType => + typeof (value as { parse?: unknown }).parse === "function"; + +export function mapAskModeResponse( + response: ApiAskModeResponse, +): AskModeResponse; +export function mapAskModeResponse( + response: ApiAskModeResponse, + outputFormat: S, +): ParsedAskModeResponse>; +export function mapAskModeResponse( response: ApiAskModeResponse, -): AskModeResponse => { + outputFormat: Record, +): ParsedAskModeResponse>; +export function mapAskModeResponse( + response: ApiAskModeResponse, + outputFormat?: OutputFormat, +): AskModeResponse | ParsedAskModeResponse { const properties: AskModeResponseProperties = { outputType: "finalState", searches: mapSearches(response.searches), @@ -60,11 +84,28 @@ export const mapAskModeResponse = ( sources: response.sources ? mapSources(response.sources) : undefined, }; - return { + if (outputFormat === undefined) { + return { + ...properties, + display: () => display(properties), + }; + } + + // A Zod schema parses and validates; a raw JSON Schema only parses the JSON. + const finalAnswerParsed: unknown = isZodSchema(outputFormat) + ? outputFormat.parse(JSON.parse(response.final_answer)) + : (JSON.parse(response.final_answer) as Record); + + const parsedProperties: ParsedAskModeResponseProperties = { ...properties, - display: () => display(properties), + finalAnswerParsed, }; -}; + + return { + ...parsedProperties, + display: () => display(parsedProperties), + }; +} const mapSearches = (searches: ApiSearch[]): Search[] => searches.map((search) => ({ @@ -312,11 +353,20 @@ const mapSources = (sources: ApiSource[]): Source[] => collection: source.collection, })); -const display = (response: AskModeResponseProperties | ResponseProperties) => { +const display = ( + response: + | AskModeResponseProperties + | ParsedAskModeResponseProperties + | ResponseProperties, +) => { console.log(JSON.stringify(response, undefined, 2)); }; type AskModeResponseProperties = Omit; +type ParsedAskModeResponseProperties = Omit< + ParsedAskModeResponse, + "display" +>; type ResponseProperties = Omit; type ProgressMessageJSON = Omit & { diff --git a/src/query/response/response.ts b/src/query/response/response.ts index 462f655..1cd1d2f 100644 --- a/src/query/response/response.ts +++ b/src/query/response/response.ts @@ -1,4 +1,5 @@ import { WeaviateReturn, WeaviateObject } from "weaviate-client"; +import { z } from "zod"; export type AskModeResponse = { outputType: "finalState"; @@ -13,6 +14,28 @@ export type AskModeResponse = { display(): void; }; +/** + * A structured output format for the ask mode. + * + * Either a Zod schema (the equivalent of a Pydantic model — the final answer is + * parsed and validated into the inferred type), or a plain object holding a + * Draft 2020-12 JSON Schema (the final answer is parsed as JSON). Pass nothing + * for no structured output. + */ +export type OutputFormat = z.ZodType | Record; + +/** + * An {@link AskModeResponse} where the final answer has been parsed into the + * requested structured output format. + * + * `T` is the inferred type of the Zod schema (`z.infer`) when a + * schema was passed, or `Record` when a raw JSON Schema was. + */ +export type ParsedAskModeResponse = AskModeResponse & { + /** The final answer parsed into the requested output format. */ + finalAnswerParsed: T; +}; + export type Search = { query?: string; filters?: PropertyFilter | FilterAndOr; From 12682627149ff7cae34542e3a67c741d243546f6 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Mon, 22 Jun 2026 16:21:24 +0100 Subject: [PATCH 2/6] merge agent.ts --- src/query/agent.ts | 186 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 183 insertions(+), 3 deletions(-) diff --git a/src/query/agent.ts b/src/query/agent.ts index ef450c2..226d17c 100644 --- a/src/query/agent.ts +++ b/src/query/agent.ts @@ -112,6 +112,16 @@ export class QueryAgent { return mapResponse(await response.json()); } + /** @hidden */ + ask( + query: QueryAgentQuery, + options: QueryAgentAskOptions & { outputFormat: S }, + ): Promise>>; + /** @hidden */ + ask( + query: QueryAgentQuery, + options: QueryAgentAskOptions & { outputFormat: Record }, + ): Promise>>; /** * Run the Query Agent ask mode. * @@ -127,8 +137,13 @@ export class QueryAgent { * `isPartialAnswer` fields of the response. * If `"none"`, the result will not be evaluated, and the sources will not be filtered. * Defaults to `"none"`. - * @returns An {@link AskModeResponse} which contains the final answer, sources, and other - * metadata such as the searches performed, usage and total time. + * @param options.outputFormat - The structured output format to return. Either a Zod schema, a JSON Schema object, or undefined. + * If a Zod schema is provided, the final answer will be parsed and validated into the schema, and the response will be returned as a {@link ParsedAskModeResponse} whose `finalAnswerParsed` is typed as `z.infer`. + * If a JSON Schema object (Draft 2020-12) is provided, the final answer will be parsed as JSON, and the response will be returned as a {@link ParsedAskModeResponse} whose `finalAnswerParsed` is typed as `Record`. + * If undefined (the default, no structured output), the answer is plain text on `finalAnswer` and an {@link AskModeResponse} is returned. + * @returns An {@link AskModeResponse} (or {@link ParsedAskModeResponse} when `outputFormat` is set) + * which contains the final answer, sources, and other metadata such as the searches performed, + * usage and total time. * * @example * ```ts @@ -137,7 +152,33 @@ export class QueryAgent { * "What are the terms of the contract signed by John Smith in May 2025?", * ); * ``` + * + * @example Structured output with a Zod schema. The final answer is parsed and + * validated into the schema, and exposed (typed) on `finalAnswerParsed`. + * ```ts + * import { z } from "zod"; + * + * const CitedText = z.object({ + * text: z.string(), + * sources: z + * .array(z.string()) + * .describe("The sources that support this section of text. Can be empty."), + * }); + * const AnswerWithSources = z.object({ texts: z.array(CitedText) }); + * + * const agent = new QueryAgent(client, { collections: ["FinancialContracts"] }); + * const result = await agent.ask( + * "What contracts were signed by Jane Doe in 2024? What were they about?", + * { outputFormat: AnswerWithSources }, + * ); + * // result.finalAnswerParsed is typed as { texts: { text: string; sources: string[] }[] } + * console.log(result.finalAnswerParsed.texts); + * ``` */ + ask( + query: QueryAgentQuery, + options?: QueryAgentAskOptions, + ): Promise; async ask( query: QueryAgentQuery, { @@ -166,7 +207,15 @@ export class QueryAgent { await handleError(await response.text()); } - return mapAskModeResponse(await response.json()); + const json = await response.json(); + if (outputFormat === undefined) { + return mapAskModeResponse(json); + } + // Both arms are the same call; the branch only narrows the type + // (Zod schema vs raw JSON Schema) to select the right overload. + return isZodSchema(outputFormat) + ? mapAskModeResponse(json, outputFormat) + : mapAskModeResponse(json, outputFormat); } /** @hidden */ @@ -276,6 +325,86 @@ export class QueryAgent { } } + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: S; + includeProgress: false; + includeFinalState: false; + }, + ): AsyncGenerator; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: S; + includeProgress: false; + includeFinalState?: true; + }, + ): AsyncGenerator>>; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: S; + includeProgress?: true; + includeFinalState: false; + }, + ): AsyncGenerator; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: S; + includeProgress?: true; + includeFinalState?: true; + }, + ): AsyncGenerator< + ProgressMessage | StreamedTokens | ParsedAskModeResponse> + >; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: Record; + includeProgress: false; + includeFinalState: false; + }, + ): AsyncGenerator; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: Record; + includeProgress: false; + includeFinalState?: true; + }, + ): AsyncGenerator< + StreamedTokens | ParsedAskModeResponse> + >; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: Record; + includeProgress?: true; + includeFinalState: false; + }, + ): AsyncGenerator; + /** @hidden */ + askStream( + query: QueryAgentQuery, + options: QueryAgentAskStreamOptions & { + outputFormat: Record; + includeProgress?: true; + includeFinalState?: true; + }, + ): AsyncGenerator< + | ProgressMessage + | StreamedTokens + | ParsedAskModeResponse> + >; /** @hidden */ askStream( query: QueryAgentQuery, @@ -320,6 +449,10 @@ export class QueryAgent { * `isPartialAnswer` fields of the response. * If `"none"`, the result will not be evaluated, and the sources will not be filtered. * Defaults to `"none"`. + * @param options.outputFormat - The structured output format to return. Either a Zod schema, a JSON Schema object, or undefined. + * If a Zod schema is provided, the final answer will be parsed and validated into the schema, and the response will be returned as a `ParsedAskModeResponse` with the inferred type. + * If a JSON Schema object is provided, the final answer will be parsed as JSON, and the response will be returned as a `ParsedAskModeResponse` with the type `Record`. + * If undefined, the final answer will be returned as a `AskModeResponse` with the type `string`. * @returns An async generator yielding any of the following: * * - {@link ProgressMessage}: informational messages about the progress of the agent's search @@ -346,6 +479,36 @@ export class QueryAgent { * } * } * ``` + * + * @example Structured output with a Zod schema. When `includeFinalState` is on + * (the default), the final-state event is a {@link ParsedAskModeResponse} whose + * `finalAnswerParsed` is typed from the schema. + * ```ts + * import { z } from "zod"; + * + * const CitedText = z.object({ + * text: z.string(), + * sources: z + * .array(z.string()) + * .describe("The sources that support this section of text. Can be empty."), + * }); + * const AnswerWithSources = z.object({ texts: z.array(CitedText) }); + * + * const agent = new QueryAgent(client, { collections: ["FinancialContracts"] }); + * for await (const event of agent.askStream( + * "What contracts were signed by Jane Doe in 2024? What were they about?", + * { outputFormat: AnswerWithSources }, + * )) { + * if ("finalAnswerParsed" in event) { + * // ParsedAskModeResponse — finalAnswerParsed is typed from the schema + * console.log(event.finalAnswerParsed.texts); + * } else if ("delta" in event) { + * process.stdout.write(event.delta); + * } else { + * console.log(event.message); + * } + * } + * ``` */ askStream( query: QueryAgentQuery, @@ -555,6 +718,23 @@ export class QueryAgent { }; } +/** + * Convert a user-supplied {@link OutputFormat} into the JSON Schema the agents + * API expects on the request body: a Zod schema is serialised to a Draft + * 2020-12 JSON Schema, a raw JSON Schema object is sent through unchanged, and + * `undefined` (no structured output) stays `undefined`. + */ +const mapOutputFormat = ( + outputFormat: OutputFormat | undefined, +): Record | undefined => { + if (outputFormat === undefined) { + return undefined; + } + return isZodSchema(outputFormat) + ? z.toJSONSchema(outputFormat, { target: "draft-2020-12" }) + : outputFormat; +}; + /** Options for constructing a {@link QueryAgent}. */ export type QueryAgentOptions = { /** From 4f159f57c7dfb401528ac8022cdd7c181eb9a070 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Wed, 24 Jun 2026 16:04:33 +0100 Subject: [PATCH 3/6] refactor: add union overload for mapAskModeResponse Add an OutputFormat union overload so callers can pass a union-typed outputFormat directly, removing the two-arms-same-call branches that only existed to narrow the union for overload resolution. Co-authored-by: Cursor --- src/query/agent.ts | 20 +++++--------------- src/query/response/response-mapping.ts | 4 ++++ 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/src/query/agent.ts b/src/query/agent.ts index 226d17c..b32fd47 100644 --- a/src/query/agent.ts +++ b/src/query/agent.ts @@ -208,13 +208,8 @@ export class QueryAgent { } const json = await response.json(); - if (outputFormat === undefined) { - return mapAskModeResponse(json); - } - // Both arms are the same call; the branch only narrows the type - // (Zod schema vs raw JSON Schema) to select the right overload. - return isZodSchema(outputFormat) - ? mapAskModeResponse(json, outputFormat) + return outputFormat === undefined + ? mapAskModeResponse(json) : mapAskModeResponse(json, outputFormat); } @@ -581,15 +576,10 @@ export class QueryAgent { output = mapStreamedTokensFromSSE(event); } else if (event.event === "final_state") { const finalState = JSON.parse(event.data); - if (outputFormat === undefined) { - output = mapAskModeResponse(finalState); - } else { - // Both arms are the same call; the branch only narrows the type - // (Zod schema vs raw JSON Schema) to select the right overload. - output = isZodSchema(outputFormat) - ? mapAskModeResponse(finalState, outputFormat) + output = + outputFormat === undefined + ? mapAskModeResponse(finalState) : mapAskModeResponse(finalState, outputFormat); - } } else { throw new Error(`Unexpected event type: ${event.event}: ${event.data}`); } diff --git a/src/query/response/response-mapping.ts b/src/query/response/response-mapping.ts index 487f668..2e040dc 100644 --- a/src/query/response/response-mapping.ts +++ b/src/query/response/response-mapping.ts @@ -63,6 +63,10 @@ export function mapAskModeResponse( response: ApiAskModeResponse, outputFormat: Record, ): ParsedAskModeResponse>; +export function mapAskModeResponse( + response: ApiAskModeResponse, + outputFormat: OutputFormat, +): ParsedAskModeResponse; export function mapAskModeResponse( response: ApiAskModeResponse, outputFormat?: OutputFormat, From 00fb14b8199ef50dfbe60e4ba6f62cb09303c04f Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Wed, 24 Jun 2026 16:06:08 +0100 Subject: [PATCH 4/6] move zod dependency to peer --- package.json | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/package.json b/package.json index bd931b1..60da78f 100644 --- a/package.json +++ b/package.json @@ -27,11 +27,9 @@ "files": [ "dist" ], - "dependencies": { - "zod": "^4.0.0" - }, "peerDependencies": { - "weaviate-client": "^3.5.4" + "weaviate-client": "^3.5.4", + "zod": "^4.0.0" }, "devDependencies": { "@eslint/js": "^9.26.0", @@ -47,4 +45,4 @@ "typescript-eslint": "^8.32.1", "zod": "^4.4.3" } -} +} \ No newline at end of file From a3c71a0847d7967db26defe24b8f9345e7284f7b Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Wed, 24 Jun 2026 16:18:56 +0100 Subject: [PATCH 5/6] lock --- pnpm-lock.yaml | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index eb40d40..e88f0ad 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -11,9 +11,6 @@ importers: weaviate-client: specifier: ^3.5.4 version: 3.5.4 - zod: - specifier: ^4.0.0 - version: 4.4.3 devDependencies: '@eslint/js': specifier: ^9.26.0 @@ -48,6 +45,9 @@ importers: typescript-eslint: specifier: ^8.32.1 version: 8.32.1(eslint@9.26.0)(typescript@5.8.3) + zod: + specifier: ^4.4.3 + version: 4.4.3 packages: @@ -595,6 +595,7 @@ packages: '@ungap/structured-clone@1.3.0': resolution: {integrity: sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g==} + deprecated: Potential CWE-502 - Update to 1.3.1 or higher '@unrs/resolver-binding-android-arm-eabi@1.11.1': resolution: {integrity: sha512-ppLRUgHVaGRWUx0R0Ut06Mjo9gBaBkg3v/8AxusGLhsIotbBLuRk51rAzqLC8gq6NyyAojEXglNjzf6R948DNw==} @@ -1199,11 +1200,12 @@ packages: glob@10.4.5: resolution: {integrity: sha512-7Bv8RF0k6xjo7d4A/PxYLbUCfb6c+Vpd2/mB2yRDlew7Jb5hEXiCD9ibfO7wpk8i4sevK6DFny9h7EYbM3/sHg==} + deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me hasBin: true glob@7.2.3: resolution: {integrity: sha512-nFR0zLpU2YCaRxwoCJvL6UvCH2JFyFVIvwTLsIf21AuHlMskA1hhTdk+LlYJtOlYt9v6dvszD2BGRqBL+iQK9Q==} - deprecated: Glob versions prior to v9 are no longer supported + deprecated: Old versions of glob are not supported, and contain widely publicized security vulnerabilities, which have been fixed in the current version. Please update. Support for old versions may be purchased (at exorbitant rates) by contacting i@izs.me globals@14.0.0: resolution: {integrity: sha512-oahGvuMGQlPw/ivIYBjVSrWAfWLBeku5tpPE2fOPLi+WHffIWbuh2tCjhyQhTBPMf5E9jDEH4FOmTYgYwbKwtQ==} @@ -2103,6 +2105,7 @@ packages: uuid@9.0.1: resolution: {integrity: sha512-b+1eJOlsR9K8HJpow9Ok3fiWOWSIcIzXodvv0rQjVoOVNpWMpxf1wZNpt4y9h10odCNrqnYp1OBzRktckBe3sA==} + deprecated: uuid@10 and below is no longer supported. For ESM codebases, update to uuid@latest. For CommonJS codebases, use uuid@11 (but be aware this version will likely be deprecated in 2028). hasBin: true v8-to-istanbul@9.3.0: From e61e0bb7c02130336f949eebf2a4b61ce8ebebd4 Mon Sep 17 00:00:00 2001 From: dannyjameswilliams Date: Wed, 24 Jun 2026 16:20:51 +0100 Subject: [PATCH 6/6] prettier formatting --- package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/package.json b/package.json index 60da78f..cd8a8be 100644 --- a/package.json +++ b/package.json @@ -45,4 +45,4 @@ "typescript-eslint": "^8.32.1", "zod": "^4.4.3" } -} \ No newline at end of file +}