Added VectorStore, Retriever with implementations

* Removed some classes out of loader to their own packages

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

@@ -1,4 +1,4 @@
-package org.springframework.ai.core.loader;
+package org.springframework.ai.core.document;
 
 import org.springframework.util.StringUtils;
 
@@ -11,11 +11,11 @@ public class Document {
 	private static String DEFAULT_METADATA_TEMPLATE = "{key}: {value}";
 
 	/**
-	 * Unique ID, creates UUID by default
+	 * Unique ID
 	 */
-	private String id;
+	private String id = UUID.randomUUID().toString();
 
-	// Embedding List<Float>
+	private List<Double> embedding = new ArrayList<>();
 
 	/**
 	 * Metadata for the document. It should not be nested and values should be restricted
@@ -51,10 +51,18 @@ public class Document {
 		this.metadata = metadata;
 	}
 
+	public String getId() {
+		return id;
+	}
+
 	public String getText() {
 		return this.text;
 	}
 
+	public String getContent() {
+		return getContent(MetadataMode.ALL);
+	}
+
 	public String getContent(MetadataMode metadataMode) {
 		if (metadataMode == MetadataMode.NONE) {
 			return this.text;
@@ -111,6 +119,14 @@ public class Document {
 		return metadata;
 	}
 
+	public List<Double> getEmbedding() {
+		return embedding;
+	}
+
+	public void setEmbedding(List<Double> embedding) {
+		this.embedding = embedding;
+	}
+
 	@Override
 	public String toString() {
 		return "Document{" + "id='" + id + '\'' + ", metadata=" + metadata + ", text='" + text + '\'' + '}';

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

@@ -1,4 +1,4 @@
-package org.springframework.ai.core.loader;
+package org.springframework.ai.core.document;
 
 import java.util.List;
 import java.util.function.Function;

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

@@ -1,4 +1,4 @@
-package org.springframework.ai.core.loader;
+package org.springframework.ai.core.document;
 
 public enum MetadataMode {
 

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

@@ -1,11 +1,15 @@
 package org.springframework.ai.core.embedding;
 
+import org.springframework.ai.core.document.Document;
+
 import java.util.List;
 
 public interface EmbeddingClient {
 
 	List<Double> createEmbedding(String text);
 
+	List<Double> createEmbedding(Document document);
+
 	List<List<Double>> createEmbedding(List<String> texts);
 
 	EmbeddingResponse createEmbeddingResult(List<String> texts);

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

@@ -1,6 +1,7 @@
 package org.springframework.ai.core.loader;
 
-import org.springframework.ai.core.loader.splitter.TextSplitter;
+import org.springframework.ai.core.document.Document;
+import org.springframework.ai.core.splitter.TextSplitter;
 
 import java.util.List;
 

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

@@ -2,10 +2,10 @@ package org.springframework.ai.core.loader.impl;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import org.springframework.ai.core.loader.Document;
+import org.springframework.ai.core.document.Document;
 import org.springframework.ai.core.loader.Loader;
-import org.springframework.ai.core.loader.splitter.TextSplitter;
-import org.springframework.ai.core.loader.splitter.TokenTextSplitter;
+import org.springframework.ai.core.splitter.TextSplitter;
+import org.springframework.ai.core.splitter.TokenTextSplitter;
 import org.springframework.core.io.Resource;
 
 import java.io.IOException;

spring-ai-core/src/main/java/org/springframework/ai/core/retriever/Retriever.java

@@ -0,0 +1,16 @@
+package org.springframework.ai.core.retriever;
+
+import org.springframework.ai.core.document.Document;
+
+import java.util.List;
+
+public interface Retriever {
+
+	/**
+	 * Retrieves relevant documents however the implementation sees fit.
+	 * @param query query string
+	 * @return relevant documents
+	 */
+	List<Document> retrieve(String query);
+
+}

spring-ai-core/src/main/java/org/springframework/ai/core/retriever/impl/VectorStoreRetriever.java

