[Inference] Implement a "1 class = 1 provider<>task pair" logic to isolate provider-specific code

packages/inference/src/lib/getProviderHelper.ts

@@ -0,0 +1,94 @@
+import { BlackForestLabsTextToImageTask } from "../providers/black-forest-labs";
+import { FalAITask, FalAITextToImageTask, FalAITextToVideoTask } from "../providers/fal-ai";
+import { FireworksConversationalTask } from "../providers/fireworks-ai";
+import {
+	HyperbolicConversationalTask,
+	HyperbolicTextGenerationTask,
+	HyperbolicTextToImageTask,
+} from "../providers/hyperbolic";
+import { NebiusTextToImageTask } from "../providers/nebius";
+import { NovitaConversationalTask, NovitaTextGenerationTask } from "../providers/novita";
+import { OpenAIConversationalTask } from "../providers/openai";
+import type { TaskProviderHelper } from "../providers/providerHelper";
+import { BaseConversationalTask, BaseTextGenerationTask } from "../providers/providerHelper";
+import { ReplicateTextToImageTask, ReplicateTextToSpeechTask, ReplicateTextToVideoTask } from "../providers/replicate";
+import { TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask } from "../providers/together";
+import type { InferenceProvider, InferenceTask } from "../types";
+
+export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask, TaskProviderHelper>>> = {
+	"black-forest-labs": {
+		"text-to-image": new BlackForestLabsTextToImageTask(),
+	},
+	cerebras: {
+		conversational: new BaseConversationalTask("cerebras", "https://api.cerebras.ai"),
+	},
+	cohere: {
+		conversational: new BaseConversationalTask("cohere", "https://api.cohere.ai"),
+	},
+	"fal-ai": {
+		// TODO: Add automatic-speech-recognition task helper
+		// "automatic-speech-recognition": new FalAIAutomaticSpeechRecognitionTask(),
+		"text-to-image": new FalAITextToImageTask(),
+		"text-to-speech": new FalAITask("text-to-speech"),
+		"text-to-video": new FalAITextToVideoTask(),
+	},
+	"fireworks-ai": {
+		conversational: new FireworksConversationalTask(),
+	},
+	"hf-inference": {
+		//TODO: Add the correct provider helpers for hf-inference
+	},
+	hyperbolic: {
+		"text-to-image": new HyperbolicTextToImageTask(),
+		conversational: new HyperbolicConversationalTask(),
+		"text-generation": new HyperbolicTextGenerationTask(),
+	},
+	nebius: {
+		"text-to-image": new NebiusTextToImageTask(),
+		conversational: new BaseConversationalTask("nebius", "https://api.nebius.ai"),
+		"text-generation": new BaseTextGenerationTask("nebius", "https://api.nebius.ai"),
+	},
+	novita: {
+		"text-generation": new NovitaTextGenerationTask(),
+		conversational: new NovitaConversationalTask(),
+	},
+	openai: {
+		conversational: new OpenAIConversationalTask(),
+	},
+	replicate: {
+		"text-to-image": new ReplicateTextToImageTask(),
+		"text-to-speech": new ReplicateTextToSpeechTask(),
+		"text-to-video": new ReplicateTextToVideoTask(),
+	},
+	sambanova: {
+		conversational: new BaseConversationalTask("sambanova", "https://api.sambanova.ai"),
+	},
+	together: {
+		"text-to-image": new TogetherTextToImageTask(),
+		"text-generation": new TogetherTextGenerationTask(),
+		conversational: new TogetherConversationalTask(),
+	},
+};
+
+/**
+ * Get provider helper instance by name and task
+ */
+export function getProviderHelper(provider: InferenceProvider, task: InferenceTask | undefined): TaskProviderHelper {
+	if (!task) {
+		throw new Error("you need to provide a task name, e.g. 'text-to-image'");
+	}
+	if (!(provider in PROVIDERS)) {
+		throw new Error(`Provider '${provider}' not supported. Available providers: ${Object.keys(PROVIDERS)}`);
+	}
+	const providerTasks = PROVIDERS[provider];
+	if (!providerTasks || !(task in providerTasks)) {
+		throw new Error(
+			`Task '${task}' not supported for provider '${provider}'. Available tasks: ${Object.keys(providerTasks ?? {})}`
+		);
+	}
+	const helper = providerTasks[task];
+	if (!helper) {
+		throw new Error(`Internal error: Helper for task '${task}' and provider '${provider}' resolved to undefined.`);
+	}
+	return helper;
+}

