Astra DB Serverless quickstart for collections (Python)

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")
    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

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


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

    collection = database.create_collection(
        "quickstart_collection",
        definition=CollectionDefinition(
            vector=CollectionVectorOptions(
                metric=VectorMetric.COSINE,
                service=VectorServiceOptions(
                    provider="nvidia",
                    model_name="nvidia/nv-embedqa-e5-v5",
                ),
            )
        ),
    )

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


if __name__ == "__main__":
    main()

import json

from astrapy.data_types import DataAPIDate
from quickstart_connect import connect_to_database


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

    collection = database.get_collection("quickstart_collection")

    data_file_path = "PATH_TO_DATA_FILE"

    # 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
    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()

from quickstart_connect import connect_to_database


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

    collection = database.get_collection("quickstart_collection")

    # 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()