Improve extensibility of SimpleVectoreStore

Additionally:

* Throw IllegalArgumentException instead of NullPointerException from constructor for null EmbeddingModel.
* Delegate Documentation embedding generation to embed(:Document) method called by doAdd(:Document).
* Break doSimilaritySearch(:SearchRequest) method into discrete, overridable operations.
* Implement load(:File) in terms of load(:Resource) using a FileSystemResource.
* Hide implementation details of EmbeddedMath dotProduct(..) and norm(..) methods.
* Use more meaningful and descriptive variable names.
* Add whitespace to improve readability.
* Fix compiler warnings.

Author: jxblum

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/SimpleVectorStore.java

@@ -26,17 +26,14 @@
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
 import java.util.Comparator;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.ObjectWriter;
 import com.fasterxml.jackson.databind.json.JsonMapper;
 import io.micrometer.observation.ObservationRegistry;
 import org.slf4j.Logger;
@@ -50,14 +47,16 @@
 import org.springframework.ai.vectorstore.observation.AbstractObservationVectorStore;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
 import org.springframework.ai.vectorstore.observation.VectorStoreObservationConvention;
+import org.springframework.core.io.FileSystemResource;
 import org.springframework.core.io.Resource;
+import org.springframework.util.Assert;
 
 /**
- * SimpleVectorStore is a simple implementation of the VectorStore interface.
- *
+ * Simple, in-memory implementation of the {@link VectorStore} interface.
+ * <p/>
  * It also provides methods to save the current state of the vectors to a file, and to
  * load vectors from a file.
- *
+ * <p/>
  * For a deeper understanding of the mathematical concepts and computations involved in
  * calculating similarity scores among vectors, refer to this
  * [resource](https://docs.spring.io/spring-ai/reference/api/vectordbs.html#_understanding_vectors).
@@ -67,6 +66,8 @@
  * @author Mark Pollack
  * @author Christian Tzolov
  * @author Sebastien Deleuze
+ * @author John Blum
+ * @see VectorStore
  */
 public class SimpleVectorStore extends AbstractObservationVectorStore {
 
@@ -87,54 +88,72 @@ public SimpleVectorStore(EmbeddingModel embeddingModel, ObservationRegistry obse
 
 		super(observationRegistry, customObservationConvention);
 
-		Objects.requireNonNull(embeddingModel, "EmbeddingModel must not be null");
+		Assert.notNull(embeddingModel, "EmbeddingModel must not be null");
+
 		this.embeddingModel = embeddingModel;
 		this.objectMapper = JsonMapper.builder().addModules(JacksonUtils.instantiateAvailableModules()).build();
 	}
 
 	@Override
 	public void doAdd(List<Document> documents) {
 		for (Document document : documents) {
-			logger.info("Calling EmbeddingModel for document id = {}", document.getId());
-			float[] embedding = this.embeddingModel.embed(document);
-			document.setEmbedding(embedding);
+			logger.info("Calling EmbeddingModel for Document id = {}", document.getId());
+			document = embed(document);
 			this.store.put(document.getId(), document);
 		}
 	}
 
+	protected Document embed(Document document) {
+		float[] documentEmbedding = this.embeddingModel.embed(document);
+		document.setEmbedding(documentEmbedding);
+		return document;
+	}
+
 	@Override
 	public Optional<Boolean> doDelete(List<String> idList) {
-		for (String id : idList) {
-			this.store.remove(id);
-		}
+		idList.forEach(this.store::remove);
 		return Optional.of(true);
 	}
 
 	@Override
 	public List<Document> doSimilaritySearch(SearchRequest request) {
+
 		if (request.getFilterExpression() != null) {
 			throw new UnsupportedOperationException(
-					"The [" + this.getClass() + "] doesn't support metadata filtering!");
+					"[%s] doesn't support metadata filtering".formatted(getClass().getName()));
 		}
 
-		float[] userQueryEmbedding = getUserQueryEmbedding(request.getQuery());
-		return this.store.values()
-			.stream()
-			.map(entry -> new Similarity(entry.getId(),
-					EmbeddingMath.cosineSimilarity(userQueryEmbedding, entry.getEmbedding())))
-			.filter(s -> s.score >= request.getSimilarityThreshold())
-			.sorted(Comparator.<Similarity>comparingDouble(s -> s.score).reversed())
+		// @formatter:off
+		return this.store.values().stream()
+			.map(document -> computeSimilarity(request, document))
+			.filter(similarity -> similarity.score >= request.getSimilarityThreshold())
+			.sorted(Comparator.<Similarity>comparingDouble(similarity -> similarity.score).reversed())
 			.limit(request.getTopK())
-			.map(s -> this.store.get(s.key))
+			.map(similarity -> this.store.get(similarity.key))
 			.toList();
+		// @formatter:on
+	}
+
+	protected Similarity computeSimilarity(SearchRequest request, Document document) {
+
+		float[] userQueryEmbedding = getUserQueryEmbedding(request);
+		float[] documentEmbedding = document.getEmbedding();
+
+		double score = computeCosineSimilarity(userQueryEmbedding, documentEmbedding);
+
+		return new Similarity(document.getId(), score);
+	}
+
+	protected double computeCosineSimilarity(float[] userQueryEmbedding, float[] storedDocumentEmbedding) {
+		return EmbeddingMath.cosineSimilarity(userQueryEmbedding, storedDocumentEmbedding);
 	}
 
 	/**
 	 * Serialize the vector store content into a file in JSON format.
 	 * @param file the file to save the vector store content
 	 */
 	public void save(File file) {
-		String json = getVectorDbAsJson();
+
 		try {
 			if (!file.exists()) {
 				logger.info("Creating new vector store file: {}", file);
@@ -145,28 +164,30 @@ public void save(File file) {
 					throw new RuntimeException("File already exists: " + file, e);
 				}
 				catch (IOException e) {
-					throw new RuntimeException("Failed to create new file: " + file + ". Reason: " + e.getMessage(), e);
+					throw new RuntimeException("Failed to create new file: " + file + "; Reason: " + e.getMessage(), e);
 				}
 			}
 			else {
 				logger.info("Overwriting existing vector store file: {}", file);
 			}
+
 			try (OutputStream stream = new FileOutputStream(file);
 					Writer writer = new OutputStreamWriter(stream, StandardCharsets.UTF_8)) {
+				String json = getVectorDbAsJson();
 				writer.write(json);
 				writer.flush();
 			}
 		}
 		catch (IOException ex) {
-			logger.error("IOException occurred while saving vector store file.", ex);
+			logger.error("IOException occurred while saving vector store file", ex);
 			throw new RuntimeException(ex);
 		}
 		catch (SecurityException ex) {
-			logger.error("SecurityException occurred while saving vector store file.", ex);
+			logger.error("SecurityException occurred while saving vector store file", ex);
 			throw new RuntimeException(ex);
 		}
 		catch (NullPointerException ex) {
-			logger.error("NullPointerException occurred while saving vector store file.", ex);
+			logger.error("NullPointerException occurred while saving vector store file", ex);
 			throw new RuntimeException(ex);
 		}
 	}
