Replicate Nodes Migration — Implementation Plan

// tests/replicate-base.test.ts
import { describe, it, expect } from "vitest";
import { getReplicateApiKey, removeNulls } from "../src/replicate-base.js";

describe("getReplicateApiKey", () => {
  it("reads from _secrets", () => {
    const key = getReplicateApiKey({
      _secrets: { REPLICATE_API_TOKEN: "test-key" },
    });
    expect(key).toBe("test-key");
  });

  it("throws when no key configured", () => {
    expect(() => getReplicateApiKey({})).toThrow("REPLICATE_API_TOKEN");
  });
});

describe("removeNulls", () => {
  it("removes null and undefined values", () => {
    const obj = { a: 1, b: null, c: undefined, d: "ok" };
    removeNulls(obj);
    expect(obj).toEqual({ a: 1, d: "ok" });
  });
});

// src/replicate-base.ts
/**
 * Shared Replicate API utilities.
 * Uses plain fetch() — no SDK dependency.
 */

const REPLICATE_API_BASE = "https://api.replicate.com/v1";

// ---------------------------------------------------------------------------
// API Key extraction
// ---------------------------------------------------------------------------

export function getReplicateApiKey(inputs: Record<string, unknown>): string {
  const key =
    (inputs._secrets as Record<string, string>)?.REPLICATE_API_TOKEN ||
    process.env.REPLICATE_API_TOKEN ||
    "";
  if (!key) throw new Error("REPLICATE_API_TOKEN is not configured");
  return key;
}

// ---------------------------------------------------------------------------
// Remove null/undefined from args before sending to API
// ---------------------------------------------------------------------------

export function removeNulls(obj: Record<string, unknown>): void {
  for (const key of Object.keys(obj)) {
    if (obj[key] === null || obj[key] === undefined) {
      delete obj[key];
    }
  }
}

// ---------------------------------------------------------------------------
// Check if an asset ref has actual content
// ---------------------------------------------------------------------------

export function isRefSet(
  ref: Record<string, unknown> | null | undefined
): boolean {
  if (!ref || typeof ref !== "object") return false;
  return Boolean(ref.uri || ref.data);
}

// ---------------------------------------------------------------------------
// Convert an asset ref to a URL string for the Replicate API
// ---------------------------------------------------------------------------

export async function assetToUrl(
  ref: Record<string, unknown>
): Promise<string | null> {
  if (!ref) return null;
  if (typeof ref.uri === "string" && ref.uri) return ref.uri;
  if (typeof ref.data === "string" && ref.data.startsWith("data:"))
    return ref.data;
  return null;
}

// ---------------------------------------------------------------------------
// Submit prediction and poll until complete
// ---------------------------------------------------------------------------

export interface ReplicatePrediction {
  id: string;
  status: string;
  output: unknown;
  error: string | null;
  logs: string;
  metrics?: Record<string, number>;
}

export async function replicateSubmit(
  apiKey: string,
  modelVersion: string,
  input: Record<string, unknown>
): Promise<ReplicatePrediction> {
  // Create prediction
  const createRes = await fetch(`${REPLICATE_API_BASE}/predictions`, {
    method: "POST",
    headers: {
      Authorization: `Bearer ${apiKey}`,
      "Content-Type": "application/json",
      Prefer: "wait",
    },
    body: JSON.stringify({ version: modelVersion, input }),
  });

  if (!createRes.ok) {
    const text = await createRes.text();
    throw new Error(
      `Replicate API error (${createRes.status}): ${text}`
    );
  }

  let prediction = (await createRes.json()) as ReplicatePrediction;

  // If "Prefer: wait" returned a completed prediction, return immediately
  if (
    prediction.status === "succeeded" ||
    prediction.status === "failed" ||
    prediction.status === "canceled"
  ) {
    if (prediction.status === "failed") {
      throw new Error(prediction.error || "Prediction failed");
    }
    return prediction;
  }

  // Poll for completion (max 30 minutes, 2s interval)
  for (let i = 0; i < 900; i++) {
    await new Promise((r) => setTimeout(r, 2000));

    const pollRes = await fetch(
      `${REPLICATE_API_BASE}/predictions/${prediction.id}`,
      {
        headers: { Authorization: `Bearer ${apiKey}` },
      }
    );
    if (!pollRes.ok) continue;

    prediction = (await pollRes.json()) as ReplicatePrediction;

    if (prediction.status === "failed" || prediction.status === "canceled") {
      throw new Error(prediction.error || "Prediction failed");
    }
    if (prediction.status === "succeeded") {
      return prediction;
    }
  }

  throw new Error("Replicate prediction timed out after 30 minutes");
}

