diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java index 58038fa31..50220c3d8 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClient.java @@ -16,9 +16,12 @@ package org.springframework.ai.bedrock.titan; +import java.util.ArrayList; import java.util.List; import java.util.concurrent.atomic.AtomicInteger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi; import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest; import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse; @@ -41,6 +44,8 @@ import org.springframework.util.Assert; */ public class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient { + private final Logger logger = LoggerFactory.getLogger(getClass()); + private final TitanEmbeddingBedrockApi embeddingApi; public enum InputType { @@ -90,17 +95,20 @@ public class BedrockTitanEmbeddingClient extends AbstractEmbeddingClient { @Override public List> embed(List inputContents) { Assert.notEmpty(inputContents, "At least one text is required!"); - Assert.isTrue(inputContents.size() == 1, "Titan Embedding does not support batch embedding!"); + if (inputContents.size() != 1) { + logger.warn( + "Titan Embedding does not support batch embedding. Will make multiple API calls to embed(Document)"); + } - String inputContent = inputContents.iterator().next(); - - var request = (this.inputType == InputType.IMAGE) - ? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build() - : new TitanEmbeddingRequest.Builder().withInputText(inputContent).build(); - - TitanEmbeddingResponse response = this.embeddingApi.embedding(request); - - return List.of(response.embedding()); + List> embeddingList = new ArrayList<>(); + for (String inputContent : inputContents) { + var request = (this.inputType == InputType.IMAGE) + ? new TitanEmbeddingRequest.Builder().withInputImage(inputContent).build() + : new TitanEmbeddingRequest.Builder().withInputText(inputContent).build(); + TitanEmbeddingResponse response = this.embeddingApi.embedding(request); + embeddingList.add(response.embedding()); + } + return embeddingList; } @Override diff --git a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java index 267d0eabe..478c01289 100644 --- a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java +++ b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingClientIT.java @@ -37,14 +37,6 @@ class BedrockTitanEmbeddingClientIT { assertThat(embeddingClient.dimensions()).isEqualTo(1024); } - @Test - void batchEmbedding() { - assertThatThrownBy( - () -> embeddingClient.embedForResponse(List.of("Hello World", "World is big and salvation is near"))) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Titan Embedding does not support batch embedding!"); - } - @Test void imageEmbedding() throws IOException { diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java index 32328e142..05c34aacf 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaEmbeddingClient.java @@ -16,10 +16,13 @@ package org.springframework.ai.ollama; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.AbstractEmbeddingClient; import org.springframework.ai.embedding.Embedding; @@ -49,6 +52,8 @@ import org.springframework.util.Assert; */ public class OllamaEmbeddingClient extends AbstractEmbeddingClient { + private final Logger logger = LoggerFactory.getLogger(getClass()); + private final OllamaApi ollamaApi; private String model = "orca-mini"; @@ -87,14 +92,18 @@ public class OllamaEmbeddingClient extends AbstractEmbeddingClient { @Override public List> embed(List texts) { Assert.notEmpty(texts, "At least one text is required!"); - Assert.isTrue(texts.size() == 1, "Ollama Embedding does not support batch embedding!"); + if (texts.size() != 1) { + logger.warn( + "Ollama Embedding does not support batch embedding. Will make multiple API calls to embed(Document)"); + } - String inputContent = texts.iterator().next(); - - OllamaApi.EmbeddingResponse response = this.ollamaApi - .embeddings(new EmbeddingRequest(this.model, inputContent, this.clientOptions)); - - return List.of(response.embedding()); + List> embeddingList = new ArrayList<>(); + for (String inputContent : texts) { + OllamaApi.EmbeddingResponse response = this.ollamaApi + .embeddings(new EmbeddingRequest(this.model, inputContent, this.clientOptions)); + embeddingList.add(response.embedding()); + } + return embeddingList; } @Override diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingClientIT.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingClientIT.java index 8016b0c07..c11175aa7 100644 --- a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingClientIT.java +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaEmbeddingClientIT.java @@ -56,14 +56,6 @@ class OllamaEmbeddingClientIT { assertThat(embeddingClient.dimensions()).isEqualTo(3200); } - @Test - void batchEmbedding() { - assertThatThrownBy( - () -> embeddingClient.embedForResponse(List.of("Hello World", "World is big and salvation is near"))) - .isInstanceOf(IllegalArgumentException.class) - .hasMessageContaining("Ollama Embedding does not support batch embedding!"); - } - @SpringBootConfiguration public static class TestConfiguration {