|
| 1 | +import { describe, expect, it, beforeEach, afterEach, vi } from "vitest"; |
| 2 | +import { BaseProvider } from "@nodetool-ai/runtime"; |
| 3 | +import type { ProcessingContext, ProviderId, VideoModel } from "@nodetool-ai/runtime"; |
| 4 | +import { Asset, ModelObserver, Storyboard, initTestDb } from "@nodetool-ai/models"; |
| 5 | +import type { Shot } from "@nodetool-ai/protocol"; |
| 6 | +import { toolForCapabilityName } from "../src/capabilities/lazy-tool.js"; |
| 7 | +import { createCapabilityRun, UNGATED } from "../src/capabilities/invoke.js"; |
| 8 | +import { withGenerationSeam } from "./_helpers/generation-seam.js"; |
| 9 | + |
| 10 | +class ReferenceVideoProvider extends BaseProvider { |
| 11 | + constructor(private readonly models: VideoModel[]) { |
| 12 | + super("fal_ai" as ProviderId); |
| 13 | + } |
| 14 | + override async referenceToVideo(): Promise<Uint8Array> { return new Uint8Array(); } |
| 15 | + override async getAvailableVideoModels(): Promise<VideoModel[]> { |
| 16 | + return this.models; |
| 17 | + } |
| 18 | +} |
| 19 | + |
| 20 | +const ctx = { userId: "u1" } as ProcessingContext; |
| 21 | +const shot = (overrides: Partial<Shot> & { id: string; index: number }): Shot => ({ |
| 22 | + type: "shot", |
| 23 | + action: `action ${overrides.index}`, |
| 24 | + status: "planned", |
| 25 | + ...overrides |
| 26 | +}); |
| 27 | + |
| 28 | +function renderContext(referenceIds: [string, string]) { |
| 29 | + return withGenerationSeam({ |
| 30 | + userId: "u1", |
| 31 | + runProviderPrediction: vi.fn(async () => new Uint8Array([0, 0, 0, 24, 102, 116, 121, 112])), |
| 32 | + hasModelInterface: () => true, |
| 33 | + createAsset: vi.fn(async (args: { name: string; contentType: string; content: Uint8Array }) => { |
| 34 | + return Asset.create<Asset>({ user_id: "u1", name: args.name, content_type: args.contentType }); |
| 35 | + }), |
| 36 | + resolveAssetBytes: vi.fn(async (uri: string) => ({ |
| 37 | + bytes: uri.includes(referenceIds[0]) ? new Uint8Array([1, 2]) : new Uint8Array([3, 4]) |
| 38 | + })) |
| 39 | + }) as ProcessingContext & { runProviderPrediction: ReturnType<typeof vi.fn> }; |
| 40 | +} |
| 41 | + |
| 42 | +describe("reference_to_video capability contracts", () => { |
| 43 | + beforeEach(() => initTestDb()); |
| 44 | + afterEach(() => ModelObserver.clear()); |
| 45 | + |
| 46 | + it("find_model filters video models by reference_to_video", async () => { |
| 47 | + const provider = new ReferenceVideoProvider([ |
| 48 | + { id: "image-model", name: "Image model", provider: "fal_ai", supportedTasks: ["image_to_video"] }, |
| 49 | + { id: "reference-model", name: "Reference model", provider: "fal_ai", supportedTasks: ["reference_to_video"] } |
| 50 | + ]); |
| 51 | + const tool = toolForCapabilityName("find_model", (context) => |
| 52 | + createCapabilityRun({ context, gate: UNGATED, providers: { fal_ai: provider } }) |
| 53 | + ); |
| 54 | + const result = (await tool.process(ctx, { capability: "reference_to_video" })) as { |
| 55 | + results: Array<{ model_id: string }>; |
| 56 | + }; |
| 57 | + expect(result.results.map((model) => model.model_id)).toEqual(["reference-model"]); |
| 58 | + }); |
| 59 | + |
| 60 | + it("render_storyboard_clips dispatches a reference shot through reference_to_video", async () => { |
| 61 | + const first = await Asset.create<Asset>({ |
| 62 | + user_id: "u1", name: "ref-a.png", content_type: "image/png", |
| 63 | + metadata: { nodetool_entity: { kind: "character", name: "A", descriptor: "first" } } |
| 64 | + }); |
| 65 | + const second = await Asset.create<Asset>({ |
| 66 | + user_id: "u1", name: "ref-b.png", content_type: "image/png", |
| 67 | + metadata: { nodetool_entity: { kind: "character", name: "B", descriptor: "second" } } |
| 68 | + }); |
| 69 | + const board = await Storyboard.create<Storyboard>({ |
| 70 | + user_id: "u1", project_id: "default", name: "Reference board", |
| 71 | + document: JSON.stringify({ |
| 72 | + screenplay: null, |
| 73 | + shots: [shot({ id: "s1", index: 0, render_mode: "reference", entity_ids: [first.id, second.id] })], |
| 74 | + brief: "", style: "", entityIds: [first.id, second.id], aspectRatio: "16:9", |
| 75 | + directorModel: null, imageModel: null, |
| 76 | + videoModel: { type: "video_model", id: "reference-model", provider: "fal_ai" } |
| 77 | + }) |
| 78 | + }); |
| 79 | + const context = renderContext([first.id, second.id]); |
| 80 | + const result = (await toolForCapabilityName("render_storyboard_clips").process(context, { |
| 81 | + storyboard_id: board.id |
| 82 | + })) as { rendered: number }; |
| 83 | + expect(result.rendered).toBe(1); |
| 84 | + expect(context.runProviderPrediction.mock.calls[0][0]).toMatchObject({ |
| 85 | + capability: "reference_to_video", |
| 86 | + model: "reference-model", |
| 87 | + params: { reference_images: [new Uint8Array([1, 2]), new Uint8Array([3, 4])] } |
| 88 | + }); |
| 89 | + }); |
| 90 | +}); |
0 commit comments