// ---------------------------------------------------------------------------
// Extract version hash from model ID "owner/name:version"
// ---------------------------------------------------------------------------

export function extractVersion(modelId: string): string {
  const colonIdx = modelId.indexOf(":");
  if (colonIdx < 0) {
    throw new Error(
      `Invalid Replicate model ID: ${modelId}. Expected format: owner/name:version`
    );
  }
  return modelId.slice(colonIdx + 1);
}

// ---------------------------------------------------------------------------
// Convert Replicate output to asset ref
// ---------------------------------------------------------------------------

export function outputToImageRef(output: unknown): Record<string, unknown> {
  if (typeof output === "string") return { type: "image", uri: output };
  if (Array.isArray(output) && output.length > 0)
    return { type: "image", uri: String(output[0]) };
  return { type: "image", uri: "" };
}

export function outputToVideoRef(output: unknown): Record<string, unknown> {
  if (typeof output === "string") return { type: "video", uri: output };
  if (Array.isArray(output) && output.length > 0)
    return { type: "video", uri: String(output[0]) };
  return { type: "video", uri: "" };
}

export function outputToAudioRef(output: unknown): Record<string, unknown> {
  if (typeof output === "string") return { type: "audio", uri: output };
  if (Array.isArray(output) && output.length > 0)
    return { type: "audio", uri: String(output[0]) };
  return { type: "audio", uri: "" };
}

export function outputToString(output: unknown): string {
  if (typeof output === "string") return output;
  if (Array.isArray(output)) return output.join("");
  return String(output ?? "");
}

// src/index.ts
import type { NodeClass } from "@nodetool/node-sdk";
import { NodeRegistry } from "@nodetool/node-sdk";

// Generated node arrays will be imported here after codegen runs.
// For now, export an empty array.
export const REPLICATE_NODES: readonly NodeClass[] = [];

export function registerReplicateNodes(registry: NodeRegistry): void {
  for (const nodeClass of REPLICATE_NODES) {
    registry.register(nodeClass);
  }
}

// src/schema-fetcher.ts
import { readFile, writeFile, mkdir } from "node:fs/promises";
import { existsSync } from "node:fs";
import { join } from "node:path";

const REPLICATE_API = "https://api.replicate.com/v1";

export interface ReplicateSchema {
  modelId: string;       // "owner/name:version"
  owner: string;
  name: string;
  version: string;
  description: string;
  inputSchema: Record<string, unknown>;
  outputSchema: Record<string, unknown>;
  hardware?: string;
}

export class SchemaFetcher {
  private cacheDir = join(process.cwd(), ".schema-cache", "replicate");

  async fetchSchema(
    modelId: string,
    useCache = true
  ): Promise<ReplicateSchema> {
    const cacheKey = modelId.replace(/[/:]/g, "_");
    const cachePath = join(this.cacheDir, `${cacheKey}.json`);

    if (useCache && existsSync(cachePath)) {
      const raw = await readFile(cachePath, "utf8");
      return JSON.parse(raw) as ReplicateSchema;
    }

    const colonIdx = modelId.indexOf(":");
    const ownerName = colonIdx > 0 ? modelId.slice(0, colonIdx) : modelId;
    const [owner, name] = ownerName.split("/");

    // Fetch latest version if no version specified
    const token = process.env.REPLICATE_API_TOKEN ?? "";
    const headers: Record<string, string> = {};
    if (token) headers.Authorization = `Bearer ${token}`;

    let version: string;
    let description = "";
    let inputSchema: Record<string, unknown> = {};
    let outputSchema: Record<string, unknown> = {};

    if (colonIdx > 0) {
      version = modelId.slice(colonIdx + 1);
    } else {
      // Fetch model to get latest version
      const modelRes = await fetch(
        `${REPLICATE_API}/models/${owner}/${name}`,
        { headers }
      );
      if (!modelRes.ok) throw new Error(`Failed to fetch model ${modelId}`);
      const modelData = (await modelRes.json()) as Record<string, unknown>;
      const latestVersion = modelData.latest_version as Record<string, unknown>;
      version = latestVersion?.id as string ?? "";
      description = (modelData.description as string) ?? "";
    }

    // Fetch version details with schema
    const versionRes = await fetch(
      `${REPLICATE_API}/models/${owner}/${name}/versions/${version}`,
      { headers }
    );
    if (!versionRes.ok)
      throw new Error(`Failed to fetch version for ${modelId}`);
    const versionData = (await versionRes.json()) as Record<string, unknown>;

    const openApiSchema = versionData.openapi_schema as Record<string, unknown>;
    if (openApiSchema) {
      const components = openApiSchema.components as Record<string, unknown>;
      const schemas = components?.schemas as Record<string, unknown>;
      if (schemas) {
        inputSchema = (schemas.Input as Record<string, unknown>) ?? {};
        outputSchema = (schemas.Output as Record<string, unknown>) ?? {};
      }
    }

    const result: ReplicateSchema = {
      modelId: `${owner}/${name}:${version}`,
      owner,
      name,
      version,
      description: description || (versionData.description as string) ?? "",
      inputSchema,
      outputSchema,
    };

    // Cache
    if (useCache) {
      await mkdir(this.cacheDir, { recursive: true });
      await writeFile(cachePath, JSON.stringify(result, null, 2));
    }

    return result;
  }
}

