Add base64 import when required

Author: Wauplin

packages/tasks/src/snippets/python.spec.ts

@@ -155,7 +155,8 @@ client = InferenceClient("impira/layoutlm-invoices", token="api_token")
 output = client.document_question_answering(cat.png, question=What is in this image?)`);
 
 		expect(snippets[1].client).toEqual("requests");
-		expect(snippets[1].content).toEqual(`import requests
+		expect(snippets[1].content).toEqual(`import base64
+import requests
 
 API_URL = "https://api-inference.huggingface.co/models/impira/layoutlm-invoices"
 headers = {"Authorization": "Bearer api_token"}
@@ -194,7 +195,8 @@ client = InferenceClient("stabilityai/stable-diffusion-xl-refiner-1.0", token="a
 image = client.image_to_image("cat.png", prompt="Turn the cat into a tiger.")`);
 
 		expect(snippets[1].client).toEqual("requests");
-		expect(snippets[1].content).toEqual(`import requests
+		expect(snippets[1].content).toEqual(`import base64
+import requests
 
 API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-refiner-1.0"
 headers = {"Authorization": "Bearer api_token"}

packages/tasks/src/snippets/python.ts

@@ -11,6 +11,22 @@ const snippetImportInferenceClient = (model: ModelDataMinimal, accessToken: stri
 client = InferenceClient("${model.id}", token="${accessToken || "{API_TOKEN}"}")
 `;
 
+const addImportsToSnippet = (snippet: string, model: ModelDataMinimal, accessToken: string): string => {
+	if (snippet.includes("requests")) {
+		snippet = `import requests
+
+API_URL = "https://api-inference.huggingface.co/models/${model.id}"
+headers = {"Authorization": ${accessToken ? `"Bearer ${accessToken}"` : `f"Bearer {API_TOKEN}"`}}
+
+${snippet}`;
+	}
+	if (snippet.includes("base64")) {
+		snippet = `import base64
+${snippet}`;
+	}
+	return snippet;
+};
+
 const snippetBasic = (model: ModelDataMinimal): InferenceSnippet => ({
 	content: `def query(payload):
 	response = requests.post(API_URL, headers=headers, json=payload)
@@ -368,14 +384,7 @@ export function getPythonInferenceSnippet(
 		return snippets.map((snippet) => {
 			return {
 				...snippet,
-				content: snippet.content.includes("requests")
-					? `import requests
-
-API_URL = "https://api-inference.huggingface.co/models/${model.id}"
-headers = {"Authorization": ${accessToken ? `"Bearer ${accessToken}"` : `f"Bearer {API_TOKEN}"`}}
-
-${snippet.content}`
-					: snippet.content,
+				content: addImportsToSnippet(snippet.content, model, accessToken),
 			};
 		});
 	}