Astra DB Serverless quickstart for collections

import os
from astrapy import DataAPIClient, Database


def connect_to_database() -> Database:
    """
    Connects to a DataStax Astra database.
    This function retrieves the database endpoint and application token from the
    environment variables `API_ENDPOINT` and `APPLICATION_TOKEN`.

    Returns:
        Database: An instance of the connected database.

    Raises:
        RuntimeError: If the environment variables `API_ENDPOINT` or
        `APPLICATION_TOKEN` are not defined.
    """
    endpoint = os.environ.get("API_ENDPOINT")  (1)
    token = os.environ.get("APPLICATION_TOKEN")

    if not token or not endpoint:
        raise RuntimeError(
            "Environment variables API_ENDPOINT and APPLICATION_TOKEN must be defined"
        )

    # Create an instance of the `DataAPIClient` class
    client = DataAPIClient()

    # Get the database specified by your endpoint and provide the token
    database = client.get_database(endpoint, token=token)

    print(f"Connected to database {database.info().name}")

    return database

import { DataAPIClient, Db } from "@datastax/astra-db-ts";

/**
 * Connects to a DataStax Astra database.
 * This function retrieves the database endpoint and application token from the
 * environment variables `API_ENDPOINT` and `APPLICATION_TOKEN`.
 *
 * @returns An instance of the connected database.
 * @throws Will throw an error if the environment variables
 * `API_ENDPOINT` or `APPLICATION_TOKEN` are not defined.
 */
export function connectToDatabase(): Db {
  const { API_ENDPOINT: endpoint, APPLICATION_TOKEN: token } =
    process.env; (1)

  if (!token || !endpoint) {
    throw new Error(
      "Environment variables API_ENDPOINT and APPLICATION_TOKEN must be defined.",
    );
  }

  // Create an instance of the `DataAPIClient` class
  const client = new DataAPIClient();

  // Get the database specified by your endpoint and provide the token
  const database = client.db(endpoint, { token });

  console.log(`Connected to database ${database.id}`);

  return database;
}

package com.quickstart;

import com.datastax.astra.client.DataAPIClient;
import com.datastax.astra.client.databases.Database;

public class QuickstartConnect {

  /**
   * Connects to a DataStax Astra database. This function retrieves the database endpoint and
   * application token from the environment variables `API_ENDPOINT` and
   * `APPLICATION_TOKEN`.
   *
   * @return an instance of the connected database
   * @throws IllegalStateException if the environment variables `API_ENDPOINT` or
   *     `APPLICATION_TOKEN` are not defined
   */
  public static Database connectToDatabase() {
    String endpoint = System.getenv("API_ENDPOINT"); (1)
    String token = System.getenv("APPLICATION_TOKEN");

    if (endpoint == null || token == null) {
      throw new IllegalStateException(
          "Environment variables API_ENDPOINT and APPLICATION_TOKEN must be defined");
    }

    // Create an instance of `DataAPIClient` with your token.
    DataAPIClient client = new DataAPIClient(token);

    // Get the database specified by your endpoint.
    Database database = client.getDatabase(endpoint);

    System.out.println("Connected to database.");

    return database;
  }
}

from quickstart_connect import connect_to_database  (1)
from astrapy.constants import VectorMetric
from astrapy.info import (
    CollectionDefinition,
    CollectionVectorOptions,
    VectorServiceOptions,
)


def main() -> None:
    database = connect_to_database()

    collection = database.create_collection(
        "quickstart_collection",  (2)
        definition=CollectionDefinition(
            vector=CollectionVectorOptions(
                metric=VectorMetric.COSINE,
                service=VectorServiceOptions(
                    provider="nvidia",  (3)
                    model_name="NV-Embed-QA",
                ),
            )
        ),
    )

    print(f"Created collection {collection.full_name}")


if __name__ == "__main__":
    main()

import { connectToDatabase } from "./quickstart-connect"; (1)

(async function () {
  const database = connectToDatabase();

  const collection = await database.createCollection(
    "quickstart_collection", (2)
    {
      vector: {
        service: {
          provider: "nvidia", (3)
          modelName: "NV-Embed-QA",
        },
      },
    },
  );

  console.log(
    `Created collection ${collection.keyspace}.${collection.name}`,
  );
})();

package com.quickstart;

import com.datastax.astra.client.collections.Collection;
import com.datastax.astra.client.collections.definition.CollectionDefinition;
import com.datastax.astra.client.collections.definition.documents.Document;
import com.datastax.astra.client.core.vector.SimilarityMetric;
import com.datastax.astra.client.databases.Database;

