Updates for moving acme-assist fully on Spring AI

* Add SimplePersistentVectorStore

spring-ai-core/src/main/java/org/springframework/ai/document/Document.java

@@ -1,5 +1,8 @@
 package org.springframework.ai.document;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import org.springframework.util.StringUtils;
 
 import java.util.*;
@@ -11,26 +14,22 @@ public class Document {
 	 */
 	private final String id;
 
+	@JsonProperty(index = 100)
 	private List<Double> embedding = new ArrayList<>();
 
 	/**
 	 * Metadata for the document. It should not be nested and values should be restricted
 	 * to string, int, float, boolean for simple use with Vector Dbs.
 	 */
-	private Map<String, Object> metadata = new HashMap<>();
+	private Map<String, Object> metadata;
 
 	// Type; introduce when support images, now only text.
 
-	// TODO: Rename to `content` instead.
 	private String text;
 
-	public Document(String text) {
-		this(UUID.randomUUID().toString(), text);
-	}
-
-	public Document(String id, String text) {
-		this.id = id;
-		this.text = text;
+	@JsonCreator(mode = JsonCreator.Mode.PROPERTIES)
+	public Document(@JsonProperty("text") String text) {
+		this(text, new HashMap<>());
 	}
 
 	public Document(String text, Map<String, Object> metadata) {
@@ -68,26 +67,21 @@ public class Document {
 		return "Document{" + "id='" + id + '\'' + ", metadata=" + metadata + ", text='" + text + '\'' + '}';
 	}
 
-	// TODO: Consider moving the following methods & fields in a seprarate
-	// dedicated class. (e.g. DocumentService, DocumentUtil or alike)¬
-
-	// private List<String> excludedMetadataKeysForEmbedding;
-	// private List<String> relatedIds;
-
 	private static String DEFAULT_TEXT_TEMPLATE = "{metadata_string}\n\n{text}";
 
 	private static String DEFAULT_METADATA_TEMPLATE = "{key}: {value}";
 
-	private final String textTemplate = DEFAULT_TEXT_TEMPLATE;
+	private String textTemplate = DEFAULT_TEXT_TEMPLATE;
 
-	private final String metadataTemplate = DEFAULT_METADATA_TEMPLATE;
+	private String metadataTemplate = DEFAULT_METADATA_TEMPLATE;
 
-	private final String metadataSeparator = "\n";
+	private String metadataSeparator = "\n";
 
 	private MetadataMode metadataMode = MetadataMode.NONE;
 
 	private List<String> excludedMetadataKeysForLlm;
 
+	@JsonIgnore
 	public String getContent() {
 		return getContent(MetadataMode.ALL);
 	}
@@ -103,6 +97,7 @@ public class Document {
 		return getTextTemplate().replace("{metadata_string}", metadataString).replace("{text}", text);
 	}
 
+	@JsonIgnore
 	public String getMetadataString() {
 		return getMetadataString(metadataMode);
 	}
@@ -144,16 +139,16 @@ public class Document {
 		return metadataSeparator;
 	}
 
-	// public void setTextTemplate(String textTemplate) {
-	// this.textTemplate = textTemplate;
-	// }
+	public void setTextTemplate(String textTemplate) {
+		this.textTemplate = textTemplate;
+	}
 
-	// public void setMetadataTemplate(String metadataTemplate) {
-	// this.metadataTemplate = metadataTemplate;
-	// }
+	public void setMetadataTemplate(String metadataTemplate) {
+		this.metadataTemplate = metadataTemplate;
+	}
 
-	// public void setMetadataSeparator(String metadataSeparator) {
-	// this.metadataSeparator = metadataSeparator;
-	// }
+	public void setMetadataSeparator(String metadataSeparator) {
+		this.metadataSeparator = metadataSeparator;
+	}
 
 }

spring-ai-core/src/main/java/org/springframework/ai/loader/impl/EmptyJsonMetadataGenerator.java

@@ -0,0 +1,15 @@
+package org.springframework.ai.loader.impl;
+
+import java.util.Collections;
+import java.util.Map;
+
+public class EmptyJsonMetadataGenerator implements JsonMetadataGenerator {
+
+	private static final Map<String, Object> EMPTY_MAP = Collections.emptyMap();
+
+	@Override
+	public Map<String, Object> generate(Map<String, Object> jsonMap) {
+		return EMPTY_MAP;
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/loader/impl/JsonLoader.java

@@ -13,23 +13,30 @@ import java.util.*;
 
 public class JsonLoader implements Loader {
 
+	private Resource resource;
+
+	private JsonMetadataGenerator jsonMetadataGenerator;
+
 	/**
 	 * The key from the JSON that we will use as the text to parse into the Document text
 	 */
-	private List<String> jsonKeysToUse = new ArrayList<>();
-
-	private Resource resource;
+	private List<String> jsonKeysToUse;
 
 	public JsonLoader(Resource resource) {
-		Objects.requireNonNull(resource, "The Spring Resource must not be null");
-		this.resource = resource;
+		this(resource, new ArrayList<>().toArray(new String[0]));
 	}
 
 	public JsonLoader(Resource resource, String... jsonKeysToUse) {
+		this(resource, new EmptyJsonMetadataGenerator(), jsonKeysToUse);
+	}
+
+	public JsonLoader(Resource resource, JsonMetadataGenerator jsonMetadataGenerator, String... jsonKeysToUse) {
 		Objects.requireNonNull(jsonKeysToUse, "keys must not be null");
+		Objects.requireNonNull(jsonMetadataGenerator, "jsonMetadataGenerator must not be null");
 		Objects.requireNonNull(resource, "The Spring Resource must not be null");
-		this.jsonKeysToUse = List.of(jsonKeysToUse);
 		this.resource = resource;
+		this.jsonMetadataGenerator = jsonMetadataGenerator;
+		this.jsonKeysToUse = List.of(jsonKeysToUse);
 	}
 
 	@Override
@@ -41,9 +48,6 @@ public class JsonLoader implements Loader {
 	public List<Document> load(TextSplitter textSplitter) {
 
 		ObjectMapper objectMapper = new ObjectMapper();
-		TypeReference<HashMap<String, Object>> typeRef = new TypeReference<HashMap<String, Object>>() {
-		};
-
 		List<Document> documents = new ArrayList<>();
 		try {
 			// TODO, not all json will be an array
@@ -61,21 +65,30 @@ public class JsonLoader implements Loader {
 					}
 				}
 
+				Map<String, Object> metadata = this.jsonMetadataGenerator.generate(item);
+
+				Document document;
 				if (!sb.isEmpty()) {
-					Document document = new Document(sb.toString());
-					documents.add(document);
+					document = new Document(sb.toString(), metadata);
 				}
 				else {
-					Document document = new Document(item.toString());
-					documents.add(document);
+					document = new Document(item.toString(), metadata);
 				}
+				// Splitting at the item level is good when the size of the json per
+				// element is large
+				// as is the case with a catalog of product, as the metadata applies
+				// across all split documents
+				// This may not be good when the size of the json element is small as it
+				// can create too many individual
+				// documents.
+				List<Document> splitDocuments = textSplitter.apply(List.of(document));
+				documents.addAll(splitDocuments);
 			}
 		}
 		catch (IOException e) {
 			throw new RuntimeException(e);
 		}
-		List<Document> splitDocuments = textSplitter.apply(documents);
-		return splitDocuments;
+		return documents;
 	}
 
 }

spring-ai-core/src/main/java/org/springframework/ai/loader/impl/JsonMetadataGenerator.java

@@ -0,0 +1,14 @@
+package org.springframework.ai.loader.impl;
+
+import java.util.Map;
+
+@FunctionalInterface
+public interface JsonMetadataGenerator {
+
+	/**
+	 * The input is the JSON document represented as a map, the output are the fields
+	 * extracted from the input map that will be used as metadata.
+	 */
+	Map<String, Object> generate(Map<String, Object> jsonMap);
+
+}

spring-ai-core/src/main/java/org/springframework/ai/splitter/TextSplitter.java

@@ -41,7 +41,7 @@ public abstract class TextSplitter implements DocumentTransformer {
 			String text = texts.get(i);
 			List<String> chunks = splitText(text);
 			if (chunks.size() > 1) {
-				logger.info("Broke up document " + i + " into " + chunks.size() + " chunks.");
+				logger.info("Splitting up document into " + chunks.size() + " chunks.");
 			}
 			for (String chunk : chunks) {
 				// only primitive values are in here -

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

@@ -1,5 +1,7 @@
 package org.springframework.ai.vectorstore.impl;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.ai.document.Document;
 import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.vectorstore.VectorStore;
@@ -14,9 +16,11 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class InMemoryVectorStore implements VectorStore {
 
-	private Map<String, Document> store = new ConcurrentHashMap<>();
+	private static final Logger logger = LoggerFactory.getLogger(InMemoryVectorStore.class);
 
-	private EmbeddingClient embeddingClient;
+	protected Map<String, Document> store = new ConcurrentHashMap<>();
+
+	protected EmbeddingClient embeddingClient;
 
 	public InMemoryVectorStore(EmbeddingClient embeddingClient) {
 		Objects.requireNonNull(embeddingClient, "EmbeddingClient must not be null");
@@ -26,6 +30,7 @@ public class InMemoryVectorStore implements VectorStore {
 	@Override
 	public void add(List<Document> documents) {
 		for (Document document : documents) {
+			logger.info("Calling EmbeddingClient for document id = " + document.getId());
 			List<Double> embedding = this.embeddingClient.embed(document);
 			document.setEmbedding(embedding);
 			this.store.put(document.getId(), document);

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

@@ -0,0 +1,96 @@
+package org.springframework.ai.vectorstore.impl;
+
+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 org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.core.io.Resource;
+import org.springframework.util.StreamUtils;
+
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.Charset;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Adds simple serialization/deserialization to the data stored in the InMemoryVectorStore
+ */
+public class SimplePersistentVectorStore extends InMemoryVectorStore {
+
+	private static final Logger logger = LoggerFactory.getLogger(SimplePersistentVectorStore.class);
+
+	public SimplePersistentVectorStore(EmbeddingClient embeddingClient) {
+		super(embeddingClient);
+	}
+
+	public void save(File file) {
+		String json = getVectorDbAsJson();
+		try {
+			if (!file.exists()) {
+				logger.info("Creating new vector store file: " + file);
+				file.createNewFile();
+			}
+			else {
+				logger.info("Replacing existing vector store file: " + file);
+				file.delete();
+				file.createNewFile();
+			}
+		}
+		catch (IOException ex) {
+			throw new RuntimeException(ex);
+		}
+		try (OutputStream stream = new FileOutputStream(file)) {
+			StreamUtils.copy(json, Charset.forName("UTF-8"), stream);
+		}
+		catch (IOException e) {
+			throw new RuntimeException(e);
+		}
+	}
+
+	public void load(File file) {
+		TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
+		};
+		ObjectMapper objectMapper = new ObjectMapper();
+		try {
+			Map<String, Document> deserializedMap = objectMapper.readValue(file, typeRef);
+			this.store = deserializedMap;
+		}
+		catch (IOException ex) {
+			throw new RuntimeException(ex);
+		}
+	}
+
+	public void load(Resource resource) {
+		TypeReference<HashMap<String, Document>> typeRef = new TypeReference<>() {
+		};
+		ObjectMapper objectMapper = new ObjectMapper();
+		try {
+			Map<String, Document> deserializedMap = objectMapper.readValue(resource.getInputStream(), typeRef);
+			this.store = deserializedMap;
+		}
+		catch (IOException ex) {
+			throw new RuntimeException(ex);
+		}
+	}
+
+	private String getVectorDbAsJson() {
+		ObjectMapper objectMapper = new ObjectMapper();
+		ObjectWriter objectWriter = objectMapper.writerWithDefaultPrettyPrinter();
+		String json;
+		try {
+			json = objectWriter.writeValueAsString(this.store);
+		}
+		catch (JsonProcessingException e) {
+			throw new RuntimeException("Error serializing documentMap to JSON.", e);
+		}
+		return json;
+	}
+
+}

spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiTestConfiguration.java

@@ -1,8 +1,9 @@
 package org.springframework.ai.openai;
 
 import com.theokanning.openai.service.OpenAiService;
-import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
+import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.openai.client.OpenAiClient;
+import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.util.StringUtils;
@@ -32,7 +33,7 @@ public class OpenAiTestConfiguration {
 	}
 
 	@Bean
-	public OpenAiEmbeddingClient openAiEmbeddingClient(OpenAiService theoOpenAiService) {
+	public EmbeddingClient openAiEmbeddingClient(OpenAiService theoOpenAiService) {
 		return new OpenAiEmbeddingClient(theoOpenAiService);
 	}
 

spring-ai-openai/src/test/java/org/springframework/ai/openai/vectorstore/SimplePersistentVectorStoreIT.java

@@ -0,0 +1,62 @@
+package org.springframework.ai.openai.vectorstore;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.CleanupMode;
+import org.junit.jupiter.api.io.TempDir;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.ai.loader.impl.JsonLoader;
+import org.springframework.ai.loader.impl.JsonMetadataGenerator;
+import org.springframework.ai.vectorstore.impl.SimplePersistentVectorStore;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.core.io.Resource;
+
+import java.io.File;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@SpringBootTest
+public class SimplePersistentVectorStoreIT {
+
+	@Value("classpath:/data/acme/bikes.json")
+	private Resource bikesJsonResource;
+
+	@Autowired
+	private EmbeddingClient embeddingClient;
+
+	@Test
+	void persist(@TempDir(cleanup = CleanupMode.ON_SUCCESS) Path workingDir) {
+		JsonLoader jsonLoader = new JsonLoader(bikesJsonResource, new ProductMetadataGenerator(), "price", "name",
+				"shortDescription", "description", "tags");
+		List<Document> documents = jsonLoader.load();
+		SimplePersistentVectorStore vectorStore = new SimplePersistentVectorStore(this.embeddingClient);
+		vectorStore.add(documents);
+
+		File tempFile = new File(workingDir.toFile(), "temp.txt");
+		vectorStore.save(tempFile);
+		assertThat(tempFile).isNotEmpty();
+		assertThat(tempFile).content().contains("Velo 99 XR1 AXS");
+		SimplePersistentVectorStore vectorStore2 = new SimplePersistentVectorStore(this.embeddingClient);
+
+		vectorStore2.load(tempFile);
+		List<Document> similaritySearch = vectorStore2.similaritySearch("Velo 99 XR1 AXS");
+		assertThat(similaritySearch).isNotEmpty();
+		assertThat(similaritySearch.get(0).getMetadata()).containsEntry("name", "Velo 99 XR1 AXS");
+
+	}
+
+	public class ProductMetadataGenerator implements JsonMetadataGenerator {
+
+		@Override
+		public Map<String, Object> generate(Map<String, Object> jsonMap) {
+			return Map.of("name", jsonMap.get("name"));
+		}
+
+	}
+
+}