diff --git a/ts/packages/dispatcher/dispatcher/src/translation/actionTemplate.ts b/ts/packages/dispatcher/dispatcher/src/translation/actionTemplate.ts index f6056173a..e8f9d6e24 100644 --- a/ts/packages/dispatcher/dispatcher/src/translation/actionTemplate.ts +++ b/ts/packages/dispatcher/dispatcher/src/translation/actionTemplate.ts @@ -20,6 +20,7 @@ import { ActionParamArray, ActionParamObject, ActionParamType, + resolveTypeReference, } from "@typeagent/action-schema"; import { getActionParamCompletion } from "./requestCompletion.js"; @@ -43,22 +44,86 @@ function getDefaultActionTemplate( return template; } +function resolveObjectType( + type: ActionParamType, +): ActionParamObject | undefined { + const resolved = resolveTypeReference(type); + return resolved?.type === "object" ? resolved : undefined; +} + +// Discriminated object union: shared single-value kind field. +function getObjectUnionDiscriminator( + types: readonly ActionParamType[], +): + | { fieldName: string; values: string[]; arms: ActionParamObject[] } + | undefined { + if (types.length < 2) { + return undefined; + } + const arms: ActionParamObject[] = []; + for (const t of types) { + const obj = resolveObjectType(t); + if (obj === undefined) { + return undefined; + } + arms.push(obj); + } + + const firstFields = Object.keys(arms[0].fields); + for (const fieldName of firstFields) { + const values: string[] = []; + let ok = true; + for (const arm of arms) { + const field = arm.fields[fieldName]; + if (field === undefined || field.optional) { + ok = false; + break; + } + const ft = resolveTypeReference(field.type) ?? field.type; + if (ft.type !== "string-union" || ft.typeEnum.length !== 1) { + ok = false; + break; + } + values.push(ft.typeEnum[0]); + } + if (!ok) { + continue; + } + // Kind values must be unique per arm. + if (new Set(values).size !== values.length) { + continue; + } + return { fieldName, values, arms }; + } + return undefined; +} + function toTemplateTypeObject( type: ActionParamObject, visited: ReadonlySet, + data: unknown, ) { const templateType: TemplateFieldObject = { type: "object", fields: {}, }; + const dataObj = + data !== null && typeof data === "object" && !Array.isArray(data) + ? (data as Record) + : undefined; + for (const [key, field] of Object.entries(type.fields)) { - const type = toTemplateType(field.type, visited); - if (type === undefined) { + const fieldData = dataObj !== undefined ? dataObj[key] : undefined; + const fieldType = toTemplateType(field.type, visited, fieldData); + if (fieldType === undefined) { // Skip undefined fields. continue; } - templateType.fields[key] = { optional: field.optional, type }; + templateType.fields[key] = { + optional: field.optional, + type: fieldType, + }; } return templateType; } @@ -66,8 +131,12 @@ function toTemplateTypeObject( function toTemplateTypeArray( type: ActionParamArray, visited: ReadonlySet, + data: unknown, ) { - const elementType = toTemplateType(type.elementType, visited); + // First array element shapes nested templates. + const elementData = + Array.isArray(data) && data.length > 0 ? data[0] : undefined; + const elementType = toTemplateType(type.elementType, visited, elementData); if (elementType === undefined) { // Skip undefined fields. return undefined; @@ -79,14 +148,92 @@ function toTemplateTypeArray( return templateType; } +// Expand discriminated unions; else first-arm fallback. +function toTemplateTypeUnion( + types: readonly ActionParamType[], + visited: ReadonlySet, + data: unknown, +): TemplateType | undefined { + const disc = getObjectUnionDiscriminator(types); + if (disc !== undefined) { + const { fieldName, values, arms } = disc; + let selectedIndex = 0; + if ( + data !== null && + typeof data === "object" && + !Array.isArray(data) && + fieldName in (data as object) + ) { + const current = (data as Record)[fieldName]; + if (typeof current === "string") { + const idx = values.indexOf(current); + if (idx >= 0) { + selectedIndex = idx; + } + } + } + const selectedArm = arms[selectedIndex]; + const template = toTemplateTypeObject(selectedArm, visited, data); + // Full kind enum; arm fields from data. + template.fields[fieldName] = { + optional: false, + type: { + type: "string-union", + typeEnum: values, + discriminator: values[selectedIndex], + }, + }; + return template; + } + + // Prefer arm that matches current data shape. + if (data !== undefined) { + for (const t of types) { + const resolved = resolveTypeReference(t) ?? t; + try { + if ( + resolved.type === "object" && + data !== null && + typeof data === "object" && + !Array.isArray(data) + ) { + return toTemplateType(t, visited, data); + } + if (resolved.type === "array" && Array.isArray(data)) { + return toTemplateType(t, visited, data); + } + if ( + (resolved.type === "string" || + resolved.type === "number" || + resolved.type === "boolean" || + resolved.type === "string-union") && + typeof data === resolved.type + ) { + return toTemplateType(t, visited, data); + } + if ( + resolved.type === "string-union" && + typeof data === "string" + ) { + return toTemplateType(t, visited, data); + } + } catch { + // try next arm + } + } + } + // Historical fallback: first arm. + return toTemplateType(types[0], visited, data); +} + function toTemplateType( type: ActionParamType, visited: ReadonlySet = new Set(), + data: unknown = undefined, ): TemplateType | undefined { switch (type.type) { case "type-union": - // TODO: smarter about type unions. - return toTemplateType(type.types[0], visited); + return toTemplateTypeUnion(type.types, visited, data); case "type-reference": if (type.definition === undefined) { throw new Error(`Unresolved type reference: ${type.name}`); @@ -98,11 +245,12 @@ function toTemplateType( return toTemplateType( type.definition.type, new Set([...visited, type.name]), + data, ); case "object": - return toTemplateTypeObject(type, visited); + return toTemplateTypeObject(type, visited, data); case "array": - return toTemplateTypeArray(type, visited); + return toTemplateTypeArray(type, visited, data); case "string-union": return type as TemplateFieldStringUnion; case "string": @@ -153,7 +301,12 @@ function toTemplate( const actionParametersType = actionSchema.type.fields.parameters?.type; if (actionParametersType) { - const type = toTemplateType(actionParametersType); + // Pass data so union arms resolve correctly. + const type = toTemplateType( + actionParametersType, + new Set(), + action.parameters, + ); if (type !== undefined) { template.fields.parameters = { // ActionParam types are compatible with TemplateFields diff --git a/ts/packages/dispatcher/dispatcher/test/actionTemplate.spec.ts b/ts/packages/dispatcher/dispatcher/test/actionTemplate.spec.ts index 23ada756d..002d0f368 100644 --- a/ts/packages/dispatcher/dispatcher/test/actionTemplate.spec.ts +++ b/ts/packages/dispatcher/dispatcher/test/actionTemplate.spec.ts @@ -92,3 +92,134 @@ describe("getSystemTemplateSchema toTemplateType", () => { expect(Object.keys(parameters.fields)).toEqual(["keep"]); }); }); + +describe("getSystemTemplateSchema discriminated type-union", () => { + function musicTargetUnion(): ActionParamType { + const track = { + type: "object", + fields: { + kind: { + type: { type: "string-union", typeEnum: ["track"] }, + }, + trackName: { type: { type: "string" } }, + artists: { + optional: true, + type: { + type: "array", + elementType: { type: "string" }, + }, + }, + }, + }; + const artist = { + type: "object", + fields: { + kind: { + type: { type: "string-union", typeEnum: ["artist"] }, + }, + artistName: { type: { type: "string" } }, + }, + }; + const album = { + type: "object", + fields: { + kind: { + type: { type: "string-union", typeEnum: ["album"] }, + }, + albumName: { type: { type: "string" } }, + }, + }; + const playlist = { + type: "object", + fields: { + kind: { + type: { type: "string-union", typeEnum: ["playlist"] }, + }, + playlistName: { type: { type: "string" } }, + }, + }; + return { + type: "object", + fields: { + target: { + type: { + type: "type-union", + types: [track, artist, album, playlist], + }, + }, + }, + } as unknown as ActionParamType; + } + + it("exposes full kind enum for MusicTarget-style discriminated unions", async () => { + const template = await getSystemTemplateSchema( + "action", + { + schemaName, + actionName, + parameters: { + target: { kind: "track", trackName: "Bohemian Rhapsody" }, + }, + }, + makeContext(musicTargetUnion()), + ); + + const parameters = template.fields.parameters?.type as any; + const target = parameters.fields.target.type; + expect(target.type).toBe("object"); + expect(target.fields.kind.type).toEqual({ + type: "string-union", + typeEnum: ["track", "artist", "album", "playlist"], + discriminator: "track", + }); + expect(target.fields.trackName.type).toEqual({ type: "string" }); + expect(target.fields.artistName).toBeUndefined(); + }); + + it("selects album arm fields when data.kind is album", async () => { + const template = await getSystemTemplateSchema( + "action", + { + schemaName, + actionName, + parameters: { + target: { + kind: "album", + albumName: "A Night at the Opera", + }, + }, + }, + makeContext(musicTargetUnion()), + ); + + const target = (template.fields.parameters?.type as any).fields.target + .type; + expect(target.fields.kind.type.typeEnum).toEqual([ + "track", + "artist", + "album", + "playlist", + ]); + expect(target.fields.kind.type.discriminator).toBe("album"); + expect(target.fields.albumName.type).toEqual({ type: "string" }); + expect(target.fields.trackName).toBeUndefined(); + expect(target.fields.playlistName).toBeUndefined(); + }); + + it("keeps non-union object parameters unchanged", async () => { + const parametersType: ActionParamType = { + type: "object", + fields: { + query: { type: { type: "string" } }, + }, + } as unknown as ActionParamType; + + const template = await getSystemTemplateSchema( + "action", + { ...data, parameters: { query: "hello" } }, + makeContext(parametersType), + ); + const parameters = template.fields.parameters?.type as any; + expect(parameters.fields.query.type).toEqual({ type: "string" }); + }); +});