public class QuickstartCreateCollectionDemo {

  public static void main(String[] args) {
    Database database = QuickstartConnect.connectToDatabase(); (1)

    Collection<Document> collection =
        database.createCollection(
            "quickstart_collection", (2)
            new CollectionDefinition()
                .vectorSimilarity(SimilarityMetric.COSINE)
                .vectorize("nvidia", "NV-Embed-QA")); (3)

    System.out.println("Created collection " + collection.getCollectionName());
  }
}

import json
from quickstart_connect import connect_to_database  (1)
from astrapy.data_types import DataAPIDate


def main() -> None:
    database = connect_to_database()

    collection = database.get_collection("quickstart_collection")  (2)

    data_file_path = "PATH_TO_DATA_FILE"  (3)

    # Read the JSON file and parse it into a JSON array
    with open(data_file_path, "r", encoding="utf8") as file:
        json_data = json.load(file)

    # Assemble the documents to insert:
    # - Convert the date string into a DataAPIDate
    # - Add a $vectorize field (4)
    documents = [
        {
            **data,
            "due_date": (
                DataAPIDate.from_string(data["due_date"])
                if data.get("due_date")
                else None
            ),
            "$vectorize": (
                f"summary: {data['summary']} | genres: {', '.join(data['genres'])}"
            ),
        }
        for data in json_data
    ]

    # Insert the data
    inserted = collection.insert_many(documents)

    print(f"Inserted {len(inserted.inserted_ids)} documents.")


if __name__ == "__main__":
    main()

import { connectToDatabase } from "./quickstart-connect"; (1)
import fs from "fs";

(async function () {
  const database = connectToDatabase();

  const collection = database.collection("quickstart_collection"); (2)

  const dataFilePath = "PATH_TO_DATA_FILE"; (3)

  // Read the JSON file and parse it into a JSON array
  const rawData = fs.readFileSync(dataFilePath, "utf8");
  const jsonData = JSON.parse(rawData);

  // Assemble the documents to insert:
  // - Convert the date string into a Date
  // - Add a $vectorize field (4)
  const documents = jsonData.map((data: any) => ({
    ...data,
    due_date: data.due_date ? new Date(data.due_date) : null,
    $vectorize: `summary: ${data["summary"]} | genres: ${data["genres"].join(", ")}`,
  }));

  // Insert the data
  const inserted = await collection.insertMany(documents);

  console.log(`Inserted ${inserted.insertedCount} documents.`);
})();

package com.quickstart;

import com.datastax.astra.client.collections.Collection;
import com.datastax.astra.client.collections.commands.results.CollectionInsertManyResult;
import com.datastax.astra.client.collections.definition.documents.Document;
import com.datastax.astra.client.databases.Database;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.StreamSupport;

public class QuickstartInsertToCollectionDemo {

  public static void main(String[] args) throws IOException {
    Database database = QuickstartConnect.connectToDatabase(); (1)

    Collection<Document> collection = database.getCollection("quickstart_collection"); (2)

    // Initialize Jackson ObjectMapper
    ObjectMapper objectMapper = new ObjectMapper();

    // Read the JSON file and parse it into a JSON array (ArrayNode). (3)
    String rawData =
        Files.readString(Paths.get("PATH_TO_DATA_FILE"), StandardCharsets.UTF_8);

    ArrayNode jsonData = (ArrayNode) objectMapper.readTree(rawData);

    // Convert the data to a list of Documents, and
    // add a $vectorize field to each piece of data. (4)
    List<Document> documents = new ArrayList<>();

    for (JsonNode data : jsonData) {
      if (data instanceof ObjectNode obj) {
        String summary = obj.path("summary").asText("");
        String genres =
            String.join(
                ", ",
                StreamSupport.stream(obj.withArray("genres").spliterator(), false)
                    .map(JsonNode::asText)
                    .toList());
        String summary_genres = String.format("summary: %s | genres: %s", summary, genres);
        obj.put("$vectorize", summary_genres);
      }

      documents.add(Document.parse(data.toString()));
    }

    CollectionInsertManyResult result = collection.insertMany(documents);

    System.out.println("Inserted " + result.getInsertedIds().size() + " documents.");
  }
}

from quickstart_connect import connect_to_database  (1)

