[InferenceSnippets] move as snippet option + more accurate typing (#1360)

Related to
#1349 (review)
(requested changes after initial approval, sorry about that 😬)

cc @kefranabg 

This PR:
- moves `billTo` inside `opts` => doesn't break the method signature +
feel more appropriate
- adds some typing to `opts` => not perfect but better than nothing

Author: Wauplin

packages/inference/src/snippets/getInferenceSnippets.ts

@@ -14,6 +14,8 @@ import { templates } from "./templates.exported";
 import type { InferenceProviderModelMapping } from "../lib/getInferenceProviderMapping";
 import { getProviderHelper } from "../lib/getProviderHelper";
 
+export type InferenceSnippetOptions = { streaming?: boolean; billTo?: string } & Record<string, unknown>;
+
 const PYTHON_CLIENTS = ["huggingface_hub", "fal_client", "requests", "openai"] as const;
 const JS_CLIENTS = ["fetch", "huggingface.js", "openai"] as const;
 const SH_CLIENTS = ["curl"] as const;
@@ -120,8 +122,7 @@ const snippetGenerator = (templateName: string, inputPreparationFn?: InputPrepar
 		accessToken: string,
 		provider: InferenceProvider,
 		inferenceProviderMapping?: InferenceProviderModelMapping,
-		billTo?: string,
-		opts?: Record<string, unknown>
+		opts?: InferenceSnippetOptions
 	): InferenceSnippet[] => {
 		const providerModelId = inferenceProviderMapping?.providerId ?? model.id;
 		/// Hacky: hard-code conversational templates here
@@ -155,7 +156,7 @@ const snippetGenerator = (templateName: string, inputPreparationFn?: InputPrepar
 			inferenceProviderMapping,
 			{
 				task,
-				billTo,
+				billTo: opts?.billTo,
 			}
 		);
 
@@ -194,7 +195,7 @@ const snippetGenerator = (templateName: string, inputPreparationFn?: InputPrepar
 			model,
 			provider,
 			providerModelId: providerModelId ?? model.id,
-			billTo,
+			billTo: opts?.billTo,
 		};
 
 		/// Iterate over clients => check if a snippet exists => generate
@@ -283,8 +284,7 @@ const snippets: Partial<
 			accessToken: string,
 			provider: InferenceProvider,
 			inferenceProviderMapping?: InferenceProviderModelMapping,
-			billTo?: string,
-			opts?: Record<string, unknown>
+			opts?: InferenceSnippetOptions
 		) => InferenceSnippet[]
 	>
 > = {
@@ -324,11 +324,10 @@ export function getInferenceSnippets(
 	accessToken: string,
 	provider: InferenceProvider,
 	inferenceProviderMapping?: InferenceProviderModelMapping,
-	billTo?: string,
 	opts?: Record<string, unknown>
 ): InferenceSnippet[] {
 	return model.pipeline_tag && model.pipeline_tag in snippets
-		? snippets[model.pipeline_tag]?.(model, accessToken, provider, inferenceProviderMapping, billTo, opts) ?? []
+		? snippets[model.pipeline_tag]?.(model, accessToken, provider, inferenceProviderMapping, opts) ?? []
 		: [];
 }
 

packages/inference/src/snippets/index.ts

@@ -1 +1 @@
-export { getInferenceSnippets } from "./getInferenceSnippets.js";
+export { getInferenceSnippets, type InferenceSnippetOptions } from "./getInferenceSnippets.js";

packages/tasks-gen/scripts/generate-snippets-fixtures.ts

@@ -33,8 +33,7 @@ const TEST_CASES: {
 	model: ModelDataMinimal;
 	providers: SnippetInferenceProvider[];
 	lora?: boolean;
-	billTo?: string;
-	opts?: Record<string, unknown>;
+	opts?: snippets.InferenceSnippetOptions;
 }[] = [
 	{
 		testName: "automatic-speech-recognition",
@@ -237,8 +236,8 @@ const TEST_CASES: {
 			tags: ["conversational"],
 			inference: "",
 		},
-		billTo: "huggingface",
 		providers: ["hf-inference"],
+		opts: { billTo: "huggingface" },
 	},
 ] as const;
 
@@ -266,7 +265,6 @@ function generateInferenceSnippet(
 	provider: SnippetInferenceProvider,
 	task: WidgetType,
 	lora: boolean = false,
-	billTo?: string,
 	opts?: Record<string, unknown>
 ): InferenceSnippet[] {
 	const allSnippets = snippets.getInferenceSnippets(
@@ -285,7 +283,6 @@ function generateInferenceSnippet(
 				  }
 				: {}),
 		},
-		billTo,
 		opts
 	);
 	return allSnippets
@@ -341,12 +338,12 @@ if (import.meta.vitest) {
 	const { describe, expect, it } = import.meta.vitest;
 
 	describe("inference API snippets", () => {
-		TEST_CASES.forEach(({ testName, task, model, providers, lora, billTo, opts }) => {
+		TEST_CASES.forEach(({ testName, task, model, providers, lora, opts }) => {
 			describe(testName, () => {
 				inferenceSnippetLanguages.forEach((language) => {
 					providers.forEach((provider) => {
 						it(language, async () => {
-							const generatedSnippets = generateInferenceSnippet(model, language, provider, task, lora, billTo, opts);
+							const generatedSnippets = generateInferenceSnippet(model, language, provider, task, lora, opts);
 							const expectedSnippets = await getExpectedInferenceSnippet(testName, language, provider);
 							expect(generatedSnippets).toEqual(expectedSnippets);
 						});
@@ -362,11 +359,11 @@ if (import.meta.vitest) {
 	await fs.rm(path.join(rootDirFinder(), "snippets-fixtures"), { recursive: true, force: true });
 
 	console.debug("  🏭 Generating new fixtures...");
-	TEST_CASES.forEach(({ testName, task, model, providers, lora, billTo, opts }) => {
+	TEST_CASES.forEach(({ testName, task, model, providers, lora, opts }) => {
 		console.debug(`      ${testName} (${providers.join(", ")})`);
 		inferenceSnippetLanguages.forEach(async (language) => {
 			providers.forEach(async (provider) => {
-				const generatedSnippets = generateInferenceSnippet(model, language, provider, task, lora, billTo, opts);
+				const generatedSnippets = generateInferenceSnippet(model, language, provider, task, lora, opts);
 				await saveExpectedInferenceSnippet(testName, language, provider, generatedSnippets);
 			});
 		});