Skip to content

Commit

Permalink
iter
Browse files Browse the repository at this point in the history
  • Loading branch information
pmpailis committed Oct 6, 2024
1 parent 88b9384 commit f7e4802
Show file tree
Hide file tree
Showing 5 changed files with 10 additions and 151 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
{
Request request = new Request("GET", INDEX_NAME + "/_search");
SearchSourceBuilder source = new SearchSourceBuilder();
source.retriever(new KnnRetrieverBuilder("vector", new float[]{1.0f}, null, 10, 15, null));
source.retriever(new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, null));
request.setJsonEntity(Strings.toString(source));
getRestClient().performRequest(request);
}
Expand All @@ -100,7 +100,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
{
Request request = new Request("GET", INDEX_NAME + "/_search");
SearchSourceBuilder source = new SearchSourceBuilder();
source.retriever(new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[]{1.0f}, 10, 15, null)));
source.retriever(new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, null)));
request.setJsonEntity(Strings.toString(source));
getRestClient().performRequest(request);
}
Expand All @@ -119,7 +119,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
{
Request request = new Request("GET", INDEX_NAME + "/_search");
SearchSourceBuilder source = new SearchSourceBuilder();
source.knnSearch(List.of(new KnnSearchBuilder("vector", new float[]{1.0f}, 10, 15, null)));
source.knnSearch(List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, null)));
request.setJsonEntity(Strings.toString(source));
getRestClient().performRequest(request);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,6 @@
import org.elasticsearch.search.collapse.CollapseBuilderTests;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.rescore.QueryRescorerBuilder;
import org.elasticsearch.search.retriever.KnnRetrieverBuilder;
import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
import org.elasticsearch.search.slice.SliceBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.ScoreSortBuilder;
Expand Down Expand Up @@ -602,61 +600,6 @@ public void testNegativeTrackTotalHits() throws IOException {
}
}

public void testStandardRetrieverParsing() throws IOException {
String restContent = "{"
+ " \"retriever\": {"
+ " \"standard\": {"
+ " \"query\": {"
+ " \"match_all\": {}"
+ " }"
+ " }"
+ " }"
+ "}";
SearchUsageHolder searchUsageHolder = new UsageService().getSearchUsageHolder();
try (XContentParser jsonParser = createParser(JsonXContent.jsonXContent, restContent)) {
SearchSourceBuilder source = new SearchSourceBuilder().parseXContent(jsonParser, true, searchUsageHolder, nf -> true);
assertThat(source.retriever(), instanceOf(StandardRetrieverBuilder.class));
try (XContentParser parseSerialized = createParser(JsonXContent.jsonXContent, Strings.toString(source))) {
SearchSourceBuilder deserializedSource = new SearchSourceBuilder().parseXContent(
parseSerialized,
true,
searchUsageHolder,
nf -> true
);
assertThat(deserializedSource.retriever(), instanceOf(StandardRetrieverBuilder.class));
}
}
}

public void testKnnRetrieverParsing() throws IOException {
String restContent = "{"
+ " \"retriever\": {"
+ " \"knn\": {"
+ " \"query_vector\": ["
+ " 3"
+ " ],"
+ " \"field\": \"vector\","
+ " \"k\": 10,"
+ " \"num_candidates\": 15"
+ " }"
+ " }"
+ "}";
SearchUsageHolder searchUsageHolder = new UsageService().getSearchUsageHolder();
try (XContentParser jsonParser = createParser(JsonXContent.jsonXContent, restContent)) {
SearchSourceBuilder source = new SearchSourceBuilder().parseXContent(jsonParser, true, searchUsageHolder, nf -> true);
assertThat(source.retriever(), instanceOf(KnnRetrieverBuilder.class));
try (XContentParser parseSerialized = createParser(JsonXContent.jsonXContent, Strings.toString(source))) {
SearchSourceBuilder deserializedSource = new SearchSourceBuilder().parseXContent(
parseSerialized,
true,
searchUsageHolder,
nf -> true
);
assertThat(deserializedSource.retriever(), instanceOf(KnnRetrieverBuilder.class));
}
}
}

public void testStoredFieldsUsage() throws IOException {
Set<String> storedFieldRestVariations = Set.of(
"{\"stored_fields\" : [\"_none_\"]}",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,6 @@
import org.elasticsearch.xpack.core.XPackPlugin;
import org.elasticsearch.xpack.inference.InferencePlugin;
import org.junit.Before;
import org.junit.Ignore;

import java.io.IOException;
import java.util.Collection;
Expand Down Expand Up @@ -92,7 +91,7 @@ public void setup() throws IOException {

public void testTelemetryForRRFRetriever() throws IOException {

if(false == isRetrieverTelemetryEnabled()) {
if (false == isRetrieverTelemetryEnabled()) {
return;
}

Expand Down Expand Up @@ -135,7 +134,8 @@ public void testTelemetryForRRFRetriever() throws IOException {
getRestClient().performRequest(request);
}

// search#5 - this will record 1 entry for "retriever" in `sections`, and 1 for "text_similarity_reranker" under `retrievers`, as well as
// search#5 - this will record 1 entry for "retriever" in `sections`, and 1 for "text_similarity_reranker" under `retrievers`, as
// well as
// 1 "standard" under `retrievers`, and eventually 1 for "match" under `queries`
{
Request request = new Request("GET", INDEX_NAME + "/_search");
Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
{
Request request = new Request("GET", INDEX_NAME + "/_search");
SearchSourceBuilder source = new SearchSourceBuilder();
source.retriever(new KnnRetrieverBuilder("vector", new float[]{1.0f}, null, 10, 15, null));
source.retriever(new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, null));
request.setJsonEntity(Strings.toString(source));
getRestClient().performRequest(request);
}
Expand All @@ -119,7 +119,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
{
Request request = new Request("GET", INDEX_NAME + "/_search");
SearchSourceBuilder source = new SearchSourceBuilder();
source.retriever(new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[]{1.0f}, 10, 15, null)));
source.retriever(new StandardRetrieverBuilder(new KnnVectorQueryBuilder("vector", new float[] { 1.0f }, 10, 15, null)));
request.setJsonEntity(Strings.toString(source));
getRestClient().performRequest(request);
}
Expand All @@ -143,7 +143,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
new RRFRetrieverBuilder(
Arrays.asList(
new CompoundRetrieverBuilder.RetrieverSource(
new KnnRetrieverBuilder("vector", new float[]{1.0f}, null, 10, 15, null),
new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, null),
null
),
new CompoundRetrieverBuilder.RetrieverSource(
Expand All @@ -163,7 +163,7 @@ public void testTelemetryForRRFRetriever() throws IOException {
{
Request request = new Request("GET", INDEX_NAME + "/_search");
SearchSourceBuilder source = new SearchSourceBuilder();
source.knnSearch(List.of(new KnnSearchBuilder("vector", new float[]{1.0f}, 10, 15, null)));
source.knnSearch(List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, null)));
request.setJsonEntity(Strings.toString(source));
getRestClient().performRequest(request);
}
Expand Down

0 comments on commit f7e4802

Please sign in to comment.