packages/inference/src/lib/makeRequestOptions.ts

@@ -1,21 +1,9 @@
+import { name as packageName, version as packageVersion } from "../../package.json";
 import { HF_HUB_URL, HF_ROUTER_URL } from "../config";
-import { BLACK_FOREST_LABS_CONFIG } from "../providers/black-forest-labs";
-import { CEREBRAS_CONFIG } from "../providers/cerebras";
-import { COHERE_CONFIG } from "../providers/cohere";
-import { FAL_AI_CONFIG } from "../providers/fal-ai";
-import { FIREWORKS_AI_CONFIG } from "../providers/fireworks-ai";
-import { HF_INFERENCE_CONFIG } from "../providers/hf-inference";
-import { HYPERBOLIC_CONFIG } from "../providers/hyperbolic";
-import { NEBIUS_CONFIG } from "../providers/nebius";
-import { NOVITA_CONFIG } from "../providers/novita";
-import { REPLICATE_CONFIG } from "../providers/replicate";
-import { SAMBANOVA_CONFIG } from "../providers/sambanova";
-import { TOGETHER_CONFIG } from "../providers/together";
-import { OPENAI_CONFIG } from "../providers/openai";
-import type { InferenceProvider, InferenceTask, Options, ProviderConfig, RequestArgs } from "../types";
-import { isUrl } from "./isUrl";
-import { version as packageVersion, name as packageName } from "../../package.json";
+import type { InferenceTask, Options, RequestArgs } from "../types";
+import { getProviderHelper } from "./getProviderHelper";
 import { getProviderModelId } from "./getProviderModelId";
+import { isUrl } from "./isUrl";
 
 const HF_HUB_INFERENCE_PROXY_TEMPLATE = `${HF_ROUTER_URL}/{{PROVIDER}}`;
 
@@ -25,25 +13,6 @@ const HF_HUB_INFERENCE_PROXY_TEMPLATE = `${HF_ROUTER_URL}/{{PROVIDER}}`;
  */
 let tasks: Record<string, { models: { id: string }[] }> | null = null;
 