@@ -0,0 +1,58 @@
+package org.springframework.ai.core.retriever.impl;
+
+import org.springframework.ai.core.document.Document;
+import org.springframework.ai.core.retriever.Retriever;
+import org.springframework.ai.core.vectorstore.VectorStore;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+
+public class VectorStoreRetriever implements Retriever {
+
+	private VectorStore vectorStore;
+
+	int k;
+
+	Optional<Double> threshold = Optional.empty();
+
+	public VectorStoreRetriever(VectorStore vectorStore) {
+		this(vectorStore, 4);
+	}
+
+	public VectorStoreRetriever(VectorStore vectorStore, int k) {
+		Objects.requireNonNull(vectorStore, "VectorStore must not be null");
+		this.vectorStore = vectorStore;
+		this.k = k;
+	}
+
+	public VectorStoreRetriever(VectorStore vectorStore, int k, double threshold) {
+		Objects.requireNonNull(vectorStore, "VectorStore must not be null");
+		this.vectorStore = vectorStore;
+		this.k = k;
+		this.threshold = Optional.of(threshold);
+	}
+
+	public VectorStore getVectorStore() {
+		return vectorStore;
+	}
+
+	public int getK() {
+		return k;
+	}
+
+	public Optional<Double> getThreshold() {
+		return threshold;
+	}
+
+	@Override
+	public List<Document> retrieve(String query) {
+		if (threshold.isPresent()) {
+			return this.vectorStore.similaritySearch(query, this.k, this.threshold.get());
+		}
+		else {
+			return this.vectorStore.similaritySearch(query, this.k);
+		}
+	}
+
+}

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

@@ -1,9 +1,9 @@
-package org.springframework.ai.core.loader.splitter;
+package org.springframework.ai.core.splitter;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.springframework.ai.core.loader.Document;
-import org.springframework.ai.core.loader.DocumentTransformer;
+import org.springframework.ai.core.document.Document;
+import org.springframework.ai.core.document.DocumentTransformer;
 
 import java.util.ArrayList;
 import java.util.HashMap;

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

@@ -1,4 +1,4 @@
-package org.springframework.ai.core.loader.splitter;
+package org.springframework.ai.core.splitter;
 
 import com.knuddels.jtokkit.Encodings;
 import com.knuddels.jtokkit.api.Encoding;

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

@@ -0,0 +1,33 @@
+package org.springframework.ai.core.vectorstore;
+
+import org.springframework.ai.core.document.Document;
+import org.springframework.ai.core.embedding.EmbeddingClient;
+
+import java.util.List;
+import java.util.Optional;
+
+public interface VectorStore {
+
+	/**
+	 * Adds Documents to the vector store.
+	 * @param documents the list of documents to store Will throw an exception if the
+	 * underlying provider checks for duplicate IDs on add
+	 */
+	void add(List<Document> documents);
+
+	Optional<Boolean> delete(List<String> idList);
+
+	List<Document> similaritySearch(String query);
+
+	List<Document> similaritySearch(String query, int k);
+
+	/**
+	 * @param query The query to send, it will be converted to an embeddeing based on the
+	 * configuration of the vector store.
+	 * @param k the top 'k' similar results
+	 * @param threshold the lower bound of the similarity score
+	 * @return similar documents
+	 */
+	List<Document> similaritySearch(String query, int k, double threshold);
+
+}

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