def main() -> None:
    database = connect_to_database()

    collection = database.get_collection("quickstart_collection")  (2)

    # Find documents that match a filter
    print("\nFinding books with rating greater than 4.7...")

    rating_cursor = collection.find({"rating": {"$gt": 4.7}})

    for document in rating_cursor:
        print(f"{document['title']} is rated {document['rating']}")

    # Perform a vector search to find the closest match to a search string
    print("\nUsing vector search to find a single scary novel...")

    single_vector_match = collection.find_one(sort={"$vectorize": "A scary novel"})

    print(f"{single_vector_match['title']} is a scary novel")

    # Combine a filter, vector search, and projection to find the 3 books with
    # more than 400 pages that are the closest matches to a search string,
    # and just return the title and author
    print("\nUsing filters and vector search to find 3 books with more than 400 pages that are set in the arctic, returning just the title and author...")

    vector_cursor = collection.find(
        {"number_of_pages": {"$gt": 400}},
        sort={"$vectorize": "A book set in the arctic"},
        limit=3,
        projection={"title": True, "author": True}
    )

    for document in vector_cursor:
        print(document)


if __name__ == "__main__":
    main()

import { connectToDatabase } from "./quickstart-connect"; (1)

(async function () {
  const database = connectToDatabase();

  const collection = database.collection("quickstart_collection"); (2)

  // Find documents that match a filter
  console.log("\nFinding books with rating greater than 4.7...");

  const ratingCursor = collection.find(
    { rating: { $gt: 4.7 } },
    { limit: 10 },
  );

  for await (const document of ratingCursor) {
    console.log(`${document.title} is rated ${document.rating}`);
  }

  // Perform a vector search to find the closest match to a search string
  console.log("\nUsing vector search to find a single scary novel...");

  const singleVectorMatch = await collection.findOne(
    {},
    { sort: { $vectorize: "A scary novel" } },
  );

  console.log(`${singleVectorMatch?.title} is a scary novel`);

  // Combine a filter, vector search, and projection to find the 3 books with
  // more than 400 pages that are the closest matches to a search string,
  // and just return the title and author
  console.log("\nUsing filters and vector search to find 3 books with more than 400 pages that are set in the arctic, returning just the title and author...");

  const vectorCursor = collection.find(
    { number_of_pages: { $gt: 400 } },
    {
      sort: { $vectorize: "A book set in the arctic" },
      limit: 3,
      projection: { title: true, author: true },
    },
  );

  for await (const document of vectorCursor) {
    console.log(document);
  }
})();

package com.quickstart;

import static com.datastax.astra.client.core.query.Projection.include;

import com.datastax.astra.client.collections.Collection;
import com.datastax.astra.client.databases.Database;
import com.datastax.astra.client.collections.definition.documents.Document;
import com.datastax.astra.client.core.query.Filter;
import com.datastax.astra.client.core.query.Filters;
import com.datastax.astra.client.core.query.Sort;
import com.datastax.astra.client.collections.commands.options.CollectionFindOneOptions;
import com.datastax.astra.client.collections.commands.options.CollectionFindOptions;

public class QuickstartFindDemo {

  public static void main(String[] args) {

    Database database = QuickstartConnect.connectToDatabase(); (1)

    Collection<Document> collection = database.getCollection(
      "quickstart_collection" (2)
    );

    // Find documents that match a filter
    System.out.println("\nFinding books with rating greater than 4.7...");

    Filter filter = Filters.gt("rating", 4.7);

    CollectionFindOptions options = new CollectionFindOptions().limit(10);

    collection
        .find(filter, options)
        .forEach(
            document -> {
              System.out.println(
                  document.getString("title") + " is rated " + document.get("rating"));
            });

    // Perform a vector search to find the closest match to a search string
    System.out.println("\nUsing vector search to find a single scary novel...");

    CollectionFindOneOptions options2 = new CollectionFindOneOptions().sort(Sort.vectorize("A scary novel"));

    collection
        .findOne(options2)
        .ifPresent(
            document -> {
              System.out.println(document.getString("title") + " is a scary novel");
            });

    // Combine a filter, vector search, and projection to find the 3 books with
    // more than 400 pages that are the closest matches to a search string,
    // and just return the title and author
    System.out.println("\nUsing filters and vector search to find 3 books with more than 400 pages that are set in the arctic, returning just the title and author...");

    Filter filter3 = Filters.gt("number_of_pages", 400);

    CollectionFindOptions options3 =
        new CollectionFindOptions()
            .limit(3)
            .sort(Sort.vectorize("A book set in the arctic"))
            .projection(include("title", "author"));

    collection
        .find(filter3, options3)
        .forEach(
            document -> {
              System.out.println(document);
            });
  }
}