diff --git a/CONFIGURATION.md b/CONFIGURATION.md index b705398cda..c6287a4115 100644 --- a/CONFIGURATION.md +++ b/CONFIGURATION.md @@ -53,6 +53,7 @@ Other Quarkus properties that are specifically relevant for the service: | `stargate.jsonapi.operations.database-config.ddl-retry-delay-millis` | `int` | `1000` | Delay time in seconds for DDL timeout. | | `stargate.jsonapi.operations.database-config.ddl-delay-millis` | `int` | `2000` | Delay between create table and create index to get the schema sync. | | `stargate.jsonapi.operations.vectorize-enabled` | `boolean` | `false` | Flag to enable server side vectorization. | +| `stargate.jsonapi.operations.return-deprecated-models` | `boolean` | `false` | Flag to indicate whether to return deprecated models in the FindEmbeddingProvidersCommand's response or not. | ## Jsonapi metering configuration diff --git a/src/main/java/io/stargate/sgv2/jsonapi/config/OperationsConfig.java b/src/main/java/io/stargate/sgv2/jsonapi/config/OperationsConfig.java index 0f502ad38c..26e28d5e6a 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/config/OperationsConfig.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/config/OperationsConfig.java @@ -296,6 +296,13 @@ interface ConsistencyConfig { @WithDefault("false") boolean vectorizeEnabled(); + /** + * @return Flag to indicate whether to return deprecated models in the + * FindEmbeddingProvidersCommand's response. + */ + @WithDefault("false") + boolean returnDeprecatedModels(); + /** Offline mode configuration. */ @NotNull @Valid diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperation.java b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperation.java index 2bdb65eeac..3b701ff840 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperation.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperation.java @@ -15,8 +15,8 @@ * Operation that list all available and enabled vector providers into the {@link * CommandStatus#EXISTING_VECTOR_PROVIDERS} command status. */ -public record FindEmbeddingProvidersOperation(EmbeddingProvidersConfig config) - implements Operation { +public record FindEmbeddingProvidersOperation( + EmbeddingProvidersConfig config, boolean returnDeprecatedModels) implements Operation { @Override public Uni> execute( DataApiRequestInfo dataApiRequestInfo, QueryExecutor queryExecutor) { @@ -29,7 +29,9 @@ public Uni> execute( .collect( Collectors.toMap( Map.Entry::getKey, - entry -> EmbeddingProviderResponse.provider(entry.getValue()))); + entry -> + EmbeddingProviderResponse.provider( + entry.getValue(), returnDeprecatedModels))); return new Result(embeddingProviders); }); } @@ -57,7 +59,7 @@ public CommandResult get() { * @param parameters Customizable parameters for the provider's service. * @param models Model configurations available from the provider. */ - private record EmbeddingProviderResponse( + protected record EmbeddingProviderResponse( String displayName, Optional url, Map< @@ -67,10 +69,16 @@ private record EmbeddingProviderResponse( List parameters, List models) { private static EmbeddingProviderResponse provider( - EmbeddingProvidersConfig.EmbeddingProviderConfig embeddingProviderConfig) { + EmbeddingProvidersConfig.EmbeddingProviderConfig embeddingProviderConfig, + boolean returnDeprecatedModels) { ArrayList modelsRemoveProperties = new ArrayList<>(); for (EmbeddingProvidersConfig.EmbeddingProviderConfig.ModelConfig model : embeddingProviderConfig.models()) { + if (!returnDeprecatedModels) { + if (model.deprecated().orElse(false)) { + continue; + } + } ModelConfigResponse returnModel = ModelConfigResponse.returnModelConfigResponse( model.name(), model.vectorDimension(), model.parameters()); @@ -94,7 +102,7 @@ private static EmbeddingProviderResponse provider( * @param vectorDimension vector dimension of the model. * @param parameters Parameters for customizing the model. */ - private record ModelConfigResponse( + protected record ModelConfigResponse( String name, Optional vectorDimension, List parameters) { private static ModelConfigResponse returnModelConfigResponse( String name, diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/FindEmbeddingProvidersCommandResolver.java b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/FindEmbeddingProvidersCommandResolver.java index 43799a8794..2d2967cfa7 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/FindEmbeddingProvidersCommandResolver.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/service/resolver/FindEmbeddingProvidersCommandResolver.java @@ -32,6 +32,7 @@ public Operation resolveDatabaseCommand( if (!operationsConfig.vectorizeEnabled()) { throw ErrorCodeV1.VECTORIZE_FEATURE_NOT_AVAILABLE.toApiException(); } - return new FindEmbeddingProvidersOperation(embeddingProvidersConfig); + return new FindEmbeddingProvidersOperation( + embeddingProvidersConfig, operationsConfig.returnDeprecatedModels()); } } diff --git a/src/test/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperationTest.java new file mode 100644 index 0000000000..7bbaffdea7 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/service/operation/embeddings/FindEmbeddingProvidersOperationTest.java @@ -0,0 +1,101 @@ +package io.stargate.sgv2.jsonapi.service.operation.embeddings; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.fasterxml.jackson.databind.ObjectMapper; +import io.quarkus.test.junit.QuarkusTest; +import io.quarkus.test.junit.TestProfile; +import io.smallrye.mutiny.helpers.test.UniAssertSubscriber; +import io.stargate.sgv2.jsonapi.api.model.command.CommandResult; +import io.stargate.sgv2.jsonapi.api.model.command.CommandStatus; +import io.stargate.sgv2.jsonapi.service.embedding.configuration.EmbeddingProvidersConfig; +import io.stargate.sgv2.jsonapi.service.embedding.configuration.EmbeddingProvidersConfigImpl; +import io.stargate.sgv2.jsonapi.service.operation.collections.OperationTestBase; +import io.stargate.sgv2.jsonapi.testresource.NoGlobalResourcesTestProfile; +import jakarta.inject.Inject; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.Supplier; +import javax.annotation.Nullable; +import org.junit.jupiter.api.Test; + +@QuarkusTest +@TestProfile(NoGlobalResourcesTestProfile.Impl.class) +public class FindEmbeddingProvidersOperationTest extends OperationTestBase { + @Inject ObjectMapper objectMapper; + + @Test + public void testDoNotReturnDeprecatedModels() { + CommandResult result = executeFindEmbeddingProvidersOperation(false); + Map configMap = + (Map) + result.status().get(CommandStatus.EXISTING_VECTOR_PROVIDERS); + assertThat(configMap).hasSize(1); + FindEmbeddingProvidersOperation.EmbeddingProviderResponse providerConfig = + configMap.get("provider1"); + assertThat(providerConfig.models()).hasSize(1); + assertThat(providerConfig.models().getFirst().name()).isEqualTo("supported-model"); + } + + @Test + public void testReturnDeprecatedModels() { + CommandResult result = executeFindEmbeddingProvidersOperation(true); + Map configMap = + (Map) + result.status().get(CommandStatus.EXISTING_VECTOR_PROVIDERS); + assertThat(configMap).hasSize(1); + FindEmbeddingProvidersOperation.EmbeddingProviderResponse providerConfig = + configMap.get("provider1"); + assertThat(providerConfig.models()).hasSize(2); + assertThat(providerConfig.models().getFirst().name()).isEqualTo("supported-model"); + assertThat(providerConfig.models().getLast().name()).isEqualTo("deprecated-model"); + } + + private CommandResult executeFindEmbeddingProvidersOperation(boolean returnDeprecatedModels) { + FindEmbeddingProvidersOperation findEmbeddingProvidersOperation = + new FindEmbeddingProvidersOperation( + getTestEmbeddingProvidersConfig(), returnDeprecatedModels); + Supplier execute = + findEmbeddingProvidersOperation + .execute(null, null) + .subscribe() + .withSubscriber(UniAssertSubscriber.create()) + .awaitItem() + .getItem(); + return execute.get(); + } + + private EmbeddingProvidersConfig getTestEmbeddingProvidersConfig() { + return new EmbeddingProvidersConfig() { + @Override + public Map providers() { + List modelConfigs = + List.of( + new EmbeddingProvidersConfigImpl.EmbeddingProviderConfigImpl.ModelConfigImpl( + "supported-model", + Optional.of(1), + List.of(), + Map.of(), + Optional.of(false), + Optional.of("url1")), + new EmbeddingProvidersConfigImpl.EmbeddingProviderConfigImpl.ModelConfigImpl( + "deprecated-model", + Optional.of(1), + List.of(), + Map.of(), + Optional.of(true), + Optional.of("url1"))); + EmbeddingProviderConfig providerConfig = + new EmbeddingProvidersConfigImpl.EmbeddingProviderConfigImpl( + "provider1", true, Optional.of("url1"), Map.of(), List.of(), null, modelConfigs); + return Map.of("provider1", providerConfig); + } + + @Override + public @Nullable CustomConfig custom() { + return null; + } + }; + } +}