Expose source_model parameter for vector-enabled collections

src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/CreateCollectionCommand.java

@@ -95,6 +95,20 @@ public record VectorSearchConfig(
             @JsonProperty("metric")
             @JsonAlias("function") // old name
             String metric,
+        @Nullable
+            @Pattern(
+                regexp =
+                    "(openai-v3-large|openai-v3-small|ada002|gecko|nv-qa-4|cohere-v3|bert|other)",
+                message =
+                    "sourceModel options are 'openai-v3-large', 'openai-v3-small', 'ada002', 'gecko', 'nv-qa-4', 'cohere-v3', 'bert', and 'other'")
+            @Schema(
+                description =
+                    "The 'sourceModel' option configures the index with the fastest settings for a given source of embeddings vectors",
+                defaultValue = "other",
+                type = SchemaType.STRING,
+                implementation = String.class)
+            @JsonProperty("sourceModel")
+            String sourceModel,
         @Valid
             @Nullable
             @JsonInclude(JsonInclude.Include.NON_NULL)
@@ -105,9 +119,11 @@ public record VectorSearchConfig(
             @JsonProperty("service")
             VectorizeConfig vectorizeConfig) {
 
-      public VectorSearchConfig(Integer dimension, String metric, VectorizeConfig vectorizeConfig) {
+      public VectorSearchConfig(
+          Integer dimension, String metric, String sourceModel, VectorizeConfig vectorizeConfig) {
         this.dimension = dimension;
         this.metric = metric == null ? "cosine" : metric;
+        this.sourceModel = sourceModel == null ? "other" : sourceModel;
         this.vectorizeConfig = vectorizeConfig;
       }
     }

src/main/java/io/stargate/sgv2/jsonapi/config/constants/DocumentConstants.java

@@ -30,6 +30,9 @@ interface Fields {
     /** Key for vector function name definition in cql index. */
     String VECTOR_INDEX_FUNCTION_NAME = "similarity_function";
 
+    /** Key for vector source model name definition in cql index. */
+    String VECTOR_INDEX_SOURCE_MODEL_NAME = "source_model";
+
     /** Field name used in projection clause to get similarity score in response. */
     String VECTOR_FUNCTION_SIMILARITY_FIELD = "$similarity";
 

src/main/java/io/stargate/sgv2/jsonapi/exception/ErrorCodeV1.java

@@ -141,6 +141,8 @@ public enum ErrorCodeV1 {
 
   VECTOR_SEARCH_INVALID_FUNCTION_NAME("Invalid vector search function name"),
 
+  VECTOR_SEARCH_INVALID_SOURCE_MODEL_NAME("Invalid vector search source model name"),
+
   VECTOR_SEARCH_TOO_BIG_VALUE("Vector embedding property '$vector' length too big"),
   VECTOR_SIZE_MISMATCH("Length of vector parameter different from declared '$vector' dimension"),
 

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/TableSchemaObject.java

@@ -6,8 +6,10 @@
 import com.datastax.oss.driver.api.core.metadata.schema.TableMetadata;
 import com.datastax.oss.driver.api.core.type.VectorType;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
 import io.stargate.sgv2.jsonapi.config.constants.VectorConstant;
 import io.stargate.sgv2.jsonapi.service.schema.SimilarityFunction;
+import io.stargate.sgv2.jsonapi.service.schema.SourceModel;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -50,15 +52,19 @@ public static TableSchemaObject from(TableMetadata tableMetadata, ObjectMapper o
                     indexMetadata -> indexMetadata.getTarget().equals(column.getKey().asCql(true)))
                 .findFirst();
         SimilarityFunction similarityFunction = SimilarityFunction.COSINE;
+        SourceModel sourceModel = SourceModel.OTHER;
         if (index.isPresent()) {
           final IndexMetadata indexMetadata = index.get();
           final Map<String, String> indexOptions = indexMetadata.getOptions();
-          final String sourceModel = indexOptions.get("source_model");
-          final String similarityFunctionValue = indexOptions.get("similarity_function");
+          final String sourceModelValue =
+              indexOptions.get(DocumentConstants.Fields.VECTOR_INDEX_SOURCE_MODEL_NAME);
+          final String similarityFunctionValue =
+              indexOptions.get(DocumentConstants.Fields.VECTOR_INDEX_FUNCTION_NAME);
           if (similarityFunctionValue != null) {
             similarityFunction = SimilarityFunction.fromString(similarityFunctionValue);
-          } else if (sourceModel != null) {
-            similarityFunction = VectorConstant.SUPPORTED_SOURCES.get(sourceModel);
+          } else if (sourceModelValue != null) {
+            similarityFunction = VectorConstant.SUPPORTED_SOURCES.get(sourceModelValue);
+            sourceModel = SourceModel.fromString(sourceModelValue);
           }
         }
         int dimension = vectorType.getDimensions();
@@ -67,6 +73,7 @@ public static TableSchemaObject from(TableMetadata tableMetadata, ObjectMapper o
                 column.getKey().asInternal(),
                 dimension,
                 similarityFunction,
+                sourceModel,
                 vectorizeConfigMap.get(column.getKey().asInternal()));
         columnVectorDefinitions.add(columnVectorDefinition);
       }

src/main/java/io/stargate/sgv2/jsonapi/service/cqldriver/executor/VectorConfig.java

@@ -4,6 +4,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
 import io.stargate.sgv2.jsonapi.service.schema.SimilarityFunction;
+import io.stargate.sgv2.jsonapi.service.schema.SourceModel;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -72,6 +73,7 @@ public record ColumnVectorDefinition(
       String fieldName,
       int vectorSize,
       SimilarityFunction similarityFunction,
+      SourceModel sourceModel,
       VectorizeConfig vectorizeConfig) {
 
     // convert a vector jsonNode from comment option to vectorConfig, used for collection
@@ -80,11 +82,15 @@ public static ColumnVectorDefinition fromJson(JsonNode jsonNode, ObjectMapper ob
       int dimension = jsonNode.get("dimension").asInt();
       SimilarityFunction similarityFunction =
           SimilarityFunction.fromString(jsonNode.get("metric").asText());
+      // sourceModel doesn't exist if the collection was created before supporting sourceModel; if
+      // missing, it will be an empty string and sourceModel becomes OTHER.
+      SourceModel sourceModel = SourceModel.fromString(jsonNode.path("sourceModel").asText());
 
       return fromJson(
           DocumentConstants.Fields.VECTOR_EMBEDDING_TEXT_FIELD,
           dimension,
           similarityFunction,
+          sourceModel,
           jsonNode,
           objectMapper);
     }
@@ -94,6 +100,7 @@ public static ColumnVectorDefinition fromJson(
         String fieldName,
         int dimension,
         SimilarityFunction similarityFunction,
+        SourceModel sourceModel,
         JsonNode jsonNode,
         ObjectMapper objectMapper) {
       VectorizeConfig vectorizeConfig = null;
@@ -102,7 +109,8 @@ public static ColumnVectorDefinition fromJson(
       if (vectorizeServiceNode != null) {
         vectorizeConfig = VectorizeConfig.fromJson(vectorizeServiceNode, objectMapper);
       }
-      return new ColumnVectorDefinition(fieldName, dimension, similarityFunction, vectorizeConfig);
+      return new ColumnVectorDefinition(
+          fieldName, dimension, similarityFunction, sourceModel, vectorizeConfig);
     }
 
     /**

src/main/java/io/stargate/sgv2/jsonapi/service/operation/collections/CreateCollectionOperation.java

@@ -20,6 +20,7 @@
 import io.stargate.sgv2.jsonapi.service.cqldriver.executor.QueryExecutor;
 import io.stargate.sgv2.jsonapi.service.operation.Operation;
 import io.stargate.sgv2.jsonapi.service.schema.SimilarityFunction;
+import io.stargate.sgv2.jsonapi.service.schema.SourceModel;
 import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionSchemaObject;
 import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionTableMatcher;
 import java.time.Duration;
@@ -41,6 +42,7 @@ public record CreateCollectionOperation(
     boolean vectorSearch,
     int vectorSize,
     String vectorFunction,
+    String sourceModel,
     String comment,
     int ddlDelayMillis,
     boolean tooManyIndexesRollbackEnabled,
@@ -60,6 +62,7 @@ public static CreateCollectionOperation withVectorSearch(
       String name,
       int vectorSize,
       String vectorFunction,
+      String sourceModel,
       String comment,
       int ddlDelayMillis,
       boolean tooManyIndexesRollbackEnabled,
@@ -73,6 +76,7 @@ public static CreateCollectionOperation withVectorSearch(
         true,
         vectorSize,
         vectorFunction,
+        sourceModel,
         comment,
         ddlDelayMillis,
         tooManyIndexesRollbackEnabled,
@@ -98,6 +102,7 @@ public static CreateCollectionOperation withoutVectorSearch(
         false,
         0,
         null,
+        null,
         comment,
         ddlDelayMillis,
         tooManyIndexesRollbackEnabled,
@@ -142,6 +147,7 @@ public Uni<Supplier<CommandResult>> execute(
             vectorSearch,
             vectorSize,
             SimilarityFunction.fromString(vectorFunction),
+            SourceModel.fromString(sourceModel),
             comment,
             objectMapper);
     // if table exists we have to choices:
@@ -492,6 +498,8 @@ public List<SimpleStatement> getIndexStatements(
           appender
               + " \"%s_query_vector_value\" ON \"%s\".\"%s\" (query_vector_value) USING 'StorageAttachedIndex' WITH OPTIONS = { 'similarity_function': '"
               + vectorFunction()
+              + "', 'source_model': '"
+              + sourceModel()
               + "'}";
       statements.add(
           SimpleStatement.newInstance(String.format(vectorSearch, table, keyspace, table)));

src/main/java/io/stargate/sgv2/jsonapi/service/operation/tables/CreateIndexAttempt.java

@@ -13,6 +13,7 @@
 import com.datastax.oss.driver.api.querybuilder.schema.CreateIndexOnTable;
 import com.datastax.oss.driver.api.querybuilder.schema.CreateIndexStart;
 import com.datastax.oss.driver.internal.querybuilder.schema.DefaultCreateIndex;
+import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
 import io.stargate.sgv2.jsonapi.service.cqldriver.executor.TableSchemaObject;
 import io.stargate.sgv2.jsonapi.service.cqldriver.override.ExtendedCreateIndex;
 import io.stargate.sgv2.jsonapi.service.operation.SchemaAttempt;
@@ -90,10 +91,11 @@ public record VectorIndexOptions(SimilarityFunction similarityFunction, String s
     public Map<String, Object> getOptions() {
       Map<String, Object> options = new HashMap<>();
       if (similarityFunction != null) {
-        options.put("similarity_function", similarityFunction.getMetric());
+        options.put(
+            DocumentConstants.Fields.VECTOR_INDEX_FUNCTION_NAME, similarityFunction.getMetric());
       }
       if (sourceModel != null) {
-        options.put("source_model", sourceModel);
+        options.put(DocumentConstants.Fields.VECTOR_INDEX_SOURCE_MODEL_NAME, sourceModel);
       }
       return options;
     }

src/main/java/io/stargate/sgv2/jsonapi/service/resolver/CreateCollectionCommandResolver.java

@@ -106,6 +106,7 @@ public Operation resolveKeyspaceCommand(
           command.name(),
           vector.dimension(),
           vector.metric(),
+          vector.sourceModel(),
           comment,
           operationsConfig.databaseConfig().ddlDelayMillis(),
           operationsConfig.tooManyIndexesRollbackEnabled(),
@@ -198,7 +199,7 @@ private CreateCollectionCommand.Options.VectorSearchConfig validateVectorOptions
       vectorDimension = validateVectorize.validateService(service, vectorDimension);
       vector =
           new CreateCollectionCommand.Options.VectorSearchConfig(
-              vectorDimension, vector.metric(), vector.vectorizeConfig());
+              vectorDimension, vector.metric(), vector.sourceModel(), vector.vectorizeConfig());
     } else {
       // Ensure vector dimension is provided when service configuration is absent.
       if (vectorDimension == null) {

src/main/java/io/stargate/sgv2/jsonapi/service/schema/SourceModel.java

@@ -0,0 +1,45 @@
+package io.stargate.sgv2.jsonapi.service.schema;
+
+import io.stargate.sgv2.jsonapi.exception.ErrorCodeV1;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * The source model used for the vector index. This is only applicable if the vector index is
+ * enabled.
+ */
+public enum SourceModel {
+  ADA002("ada002"),
+  BERT("bert"),
+  COHERE_V3("cohere-v3"),
+  GECKO("gecko"),
+  NV_QA_4("nv-qa-4"),
+  OPENAI_V3_LARGE("openai-v3-large"),
+  OPENAI_V3_SMALL("openai-v3-small"),
+  OTHER("other"),
+  UNDEFINED("undefined");
+
+  private final String sourceModel;
+  private static final java.util.Map<String, SourceModel> FUNCTIONS_MAP =
+      Stream.of(SourceModel.values())
+          .collect(Collectors.toMap(SourceModel::getSourceModel, sourceModel -> sourceModel));
+
+  SourceModel(String sourceModel) {
+    this.sourceModel = sourceModel;
+  }
+
+  public String getSourceModel() {
+    return sourceModel;
+  }
+
+  public static SourceModel fromString(String sourceModel) {
+    if (sourceModel == null) return UNDEFINED;
+    // The string may be empty if the collection was created before supporting source models
+    if (sourceModel.isEmpty()) return OTHER;
+    SourceModel model = FUNCTIONS_MAP.get(sourceModel);
+    if (model == null) {
+      throw ErrorCodeV1.VECTOR_SEARCH_INVALID_SOURCE_MODEL_NAME.toApiException("'%s'", sourceModel);
+    }
+    return model;
+  }
+}