ai-chatbot/lib/ai/vip98k-task-invocation.ts

export type Vip98kTaskInvocation = {
  kind: string;
  method: string;
  path: string;
  contentType: string;
  requestExample: Record<string, unknown>;
  resultMethod: string;
  resultPath: string;
};

type KnownImageInput = "images" | "video_clips";

const IMAGE_ARRAY_INPUT_KEYS = [
  "file_urls",
  "images",
  "imageUrls",
  "referenceImages",
  "reference_images",
] as const;

const IMAGE_SINGLE_INPUT_KEYS = [
  "image",
  "imageUrl",
  "image_url",
  "firstFrameImage",
  "first_frame_image",
  "firstFrameImageUrl",
  "first_frame_image_url",
] as const;

const KNOWN_IMAGE_TOOL_INPUTS: Record<string, KnownImageInput> = {
  "google/nano-banana-2/reference-to-image": "video_clips",
  "google/nano-banana-pro/edit": "images",
  "openai/gpt-image-2/edit": "images",
};

const KNOWN_IMAGE_PROCESSING_TOOL_REQUESTS: Record<
  string,
  Record<string, unknown>
> = {
  "remove-background": { image_url: "" },
  enhance: { image_url: "", scale_factor: 4 },
};

function getInvocationToolId(invocation: Vip98kTaskInvocation | null) {
  if (!invocation || invocation.kind !== "tool") return null;
  const query = invocation.path.split("?", 2)[1] || "";
  return new URLSearchParams(query).get("tool_id");
}

function invocationTargetsModel(
  invocation: Vip98kTaskInvocation | null,
  modelId: string,
) {
  if (!invocation) return false;
  const expected = modelId.trim().toLowerCase();
  if (!expected) return false;
  const requestModel = String(invocation.requestExample.model || "")
    .trim()
    .toLowerCase();
  if (requestModel) return requestModel === expected;
  return getInvocationToolId(invocation)?.toLowerCase() === expected;
}

export function getVip98kImageBindingKeys(invocation: Vip98kTaskInvocation | null) {
  if (!invocation) return [];
  const body = invocation.requestExample || {};
  return [
    ...IMAGE_ARRAY_INPUT_KEYS.filter((key) => key in body),
    ...IMAGE_SINGLE_INPUT_KEYS.filter((key) => key in body),
    ...(Array.isArray(body.video_clips) ? ["video_clips" as const] : []),
  ];
}

export function invocationSupportsVip98kImages(
  invocation: Vip98kTaskInvocation | null,
) {
  return getVip98kImageBindingKeys(invocation).length > 0;
}

export function bindVip98kImages(
  body: Record<string, unknown>,
  imageUrls: string[],
) {
  const bindingKeys: string[] = [];
  if (imageUrls.length === 0) return bindingKeys;

  for (const key of IMAGE_ARRAY_INPUT_KEYS) {
    if (!(key in body)) continue;
    body[key] = imageUrls;
    bindingKeys.push(key);
  }
  for (const key of IMAGE_SINGLE_INPUT_KEYS) {
    if (!(key in body)) continue;
    body[key] = imageUrls[0];
    bindingKeys.push(key);
  }
  if (Array.isArray(body.video_clips)) {
    const first = body.video_clips[0];
    const template =
      first && typeof first === "object" && !Array.isArray(first)
        ? (first as Record<string, unknown>)
        : {};
    body.video_clips = imageUrls.map((url) => ({ ...template, url }));
    bindingKeys.push("video_clips");
  }

  const secondImageBindings = [
    "lastFrameImage",
    "last_frame_image",
    "endImageUrl",
    "end_image_url",
  ] as const;
  if (imageUrls[1]) {
    for (const key of secondImageBindings) {
      if (!(key in body)) continue;
      body[key] = imageUrls[1];
      bindingKeys.push(key);
    }
  }
  return bindingKeys;
}

