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/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java
index d6e9649f401e2..40fb875d91cae 100644
--- a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java
+++ b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java
@@ -103,7 +103,7 @@ protected String getRequirements() {
@Override
protected String getPythonScript() {
- return loadPythonScript("sentence_embeddings.py", config.getDevice(), config.getModelId());
+ return loadPythonScript("sentence_embeddings.py", config.getDevice(), config.getModelId(), config.getRevision());
}
@Override
diff --git a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java
index a928854acd1ac..561e7cd0d381b 100644
--- a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java
+++ b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java
@@ -22,6 +22,7 @@
import ai.djl.modality.Output;
import org.apache.camel.Exchange;
import org.apache.camel.RuntimeCamelException;
+import org.apache.camel.component.huggingface.HuggingFaceConstants;
import org.apache.camel.component.huggingface.HuggingFaceEndpoint;
/**
@@ -92,7 +93,7 @@ protected String getRequirements() {
@Override
protected String getPythonScript() {
- return loadPythonScript("text_to_image.py", config.getModelId(), config.getDevice());
+ return loadPythonScript("text_to_image.py", config.getModelId(), config.getRevision(), config.getDevice());
}
@Override
@@ -117,5 +118,6 @@ protected void processOutput(Exchange exchange, Output output) throws Exception
}
exchange.getMessage().setBody(imageBytes);
exchange.getMessage().setHeader("Content-Type", "image/png");
+ exchange.getMessage().setHeader(HuggingFaceConstants.OUTPUT, imageBytes);
}
}
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/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py
index 6fe310f35f245..9ee585c2dde64 100644
--- a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py
+++ b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py
@@ -31,7 +31,7 @@ def handle(inputs: Input):
device = '%s'
if device == 'auto':
device = 'cuda' if torch.cuda.is_available() else 'cpu'
- model = SentenceTransformer('%s', device=device)
+ model = SentenceTransformer('%s', device=device, revision='%s')
logging.debug("Model initialized")
if inputs.content.size() == 0:
diff --git a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py
index e8b41788c2b52..aa7eae92638cc 100644
--- a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py
+++ b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py
@@ -31,6 +31,7 @@ def handle(inputs: Input):
logging.debug("Initializing pipeline")
pipe = StableDiffusionPipeline.from_pretrained(
'%s',
+ revision='%s',
torch_dtype=torch.float32, # CPU-safe
safety_checker=None
)
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)");
+ }
+}
diff --git a/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/RevisionAndOutputHeaderTest.java b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/RevisionAndOutputHeaderTest.java
new file mode 100644
index 0000000000000..535a7fa69737d
--- /dev/null
+++ b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/RevisionAndOutputHeaderTest.java
@@ -0,0 +1,90 @@
+/*
+ * 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 ai.djl.modality.Output;
+import org.apache.camel.Exchange;
+import org.apache.camel.component.huggingface.HuggingFaceConfiguration;
+import org.apache.camel.component.huggingface.HuggingFaceConstants;
+import org.apache.camel.component.huggingface.HuggingFaceEndpoint;
+import org.apache.camel.impl.DefaultCamelContext;
+import org.apache.camel.support.DefaultExchange;
+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.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * The sentence-embeddings and text-to-image tasks must honour the configured model revision (they were the only two
+ * that dropped it), and the text-to-image task must publish its result on the OUTPUT header as its Javadoc promises.
+ */
+class RevisionAndOutputHeaderTest {
+
+ private DefaultCamelContext context;
+
+ @BeforeEach
+ void setUp() {
+ context = new DefaultCamelContext();
+ }
+
+ @AfterEach
+ void tearDown() {
+ context.stop();
+ }
+
+ private HuggingFaceEndpoint endpoint(HuggingFaceConfiguration config) {
+ HuggingFaceEndpoint endpoint = new HuggingFaceEndpoint(null, null, config);
+ endpoint.setCamelContext(context);
+ return endpoint;
+ }
+
+ @Test
+ void sentenceEmbeddingsScriptPinsRevision() {
+ HuggingFaceConfiguration config = new HuggingFaceConfiguration();
+ config.setModelId("sentence-transformers/all-MiniLM-L6-v2");
+ config.setRevision("v1.5");
+ SentenceEmbeddingsPredictor predictor = new SentenceEmbeddingsPredictor(endpoint(config));
+ assertTrue(predictor.getPythonScript().contains("revision='v1.5'"),
+ "the generated script must pin the configured revision");
+ }
+
+ @Test
+ void textToImageScriptPinsRevision() {
+ HuggingFaceConfiguration config = new HuggingFaceConfiguration();
+ config.setModelId("stabilityai/stable-diffusion");
+ config.setRevision("fp16");
+ TextToImagePredictor predictor = new TextToImagePredictor(endpoint(config));
+ assertTrue(predictor.getPythonScript().contains("revision='fp16'"),
+ "the generated script must pin the configured revision");
+ }
+
+ @Test
+ void textToImagePublishesTheImageOnTheOutputHeader() throws Exception {
+ HuggingFaceConfiguration config = new HuggingFaceConfiguration();
+ TextToImagePredictor predictor = new TextToImagePredictor(endpoint(config));
+ Exchange exchange = new DefaultExchange(context);
+ Output output = new Output();
+ byte[] image = { 1, 2, 3, 4 };
+ output.add("data", image);
+
+ predictor.processOutput(exchange, output);
+
+ assertArrayEquals(image, exchange.getMessage().getHeader(HuggingFaceConstants.OUTPUT, byte[].class));
+ }
+}
diff --git a/components/camel-ai/camel-weaviate/src/main/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformer.java b/components/camel-ai/camel-weaviate/src/main/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformer.java
index f9c33054ecc16..871acea76253d 100644
--- a/components/camel-ai/camel-weaviate/src/main/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformer.java
+++ b/components/camel-ai/camel-weaviate/src/main/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformer.java
@@ -19,6 +19,7 @@
import java.util.HashMap;
import java.util.List;
+import java.util.Map;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
@@ -58,7 +59,8 @@ public void transform(Message message, DataType fromType, DataType toType) {
case QUERY ->
queryEmbeddingOperation(message, embedding, vectorFieldName, textFieldName, text, collectionName, keyValue,
keyName);
- default -> throw new IllegalStateException("The only operations supported are create and updatebyid");
+ default ->
+ throw new IllegalStateException("The only operations supported are create, updatebyid and query");
}
}
@@ -66,22 +68,33 @@ private static void createEmbeddingOperation(
Message message, Embedding embedding, String vectorFieldName, String textFieldName, TextSegment text,
String collectionName, Object keyValue, String keyName) {
message.setBody(embedding.vectorAsList(), List.class);
-
- if (ObjectHelper.isNotEmpty(keyValue) && ObjectHelper.isNotEmpty(keyName)) {
- HashMap maps = new HashMap();
- maps.put(keyName, keyValue);
- message.setHeader(WeaviateVectorDbHeaders.PROPERTIES, maps);
- }
+ setProperties(message, textFieldName, text, keyValue, keyName);
}
private static void updateEmbeddingOperation(
Message message, Embedding embedding, String vectorFieldName, String textFieldName, TextSegment text,
String collectionName, Object keyValue, String keyName) {
message.setBody(embedding.vectorAsList(), List.class);
+ setProperties(message, textFieldName, text, keyValue, keyName);
+ }
+ /**
+ * Writes the object properties for a CREATE / UPDATE_BY_ID operation. The embedded text is stored under
+ * textFieldName so that the source passage can be retrieved later; without it only the vector (and optional id) was
+ * persisted and the original text was lost. Mirrors the Milvus transformer. Any PROPERTIES header the caller
+ * already set is preserved: the text (and optional id) are merged into a copy of it rather than replacing it.
+ */
+ private static void setProperties(
+ Message message, String textFieldName, TextSegment text, Object keyValue, String keyName) {
+ Map existing = message.getHeader(WeaviateVectorDbHeaders.PROPERTIES, Map.class);
+ HashMap maps = existing != null ? new HashMap<>(existing) : new HashMap<>();
+ if (text != null && text.text() != null) {
+ maps.put(textFieldName, text.text());
+ }
if (ObjectHelper.isNotEmpty(keyValue) && ObjectHelper.isNotEmpty(keyName)) {
- HashMap maps = new HashMap();
maps.put(keyName, keyValue);
+ }
+ if (!maps.isEmpty()) {
message.setHeader(WeaviateVectorDbHeaders.PROPERTIES, maps);
}
}
diff --git a/components/camel-ai/camel-weaviate/src/test/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformerTest.java b/components/camel-ai/camel-weaviate/src/test/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformerTest.java
new file mode 100644
index 0000000000000..d454dcb6b73bd
--- /dev/null
+++ b/components/camel-ai/camel-weaviate/src/test/java/org/apache/camel/component/weaviate/transform/WeaviateEmbeddingsDataTypeTransformerTest.java
@@ -0,0 +1,102 @@
+/*
+ * 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.weaviate.transform;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import dev.langchain4j.data.embedding.Embedding;
+import dev.langchain4j.data.segment.TextSegment;
+import org.apache.camel.Message;
+import org.apache.camel.ai.CamelLangchain4jAttributes;
+import org.apache.camel.component.weaviate.WeaviateVectorDbAction;
+import org.apache.camel.component.weaviate.WeaviateVectorDbHeaders;
+import org.apache.camel.impl.DefaultCamelContext;
+import org.apache.camel.spi.DataType;
+import org.apache.camel.support.DefaultExchange;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class WeaviateEmbeddingsDataTypeTransformerTest {
+
+ @SuppressWarnings("unchecked")
+ private Map transformProperties(WeaviateVectorDbAction action) throws Exception {
+ Embedding embedding = new Embedding(new float[] { 0.1f, 0.2f, 0.3f });
+ TextSegment segment = TextSegment.from("the source passage");
+
+ try (DefaultCamelContext context = new DefaultCamelContext()) {
+ context.start();
+ Message in = new DefaultExchange(context).getMessage();
+ in.setHeader(CamelLangchain4jAttributes.CAMEL_LANGCHAIN4J_EMBEDDING_VECTOR, embedding);
+ in.setHeader(WeaviateVectorDbHeaders.ACTION, action);
+ in.setHeader(WeaviateVectorDbHeaders.KEY_NAME, "id");
+ in.setHeader(WeaviateVectorDbHeaders.KEY_VALUE, "doc-1");
+ in.setBody(segment);
+
+ new WeaviateEmbeddingsDataTypeTransformer().transform(in, DataType.ANY, DataType.ANY);
+
+ return in.getHeader(WeaviateVectorDbHeaders.PROPERTIES, Map.class);
+ }
+ }
+
+ @Test
+ void createStoresTheDocumentTextInProperties() throws Exception {
+ Map props = transformProperties(WeaviateVectorDbAction.CREATE);
+ assertThat(props)
+ .isNotNull()
+ .containsEntry("text", "the source passage")
+ .containsEntry("id", "doc-1");
+ }
+
+ @Test
+ void updateStoresTheDocumentTextInProperties() throws Exception {
+ Map props = transformProperties(WeaviateVectorDbAction.UPDATE_BY_ID);
+ assertThat(props)
+ .isNotNull()
+ .containsEntry("text", "the source passage")
+ .containsEntry("id", "doc-1");
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ void createMergesTextIntoAnExistingPropertiesHeader() throws Exception {
+ Embedding embedding = new Embedding(new float[] { 0.1f, 0.2f, 0.3f });
+ TextSegment segment = TextSegment.from("the source passage");
+
+ try (DefaultCamelContext context = new DefaultCamelContext()) {
+ context.start();
+ Message in = new DefaultExchange(context).getMessage();
+ in.setHeader(CamelLangchain4jAttributes.CAMEL_LANGCHAIN4J_EMBEDDING_VECTOR, embedding);
+ in.setHeader(WeaviateVectorDbHeaders.ACTION, WeaviateVectorDbAction.CREATE);
+ // Properties the caller populated before the transformer runs must survive.
+ Map callerProperties = new HashMap<>();
+ callerProperties.put("sky", "blue");
+ callerProperties.put("age", "34");
+ in.setHeader(WeaviateVectorDbHeaders.PROPERTIES, callerProperties);
+ in.setBody(segment);
+
+ new WeaviateEmbeddingsDataTypeTransformer().transform(in, DataType.ANY, DataType.ANY);
+
+ Map props = in.getHeader(WeaviateVectorDbHeaders.PROPERTIES, Map.class);
+ assertThat(props)
+ .containsEntry("sky", "blue")
+ .containsEntry("age", "34")
+ .containsEntry("text", "the source passage");
+ }
+ }
+}