@@ -0,0 +1,132 @@
+package org.springframework.ai.core.vectorstore.impl;
+
+import org.springframework.ai.core.document.Document;
+import org.springframework.ai.core.embedding.EmbeddingClient;
+import org.springframework.ai.core.vectorstore.VectorStore;
+
+import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
+
+/***
+ * @author Raphael Yu
+ * @author Dingmeng Xue
+ * @author Mark Pollack
+ */
+public class InMemoryVectorStore implements VectorStore {
+
+	private Map<String, Document> store = new ConcurrentHashMap<>();
+
+	private EmbeddingClient embeddingClient;
+
+	public InMemoryVectorStore(EmbeddingClient embeddingClient) {
+		Objects.requireNonNull(embeddingClient, "EmbeddingClient must not be null");
+		this.embeddingClient = embeddingClient;
+	}
+
+	@Override
+	public void add(List<Document> documents) {
+		for (Document document : documents) {
+			List<Double> embedding = this.embeddingClient.createEmbedding(document);
+			document.setEmbedding(embedding);
+			this.store.put(document.getId(), document);
+		}
+	}
+
+	@Override
+	public Optional<Boolean> delete(List<String> idList) {
+		for (String id : idList) {
+			this.store.remove(id);
+		}
+		return Optional.of(true);
+	}
+
+	@Override
+	public List<Document> similaritySearch(String query) {
+		return similaritySearch(query, 4);
+	}
+
+	@Override
+	public List<Document> similaritySearch(String query, int k) {
+		List<Double> userQueryEmbedding = getUserQueryEmbedding(query);
+		var similarities = this.store.values()
+			.stream()
+			.map(entry -> new Similarity(entry.getId(),
+					EmbeddingMath.cosineSimilarity(userQueryEmbedding, entry.getEmbedding())))
+			.sorted(Comparator.<Similarity>comparingDouble(s -> s.similarity).reversed())
+			.limit(k)
+			.map(s -> store.get(s.key))
+			.toList();
+		return similarities;
+	}
+
+	@Override
+	public List<Document> similaritySearch(String query, int k, double threshold) {
+		List<Double> userQueryEmbedding = getUserQueryEmbedding(query);
+		var similarities = this.store.values()
+			.stream()
+			.map(entry -> new Similarity(entry.getId(),
+					EmbeddingMath.cosineSimilarity(userQueryEmbedding, entry.getEmbedding())))
+			.filter(s -> s.similarity >= threshold)
+			.sorted(Comparator.<Similarity>comparingDouble(s -> s.similarity).reversed())
+			.limit(k)
+			.map(s -> store.get(s.key))
+			.toList();
+		return similarities;
+	}
+
+	private List<Double> getUserQueryEmbedding(String query) {
+		List<Double> userQueryEmbedding = this.embeddingClient.createEmbedding(query);
+		return userQueryEmbedding;
+	}
+
+	public static class Similarity {
+
+		private String key;
+
+		private double similarity;
+
+		public Similarity(String key, double similarity) {
+			this.key = key;
+			this.similarity = similarity;
+		}
+
+	}
+
+	public class EmbeddingMath {
+
+		public static double cosineSimilarity(List<Double> vectorX, List<Double> vectorY) {
+			if (vectorX.size() != vectorY.size()) {
+				throw new IllegalArgumentException("Vectors lengths must be equal");
+			}
+
+			double dotProduct = dotProduct(vectorX, vectorY);
+			double normX = norm(vectorX);
+			double normY = norm(vectorY);
+
+			if (normX == 0 || normY == 0) {
+				throw new IllegalArgumentException("Vectors cannot have zero norm");
+			}
+
+			return dotProduct / (Math.sqrt(normX) * Math.sqrt(normY));
+		}
+
+		public static double dotProduct(List<Double> vectorX, List<Double> vectorY) {
+			if (vectorX.size() != vectorY.size()) {
+				throw new IllegalArgumentException("Vectors lengths must be equal");
+			}
+
+			double result = 0;
+			for (int i = 0; i < vectorX.size(); ++i) {
+				result += vectorX.get(i) * vectorY.get(i);
+			}
+
+			return result;
+		}
+
+		public static double norm(List<Double> vector) {
+			return dotProduct(vector, vector);
+		}
+
+	}
+
+}

spring-ai-core/src/test/java/org/springframework/ai/core/loader/LoaderTests.java

@@ -1,6 +1,7 @@
 package org.springframework.ai.core.loader;
 
 import org.junit.jupiter.api.Test;
+import org.springframework.ai.core.document.Document;
 import org.springframework.ai.core.loader.impl.JsonLoader;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.test.context.SpringBootTest;
@@ -23,7 +24,7 @@ public class LoaderTests {
 		List<Document> documents = jsonLoader.load();
 		assertThat(documents).isNotEmpty();
 		for (Document document : documents) {
-			System.out.println(document);
+			assertThat(document.getText()).isNotEmpty();
 		}
 	}
 

spring-ai-openai/src/main/java/org/springframework/ai/openai/embedding/OpenAiEmbeddingClient.java

@@ -5,6 +5,7 @@ import com.theokanning.openai.embedding.EmbeddingRequest;
 import com.theokanning.openai.service.OpenAiService;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.ai.core.document.Document;
 import org.springframework.ai.core.embedding.Embedding;
 import org.springframework.ai.core.embedding.EmbeddingClient;
 import org.springframework.ai.core.embedding.EmbeddingResponse;
@@ -37,6 +38,16 @@ public class OpenAiEmbeddingClient implements EmbeddingClient {
 		return generateEmbeddingResult(nativeEmbeddingResult).getData().get(0).getEmbedding();
 	}
 
