Updates the SearchRequest class to be non-final and adds a MilvusSearchRequest subclass that includes Milvus-specific fields for native expressions and search parameters.

- Updated `SearchRequest.java` to make the class non-final.
- Added `MilvusSearchRequest` with specific Milvus parameters such as `nativeExpression` and `searchParamsJson`.
- Modified `doSimilaritySearch` method in `MilvusVectorStore` to handle these new fields from `MilvusSearchRequest`.

Add unit tests for MilvusVectorStore and MilvusSearchRequest

Introduce comprehensive unit tests to validate the functionality of MilvusVectorStore and MilvusSearchRequest, including scenarios for native and filter expressions. Refactor MilvusVectorStore to improve filter expression handling by introducing a helper method for converted expressions.

Add detailed documentation for MilvusSearchRequest usage

Introduced sections explaining MilvusSearchRequest's parameters, `nativeExpression`, and `searchParamsJson`, with examples for enhanced clarity. This update provides guidance on leveraging Milvus-specific features for precise filtering and optimal search performance.

Signed-off-by: waileong <wai_leong1015@hotmail.com>
This commit is contained in:
waileong
2025-02-24 08:30:51 +08:00
committed by Ilayaperumal Gopinathan
parent bb7c52cdb1
commit e128a39ec3
6 changed files with 453 additions and 3 deletions

View File