-/**
- * Config to define how to serialize requests for each provider
- */
-const providerConfigs: Record<InferenceProvider, ProviderConfig> = {
-	"black-forest-labs": BLACK_FOREST_LABS_CONFIG,
-	cerebras: CEREBRAS_CONFIG,
-	cohere: COHERE_CONFIG,
-	"fal-ai": FAL_AI_CONFIG,
-	"fireworks-ai": FIREWORKS_AI_CONFIG,
-	"hf-inference": HF_INFERENCE_CONFIG,
-	hyperbolic: HYPERBOLIC_CONFIG,
-	openai: OPENAI_CONFIG,
-	nebius: NEBIUS_CONFIG,
-	novita: NOVITA_CONFIG,
-	replicate: REPLICATE_CONFIG,
-	sambanova: SAMBANOVA_CONFIG,
-	together: TOGETHER_CONFIG,
-};
-
 /**
  * Helper that prepares request arguments.
  * This async version handle the model ID resolution step.
@@ -61,9 +30,8 @@ export async function makeRequestOptions(
 ): Promise<{ url: string; info: RequestInit }> {
 	const { provider: maybeProvider, model: maybeModel } = args;
 	const provider = maybeProvider ?? "hf-inference";
-	const providerConfig = providerConfigs[provider];
 	const { task, chatCompletion } = options ?? {};
-
+	const providerHelper = getProviderHelper(provider, task);
 	// Validate inputs
 	if (args.endpointUrl && provider !== "hf-inference") {
 		throw new Error(`Cannot use endpointUrl with a third-party provider.`);
@@ -74,16 +42,16 @@ export async function makeRequestOptions(
 	if (!maybeModel && !task) {
 		throw new Error("No model provided, and no task has been specified.");
 	}
-	if (!providerConfig) {
-		throw new Error(`No provider config found for provider ${provider}`);
+	if (!providerHelper) {
+		throw new Error(`No provider helper found for provider ${provider}`);
 	}
-	if (providerConfig.clientSideRoutingOnly && !maybeModel) {
+	if (providerHelper.clientSideRoutingOnly && !maybeModel) {
 		throw new Error(`Provider ${provider} requires a model ID to be passed directly.`);
 	}
 
 	// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 	const hfModel = maybeModel ?? (await loadDefaultModel(task!));
-	const resolvedModel = providerConfig.clientSideRoutingOnly
+	const resolvedModel = providerHelper.clientSideRoutingOnly
 		? // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 		  removeProviderPrefix(maybeModel!, provider)
 		: await getProviderModelId({ model: hfModel, provider }, args, {
@@ -115,12 +83,10 @@ export function makeRequestOptionsFromResolvedModel(
 	void model;
 
 	const provider = maybeProvider ?? "hf-inference";
-	const providerConfig = providerConfigs[provider];
-
 	const { includeCredentials, task, chatCompletion, signal } = options ?? {};
-
+	const providerHelper = getProviderHelper(provider, task);
 	const authMethod = (() => {
-		if (providerConfig.clientSideRoutingOnly) {
+		if (providerHelper.clientSideRoutingOnly) {
 			// Closed-source providers require an accessToken (cannot be routed).
 			if (accessToken && accessToken.startsWith("hf_")) {
 				throw new Error(`Provider ${provider} is closed-source and does not support HF tokens.`);
@@ -142,24 +108,22 @@ export function makeRequestOptionsFromResolvedModel(
 		? chatCompletion
 			? endpointUrl + `/v1/chat/completions`
 			: endpointUrl
-		: providerConfig.makeUrl({
+		: providerHelper.makeUrl({
 				authMethod,
 				baseUrl:
 					authMethod !== "provider-key"
 						? HF_HUB_INFERENCE_PROXY_TEMPLATE.replace("{{PROVIDER}}", provider)
-						: providerConfig.makeBaseUrl(task),
+						: providerHelper.makeBaseUrl(),
 				model: resolvedModel,
 				chatCompletion,
 				task,
 		  });
-
 	// Make headers
 	const binary = "data" in args && !!args.data;
-	const headers = providerConfig.makeHeaders({
+	const headers = providerHelper.prepareHeaders({
 		accessToken,
 		authMethod,
 	});
-
 	// Add content-type to headers
 	if (!binary) {
 		headers["Content-Type"] = "application/json";
@@ -177,14 +141,13 @@ export function makeRequestOptionsFromResolvedModel(
 	const body = binary
 		? args.data
 		: JSON.stringify(
-				providerConfig.makeBody({
+				providerHelper.makeBody({
 					args: remainingArgs as Record<string, unknown>,
 					model: resolvedModel,
 					task,
 					chatCompletion,
 				})
 		  );
-
 	/**
 	 * For edge runtimes, leave 'credentials' undefined, otherwise cloudflare workers will error
 	 */
@@ -202,7 +165,6 @@ export function makeRequestOptionsFromResolvedModel(
 		...(credentials ? { credentials } : undefined),
 		signal,
 	};
-
 	return { url, info };
 }
 

packages/inference/src/providers/black-forest-labs.ts

@@ -14,33 +14,77 @@
  *
  * Thanks!
  */
-import type { BodyParams, HeaderParams, ProviderConfig, UrlParams } from "../types";
+import { InferenceOutputError } from "../lib/InferenceOutputError";
+import type { BodyParams, HeaderParams, UrlParams } from "../types";
+import { delay } from "../utils/delay";
+import { omit } from "../utils/omit";
+import { TaskProviderHelper } from "./providerHelper";
 
 const BLACK_FOREST_LABS_AI_API_BASE_URL = "https://api.us1.bfl.ai";
 