export class FluxSchnell extends ReplicateNode {
  static readonly nodeType = "replicate.image_generate.FluxSchnell";
  static readonly title = "Flux Schnell";
  static readonly description = "...";
  static readonly requiredSettings = ["REPLICATE_API_TOKEN"];

  @prop({ type: "str", default: "", description: "..." })
  declare prompt: any;

  async process(inputs: Record<string, unknown>): Promise<Record<string, unknown>> {
    const apiKey = getReplicateApiKey(inputs);
    const prompt = String(inputs.prompt ?? this.prompt ?? "");
    const args: Record<string, unknown> = { prompt };
    // ... handle asset inputs ...
    removeNulls(args);
    const res = await replicateSubmit(apiKey, extractVersion("owner/name:version"), args);
    return { output: outputToImageRef(res.output) };
  }
}

"black-forest-labs/flux-schnell": {
  className: "FluxSchnell",
  returnType: "image",
  fieldOverrides: {
    image: { propType: "image" },
  },
}

// src/configs/index.ts
import type { ModuleConfig } from "../types.js";
import { imageGenerateConfig } from "./image-generate.js";
import { imageUpscaleConfig } from "./image-upscale.js";
// ... all other configs

export const allConfigs: Record<string, ModuleConfig> = {
  "image-generate": imageGenerateConfig,
  "image-upscale": imageUpscaleConfig,
  // ... all modules
};

#!/usr/bin/env node
/**
 * CLI for generating Replicate node TypeScript classes.
 *
 * Usage:
 *   npx tsx src/generate.ts --all
 *   npx tsx src/generate.ts --module image-generate
 *   npx tsx src/generate.ts --all --no-cache
 */

import { parseArgs } from "node:util";
import { mkdir, writeFile } from "node:fs/promises";
import { join } from "node:path";
import { SchemaFetcher } from "./schema-fetcher.js";
import { SchemaParser } from "./schema-parser.js";
import { NodeGenerator } from "./node-generator.js";
import { allConfigs } from "./configs/index.js";
import type { ModuleConfig } from "./types.js";

// ... same structure as fal-codegen generate.ts
// Output dir defaults to ../replicate-nodes/src/generated

import type { NodeClass } from "@nodetool/node-sdk";
import { NodeRegistry } from "@nodetool/node-sdk";

import { REPLICATE_IMAGE_GENERATE_NODES } from "./generated/image-generate.js";
import { REPLICATE_IMAGE_UPSCALE_NODES } from "./generated/image-upscale.js";
// ... all generated modules

export const REPLICATE_NODES: readonly NodeClass[] = [
  ...REPLICATE_IMAGE_GENERATE_NODES,
  ...REPLICATE_IMAGE_UPSCALE_NODES,
  // ... all
];

export function registerReplicateNodes(registry: NodeRegistry): void {
  for (const nodeClass of REPLICATE_NODES) {
    registry.register(nodeClass);
  }
}

import { registerReplicateNodes } from "@nodetool/replicate-nodes";
// ...
registerReplicateNodes(registry);

import { describe, it, expect } from "vitest";
import { REPLICATE_NODES } from "../src/index.js";

describe("Generated Replicate nodes", () => {
  it("exports at least 150 nodes", () => {
    expect(REPLICATE_NODES.length).toBeGreaterThanOrEqual(150);
  });

  it("all nodes have valid nodeType", () => {
    for (const node of REPLICATE_NODES) {
      expect(node.nodeType).toMatch(/^replicate\./);
      expect(node.title).toBeTruthy();
    }
  });

  it("all nodes have requiredSettings with REPLICATE_API_TOKEN", () => {
    for (const node of REPLICATE_NODES) {
      expect(node.requiredSettings).toContain("REPLICATE_API_TOKEN");
    }
  });

  it("all nodes can be instantiated", () => {
    for (const NodeClass of REPLICATE_NODES) {
      const instance = new NodeClass({});
      expect(instance).toBeDefined();
      expect(typeof instance.process).toBe("function");
    }
  });
});