Add operations retry capabilities for the AiClient, EmbeddingClient and VectorStore clients

- Refactor the EmbeddingResponse and Embedding from “class” into “record” types. Fix the affected code get calls.
- Refactor the TransformersEmbeddingClient to avoid telescopic public methods call anti-patter.
- Add spring-retry dependency to spring-ai-core.
- Implement RetryAiClient, RetryEmbeddingClient and RetryVectorStore as decorators around AiClient, EmbeddingClient and VectorStore instance. The retry decorators use RetryTemplate to wrap all delegate methods calls.  Default RetryTemplate is provided with each instance.
- Add Tests for all Retry decorators.
- Add missing javadocs to existing core classes.

 Resolves #123

Author: tzolov

embedding-clients/spring-ai-postgresml-embedding-client/src/test/java/org/springframework/ai/embedding/PostgresMlEmbeddingClientIT.java

@@ -108,15 +108,15 @@ void embedForResponse(String vectorType) {
 		EmbeddingResponse embeddingResponse = embeddingClient
 			.embedForResponse(List.of("Hello World!", "Spring AI!", "LLM!"));
 		assertThat(embeddingResponse).isNotNull();
-		assertThat(embeddingResponse.getData()).hasSize(3);
-		assertThat(embeddingResponse.getMetadata()).containsExactlyEntriesOf(
+		assertThat(embeddingResponse.data()).hasSize(3);
+		assertThat(embeddingResponse.metadata()).containsExactlyEntriesOf(
 				Map.of("transformer", "distilbert-base-uncased", "vector-type", vectorType, "kwargs", "{}"));
-		assertThat(embeddingResponse.getData().get(0).getIndex()).isEqualTo(0);
-		assertThat(embeddingResponse.getData().get(0).getEmbedding()).hasSize(768);
-		assertThat(embeddingResponse.getData().get(1).getIndex()).isEqualTo(1);
-		assertThat(embeddingResponse.getData().get(1).getEmbedding()).hasSize(768);
-		assertThat(embeddingResponse.getData().get(2).getIndex()).isEqualTo(2);
-		assertThat(embeddingResponse.getData().get(2).getEmbedding()).hasSize(768);
+		assertThat(embeddingResponse.data().get(0).index()).isEqualTo(0);
+		assertThat(embeddingResponse.data().get(0).embedding()).hasSize(768);
+		assertThat(embeddingResponse.data().get(1).index()).isEqualTo(1);
+		assertThat(embeddingResponse.data().get(1).embedding()).hasSize(768);
+		assertThat(embeddingResponse.data().get(2).index()).isEqualTo(2);
+		assertThat(embeddingResponse.data().get(2).embedding()).hasSize(768);
 		// embeddingClient.dropPgmlExtension();
 	}
 

embedding-clients/transformers-embedding/src/main/java/org/springframework/ai/embedding/TransformersEmbeddingClient.java

@@ -194,18 +194,19 @@ private Resource getCachedResource(Resource resource) {
 
 	@Override
 	public List<Double> embed(String text) {
-		return embed(List.of(text)).get(0);
+		return this.internalEmbed(List.of(text)).get(0);
 	}
 
 	@Override
 	public List<Double> embed(Document document) {
-		return this.embed(document.getFormattedContent(this.metadataMode));
+		String content = document.getFormattedContent(this.metadataMode);
+		return this.internalEmbed(List.of(content)).get(0);
 	}
 
 	@Override
 	public EmbeddingResponse embedForResponse(List<String> texts) {
 		List<Embedding> data = new ArrayList<>();
-		List<List<Double>> embed = this.embed(texts);
+		List<List<Double>> embed = this.internalEmbed(texts);
 		for (int i = 0; i < embed.size(); i++) {
 			data.add(new Embedding(embed.get(i), i));
 		}
@@ -214,6 +215,10 @@ public EmbeddingResponse embedForResponse(List<String> texts) {
 
 	@Override
 	public List<List<Double>> embed(List<String> texts) {
+		return this.internalEmbed(texts);
+	}
+
+	private List<List<Double>> internalEmbed(List<String> texts) {
 
 		List<List<Double>> resultEmbeddings = new ArrayList<>();
 

embedding-clients/transformers-embedding/src/test/java/org/springframework/ai/embedding/TransformersEmbeddingClientTests.java

@@ -75,16 +75,16 @@ void embedForResponse() throws Exception {
 		TransformersEmbeddingClient embeddingClient = new TransformersEmbeddingClient();
 		embeddingClient.afterPropertiesSet();
 		EmbeddingResponse embed = embeddingClient.embedForResponse(List.of("Hello world", "World is big"));
-		assertThat(embed.getData()).hasSize(2);
-		assertThat(embed.getMetadata()).isEmpty();
+		assertThat(embed.data()).hasSize(2);
+		assertThat(embed.metadata()).isEmpty();
 
-		assertThat(embed.getData().get(0).getEmbedding()).hasSize(384);
-		assertThat(DF.format(embed.getData().get(0).getEmbedding().get(0))).isEqualTo(DF.format(-0.19744634628295898));
-		assertThat(DF.format(embed.getData().get(0).getEmbedding().get(383))).isEqualTo(DF.format(0.17298996448516846));
+		assertThat(embed.data().get(0).embedding()).hasSize(384);
+		assertThat(DF.format(embed.data().get(0).embedding().get(0))).isEqualTo(DF.format(-0.19744634628295898));
+		assertThat(DF.format(embed.data().get(0).embedding().get(383))).isEqualTo(DF.format(0.17298996448516846));
 
-		assertThat(embed.getData().get(1).getEmbedding()).hasSize(384);
-		assertThat(DF.format(embed.getData().get(1).getEmbedding().get(0))).isEqualTo(DF.format(0.4293745160102844));
-		assertThat(DF.format(embed.getData().get(1).getEmbedding().get(383))).isEqualTo(DF.format(0.05501303821802139));
+		assertThat(embed.data().get(1).embedding()).hasSize(384);
+		assertThat(DF.format(embed.data().get(1).embedding().get(0))).isEqualTo(DF.format(0.4293745160102844));
+		assertThat(DF.format(embed.data().get(1).embedding().get(383))).isEqualTo(DF.format(0.05501303821802139));
 	}
 
 	@Test

pom.xml

@@ -87,6 +87,7 @@
 		<azure-open-ai-client.version>1.0.0-beta.3</azure-open-ai-client.version>
 		<jtokkit.version>0.6.1</jtokkit.version>
 		<victools.version>4.31.1</victools.version>
+		<spring-retry.version>2.0.4</spring-retry.version>
 
 		<!-- readers/writer/stores dependencies-->
 		<pdfbox.version>3.0.0</pdfbox.version>

spring-ai-core/pom.xml

@@ -26,6 +26,16 @@
 
 	<dependencies>
 		<!-- production dependencies -->
+		<dependency>
+			<groupId>org.springframework.retry</groupId>
+			<artifactId>spring-retry</artifactId>
+			<version>${spring-retry.version}</version>
+		</dependency>
+		<!-- <dependency>
+			<groupId>org.springframework</groupId>
+			<artifactId>spring-aspects</artifactId>
+			<version>6.0.14</version>
+		</dependency> -->
 		<dependency>
 			<groupId>org.antlr</groupId>
 			<artifactId>stringtemplate</artifactId>

spring-ai-core/src/main/java/org/springframework/ai/client/AiResponse.java

@@ -15,9 +15,7 @@
  */
 package org.springframework.ai.client;
 
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 
 import org.springframework.ai.metadata.GenerationMetadata;
 import org.springframework.ai.metadata.PromptMetadata;
@@ -56,8 +54,6 @@ public Generation getGeneration() {
 	}
 
 	/**
-	 * Returns {@link GenerationMetadata} containing information about the use of the AI
-	 * provider's API.
 	 * @return {@link GenerationMetadata} containing information about the use of the AI
 	 * provider's API.
 	 */
@@ -66,8 +62,6 @@ public GenerationMetadata getGenerationMetadata() {
 	}
 
 	/**
-	 * Returns {@link PromptMetadata} containing information on prompt processing by the
-	 * AI.
 	 * @return {@link PromptMetadata} containing information on prompt processing by the
 	 * AI.
 	 */

spring-ai-core/src/main/java/org/springframework/ai/client/RetryAiClient.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.ai.client;
+
+import java.time.Duration;
+
+import org.springframework.ai.prompt.Prompt;
+import org.springframework.retry.support.RetryTemplate;
+
+/**
+ * The {@link RetryAiClient} is a {@link AiClient} decorator that automatically re-invoke
+ * the failed generate operations according to pre-configured retry policies. This is
+ * helpful transient errors such as a momentary network glitch.
+ *
+ * @author Christian Tzolov
+ */
+public class RetryAiClient implements AiClient {
+
+	private final RetryTemplate retryTemplate;
+
+	private final AiClient delegate;
+
+	public RetryAiClient(AiClient delegate) {
+		this(RetryTemplate.builder()
+			.maxAttempts(10)
+			.exponentialBackoff(Duration.ofMillis(2000), 5, Duration.ofMillis(3 * 60000))
+			.build(), delegate);
+	}
+
+	public RetryAiClient(RetryTemplate retryTemplate, AiClient delegate) {
+		this.retryTemplate = retryTemplate;
+		this.delegate = delegate;
+	}
+
+	@Override
+	public AiResponse generate(Prompt prompt) {
+		return this.retryTemplate.execute(ctx -> {
+			return this.delegate.generate(prompt);
+		});
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/embedding/Embedding.java

@@ -1,46 +1,24 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.ai.embedding;
 
 import java.util.List;
-import java.util.Objects;
-
-public class Embedding {
-
-	private List<Double> embedding;
-
-	private Integer index;
-
-	public Embedding(List<Double> embedding, Integer index) {
-		this.embedding = embedding;
-		this.index = index;
-	}
-
-	public List<Double> getEmbedding() {
-		return embedding;
-	}
-
-	public Integer getIndex() {
-		return index;
-	}
-
-	@Override
-	public boolean equals(Object o) {
-		if (this == o)
-			return true;
-		if (o == null || getClass() != o.getClass())
-			return false;
-		Embedding embedding1 = (Embedding) o;
-		return Objects.equals(embedding, embedding1.embedding) && Objects.equals(index, embedding1.index);
-	}
-
-	@Override
-	public int hashCode() {
-		return Objects.hash(embedding, index);
-	}
-
-	@Override
-	public String toString() {
-		String message = this.embedding.size() == 0 ? "<empty>" : "<has data>";
-		return "Embedding{" + "embedding=" + message + ", index=" + index + '}';
-	}
 
+/**
+ * Represents a single embedding response.
+ */
+public record Embedding(List<Double> embedding, Integer index) {
 }

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingClient.java

@@ -1,19 +1,63 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.ai.embedding;
 
 import org.springframework.ai.document.Document;
 
 import java.util.List;
 
+/**
+ * Converts a {@link Document} text into its vector representation (e.g. embedding).
+ */
 public interface EmbeddingClient {
 
+	/**
+	 * Computes embedding for provided, raw, text.
+	 * @param text Input text to commute the embedding for.
+	 * @return Returns a raw, double, list of the embedding representation of the input
+	 * text.
+	 */
 	List<Double> embed(String text);
 
+	/**
+	 * Computes embedding for provided document.
+	 * @param document Document to commute the embedding for.
+	 * @return Returns a raw, double, list of the embedding representation of the input
+	 * text.
+	 */
 	List<Double> embed(Document document);
 
+	/**
+	 * Computes embeddings for provided list of text.
+	 * @param texts list of input text to compute embeddings for.
+	 * @return Returns a list of embeddings. The order corresponds to the input text list.
+	 */
 	List<List<Double>> embed(List<String> texts);
 
+	/**
+	 * Computes embeddings for provided list of text.
+	 * @param texts list of input text to compute embeddings for.
+	 * @return Returns
+	 */
 	EmbeddingResponse embedForResponse(List<String> texts);
 
+	/**
+	 * Retrieves embedding model's dimensions.
+	 * @return Returns the vector dimensions for the configured embedding model.
+	 */
 	default int dimensions() {
 		return embed("Test String").size();
 	}

spring-ai-core/src/main/java/org/springframework/ai/embedding/EmbeddingResponse.java

@@ -1,47 +1,26 @@
+/*
+ * Copyright 2023-2023 the original author or authors.
+ *
+ * Licensed 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
+ *
+ *      https://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.springframework.ai.embedding;
 
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
-
-public class EmbeddingResponse {
-
-	private List<Embedding> data;
-
-	private Map<String, Object> metadata = new HashMap<>();
-
-	public EmbeddingResponse(List<Embedding> data, Map<String, Object> metadata) {
-		this.data = data;
-		this.metadata = metadata;
-	}
-
-	public List<Embedding> getData() {
-		return data;
-	}
-
-	public Map<String, Object> getMetadata() {
-		return metadata;
-	}
-
-	@Override
-	public boolean equals(Object o) {
-		if (this == o)
-			return true;
-		if (o == null || getClass() != o.getClass())
-			return false;
-		EmbeddingResponse that = (EmbeddingResponse) o;
-		return Objects.equals(data, that.data) && Objects.equals(metadata, that.metadata);
-	}
-
-	@Override
-	public int hashCode() {
-		return Objects.hash(data, metadata);
-	}
-
-	@Override
-	public String toString() {
-		return "EmbeddingResult{" + "data=" + data + ", metadata=" + metadata + '}';
-	}
 
+/**
+ * Structured embedding response, containing list of {@link Embedding}s data and related
+ * metadata.
+ */
+public record EmbeddingResponse(List<Embedding> data, Map<String, Object> metadata) {
 }