@@ -33,7 +33,7 @@ import org.springframework.util.Assert;
* @author Thomas Vitale
* @author Ilayaperumal Gopinathan
*/
public final class SearchRequest {
public class SearchRequest {
/**
* Similarity threshold that accepts all search scores. A threshold value of 0.0 means
@@ -71,6 +71,16 @@ public final class SearchRequest {
.filterExpression(originalSearchRequest.getFilterExpression());
}
public SearchRequest() {
}
protected SearchRequest(SearchRequest original) {
this.query = original.query;
this.topK = original.topK;
this.similarityThreshold = original.similarityThreshold;
this.filterExpression = original.filterExpression;
}
public String getQuery() {
return this.query;
}

View File

@@ -161,6 +161,58 @@ vectorStore.similaritySearch(SearchRequest.builder()
NOTE: These filter expressions are converted into the equivalent Milvus filters.
== Using MilvusSearchRequest
MilvusSearchRequest extends SearchRequest, allowing you to use Milvus-specific search parameters such as native expressions and search parameter JSON.
[source,java]
----
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder()
.query("sample query")
.topK(5)
.similarityThreshold(0.7)
.nativeExpression("metadata[\"age\"] > 30") // Overrides filterExpression if both are set
.filterExpression("age <= 30") // Ignored if nativeExpression is set
.searchParamsJson("{\"nprobe\":128}")
.build();
List results = vectorStore.similaritySearch(request);
----
This allows greater flexibility when using Milvus-specific search features.
== Importance of `nativeExpression` and `searchParamsJson` in `MilvusSearchRequest`
These two parameters enhance Milvus search precision and ensure optimal query performance:
*nativeExpression*: Enables additional filtering capabilities using Milvus' native filtering expressions.
https://milvus.io/docs/boolean.md[Milvus Filtering]
Example:
[source,java]
----
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder()
.query("sample query")
.topK(5)
.nativeExpression("metadata['category'] == 'science'")
.build();
----
*searchParamsJson*: Essential for tuning search behavior when using IVF_FLAT, Milvus' default index.
https://milvus.io/docs/index.md?tab=floating[Milvus Vector Index]
By default, `IVF_FLAT` requires `nprobe` to be set for accurate results. If not specified, `nprobe` defaults to `1`, which can lead to poor recall or even zero search results.
Example:
[source,java]
----
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder()
.query("sample query")
.topK(5)
.searchParamsJson("{\"nprobe\":128}")
.build();
----
Using `nativeExpression` ensures advanced filtering, while `searchParamsJson` prevents ineffective searches caused by a low default `nprobe` value.
[[milvus-properties]]
== Milvus VectorStore properties

View File

@@ -0,0 +1,164 @@
package org.springframework.ai.vectorstore.milvus;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.lang.Nullable;
/**
* A specialized {@link SearchRequest} for Milvus vector search, extending the base
* request with Milvus-specific parameters.
* <p>
* This class introduces two additional fields:
* <ul>
* <li>{@code nativeExpression} - A native Milvus filter expression (e.g.,
* {@code "city LIKE
* 'New%'"}).</li>
* <li>{@code searchParamsJson} - A JSON string containing search parameters (e.g.,
* {@code "{\"nprobe\":128}"}).</li>
* </ul>
* <p>
* Use the {@link MilvusBuilder} to construct instances of this class.
*
* @author waileong
*/
public final class MilvusSearchRequest extends SearchRequest {
@Nullable
private final String nativeExpression;
@Nullable
private final String searchParamsJson;
/**
* Private constructor to initialize a MilvusSearchRequest using the base request and
* builder.
* @param baseRequest The base {@link SearchRequest} containing standard search
* fields.
* @param builder The {@link MilvusBuilder} containing Milvus-specific parameters.
*/
private MilvusSearchRequest(SearchRequest baseRequest, MilvusBuilder builder) {
super(baseRequest); // Copy all standard fields
this.nativeExpression = builder.nativeExpression;
this.searchParamsJson = builder.searchParamsJson;
}
/**
* Retrieves the native Milvus filter expression.
* @return A string representing the native Milvus expression, or {@code null} if not
* set.
*/
@Nullable
public String getNativeExpression() {
return this.nativeExpression;
}
/**
* Retrieves the JSON-encoded search parameters.
* @return A JSON string containing search parameters, or {@code null} if not set.
*/
@Nullable
public String getSearchParamsJson() {
return this.searchParamsJson;
}
/**
* Creates a new {@link MilvusBuilder} for constructing a {@link MilvusSearchRequest}.
* @return A new {@link MilvusBuilder} instance.
*/
public static MilvusBuilder milvusBuilder() {
return new MilvusBuilder();
}
/**
* Builder class for constructing instances of {@link MilvusSearchRequest}.
*/
public static class MilvusBuilder {
private final SearchRequest.Builder baseBuilder = SearchRequest.builder();
@Nullable
private String nativeExpression;
@Nullable
private String searchParamsJson;
/**
* {@link Builder#query(java.lang.String)}
*/
public MilvusBuilder query(String query) {
this.baseBuilder.query(query);
return this;
}
/**
* {@link Builder#topK(int)}
*/
public MilvusBuilder topK(int topK) {
this.baseBuilder.topK(topK);
return this;
}
/**
* {@link Builder#similarityThreshold(double)}
*/
public MilvusBuilder similarityThreshold(double threshold) {
this.baseBuilder.similarityThreshold(threshold);
return this;
}
/**
* {@link Builder#similarityThresholdAll()}
*/
public MilvusBuilder similarityThresholdAll() {
this.baseBuilder.similarityThresholdAll();
return this;
}
/**
* {@link Builder#filterExpression(String)}
*/
public MilvusBuilder filterExpression(String textExpression) {
this.baseBuilder.filterExpression(textExpression);
return this;
}
/**
* {@link Builder#filterExpression(Filter.Expression)}
*/
public MilvusBuilder filterExpression(Filter.Expression expression) {
this.baseBuilder.filterExpression(expression);
return this;
}
/**
* Sets the native Milvus filter expression.
* @param nativeExpression The native Milvus expression string.
* @return This builder instance.
*/
public MilvusBuilder nativeExpression(String nativeExpression) {
this.nativeExpression = nativeExpression;
return this;
}
/**
* Sets the JSON-encoded search parameters.
* @param searchParamsJson A JSON string containing search parameters.
* @return This builder instance.
*/
public MilvusBuilder searchParamsJson(String searchParamsJson) {
this.searchParamsJson = searchParamsJson;
return this;
}
/**
* Builds and returns a {@link MilvusSearchRequest} instance.
* @return A new {@link MilvusSearchRequest} object with the specified parameters.
*/
public MilvusSearchRequest build() {
SearchRequest parentRequest = this.baseBuilder.build();
return new MilvusSearchRequest(parentRequest, this);
}
}
}

View File

@@ -323,9 +323,18 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
@Override
public List<Document> doSimilaritySearch(SearchRequest request) {
String nativeFilterExpressions = "";
String searchParamsJson = null;
if (request instanceof MilvusSearchRequest milvusReq) {
nativeFilterExpressions = StringUtils.hasText(milvusReq.getNativeExpression())
? milvusReq.getNativeExpression() : getConvertedFilterExpression(request);
String nativeFilterExpressions = (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
searchParamsJson = StringUtils.hasText(milvusReq.getSearchParamsJson()) ? milvusReq.getSearchParamsJson()
: null;
}
else {
nativeFilterExpressions = getConvertedFilterExpression(request);
}
Assert.notNull(request.getQuery(), "Query string must not be null");
List<String> outFieldNames = new ArrayList<>();
@@ -348,6 +357,10 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
searchParamBuilder.withExpr(nativeFilterExpressions);
}
if (StringUtils.hasText(searchParamsJson)) {
searchParamBuilder.withParams(searchParamsJson);
}
R<SearchResults> respSearch = this.milvusClient.search(searchParamBuilder.build());
if (respSearch.getException() != null) {
@@ -385,6 +398,11 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
.toList();
}
private String getConvertedFilterExpression(SearchRequest request) {
return (request.getFilterExpression() != null)
? this.filterExpressionConverter.convertExpression(request.getFilterExpression()) : "";
}
private float getResultSimilarity(RowRecord rowRecord) {
Float score = (Float) rowRecord.get(SIMILARITY_FIELD_NAME);
return (this.metricType == MetricType.IP || this.metricType == MetricType.COSINE) ? score : (1 - score);

View File

@@ -0,0 +1,66 @@
package org.springframework.ai.vectorstore.milvus;
import org.junit.jupiter.api.Test;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.ai.vectorstore.SearchRequest.DEFAULT_TOP_K;
import static org.springframework.ai.vectorstore.SearchRequest.SIMILARITY_THRESHOLD_ACCEPT_ALL;
/**
* Test class for verifying the functionality of the {@link MilvusSearchRequest} class.
*
* @author waileong
*/
class MilvusSearchRequestTest {
@Test
void shouldBuildMilvusSearchRequestWithNativeExpression() {
String query = "sample query";
int topK = 10;
double similarityThreshold = 0.8;
String nativeExpression = "city LIKE 'New%'";
String searchParamsJson = "{\"nprobe\":128}";
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder()
.query(query)
.topK(topK)
.similarityThreshold(similarityThreshold)
.nativeExpression(nativeExpression)
.searchParamsJson(searchParamsJson)
.build();
assertThat(request.getQuery()).isEqualTo(query);
assertThat(request.getTopK()).isEqualTo(topK);
assertThat(request.getSimilarityThreshold()).isEqualTo(similarityThreshold);
assertThat(request.getNativeExpression()).isEqualTo(nativeExpression);
assertThat(request.getSearchParamsJson()).isEqualTo(searchParamsJson);
}
@Test
void shouldBuildMilvusSearchRequestWithDefaults() {
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder().build();
assertThat(request.getQuery()).isEmpty();
assertThat(request.getTopK()).isEqualTo(DEFAULT_TOP_K);
assertThat(request.getSimilarityThreshold()).isEqualTo(SIMILARITY_THRESHOLD_ACCEPT_ALL);
assertThat(request.getNativeExpression()).isNull();
assertThat(request.getSearchParamsJson()).isNull();
}
@Test
void shouldAllowSettingNativeExpressionIndependently() {
String nativeExpression = "age > 30";
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder().nativeExpression(nativeExpression).build();
assertThat(request.getNativeExpression()).isEqualTo(nativeExpression);
}
@Test
void shouldAllowSettingSearchParamsJsonIndependently() {
String searchParamsJson = "{\"metric_type\": \"IP\"}";
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder().searchParamsJson(searchParamsJson).build();
assertThat(request.getSearchParamsJson()).isEqualTo(searchParamsJson);
}
}

View File

@@ -0,0 +1,140 @@
package org.springframework.ai.vectorstore.milvus;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.R;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockedConstruction;
import org.mockito.MockedStatic;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.model.EmbeddingUtils;
import org.springframework.ai.vectorstore.SearchRequest;
import java.util.List;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
/**
* Unit test class for {@link MilvusVectorStore}.
*
* @author waileong
*/
@ExtendWith(MockitoExtension.class)
class MilvusVectorStoreTest {
@Mock
private MilvusServiceClient milvusClient;
@Mock
private EmbeddingModel embeddingModel;
private MilvusVectorStore vectorStore;
@BeforeEach
void setUp() {
vectorStore = MilvusVectorStore.builder(milvusClient, embeddingModel).build();
}
@Test
void shouldPerformSimilaritySearchWithNativeExpression() {
try (MockedStatic<EmbeddingUtils> mockedEmbeddingUtils = mockStatic(EmbeddingUtils.class);
MockedConstruction<SearchResultsWrapper> mockedSearchResultsWrapper = mockConstruction(
SearchResultsWrapper.class,
(mock, context) -> when(mock.getRowRecords(0)).thenReturn(List.of()))) {
String query = "sample query";
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder()
.query(query)
.topK(5)
.similarityThreshold(0.7)
.nativeExpression("metadata[\"age\"] > 30") // this has higher priority
.filterExpression("age <= 30") // this will be ignored
.searchParamsJson("{\"nprobe\":128}")
.build();
SearchParam capturedParam = performSimilaritySearch(mockedEmbeddingUtils, request);
assertThat(capturedParam.getTopK()).isEqualTo(request.getTopK());
assertThat(capturedParam.getExpr()).isEqualTo(request.getNativeExpression());
assertThat(capturedParam.getParams()).isEqualTo(request.getSearchParamsJson());
}
}
@Test
void shouldPerformSimilaritySearchWithFilterExpression() {
try (MockedStatic<EmbeddingUtils> mockedEmbeddingUtils = mockStatic(EmbeddingUtils.class);
MockedConstruction<SearchResultsWrapper> mockedSearchResultsWrapper = mockConstruction(
SearchResultsWrapper.class,
(mock, context) -> when(mock.getRowRecords(0)).thenReturn(List.of()))) {
String query = "sample query";
MilvusSearchRequest request = MilvusSearchRequest.milvusBuilder()
.query(query)
.topK(5)
.similarityThreshold(0.7)
.filterExpression("age > 30")
.searchParamsJson("{\"nprobe\":128}")
.build();
SearchParam capturedParam = performSimilaritySearch(mockedEmbeddingUtils, request);
assertThat(capturedParam.getTopK()).isEqualTo(request.getTopK());
assertThat(capturedParam.getExpr()).isEqualTo("metadata[\"age\"] > 30"); // filter
assertThat(capturedParam.getParams()).isEqualTo(request.getSearchParamsJson());
}
}
@Test
void shouldPerformSimilaritySearchWithOriginalSearchRequest() {
try (MockedStatic<EmbeddingUtils> mockedEmbeddingUtils = mockStatic(EmbeddingUtils.class);
MockedConstruction<SearchResultsWrapper> mockedSearchResultsWrapper = mockConstruction(
SearchResultsWrapper.class,
(mock, context) -> when(mock.getRowRecords(0)).thenReturn(List.of()))) {
String query = "sample query";
SearchRequest request = SearchRequest.builder()
.query(query)
.topK(5)
.similarityThreshold(0.7)
.filterExpression("age > 30")
.build();
SearchParam capturedParam = performSimilaritySearch(mockedEmbeddingUtils, request);
assertThat(capturedParam.getTopK()).isEqualTo(request.getTopK());
assertThat(capturedParam.getExpr()).isEqualTo("metadata[\"age\"] > 30"); // filter
assertThat(capturedParam.getParams()).isEqualTo("{}");
}
}
private SearchParam performSimilaritySearch(MockedStatic<EmbeddingUtils> mockedEmbeddingUtils,
SearchRequest request) {
List<Float> mockVector = List.of(1.0f, 2.0f, 3.0f);
mockedEmbeddingUtils.when(() -> EmbeddingUtils.toList(any())).thenReturn(mockVector);
SearchResults mockResults = mock(SearchResults.class);
when(mockResults.getResults()).thenReturn(SearchResultData.getDefaultInstance());
R<SearchResults> mockResponse = R.success(mockResults);
when(milvusClient.search(any(SearchParam.class))).thenReturn(mockResponse);
ArgumentCaptor<SearchParam> searchParamCaptor = ArgumentCaptor.forClass(SearchParam.class);
List<Document> results = vectorStore.doSimilaritySearch(request);
assertThat(results).isNotNull();
verify(milvusClient).search(searchParamCaptor.capture());
return searchParamCaptor.getValue();
}
}