Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update createCollection Command #1059

Merged
merged 8 commits into from
May 6, 2024
Merged

Update createCollection Command #1059

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0f1e776
update createCollection auth field
Hazel-Datastax May 6, 2024
3745f10
Merge branch 'main' into hazel/create_collection
Hazel-Datastax May 6, 2024
b18e6eb
update BeginOfflineSessionCommand
Hazel-Datastax May 6, 2024
d3aece1
Fix unit test
Hazel-Datastax May 6, 2024
3c7b783
Fix typo
Hazel-Datastax May 6, 2024
23573b9
Fix unit test
Hazel-Datastax May 6, 2024
9a96697
Fix IT
Hazel-Datastax May 6, 2024
529444d
Fix code
Hazel-Datastax May 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 3 additions & 30 deletions src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/CreateCollectionCommand.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,45 +133,18 @@ public record VectorizeConfig(
@Nullable
@Schema(
description = "Authentication config for chosen embedding service",
type = SchemaType.OBJECT,
implementation = VectorizeServiceAuthentication.class)
type = SchemaType.OBJECT)
@JsonProperty("authentication")
@JsonInclude(JsonInclude.Include.NON_NULL)
VectorizeServiceAuthentication vectorizeServiceAuthentication,
Map<String, String> authentication,
@Nullable
@Schema(
description =
"Optional parameters that match the template provided for the provider",
type = SchemaType.OBJECT)
@JsonProperty("parameters")
@JsonInclude(JsonInclude.Include.NON_NULL)
Map<String, Object> vectorizeServiceParameter) {
public record VectorizeServiceAuthentication(
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
@Schema(
description =
"List of authentications that can be used when sending documents that need vectorization. One or more of \"NONE\", \"HEADER\", \"SHARED_SECRET\"",
type = SchemaType.ARRAY,
implementation = String.class)
@JsonProperty("type")
List<
@Pattern(
regexp = "(NONE|HEADER|SHARED_SECRET)",
message =
"authentication type can only be one or more of 'NONE', 'HEADER' or 'SHARED_SECRET'")
String>
type,
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
@Schema(
description =
"Secret name. when stored_secrets authentication is used must be provided with the name of a pre-registered secret",
type = SchemaType.STRING,
implementation = String.class)
@JsonProperty("secretName")
String secretName) {}
}
Map<String, Object> parameters) {}
}

public record IndexingConfig(
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/io/stargate/sgv2/jsonapi/api/v1/CollectionResource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ else if (error instanceof JsonApiException jsonApiException) {
vectorizeConfig.provider(),
vectorizeConfig.modelName(),
collectionProperty.vectorConfig().vectorSize(),
vectorizeConfig.vectorizeServiceParameter());
vectorizeConfig.parameters());
}

CommandContext commandContext =
Expand Down
65 changes: 13 additions & 52 deletions src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/CollectionSettings.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
* Refactored as seperate class that represent a collection property.*
Expand Down Expand Up @@ -134,27 +133,13 @@ public static VectorConfig fromJson(JsonNode jsonNode, ObjectMapper objectMapper
// provider, modelName, must exist
String provider = vectorizeServiceNode.get("provider").asText();
String modelName = vectorizeServiceNode.get("modelName").asText();
// construct VectorizeConfig.VectorizeServiceAuthentication, can be null
// construct VectorizeConfig.authentication, can be null
JsonNode vectorizeServiceAuthenticationNode = vectorizeServiceNode.get("authentication");
VectorizeConfig.VectorizeServiceAuthentication vectorizeServiceAuthentication = null;
if (vectorizeServiceAuthenticationNode != null) {
Set<AuthenticationType> authenticationTypeSet = new HashSet<>();
vectorizeServiceAuthenticationNode
.get("type")
.forEach(
node -> authenticationTypeSet.add(AuthenticationType.fromString(node.asText())));
// when stored_secrets authentication is used must be provided with the name of a
// pre-registered secret.

String authenticationSecretName =
vectorizeServiceAuthenticationNode.has("secretName")
? vectorizeServiceAuthenticationNode.get("secretName").asText()
: null;
vectorizeServiceAuthentication =
new VectorizeConfig.VectorizeServiceAuthentication(
authenticationTypeSet, authenticationSecretName);
}
// construct VectorizeConfig.VectorizeServiceParameter, can be null
Map<String, String> vectorizeServiceAuthentication =
vectorizeServiceAuthenticationNode == null
? null
: objectMapper.convertValue(vectorizeServiceAuthenticationNode, Map.class);
// construct VectorizeConfig.parameters, can be null
JsonNode vectorizeServiceParameterNode = vectorizeServiceNode.get("parameters");
Map<String, Object> vectorizeServiceParameter =
vectorizeServiceParameterNode == null
Expand All @@ -171,12 +156,8 @@ public static VectorConfig fromJson(JsonNode jsonNode, ObjectMapper objectMapper
public record VectorizeConfig(
String provider,
String modelName,
VectorizeServiceAuthentication vectorizeServiceAuthentication,
Map<String, Object> vectorizeServiceParameter) {

public record VectorizeServiceAuthentication(
Set<AuthenticationType> type, String secretName) {}
}
Map<String, String> authentication,
Map<String, Object> parameters) {}
}

