UHD_Media_Tool/packages/agent/test/vision.test.ts

import { afterEach, describe, expect, it } from "vitest";
import { z } from "zod";
import {
  type AgentRunHandle,
  CapabilityRegistry,
  createFakeRuntime,
  type FakeScript,
  toolContent,
} from "../src/index.ts";
import { settings, tempDir } from "./helpers.ts";

// VIS-1:能力输出里的 modelImage 以 image content 交给支持视觉的模型;不支持时降级为文字提示,图片数据不进 JSON 文本。

const PNG_1X1 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==";

const opened: AgentRunHandle[] = [];
afterEach(async () => {
  await Promise.all(opened.splice(0).map((r) => r.close()));
});

function registry(seen: Array<boolean | undefined>) {
  const r = new CapabilityRegistry();
  r.register(
    {
      id: "browser.screenshot",
      version: 1,
      title: "截图",
      description: "截取页面",
      inputSchema: z.object({ sessionId: z.string() }),
      outputSchema: z.object({
        url: z.string(),
        modelImage: z
          .object({ mimeType: z.string(), data: z.string(), width: z.number(), height: z.number() })
          .nullable(),
      }),
      kind: "read",
      scope: "read",
      riskLevel: "low",
      idempotent: true,
      supportsPreview: false,
      replay: "safe",
    },
    (_input, ctx) => {
      seen.push(ctx.modelImages);
      return {
        url: "https://site.example/login",
        modelImage: { mimeType: "image/png", data: PNG_1X1, width: 1, height: 1 },
      };
    },
  );
  return r;
}

const script: FakeScript = {
  match: /看页面/,
  steps: [{ toolCalls: [{ capability: "browser.screenshot", input: { sessionId: "s1" } }] }],
  finalText: (r) => `images=${r[0]?.images};text=${r[0]?.text}`,
};

async function run(vision: boolean) {
  const seen: Array<boolean | undefined> = [];
  const { runtime, fake } = createFakeRuntime({ registry: registry(seen), scripts: [script], fake: { vision } });
  const handle = await runtime.createRun({
    runId: `vision-${vision}`,
    goal: "看页面",
    systemPrompt: "测试",
    capabilities: ["browser.screenshot"],
    settings: settings(),
    sessionDir: await tempDir(),
    actor: { actorId: "user:1", taskId: "task:1", scopeIds: [] },
  });
  opened.push(handle);
  const outcome = await handle.prompt();
  return { handle, fake, outcome, seen };
}

describe("VIS-1 图片内容块", () => {
  it("支持视觉的模型:截图作为 image content 交给模型,JSON 文本里没有图片数据", async () => {
    const { handle, fake, outcome, seen } = await run(true);
    expect(handle.supportsImages).toBe(true);
    expect(seen).toEqual([true]);
    expect(fake.images).toEqual([{ mimeType: "image/png", bytes: PNG_1X1.length }]);
    expect(outcome.finalText).toContain("images=1");
    expect(outcome.finalText).not.toContain(PNG_1X1);
  });

  it("不支持视觉的模型:能力收到 modelImages=false,图片不发送并提示", async () => {
    const { handle, fake, outcome, seen } = await run(false);
    expect(handle.supportsImages).toBe(false);
    expect(seen).toEqual([false]);
    expect(fake.images).toEqual([]);
    expect(outcome.finalText).toContain("images=0");
    expect(outcome.finalText).not.toContain(PNG_1X1);
  });

  it("toolContent:摘出 modelImage;null 时只有文本", () => {
    const img = toolContent({ a: 1, modelImage: { mimeType: "image/jpeg", data: "QUJD" } }, true);
    expect(img).toEqual([
      { type: "text", text: '{"a":1}' },
      { type: "image", data: "QUJD", mimeType: "image/jpeg" },
    ]);
    const none = toolContent({ a: 1, modelImage: null }, true);
    expect(none).toEqual([{ type: "text", text: '{"a":1}' }]);
    const text = toolContent({ a: 1, modelImage: { mimeType: "image/jpeg", data: "QUJD" } }, false);
    expect(text).toHaveLength(2);
    expect(text[1]).toMatchObject({ type: "text" });
    expect(JSON.stringify(text)).not.toContain("QUJD");
  });
});