+	public List<Double> createEmbedding(Document document) {
+		EmbeddingRequest embeddingRequest = EmbeddingRequest.builder()
+			.input(List.of(document.getContent()))
+			.model(this.model)
+			.build();
+		com.theokanning.openai.embedding.EmbeddingResult nativeEmbeddingResult = this.openAiService
+			.createEmbeddings(embeddingRequest);
+		return generateEmbeddingResult(nativeEmbeddingResult).getData().get(0).getEmbedding();
+	}
+
 	public List<List<Double>> createEmbedding(List<String> texts) {
 		EmbeddingResponse embeddingResponse = createEmbeddingResult(texts);
 		return embeddingResponse.getData().stream().map(emb -> emb.getEmbedding()).collect(Collectors.toList());

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

@@ -8,7 +8,6 @@ import org.springframework.context.annotation.Bean;
 import org.springframework.util.StringUtils;
 
 import java.io.IOException;
-import java.util.Properties;
 
 @SpringBootConfiguration
 public class OpenAiTestConfiguration {

spring-ai-openai/src/test/java/org/springframework/ai/openai/acme/AcmeIntegrationTest.java

@@ -0,0 +1,96 @@
+package org.springframework.ai.openai.acme;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.core.document.Document;
+import org.springframework.ai.core.llm.LLMResponse;
+import org.springframework.ai.core.llm.LlmClient;
+import org.springframework.ai.core.loader.impl.JsonLoader;
+import org.springframework.ai.core.prompt.ChatPromptTemplate;
+import org.springframework.ai.core.prompt.Prompt;
+import org.springframework.ai.core.prompt.PromptTemplate;
+import org.springframework.ai.core.prompt.messages.ChatMessage;
+import org.springframework.ai.core.prompt.messages.SystemMessage;
+import org.springframework.ai.core.prompt.messages.UserMessage;
+import org.springframework.ai.core.retriever.impl.VectorStoreRetriever;
+import org.springframework.ai.core.vectorstore.VectorStore;
+import org.springframework.ai.core.vectorstore.impl.InMemoryVectorStore;
+import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
+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.util.List;
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@SpringBootTest
+public class AcmeIntegrationTest {
+
+	@Value("classpath:bikes.json")
+	private Resource resource;
+
+	@Autowired
+	private OpenAiEmbeddingClient embeddingClient;
+
+	@Autowired
+	private LlmClient llmClient;
+
+	@Test
+	void beanTest() {
+		assertThat(resource).isNotNull();
+		assertThat(embeddingClient).isNotNull();
+		assertThat(llmClient).isNotNull();
+	}
+
+	void acmeChain() {
+
+		// Step 1 - load documents
+		JsonLoader jsonLoader = new JsonLoader("description", resource);
+		List<Document> documents = jsonLoader.load();
+
+		// Step 2 - Create embeddings and save to vector store
+
+		VectorStore vectorStore = new InMemoryVectorStore(embeddingClient);
+
+		vectorStore.add(documents);
+
+		// Now user query
+
+		// This will be wrapped up in a chain
+
+		VectorStoreRetriever vectorStoreRetriever = new VectorStoreRetriever(vectorStore);
+
+		String userQuery = "What bike is good for city commuting?";
+		List<Document> similarDocuments = vectorStoreRetriever.retrieve(userQuery);
+
+		// Try the case where not product was specified, so query over whatever docs might
+		// be releveant.
+
+		SystemMessage systemMessage = getSystemMessage(similarDocuments);
+		UserMessage userMessage = new UserMessage(userQuery);
+
+		// Create the prompt ad-hoc for now, need to put in system message and user
+		// message via ChatPromptTemplate or some other message building mechanic
+		Prompt prompt = new Prompt(List.of(systemMessage, userMessage));
+		LLMResponse response = llmClient.generate(prompt);
+
+		// Chain
+		// qa = new ConversationalRetrievalChain(llmClient, vectorStore, QueryOptions)
+	}
+
+	private SystemMessage getSystemMessage(List<Document> similarDocuments) {
+
+		// Would need to figure out which of the documenta metadata fields to add, from
+		// the loader, now just the 'full description.'
+
+		String systemMessageText = similarDocuments.stream()
+			.map(entry -> entry.getContent())
+			.collect(Collectors.joining("\n"));
+
+		return new SystemMessage(systemMessageText);
+
+	}
+
+}