diff --git a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/AbstractTaskPredictor.java b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/AbstractTaskPredictor.java index 46dc72bdc9af3..ce0f6808fecde 100644 --- a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/AbstractTaskPredictor.java +++ b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/AbstractTaskPredictor.java @@ -73,12 +73,14 @@ public void loadModel() throws Exception { } Path handlerPath = tmpDir.resolve("handler.py"); String pythonScript = getPythonScript(); - Files.writeString(handlerPath, pythonScript); - Path reqPath = tmpDir.resolve("requirements.txt"); - Files.writeString(reqPath, getRequirements()); + // logged before the token is prepended: withAuthToken writes the configured token into the + // script, and this now runs for every task rather than only chat if (LOG.isDebugEnabled()) { LOG.debug("Generated Python script for task {}:\n{}", config.getTask(), pythonScript); } + Files.writeString(handlerPath, withAuthToken(pythonScript)); + Path reqPath = tmpDir.resolve("requirements.txt"); + Files.writeString(reqPath, getRequirements()); String modelUrl = "file://" + tmpDir.toAbsolutePath(); Criteria.Builder criteriaBuilder = Criteria.builder() .setTypes(Input.class, Output.class) @@ -107,6 +109,22 @@ public void setEndpoint(HuggingFaceEndpoint endpoint) { protected abstract String getPythonScript(); + /** + * Prepends the configured Hugging Face token to the generated handler as the {@code HF_TOKEN} environment variable + * so that every task can load gated or private models. {@code transformers.pipeline()} reads {@code HF_TOKEN} from + * the environment when no explicit token is passed; previously only the chat task passed a token, so the other + * tasks failed with 401 on gated models. The token comes from the {@code authToken} option or is resolved from an + * OAuth profile (see {@link org.apache.camel.component.huggingface.HuggingFaceProducer}), both surfaced through + * {@code config.getAuthToken()}. + */ + protected String withAuthToken(String pythonScript) { + String authToken = config.getAuthToken(); + if (authToken == null || authToken.isEmpty()) { + return pythonScript; + } + return "import os\nos.environ['HF_TOKEN'] = '" + authToken.replace("'", "\\'") + "'\n\n" + pythonScript; + } + protected String loadPythonScript(String resourcePath, Object... args) { InputStream is = null; try { diff --git a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/ChatPredictor.java b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/ChatPredictor.java index d7d163c958c60..0859ffc0956f1 100644 --- a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/ChatPredictor.java +++ b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/ChatPredictor.java @@ -110,10 +110,10 @@ public ChatPredictor(HuggingFaceEndpoint endpoint) { protected String getPythonScript() { String doSample = config.getTemperature() > 0 ? "True" : "False"; float temperature = config.getTemperature() > 0 ? config.getTemperature() : 1.0f; - String tokenClause = config.getAuthToken() != null ? ", token='" + config.getAuthToken() + "'" : ""; - return loadPythonScript("chat.py", config.getModelId(), config.getRevision(), config.getDevice(), tokenClause, - config.getMaxTokens(), - doSample, temperature); + // The token is applied centrally as the HF_TOKEN environment variable in + // AbstractTaskPredictor.withAuthToken, so no per-task token clause is needed here. + return loadPythonScript("chat.py", config.getModelId(), config.getRevision(), config.getDevice(), + config.getMaxTokens(), doSample, temperature); } @Override diff --git a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/chat.py b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/chat.py index 56e4ef5bb18c2..f20d5c3c35fc3 100644 --- a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/chat.py +++ b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/chat.py @@ -28,7 +28,7 @@ def handle(inputs: Input): try: if not pipe: logging.debug("Initializing pipeline") - pipe = pipeline(task='text-generation', model='%s', revision='%s', device_map='%s'%s) + pipe = pipeline(task='text-generation', model='%s', revision='%s', device_map='%s') logging.debug("Pipeline initialized") if inputs.content.size() == 0: diff --git a/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/AuthTokenInjectionTest.java b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/AuthTokenInjectionTest.java new file mode 100644 index 0000000000000..1f651614bfd89 --- /dev/null +++ b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/AuthTokenInjectionTest.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.camel.component.huggingface.tasks; + +import org.apache.camel.component.huggingface.HuggingFaceConfiguration; +import org.apache.camel.component.huggingface.HuggingFaceEndpoint; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The configured Hugging Face token must be applied for every task, not only chat, so that gated or private models can + * be loaded. The token is injected centrally as the HF_TOKEN environment variable of the generated handler. + */ +class AuthTokenInjectionTest { + + private TextGenerationPredictor predictorWithToken(String token) { + HuggingFaceConfiguration config = new HuggingFaceConfiguration(); + config.setAuthToken(token); + return new TextGenerationPredictor(new HuggingFaceEndpoint(null, null, config)); + } + + @Test + void authTokenIsExposedAsHfTokenEnvForEveryTask() { + String result = predictorWithToken("hf_secret123").withAuthToken("PIPELINE"); + assertTrue(result.contains("os.environ['HF_TOKEN'] = 'hf_secret123'"), + "generated script should export the token as HF_TOKEN"); + assertTrue(result.endsWith("PIPELINE"), "the original task script must be preserved"); + } + + @Test + void noAuthTokenLeavesTheScriptUnchanged() { + assertEquals("PIPELINE", predictorWithToken(null).withAuthToken("PIPELINE")); + assertEquals("PIPELINE", predictorWithToken("").withAuthToken("PIPELINE")); + } + + @Test + void authTokenWithASingleQuoteIsEscaped() { + String result = predictorWithToken("ab'cd").withAuthToken("PIPELINE"); + assertTrue(result.contains("os.environ['HF_TOKEN'] = 'ab\\'cd'"), + "a single quote in the token must be escaped so it cannot break the Python string literal"); + } +} diff --git a/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/ChatScriptFormatTest.java b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/ChatScriptFormatTest.java new file mode 100644 index 0000000000000..48d8369312d58 --- /dev/null +++ b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/ChatScriptFormatTest.java @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.camel.component.huggingface.tasks; + +import org.apache.camel.component.huggingface.HuggingFaceConfiguration; +import org.apache.camel.component.huggingface.HuggingFaceEndpoint; +import org.apache.camel.impl.DefaultCamelContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Guards the chat.py template's format-argument alignment after the per-task token clause was removed (the token is now + * applied centrally as HF_TOKEN). A misaligned placeholder would make getPythonScript throw. + */ +class ChatScriptFormatTest { + + private DefaultCamelContext context; + + @BeforeEach + void setUp() { + context = new DefaultCamelContext(); + } + + @AfterEach + void tearDown() { + context.stop(); + } + + @Test + void chatScriptFormatsAndCarriesNoTokenClause() { + HuggingFaceConfiguration config = new HuggingFaceConfiguration(); + config.setModelId("gpt2"); + HuggingFaceEndpoint endpoint = new HuggingFaceEndpoint(null, null, config); + endpoint.setCamelContext(context); + + String script = new ChatPredictor(endpoint).getPythonScript(); + + assertTrue(script.contains("pipeline(task='text-generation'"), "the chat pipeline call must be rendered"); + assertTrue(script.contains("model='gpt2'"), "the configured model must be interpolated"); + assertFalse(script.contains("token="), "the per-task token clause must be gone (token is applied via HF_TOKEN)"); + } +}