Integrate NVIDIA as an embedding provider

import os
from astrapy import DataAPIClient
from astrapy.constants import VectorMetric
from astrapy.info import (
    CreateTableDefinition,
    ColumnType,
    VectorServiceOptions,
    TableVectorColumnTypeDescriptor,
    TableScalarColumnTypeDescriptor,
    TablePrimaryKeyDescriptor,
    TableVectorIndexOptions
)

# Initialize the client and get a "Database" object
client = DataAPIClient(os.environ["APPLICATION_TOKEN"])
database = client.get_database(os.environ["API_ENDPOINT"])
print(f"* Database: {database.info().name}\n")

# Define the columns and primary key for the table
table_definition = CreateTableDefinition(
    columns={
        # This column will store vector embeddings.
        # The NVIDIA integration
        # will automatically generate vector embeddings
        # for any text inserted to this column.
        "example_vector": TableVectorColumnTypeDescriptor(
            dimension=1024, # optional
            service=VectorServiceOptions(
                provider="nvidia",
                model_name="NV-Embed-QA",
            ),
        ),
        # If you want to store the original text
        # in addition to the generated embeddings
        # you must create a separate column.
        "original_text": TableScalarColumnTypeDescriptor(
            column_type=ColumnType.TEXT
        ),
    },
    # Define the primary key for the table.
    # You should change the primary key definition to meet the needs of your data.
    primary_key=TablePrimaryKeyDescriptor(
        partition_by=["original_text"],
        partition_sort={}
    ),
)

table = database.create_table(
    "TABLE_NAME",
    definition=table_definition,
)

print("Created table")

# Index the vector column so that you can perform a vector search on it.
table.create_vector_index(
    "example_vector_index",
    column="example_vector",
    options=TableVectorIndexOptions(
        metric="cosine",
    ),
)

print("Indexed columns")

import { DataAPIClient, Table, InferTableSchema } from '@datastax/astra-db-ts';

const { APPLICATION_TOKEN, API_ENDPOINT } = process.env;

// Initialize the client and get a 'Db' object
const client = new DataAPIClient(APPLICATION_TOKEN);
const db = client.db(API_ENDPOINT);

console.log(`* Connected to DB ${db.id}`);

const tableDefinition = Table.schema({
  // Define all of the columns in the table
  columns: {
    // This column will store vector embeddings.
    // The NVIDIA integration
    // will automatically generate vector embeddings
    // for any text inserted to this column.
    example_vector: {
      type: "vector",
      dimension: 1024, // optional
      service: {
        provider: 'nvidia',
        modelName: 'NV-Embed-QA',
      },
    },
    // If you want to store the original text
    // in addition to the generated embeddings
    // you must create a separate column.
    original_text: "text",
  },
  // Define the primary key for the table.
  // You should change the primary key definition to meet the needs of your data.
  primaryKey: {
    partitionBy: ["original_text"],
  },
});

type TableSchema = InferTableSchema<typeof tableDefinition>;

(async function () {
  const table = await db.createTable<TableSchema>(
    'TABLE_NAME',
    { definition: tableDefinition },
  );

  console.log("Created table");

  // Index the vector column so that you can perform a vector search on it.
  await table.createVectorIndex(
    "example_vector_index",
    "example_vector",
    {
      options: {
        metric: 'cosine',
      },
    },
  );

  console.log("Indexed columns");
})();

import com.datastax.astra.client.DataAPIClient;
import com.datastax.astra.client.core.vector.SimilarityMetric;
import com.datastax.astra.client.core.vectorize.VectorServiceOptions;
import com.datastax.astra.client.databases.Database;
import com.datastax.astra.client.tables.Table;
import com.datastax.astra.client.tables.definition.TableDefinition;
import com.datastax.astra.client.tables.definition.columns.ColumnDefinitionVector;
import com.datastax.astra.client.tables.definition.indexes.TableVectorIndexDefinition;
import com.datastax.astra.client.tables.definition.rows.Row;

import java.util.HashMap;
import java.util.Map;

public class EmbeddingDemo {

  public static void main(String[] args) {
    // Loading Arguments
    String astraToken = System.getenv("APPLICATION_TOKEN");
    String astraApiEndpoint = System.getenv("API_ENDPOINT");

    // Initialize the client
    DataAPIClient client = new DataAPIClient(astraToken);
    System.out.println("Connected to AstraDB");

    Database database = client.getDatabase(astraApiEndpoint);
    System.out.println("Connected to Database.");

    // Define the columns and primary key for the table
    TableDefinition tableDefinition =
        new TableDefinition()
            // This column will store vector embeddings.
            // The NVIDIA integration
            // will automatically generate vector embeddings
            // for any text inserted to this column.
            .addColumnVector(
                "example_vector",
                new ColumnDefinitionVector()
                    .metric(SimilarityMetric.COSINE)
                    .service(
                        new VectorServiceOptions()
                            .provider("nvidia")
                            .modelName("NV-Embed-QA")
                    )
            )
            // If you want to store the original text
            // in addition to the generated embeddings
            // you must create a separate column.
            .addColumnText("original_text")
            // Define the primary key for the table.
            // You should change the primary key definition to meet the needs of your data.
            .addPartitionBy("original_text");

    Table<Row> table = database.createTable("TABLE_NAME", tableDefinition);

    System.out.println("Created table");

    // Index the vector column so that you can perform a vector search on it.
    TableVectorIndexDefinition definition =
        new TableVectorIndexDefinition().column("example_vector").metric(SimilarityMetric.COSINE);

    table.createVectorIndex("example_vector_index", definition);

    System.out.println("Indexed columns");
  }
}

curl -sS -L -X POST "$API_ENDPOINT/api/json/v1/default_keyspace" \
  --header "Token: $APPLICATION_TOKEN" \
  --header "Content-Type: application/json" \
  --data '{
  "createTable": {
    "name": "TABLE_NAME",
    "definition": {
      "columns": {
        # This column will store vector embeddings.
        # The NVIDIA integration
        # will automatically generate vector embeddings
        # for any text inserted to this column.
        "example_vector": {
          "type": "vector",
          "dimension": 1024, # optional
          "service": {
            "provider": "nvidia",
            "modelName": "NV-Embed-QA"
          }
        },
        # If you want to store the original text
        # in addition to the generated embeddings
        # you must create a separate column.
        "original_text": "text"
      },
      # Define the primary key for the table.
      # You should change the primary key definition to meet the needs of your data.
      "primaryKey": "original_text"
    }
  }
}'

curl -sS -L -X POST "$API_ENDPOINT/api/json/v1/default_keyspace" \
  --header "Token: $APPLICATION_TOKEN" \
  --header "Content-Type: application/json" \
  --data '{
  "createVectorIndex": {
    "name": "example_vector_index",
    "definition": {
      "column": "example_vector",
      "options": {
        "metric": "cosine",
        "sourceModel": "NV-Embed-QA"
      }
    }
  }
}'