Remove VectorStoreRetriever interface and retriever package

This commit is contained in:
Mark Pollack
2024-01-11 10:52:28 -05:00
parent 4dc92f0cb7
commit b35c71061e
2 changed files with 1 additions and 64 deletions

View File

@@ -20,7 +20,6 @@ import org.springframework.ai.prompt.SystemPromptTemplate;
import org.springframework.ai.prompt.messages.Message;
import org.springframework.ai.prompt.messages.UserMessage;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.retriever.VectorStoreRetriever;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
@@ -73,8 +72,6 @@ public class AcmeIT extends AbstractIT {
// Now user query
VectorStoreRetriever vectorStoreRetriever = new VectorStoreRetriever(vectorStore);
logger.info("Retrieving relevant documents");
String userQuery = "What bike is good for city commuting?";
@@ -82,7 +79,7 @@ public class AcmeIT extends AbstractIT {
// "How much does the SonicRide 8S cost?";
// Eventually include metadata in query.
List<Document> similarDocuments = vectorStoreRetriever.retrieve(userQuery);
List<Document> similarDocuments = vectorStore.similaritySearch(userQuery);
logger.info(String.format("Found %s relevant documents.", similarDocuments.size()));
// Try the case where not product was specified, so query over whatever docs might

View File

@@ -1,60 +0,0 @@
package org.springframework.ai.retriever;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
public class VectorStoreRetriever implements DocumentRetriever {
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) {
SearchRequest request = SearchRequest.query(query).withTopK(this.k);
if (threshold.isPresent()) {
request.withSimilarityThreshold(this.threshold.get());
}
return this.vectorStore.similaritySearch(request);
}
}