-const makeBaseUrl = (): string => {
-	return BLACK_FOREST_LABS_AI_API_BASE_URL;
-};
+export class BlackForestLabsTextToImageTask extends TaskProviderHelper {
+	constructor() {
+		super("black-forest-labs", BLACK_FOREST_LABS_AI_API_BASE_URL, "text-to-image");
+	}
 
-const makeBody = (params: BodyParams): Record<string, unknown> => {
-	return params.args;
-};
+	override makeBody(params: BodyParams): Record<string, unknown> {
+		return {
+			...omit(params.args, ["inputs", "parameters"]),
+			...(params.args.parameters as Record<string, unknown>),
+			prompt: params.args.inputs,
+		};
+	}
 
-const makeHeaders = (params: HeaderParams): Record<string, string> => {
-	if (params.authMethod === "provider-key") {
-		return { "X-Key": `${params.accessToken}` };
-	} else {
-		return { Authorization: `Bearer ${params.accessToken}` };
+	override prepareHeaders(params: HeaderParams): Record<string, string> {
+		if (params.authMethod !== "provider-key") {
+			return { Authorization: `Bearer ${params.accessToken}` };
+		} else {
+			return { "X-Key": `${params.accessToken}` };
+		}
 	}
-};
 
-const makeUrl = (params: UrlParams): string => {
-	return `${params.baseUrl}/v1/${params.model}`;
-};
+	override makeRoute(params?: UrlParams): string {
+		if (!params) {
+			throw new Error("Params are required");
+		}
+		return `/v1/${params.model}`;
+	}
 
-export const BLACK_FOREST_LABS_CONFIG: ProviderConfig = {
-	makeBaseUrl,
-	makeBody,
-	makeHeaders,
-	makeUrl,
-};
+	async getResponse(
+		res: Response,
+		url?: string,
+		headers?: Record<string, string>,
+		outputType?: "url" | "blob"
+	): Promise<string | Blob> {
+		const urlObj = new URL(res.url);
+		for (let step = 0; step < 5; step++) {
+			await delay(1000);
+			console.debug(`Polling Black Forest Labs API for the result... ${step + 1}/5`);
+			urlObj.searchParams.set("attempt", step.toString(10));
+			const resp = await fetch(urlObj, { headers: { "Content-Type": "application/json" } });
+			if (!resp.ok) {
+				throw new InferenceOutputError("Failed to fetch result from black forest labs API");
+			}
+			const payload = await resp.json();
+			if (
+				typeof payload === "object" &&
+				payload &&
+				"status" in payload &&
+				typeof payload.status === "string" &&
+				payload.status === "Ready" &&
+				"result" in payload &&
+				typeof payload.result === "object" &&
+				payload.result &&
+				"sample" in payload.result &&
+				typeof payload.result.sample === "string"
+			) {
+				if (outputType === "url") {
+					return payload.result.sample;
+				}
+				const image = await fetch(payload.result.sample);
+				return await image.blob();
+			}
+		}
+		throw new InferenceOutputError("Failed to fetch result from black forest labs API");
+	}
+}

packages/inference/src/providers/cerebras.ts

@@ -14,32 +14,11 @@
  *
  * Thanks!
  */
-import type { BodyParams, HeaderParams, ProviderConfig, UrlParams } from "../types";
 
-const CEREBRAS_API_BASE_URL = "https://api.cerebras.ai";
+import { BaseConversationalTask } from "./providerHelper";
 
-const makeBaseUrl = (): string => {
-	return CEREBRAS_API_BASE_URL;
-};
-
-const makeBody = (params: BodyParams): Record<string, unknown> => {
-	return {
-		...params.args,
-		model: params.model,
-	};
-};
-
-const makeHeaders = (params: HeaderParams): Record<string, string> => {
-	return { Authorization: `Bearer ${params.accessToken}` };
-};
-
-const makeUrl = (params: UrlParams): string => {
-	return `${params.baseUrl}/v1/chat/completions`;
-};
-
-export const CEREBRAS_CONFIG: ProviderConfig = {
-	makeBaseUrl,
-	makeBody,
-	makeHeaders,
-	makeUrl,
-};
+export class CerebrasConversationalTask extends BaseConversationalTask {
+	constructor() {
+		super("cerebras", "https://api.cerebras.ai");
+	}
+}