@@ -176,45 +197,40 @@ public void save(File file) {
 	 * @param file the file to load the vector store content
 	 */
 	public void load(File file) {
-		TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
-
-		};
-		try {
-			Map<String, Document> deserializedMap = this.objectMapper.readValue(file, typeRef);
-			this.store = deserializedMap;
-		}
-		catch (IOException ex) {
-			throw new RuntimeException(ex);
-		}
+		load(new FileSystemResource(file));
 	}
 
 	/**
 	 * Deserialize the vector store content from a resource in JSON format into memory.
 	 * @param resource the resource to load the vector store content
 	 */
 	public void load(Resource resource) {
-		TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
 
-		};
 		try {
-			Map<String, Document> deserializedMap = this.objectMapper.readValue(resource.getInputStream(), typeRef);
-			this.store = deserializedMap;
+			this.store = this.objectMapper.readValue(resource.getInputStream(), documentMapTypeRef());
 		}
 		catch (IOException ex) {
 			throw new RuntimeException(ex);
 		}
 	}
 
+	private TypeReference<Map<String, Document>> documentMapTypeRef() {
+		return new TypeReference<>() {
+		};
+	}
+
 	private String getVectorDbAsJson() {
-		ObjectWriter objectWriter = this.objectMapper.writerWithDefaultPrettyPrinter();
-		String json;
+
 		try {
-			json = objectWriter.writeValueAsString(this.store);
+			return this.objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(this.store);
 		}
 		catch (JsonProcessingException e) {
-			throw new RuntimeException("Error serializing documentMap to JSON.", e);
+			throw new RuntimeException("Error serializing Map of Documents to JSON", e);
 		}
-		return json;
+	}
+
+	private float[] getUserQueryEmbedding(SearchRequest request) {
+		return getUserQueryEmbedding(request.getQuery());
 	}
 
 	private float[] getUserQueryEmbedding(String query) {
@@ -232,9 +248,9 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str
 
 	public static class Similarity {
 
-		private String key;
+		private final String key;
 
-		private double score;
+		private final double score;
 
 		public Similarity(String key, double score) {
 			this.key = key;
@@ -243,16 +259,18 @@ public Similarity(String key, double score) {
 
 	}
 
-	public final class EmbeddingMath {
+	public static final class EmbeddingMath {
 
 		private EmbeddingMath() {
 			throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
 		}
 
 		public static double cosineSimilarity(float[] vectorX, float[] vectorY) {
+
 			if (vectorX == null || vectorY == null) {
-				throw new RuntimeException("Vectors must not be null");
+				throw new IllegalArgumentException("Vectors must not be null");
 			}
+
 			if (vectorX.length != vectorY.length) {
 				throw new IllegalArgumentException("Vectors lengths must be equal");
 			}
@@ -268,20 +286,22 @@ public static double cosineSimilarity(float[] vectorX, float[] vectorY) {
 			return dotProduct / (Math.sqrt(normX) * Math.sqrt(normY));
 		}
 
-		public static float dotProduct(float[] vectorX, float[] vectorY) {
+		private static float dotProduct(float[] vectorX, float[] vectorY) {
+
 			if (vectorX.length != vectorY.length) {
 				throw new IllegalArgumentException("Vectors lengths must be equal");
 			}
 
 			float result = 0;
-			for (int i = 0; i < vectorX.length; ++i) {
-				result += vectorX[i] * vectorY[i];
+
+			for (int index = 0; index < vectorX.length; ++index) {
+				result += vectorX[index] * vectorY[index];
 			}
 
 			return result;
 		}
 
-		public static float norm(float[] vector) {
+		private static float norm(float[] vector) {
 			return dotProduct(vector, vector);
 		}