function scoreReferenceInvocation(invocation: Vip98kTaskInvocation) {
  const body = invocation.requestExample || {};
  let score = 0;
  if ("file_urls" in body) score += 60;
  if (IMAGE_ARRAY_INPUT_KEYS.some((key) => key in body)) score += 40;
  if (Array.isArray(body.video_clips)) score += 35;
  if (IMAGE_SINGLE_INPUT_KEYS.some((key) => key in body)) score += 15;
  if (invocation.kind === "task") score += 10;
  if (invocation.path === "/extensions/v1/tasks") score += 10;
  return score;
}

export function selectVip98kTaskInvocation(params: {
  modelId: string;
  configuredInvocation: Vip98kTaskInvocation | null;
  configuredInvocations?: Vip98kTaskInvocation[];
  imageCount?: number;
}) {
  const candidates = [
    ...(params.configuredInvocations || []),
    ...(params.configuredInvocation ? [params.configuredInvocation] : []),
  ].filter(
    (candidate, index, items) =>
      candidate.path.startsWith("/") &&
      items.findIndex(
        (item) =>
          item.method === candidate.method &&
          item.path === candidate.path &&
          JSON.stringify(item.requestExample) === JSON.stringify(candidate.requestExample),
      ) === index,
  );

  if ((params.imageCount || 0) > 0) {
    const referenceInvocation = candidates
      .filter(invocationSupportsVip98kImages)
      .sort((a, b) => scoreReferenceInvocation(b) - scoreReferenceInvocation(a))[0];
    if (referenceInvocation) {
      return resolveVip98kTaskInvocation(params.modelId, referenceInvocation);
    }
  }

  return resolveVip98kTaskInvocation(
    params.modelId,
    params.configuredInvocation || candidates[0] || null,
  );
}

function createKnownImageToolInvocation(
  modelId: string,
  imageInput: KnownImageInput
): Vip98kTaskInvocation {
  const encodedModelId = encodeURIComponent(modelId);
  const requestExample =
    imageInput === "video_clips"
      ? {
          prompt: "",
          video_clips: [{ fps: 1, url: "", ends: 0, start: 0 }],
        }
      : { prompt: "", images: [""] };

  return {
    kind: "tool",
    method: "POST",
    path: `/api/external/tools/generate?tool_id=${encodedModelId}`,
    contentType: "application/json",
    requestExample,
    resultMethod: "GET",
    resultPath: `/api/external/tools/result?tool_id=${encodedModelId}&taskId={task_id}`,
  };
}

function createKnownProcessingToolInvocation(
  modelId: string,
  requestExample: Record<string, unknown>
): Vip98kTaskInvocation {
  const encodedModelId = encodeURIComponent(modelId);
  return {
    kind: "tool",
    method: "POST",
    path: `/api/external/tools/generate?tool_id=${encodedModelId}`,
    contentType: "application/json",
    requestExample: { ...requestExample },
    resultMethod: "GET",
    resultPath: `/api/external/tools/result?tool_id=${encodedModelId}&taskId={task_id}`,
  };
}

/**
 * A frontend alias may resolve to a different VIP98K tool at runtime. In that
 * case the selected model's invocation cannot be reused because its endpoint
 * and reference-image field belong to the alias rather than the resolved tool.
 */
export function resolveVip98kTaskInvocation(
  modelId: string,
  configuredInvocation: Vip98kTaskInvocation | null
) {
  const normalizedModelId = modelId.trim().toLowerCase();
  const processingRequest = KNOWN_IMAGE_PROCESSING_TOOL_REQUESTS[normalizedModelId];
  if (processingRequest) {
    if (invocationTargetsModel(configuredInvocation, normalizedModelId)) {
      return configuredInvocation;
    }
    return createKnownProcessingToolInvocation(modelId, processingRequest);
  }

  const knownImageInput = KNOWN_IMAGE_TOOL_INPUTS[normalizedModelId];
  if (!knownImageInput) return configuredInvocation;

  if (invocationTargetsModel(configuredInvocation, normalizedModelId)) {
    return configuredInvocation;
  }

  return createKnownImageToolInvocation(modelId, knownImageInput);
}