diff --git a/src/core/eval.tsx b/src/core/eval.tsx index 439c67f4f..d438ed26c 100644 --- a/src/core/eval.tsx +++ b/src/core/eval.tsx @@ -76,6 +76,7 @@ import { type EvaluationReferenceInput, type EvaluationResultContent, type EvaluationTarget, + type CreateABTestRequest, type CreateABTestResponse, type GetABTestResponse, type ListABTestsResponse, @@ -123,6 +124,7 @@ import type { CoreEvalClient, CreateConfigurationBundleInput, CreateConfigBundleABTestInput, + CreateTargetBasedABTestInput, CreateDatasetInput, CreateOnlineEvalInput, CreateOnlineInsightInput, @@ -481,65 +483,38 @@ export class EvalClient implements CoreEvalClient { .send(new DeleteABTestCommand({ abTestId: id })); } - async createConfigBundleABTest( - input: CreateConfigBundleABTestInput, + private async createABTest( + name: string, + gateway: string, + callerRoleArn: string | undefined, + build: (context: { + gatewayArn: string; + accountId: string; + roleArn: string; + }) => CreateABTestRequest, options: CoreOptions, ): Promise { const control = this.clients.control(toClientConfig(options)); - const gateway = await control.send(new GetGatewayCommand({ gatewayIdentifier: input.gateway })); - const gatewayArn = gateway.gatewayArn!; + const gatewayArn = (await control.send(new GetGatewayCommand({ gatewayIdentifier: gateway }))) + .gatewayArn!; const accountId = accountIdFromArn(gatewayArn); - const controlBundleArn = `arn:aws:bedrock-agentcore:${options.region}:${accountId}:configuration-bundle/${input.control.configBundle}`; - const treatmentBundleArn = `arn:aws:bedrock-agentcore:${options.region}:${accountId}:configuration-bundle/${input.treatment.configBundle}`; - const onlineEvaluationConfigArn = `arn:aws:bedrock-agentcore:${options.region}:${accountId}:online-evaluation-config/${input.onlineEval}`; - - const treatmentWeight = input.treatmentWeight ?? 50; - const variants = [ - { - name: "C", - weight: 100 - treatmentWeight, - variantConfiguration: { - configurationBundle: { - bundleArn: controlBundleArn, - bundleVersion: input.control.bundleVersion, - }, - }, - }, - { - name: "T1", - weight: treatmentWeight, - variantConfiguration: { - configurationBundle: { - bundleArn: treatmentBundleArn, - bundleVersion: input.treatment.bundleVersion, - }, - }, - }, - ]; - - let roleArn = input.roleArn; + let roleArn = callerRoleArn; let provisionedRoleArn: string | undefined; if (!roleArn) { - const iam = this.clients.iam({ region: options.region }); - const provisioned = await provisionAbTestRole(iam, input.name, gatewayArn, options.region); + const provisioned = await provisionAbTestRole( + this.clients.iam({ region: options.region }), + name, + gatewayArn, + options.region, + ); roleArn = provisioned.roleArn; if (provisioned.created) provisionedRoleArn = provisioned.roleArn; } - const command = new CreateABTestCommand({ - name: input.name, - gatewayArn, - variants, - evaluationConfig: { onlineEvaluationConfigArn }, - roleArn, - gatewayFilter: input.gatewayFilter, - enableOnCreate: input.enableOnCreate ?? true, - clientToken: randomUUID(), - }); - + const command = new CreateABTestCommand(build({ gatewayArn, accountId, roleArn })); try { - return input.roleArn + return callerRoleArn ? await this.clients.data(toClientConfig(options)).send(command) : await retryWhileRolePropagates(() => this.clients.data(toClientConfig(options)).send(command), @@ -556,6 +531,99 @@ export class EvalClient implements CoreEvalClient { } } + async createConfigBundleABTest( + input: CreateConfigBundleABTestInput, + options: CoreOptions, + ): Promise { + const treatmentWeight = input.treatmentWeight ?? 50; + return this.createABTest( + input.name, + input.gateway, + input.roleArn, + ({ gatewayArn, accountId, roleArn }) => { + const bundleArn = (id: string) => + `arn:aws:bedrock-agentcore:${options.region}:${accountId}:configuration-bundle/${id}`; + return { + name: input.name, + gatewayArn, + variants: [ + { + name: "C", + weight: 100 - treatmentWeight, + variantConfiguration: { + configurationBundle: { + bundleArn: bundleArn(input.control.configBundle), + bundleVersion: input.control.bundleVersion, + }, + }, + }, + { + name: "T1", + weight: treatmentWeight, + variantConfiguration: { + configurationBundle: { + bundleArn: bundleArn(input.treatment.configBundle), + bundleVersion: input.treatment.bundleVersion, + }, + }, + }, + ], + evaluationConfig: { + onlineEvaluationConfigArn: `arn:aws:bedrock-agentcore:${options.region}:${accountId}:online-evaluation-config/${input.onlineEval}`, + }, + roleArn, + gatewayFilter: input.gatewayFilter, + enableOnCreate: input.enableOnCreate ?? true, + clientToken: randomUUID(), + }; + }, + options, + ); + } + + async createTargetBasedABTest( + input: CreateTargetBasedABTestInput, + options: CoreOptions, + ): Promise { + const treatmentWeight = input.treatmentWeight ?? 50; + return this.createABTest( + input.name, + input.gateway, + input.roleArn, + ({ gatewayArn, accountId, roleArn }) => { + const evalArn = (id: string) => + `arn:aws:bedrock-agentcore:${options.region}:${accountId}:online-evaluation-config/${id}`; + return { + name: input.name, + gatewayArn, + variants: [ + { + name: "C", + weight: 100 - treatmentWeight, + variantConfiguration: { target: { name: input.control.gatewayTarget } }, + }, + { + name: "T1", + weight: treatmentWeight, + variantConfiguration: { target: { name: input.treatment.gatewayTarget } }, + }, + ], + evaluationConfig: { + perVariantOnlineEvaluationConfig: [ + { name: "C", onlineEvaluationConfigArn: evalArn(input.control.onlineEval) }, + { name: "T1", onlineEvaluationConfigArn: evalArn(input.treatment.onlineEval) }, + ], + }, + roleArn, + gatewayFilter: input.gatewayFilter, + enableOnCreate: input.enableOnCreate ?? true, + clientToken: randomUUID(), + }; + }, + options, + ); + } + async listBatchInsights( nextToken: string | undefined, maxResults: number | undefined, diff --git a/src/handlers/eval/ab-test/ab-test.test.tsx b/src/handlers/eval/ab-test/ab-test.test.tsx index 3f75f8647..73ac93b64 100644 --- a/src/handlers/eval/ab-test/ab-test.test.tsx +++ b/src/handlers/eval/ab-test/ab-test.test.tsx @@ -79,9 +79,12 @@ describe("eval ab-test command hierarchy", () => { "stop", "delete", "config-bundle", + "target-based", ]); const cb = group?.children().find((c) => c.name() === "config-bundle"); expect(cb?.children().map((c) => c.name())).toEqual(["run"]); + const tb = group?.children().find((c) => c.name() === "target-based"); + expect(tb?.children().map((c) => c.name())).toEqual(["run"]); }); }); @@ -264,3 +267,86 @@ describe("eval ab-test config-bundle run validation", () => { }); }); }); + +describe("eval ab-test target-based run validation", () => { + const TB_BASE = [ + "eval", + "ab-test", + "target-based", + "run", + "--name", + "orders-v2-canary", + "--gateway", + "orders-gateway-abc123", + "--control", + '{"gateway-target":"orders-prod-target","online-eval":"prod-quality"}', + "--treatment", + '{"gateway-target":"orders-v2-target","online-eval":"v2-quality"}', + "--json", + ]; + + test.each(["name", "gateway", "control", "treatment"] as const)( + "requires --%s", + async (missing) => { + const args = TB_BASE.filter( + (a, i) => a !== `--${missing}` && TB_BASE[i - 1] !== `--${missing}`, + ); + await expect(run(args)).rejects.toThrow(new RegExp(`--${missing}`)); + }, + ); + + test("rejects a mis-shaped --control object", async () => { + const args = TB_BASE.map((a) => + a === '{"gateway-target":"orders-prod-target","online-eval":"prod-quality"}' + ? '{"wrong":"shape"}' + : a, + ); + await expect(run(args)).rejects.toThrow(/--control must be/); + }); + + test("rejects identical control/treatment targets", async () => { + const same = '{"gateway-target":"t","online-eval":"e"}'; + await expect( + run([ + "eval", + "ab-test", + "target-based", + "run", + "--name", + "x", + "--gateway", + "g", + "--control", + same, + "--treatment", + same, + "--json", + ]), + ).rejects.toThrow(/different gateway targets/); + }); + + test("maps flags to a createTargetBasedABTest call", async () => { + const { core } = await run([...TB_BASE, "--treatment-weight", "20"], (c) => + c.eval.setAbTestCreateResponse({ + abTestId: "x", + abTestArn: ARN, + name: "x", + status: "CREATING", + executionStatus: "RUNNING", + createdAt: new Date("2026-08-26T10:00:00.000Z"), + }), + ); + const call = core.eval.calls.find((c) => c.method === "createTargetBasedABTest"); + expect(call).toBeDefined(); + expect(call!.args[0]).toEqual({ + name: "orders-v2-canary", + gateway: "orders-gateway-abc123", + control: { gatewayTarget: "orders-prod-target", onlineEval: "prod-quality" }, + treatment: { gatewayTarget: "orders-v2-target", onlineEval: "v2-quality" }, + treatmentWeight: 20, + gatewayFilter: undefined, + roleArn: undefined, + enableOnCreate: undefined, + }); + }); +}); diff --git a/src/handlers/eval/ab-test/index.tsx b/src/handlers/eval/ab-test/index.tsx index 4e42ada3a..42b16fd9d 100644 --- a/src/handlers/eval/ab-test/index.tsx +++ b/src/handlers/eval/ab-test/index.tsx @@ -10,6 +10,7 @@ import { createResumeAbTestHandler } from "./resume"; import { createStopAbTestHandler } from "./stop"; import { createDeleteAbTestHandler } from "./delete"; import { createConfigBundleAbTestHandler } from "./config-bundle"; +import { createTargetBasedAbTestHandler } from "./target-based"; export function createAbTestHandler(core: Core, io: AppIO): Router { return new Router("ab-test", "inspect AgentCore A/B tests") @@ -22,7 +23,8 @@ export function createAbTestHandler(core: Core, io: AppIO): Router { .handler(createResumeAbTestHandler(core)) .handler(createStopAbTestHandler(core)) .handler(createDeleteAbTestHandler(core)) - .handler(createConfigBundleAbTestHandler(core, io)); + .handler(createConfigBundleAbTestHandler(core, io)) + .handler(createTargetBasedAbTestHandler(core, io)); } export { AbTestScreen } from "./screen.tsx"; diff --git a/src/handlers/eval/ab-test/target-based/index.tsx b/src/handlers/eval/ab-test/target-based/index.tsx new file mode 100644 index 000000000..e13f20b18 --- /dev/null +++ b/src/handlers/eval/ab-test/target-based/index.tsx @@ -0,0 +1,10 @@ +import { Router } from "../../../../router"; +import type { AppIO } from "../../../../io"; +import type { Core } from "../../../types"; +import { createTargetBasedRunHandler } from "./run"; + +export function createTargetBasedAbTestHandler(core: Core, io: AppIO): Router { + return new Router("target-based", "target-based A/B tests").handler( + createTargetBasedRunHandler(core, io), + ); +} diff --git a/src/handlers/eval/ab-test/target-based/run/index.tsx b/src/handlers/eval/ab-test/target-based/run/index.tsx new file mode 100644 index 000000000..8ef811eb1 --- /dev/null +++ b/src/handlers/eval/ab-test/target-based/run/index.tsx @@ -0,0 +1,116 @@ +import z from "zod"; +import { createHandler, flag } from "../../../../../router"; +import { InputValidationError } from "../../../../../errors"; +import { JsonRendererKey } from "../../../../../tui"; +import { SourceResolver, type AppIO } from "../../../../../io"; +import type { Core } from "../../../../types"; +import type { TargetVariantRef } from "../../../types"; +import { coreOptsFromCtx, parseJsonFlag } from "../../../../utils"; + +const targetRefSchema = z + .object({ + "gateway-target": z.string().min(1), + "online-eval": z.string().min(1), + }) + .strict(); + +function toTargetRef(name: string, raw: unknown): TargetVariantRef { + const parsed = targetRefSchema.safeParse(raw); + if (!parsed.success) { + throw new InputValidationError( + `--${name} must be {"gateway-target": "", "online-eval": ""}`, + ); + } + return { gatewayTarget: parsed.data["gateway-target"], onlineEval: parsed.data["online-eval"] }; +} + +export const createTargetBasedRunHandler = (core: Core, io: AppIO) => + createHandler({ + name: "run", + description: "run an A/B test between two gateway targets and their online evaluations", + flags: [ + flag("name", "the A/B test name", z.string().optional()), + flag("gateway", "deployed gateway id", z.string().optional()), + flag( + "control", + 'control JSON {"gateway-target","online-eval"} (inline, file://, or -)', + z.string().optional(), + ), + flag( + "treatment", + 'treatment JSON {"gateway-target","online-eval"} (inline, file://, or -)', + z.string().optional(), + ), + flag( + "treatment-weight", + "1-99; control weight = 100 - this (default 50)", + z.number().int().optional(), + ), + flag( + "gateway-filter", + 'GatewayFilter JSON, e.g. {"targetPaths":["/orders"]} (inline, file://, or -)', + z.string().optional(), + ), + flag( + "role-arn", + "execution-role override (default: auto-provisioned)", + z.string().optional(), + ), + flag( + "enable-on-create", + "whether to start the test immediately (default true; pass false to create it paused)", + z.enum(["true", "false"]).optional(), + ), + ], + handle: async (ctx, flags) => { + const required = ["name", "gateway", "control", "treatment"] as const; + for (const f of required) { + if (!flags[f]) throw new InputValidationError(`required option '--${f}' not specified`); + } + + const source = new SourceResolver({ stdin: io.stdin }); + const controlRaw = parseJsonFlag( + "control", + await source.resolveText("control", flags["control"]), + ); + const treatmentRaw = parseJsonFlag( + "treatment", + await source.resolveText("treatment", flags["treatment"]), + ); + const gatewayFilter = parseJsonFlag< + import("@aws-sdk/client-bedrock-agentcore").GatewayFilter + >("gateway-filter", await source.resolveText("gateway-filter", flags["gateway-filter"])); + + const control = toTargetRef("control", controlRaw); + const treatment = toTargetRef("treatment", treatmentRaw); + if (control.gatewayTarget === treatment.gatewayTarget) { + throw new InputValidationError( + "control and treatment must reference different gateway targets", + ); + } + + const treatmentWeight = flags["treatment-weight"]; + if (treatmentWeight !== undefined && (treatmentWeight < 1 || treatmentWeight > 99)) { + throw new InputValidationError("--treatment-weight must be between 1 and 99"); + } + + const result = await core.eval.createTargetBasedABTest( + { + name: flags["name"]!, + gateway: flags["gateway"]!, + control, + treatment, + treatmentWeight, + gatewayFilter, + roleArn: flags["role-arn"], + enableOnCreate: + flags["enable-on-create"] === undefined + ? undefined + : flags["enable-on-create"] === "true", + }, + coreOptsFromCtx(ctx), + ); + + ctx.require(JsonRendererKey).renderJson(result); + }, + }); diff --git a/src/handlers/eval/types.tsx b/src/handlers/eval/types.tsx index b13c6f9cd..91a70c705 100644 --- a/src/handlers/eval/types.tsx +++ b/src/handlers/eval/types.tsx @@ -247,6 +247,19 @@ export type CreateConfigBundleABTestInput = { enableOnCreate?: boolean; }; +export type TargetVariantRef = { gatewayTarget: string; onlineEval: string }; + +export type CreateTargetBasedABTestInput = { + name: string; + gateway: string; + control: TargetVariantRef; + treatment: TargetVariantRef; + treatmentWeight?: number; + gatewayFilter?: GatewayFilter; + roleArn?: string; + enableOnCreate?: boolean; +}; + export type CreateDatasetInput = CreateDatasetRequest; export type StartRecommendationInput = { name: string; @@ -447,6 +460,10 @@ export interface CoreEvalClient { input: CreateConfigBundleABTestInput, options: CoreOptions, ): Promise; + createTargetBasedABTest( + input: CreateTargetBasedABTestInput, + options: CoreOptions, + ): Promise; // startBatchEvaluation submits an async, service-side evaluation over sessions // the service gathers from the resolved data source. Returns the durable job id // + RUNNING status; poll with getBatchEvaluation. diff --git a/src/testing/TestCoreClient.tsx b/src/testing/TestCoreClient.tsx index d547dd956..41bd264ca 100644 --- a/src/testing/TestCoreClient.tsx +++ b/src/testing/TestCoreClient.tsx @@ -139,6 +139,7 @@ import type { CoreEvalClient, CreateConfigurationBundleInput, CreateConfigBundleABTestInput, + CreateTargetBasedABTestInput, CreateDatasetInput, CreateOnlineEvalInput, CreateOnlineInsightInput, @@ -1911,6 +1912,15 @@ export class TestEvalClient implements CoreEvalClient { return this.abTestCreateResponse; } + async createTargetBasedABTest( + input: CreateTargetBasedABTestInput, + options: CoreOptions, + ): Promise { + this.calls.push({ method: "createTargetBasedABTest", args: [input, options] }); + if (this.error) throw this.error; + return this.abTestCreateResponse; + } + async startBatchEvaluation( input: StartBatchEvaluationInput, options: CoreOptions,