/**
Expand Down Expand Up @@ -369,35 +350,15 @@ public static CreateCollectionCommand collectionSettingToCreateCollectionCommand
if (collectionSetting.vectorConfig.vectorEnabled) {
CreateCollectionCommand.Options.VectorSearchConfig.VectorizeConfig vectorizeConfig = null;
if (collectionSetting.vectorConfig.vectorizeConfig != null) {
CreateCollectionCommand.Options.VectorSearchConfig.VectorizeConfig
.VectorizeServiceAuthentication
vectorizeServiceAuthentication = null;
if (collectionSetting.vectorConfig.vectorizeConfig.vectorizeServiceAuthentication != null) {
vectorizeServiceAuthentication =
new CreateCollectionCommand.Options.VectorSearchConfig.VectorizeConfig
.VectorizeServiceAuthentication(
collectionSetting
.vectorConfig
.vectorizeConfig
.vectorizeServiceAuthentication
.type
.stream()
.map(Enum::name)
.collect(Collectors.toList()),
collectionSetting
.vectorConfig
.vectorizeConfig
.vectorizeServiceAuthentication
.secretName);
}
Map<String, Object> vectorizeServiceParameter =
collectionSetting.vectorConfig.vectorizeConfig.vectorizeServiceParameter;
Map<String, String> authentication =
collectionSetting.vectorConfig.vectorizeConfig.authentication;
Map<String, Object> parameters = collectionSetting.vectorConfig.vectorizeConfig.parameters;
vectorizeConfig =
new CreateCollectionCommand.Options.VectorSearchConfig.VectorizeConfig(
collectionSetting.vectorConfig.vectorizeConfig.provider,
collectionSetting.vectorConfig.vectorizeConfig.modelName,
vectorizeServiceAuthentication,
vectorizeServiceParameter == null ? null : Map.copyOf(vectorizeServiceParameter));
authentication == null ? null : Map.copyOf(authentication),
parameters == null ? null : Map.copyOf(parameters));
}
vectorSearchConfig =
new CreateCollectionCommand.Options.VectorSearchConfig(
Expand Down
16 changes: 3 additions & 13 deletions ...io/stargate/sgv2/jsonapi/service/resolver/model/impl/CreateCollectionCommandResolver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,7 @@ private EmbeddingProvidersConfig.EmbeddingProviderConfig getAndValidateProviderC
private void validateAuthentication(
CreateCollectionCommand.Options.VectorSearchConfig.VectorizeConfig userConfig,
EmbeddingProvidersConfig.EmbeddingProviderConfig providerConfig) {
if (userConfig.vectorizeServiceAuthentication() == null) {
if (userConfig.authentication() == null) {
return;
}
// Check if user authentication type is support
Expand All @@ -280,13 +280,6 @@ private void validateAuthentication(
// throw ErrorCode.INVALID_CREATE_COLLECTION_OPTIONS.toApiException(
// "Authentication type '%s' is not supported", type);
// });
// Check if 'secretName' is provided if authentication type is 'SHARED_SECRET'
if (userConfig.vectorizeServiceAuthentication().type().contains("SHARED_SECRET")
&& (userConfig.vectorizeServiceAuthentication().secretName() == null
|| userConfig.vectorizeServiceAuthentication().secretName().isEmpty())) {
throw ErrorCode.INVALID_CREATE_COLLECTION_OPTIONS.toApiException(
"'secretName' must be provided for 'SHARED_SECRET' authentication type");
}
}

private void validateUserParameters(
Expand All @@ -296,8 +289,7 @@ private void validateUserParameters(
if (providerConfig.parameters() == null || providerConfig.parameters().isEmpty()) {
// If providerConfig.parameters() is null or empty but the user still provides parameters,
// it's an error
if (userConfig.vectorizeServiceParameter() != null
&& !userConfig.vectorizeServiceParameter().isEmpty()) {
if (userConfig.parameters() != null && !userConfig.parameters().isEmpty()) {
throw ErrorCode.INVALID_CREATE_COLLECTION_OPTIONS.toApiException(
"Parameters provided but the provider '%s' expects none", userConfig.provider());
}
Expand All @@ -310,9 +302,7 @@ private void validateUserParameters(
.collect(Collectors.toSet());

Map<String, Object> userParameters =
(userConfig.vectorizeServiceParameter() != null)
? userConfig.vectorizeServiceParameter()
: Collections.emptyMap();
(userConfig.parameters() != null) ? userConfig.parameters() : Collections.emptyMap();
// Check for unconfigured parameters provided by the user
userParameters
.keySet()
Expand Down
26 changes: 1 addition & 25 deletions ...n/offline/io/stargate/sgv2/jsonapi/api/model/command/impl/BeginOfflineSessionCommand.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,32 +129,8 @@ private CollectionSettings.VectorConfig.VectorizeConfig toCollectionSettingsVect
}
String provider = vectorize.provider();
String model = vectorize.modelName();
CollectionSettings.VectorConfig.VectorizeConfig.VectorizeServiceAuthentication
vectorizeServiceAuthentication =
toCollectionSettingsVectorizeServiceAuthentication(
vectorize.vectorizeServiceAuthentication());
return new CollectionSettings.VectorConfig.VectorizeConfig(
provider, model, vectorizeServiceAuthentication, vectorize.vectorizeServiceParameter());
}

private CollectionSettings.VectorConfig.VectorizeConfig.VectorizeServiceAuthentication
toCollectionSettingsVectorizeServiceAuthentication(
CreateCollectionCommand.Options.VectorSearchConfig.VectorizeConfig
.VectorizeServiceAuthentication
vectorizeServiceAuthentication) {
if (vectorizeServiceAuthentication == null) {
return null;
}
Set<CollectionSettings.AuthenticationType> authenticationTypes = null;
if (vectorizeServiceAuthentication.type() != null) {
authenticationTypes = new HashSet<>();
for (String authenticationType : vectorizeServiceAuthentication.type()) {
authenticationTypes.add(
CollectionSettings.AuthenticationType.fromString(authenticationType));
}
}
return new CollectionSettings.VectorConfig.VectorizeConfig.VectorizeServiceAuthentication(
authenticationTypes, vectorizeServiceAuthentication.secretName());
provider, model, vectorize.authentication(), vectorize.parameters());
}

private FileWriterParams buildFileWriterParams() {
Expand Down
42 changes: 10 additions & 32 deletions src/test/java/io/stargate/sgv2/jsonapi/api/configuration/ObjectMapperConfigurationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -469,10 +469,7 @@ public void happyPathVectorizeSearch() throws Exception {
"provider": "openai",
"modelName": "text-embedding-ada-002",
"authentication": {
"type": [
"SHARED_SECRET"
],
"secretName": "name_given_by_user"
"x-embedding-provider-key": "user_key"
},
"parameters": {
"projectId": "test project"
Expand All @@ -494,6 +491,9 @@ public void happyPathVectorizeSearch() throws Exception {
Map<String, Object> parameterMap = new HashMap<>();
parameterMap.put("projectId", "test project");

Map<String, Object> authenticationMap = new HashMap<>();
authenticationMap.put("x-embedding-provider-key", "user_key");

assertThat(result)
.isInstanceOfSatisfying(
CreateCollectionCommand.class,
Expand All @@ -509,35 +509,13 @@ public void happyPathVectorizeSearch() throws Exception {
.isEqualTo("openai");
assertThat(createCollection.options().vector().vectorizeConfig().modelName())
.isEqualTo("text-embedding-ada-002");
assertThat(
createCollection
.options()
.vector()
.vectorizeConfig()
.vectorizeServiceAuthentication()
.type())
.contains("SHARED_SECRET");
assertThat(
createCollection
.options()
.vector()
.vectorizeConfig()
.vectorizeServiceAuthentication()
.secretName())
.isEqualTo("name_given_by_user");
assertThat(
createCollection
.options()
.vector()
.vectorizeConfig()
.vectorizeServiceParameter())
assertThat(createCollection.options().vector().vectorizeConfig().authentication())
.isNotNull();
assertThat(
createCollection
.options()
.vector()
.vectorizeConfig()
.vectorizeServiceParameter())
assertThat(createCollection.options().vector().vectorizeConfig().authentication())
.isEqualTo(authenticationMap);
assertThat(createCollection.options().vector().vectorizeConfig().parameters())
.isNotNull();
assertThat(createCollection.options().vector().vectorizeConfig().parameters())
.isEqualTo(parameterMap);
assertThat(createCollection.options().indexing()).isNotNull();
assertThat(createCollection.options().indexing().allow()).isNull();
Expand Down
Loading