[InferenceSnippet] Take token from env variable if not set

Author: Wauplin

packages/inference/src/snippets/getInferenceSnippets.ts

@@ -121,6 +121,8 @@ const HF_JS_METHODS: Partial<Record<WidgetType, string>> = {
 	translation: "translation",
 };
 
+const ACCESS_TOKEN_PLACEHOLDER = "<ACCESS_TOKEN>"; // Placeholder to replace with env variable in snippets
+
 // Snippet generators
 const snippetGenerator = (templateName: string, inputPreparationFn?: InputPreparationFn) => {
 	return (
@@ -149,13 +151,15 @@ const snippetGenerator = (templateName: string, inputPreparationFn?: InputPrepar
 			console.error(`Failed to get provider helper for ${provider} (${task})`, e);
 			return [];
 		}
+		const accessTokenOrPlaceholder = accessToken == "" ? ACCESS_TOKEN_PLACEHOLDER : accessToken;
+
 		/// Prepare inputs + make request
 		const inputs = inputPreparationFn ? inputPreparationFn(model, opts) : { inputs: getModelInputSnippet(model) };
 		const request = makeRequestOptionsFromResolvedModel(
 			providerModelId,
 			providerHelper,
 			{
-				accessToken,
+				accessToken: accessTokenOrPlaceholder,
 				provider,
 				...inputs,
 			} as RequestArgs,
@@ -180,7 +184,7 @@ const snippetGenerator = (templateName: string, inputPreparationFn?: InputPrepar
 
 		/// Prepare template injection data
 		const params: TemplateParams = {
-			accessToken,
+			accessToken: accessTokenOrPlaceholder,
 			authorizationHeader: (request.info.headers as Record<string, string>)?.Authorization,
 			baseUrl: removeSuffix(request.url, "/chat/completions"),
 			fullUrl: request.url,
@@ -248,6 +252,11 @@ const snippetGenerator = (templateName: string, inputPreparationFn?: InputPrepar
 							snippet = `${importSection}\n\n${snippet}`;
 						}
 
+						/// Replace access token placeholder
+						if (snippet.includes(ACCESS_TOKEN_PLACEHOLDER)) {
+							snippet = replaceAccessTokenPlaceholder(snippet, language, provider);
+						}
+
 						/// Snippet is ready!
 						return { language, client: client as string, content: snippet };
 					})
@@ -420,3 +429,48 @@ function indentString(str: string): string {
 function removeSuffix(str: string, suffix: string) {
 	return str.endsWith(suffix) ? str.slice(0, -suffix.length) : str;
 }
+
+function replaceAccessTokenPlaceholder(
+	snippet: string,
+	language: InferenceSnippetLanguage,
+	provider: InferenceProviderOrPolicy
+): string {
+	// If "accessToken" is empty, the snippets are generated with a placeholder.
+	// Once snippets are rendered, we replace the placeholder with code to fetch the access token from an environment variable.
+
+	// Determine if HF_TOKEN or specific provider token should be used
+	const accessTokenEnvVar =
+		!snippet.includes("https://") || // no URL provided => using a client => use $HF_TOKEN
+		snippet.includes("https://router.huggingface.co") || // explicit routed request => use $HF_TOKEN
+		provider == "hf-inference" // hf-inference provider => use $HF_TOKEN
+			? "HF_TOKEN"
+			: provider.toUpperCase().replace("-", "_") + "_API_TOKEN"; // e.g. "REPLICATE_API_TOKEN"
+
+	// Replace the placeholder with the env variable
+	if (language === "sh") {
+		snippet = snippet.replace(
+			`'Authorization: Bearer ${ACCESS_TOKEN_PLACEHOLDER}'`,
+			`"Authorization: Bearer $${accessTokenEnvVar}"` // e.g. "Authorization: Bearer $HF_TOKEN"
+		);
+	} else if (language === "python") {
+		snippet = "import os\n" + snippet;
+		snippet = snippet.replace(
+			`"${ACCESS_TOKEN_PLACEHOLDER}"`,
+			`os.getenv("${accessTokenEnvVar}")` // e.g. os.getenv("HF_TOKEN")
+		);
+		snippet = snippet.replace(
+			`"Bearer <ACCESS_TOKEN>"`,
+			`f"Bearer {os.getenv('${accessTokenEnvVar}')}"` // e.g. f"Bearer {os.getenv('HF_TOKEN')}"
+		);
+		snippet = snippet.replace(
+			`"Key <ACCESS_TOKEN>"`,
+			`f"Key {os.getenv('${accessTokenEnvVar}')}"` // e.g. f"Key {os.getenv('FAL_AI_API_TOKEN')}"
+		);
+	} else if (language === "js") {
+		snippet = snippet.replace(
+			`"${ACCESS_TOKEN_PLACEHOLDER}"`,
+			`process.env.${accessTokenEnvVar}` // e.g. process.env.HF_TOKEN
+		);
+	}
+	return snippet;
+}

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

@@ -314,7 +314,7 @@ function generateInferenceSnippet(
 ): InferenceSnippet[] {
 	const allSnippets = snippets.getInferenceSnippets(
 		model,
-		"api_token",
+		"",
 		provider,
 		{
 			hfModelId: model.id,

packages/tasks-gen/snippets-fixtures/automatic-speech-recognition/js/fetch/0.hf-inference.js

@@ -3,7 +3,7 @@ async function query(data) {
 		"https://router.huggingface.co/hf-inference/models/openai/whisper-large-v3-turbo",
 		{
 			headers: {
-				Authorization: "Bearer api_token",
+				Authorization: "Bearer <ACCESS_TOKEN>",
 				"Content-Type": "audio/flac",
 			},
 			method: "POST",

packages/tasks-gen/snippets-fixtures/automatic-speech-recognition/js/huggingface.js/0.hf-inference.js

@@ -1,6 +1,6 @@
 import { InferenceClient } from "@huggingface/inference";
 
-const client = new InferenceClient("api_token");
+const client = new InferenceClient(process.env.HF_TOKEN);
 
 const data = fs.readFileSync("sample1.flac");
 

packages/tasks-gen/snippets-fixtures/automatic-speech-recognition/python/huggingface_hub/0.hf-inference.py

@@ -1,8 +1,9 @@
+import os
 from huggingface_hub import InferenceClient
 
 client = InferenceClient(
     provider="hf-inference",
-    api_key="api_token",
+    api_key=os.getenv("HF_TOKEN"),
 )
 
 output = client.automatic_speech_recognition("sample1.flac", model="openai/whisper-large-v3-turbo")

packages/tasks-gen/snippets-fixtures/automatic-speech-recognition/python/requests/0.hf-inference.py

@@ -1,8 +1,9 @@
+import os
 import requests
 
 API_URL = "https://router.huggingface.co/hf-inference/models/openai/whisper-large-v3-turbo"
 headers = {
-    "Authorization": "Bearer api_token",
+    "Authorization": f"Bearer {os.getenv('HF_TOKEN')}",
 }
 
 def query(filename):

packages/tasks-gen/snippets-fixtures/automatic-speech-recognition/sh/curl/0.hf-inference.sh

@@ -1,5 +1,5 @@
 curl https://router.huggingface.co/hf-inference/models/openai/whisper-large-v3-turbo \
     -X POST \
-    -H 'Authorization: Bearer api_token' \
+    -H "Authorization: Bearer $HF_TOKEN" \
     -H 'Content-Type: audio/flac' \
     --data-binary @"sample1.flac"

packages/tasks-gen/snippets-fixtures/basic-snippet--token-classification/js/fetch/0.hf-inference.js

@@ -3,7 +3,7 @@ async function query(data) {
 		"https://router.huggingface.co/hf-inference/models/FacebookAI/xlm-roberta-large-finetuned-conll03-english",
 		{
 			headers: {
-				Authorization: "Bearer api_token",
+				Authorization: "Bearer <ACCESS_TOKEN>",
 				"Content-Type": "application/json",
 			},
 			method: "POST",

packages/tasks-gen/snippets-fixtures/basic-snippet--token-classification/js/huggingface.js/0.hf-inference.js

@@ -1,6 +1,6 @@
 import { InferenceClient } from "@huggingface/inference";
 
-const client = new InferenceClient("api_token");
+const client = new InferenceClient(process.env.HF_TOKEN);
 
 const output = await client.tokenClassification({
 	model: "FacebookAI/xlm-roberta-large-finetuned-conll03-english",

packages/tasks-gen/snippets-fixtures/basic-snippet--token-classification/python/huggingface_hub/0.hf-inference.py

@@ -1,8 +1,9 @@
+import os
 from huggingface_hub import InferenceClient
 
 client = InferenceClient(
     provider="hf-inference",
-    api_key="api_token",
+    api_key=os.getenv("HF_TOKEN"),
 )
 
 result = client.token_classification(