@@ -122,20 +122,20 @@ static Stream<Arguments> provideFilters() {
122
122
);
123
123
}
124
124
125
- private static boolean isSortedByDistance (List <Document > docs ) {
125
+ private static boolean isSortedByScore (List <Document > docs ) {
126
126
127
- List <Float > distances = docs .stream ().map (doc -> ( Float ) doc . getMetadata (). get ( "distance" ) ).toList ();
127
+ List <Double > scores = docs .stream ().map (Document :: getScore ).toList ();
128
128
129
- if (CollectionUtils .isEmpty (distances ) || distances .size () == 1 ) {
129
+ if (CollectionUtils .isEmpty (scores ) || scores .size () == 1 ) {
130
130
return true ;
131
131
}
132
132
133
- Iterator <Float > iter = distances .iterator ();
134
- Float current ;
135
- Float previous = iter .next ();
133
+ Iterator <Double > iter = scores .iterator ();
134
+ Double current ;
135
+ Double previous = iter .next ();
136
136
while (iter .hasNext ()) {
137
137
current = iter .next ();
138
- if (previous > current ) {
138
+ if (previous < current ) {
139
139
return false ;
140
140
}
141
141
previous = current ;
@@ -166,7 +166,8 @@ public void addAndSearch(String distanceType) {
166
166
assertThat (results ).hasSize (1 );
167
167
Document resultDoc = results .get (0 );
168
168
assertThat (resultDoc .getId ()).isEqualTo (this .documents .get (2 ).getId ());
169
- assertThat (resultDoc .getMetadata ()).containsKeys ("meta2" , "distance" );
169
+ assertThat (resultDoc .getMetadata ()).containsKeys ("meta2" );
170
+ assertThat (resultDoc .getScore ()).isBetween (0.0 , 1.0 );
170
171
171
172
// Remove all documents from the store
172
173
vectorStore .delete (this .documents .stream ().map (doc -> doc .getId ()).toList ());
@@ -315,7 +316,8 @@ public void documentUpdate(String distanceType) {
315
316
Document resultDoc = results .get (0 );
316
317
assertThat (resultDoc .getId ()).isEqualTo (document .getId ());
317
318
assertThat (resultDoc .getText ()).isEqualTo ("Spring AI rocks!!" );
318
- assertThat (resultDoc .getMetadata ()).containsKeys ("meta1" , "distance" );
319
+ assertThat (resultDoc .getMetadata ()).containsKeys ("meta1" );
320
+ assertThat (resultDoc .getScore ()).isBetween (0.0 , 1.0 );
319
321
320
322
Document sameIdDocument = new Document (document .getId (),
321
323
"The World is Big and Salvation Lurks Around the Corner" ,
@@ -329,7 +331,8 @@ public void documentUpdate(String distanceType) {
329
331
resultDoc = results .get (0 );
330
332
assertThat (resultDoc .getId ()).isEqualTo (document .getId ());
331
333
assertThat (resultDoc .getText ()).isEqualTo ("The World is Big and Salvation Lurks Around the Corner" );
332
- assertThat (resultDoc .getMetadata ()).containsKeys ("meta2" , "distance" );
334
+ assertThat (resultDoc .getMetadata ()).containsKeys ("meta2" );
335
+ assertThat (resultDoc .getScore ()).isBetween (0.0 , 1.0 );
333
336
334
337
dropTable (context );
335
338
});
@@ -350,19 +353,14 @@ public void searchWithThreshold(String distanceType) {
350
353
351
354
assertThat (fullResult ).hasSize (3 );
352
355
353
- assertThat (isSortedByDistance (fullResult )).isTrue ();
356
+ assertThat (isSortedByScore (fullResult )).isTrue ();
354
357
355
- List <Float > distances = fullResult .stream ()
356
- .map (doc -> (Float ) doc .getMetadata ().get ("distance" ))
357
- .toList ();
358
+ List <Double > scores = fullResult .stream ().map (Document ::getScore ).toList ();
358
359
359
- float threshold = (distances .get (0 ) + distances .get (1 )) / 2 ;
360
+ double threshold = (scores .get (0 ) + scores .get (1 )) / 2 ;
360
361
361
- List <Document > results = vectorStore .similaritySearch (SearchRequest .builder ()
362
- .query ("Time Shelter" )
363
- .topK (5 )
364
- .similarityThreshold (1 - threshold )
365
- .build ());
362
+ List <Document > results = vectorStore .similaritySearch (
363
+ SearchRequest .builder ().query ("Time Shelter" ).topK (5 ).similarityThreshold (threshold ).build ());
366
364
367
365
assertThat (results ).hasSize (1 );
368
366
Document resultDoc = results .get (0 );
0 commit comments