Quickstart

This quickstart demonstrates a basic RAG pattern using RAGStack and the vector-enabled Astra DB Serverless database to retrieve context and pass it to a language model for generation.

The Complete code example is available at the bottom of the page.

Setup

RAGStack includes all the libraries you need for the RAG pattern, including the vector database, embeddings pipeline, and retrieval.

This quickstart also uses the HuggingFace datasets library to load a small dataset of philosophical quotes.

  1. Install RAGStack and the datasets library:

    pip3 install ragstack-ai datasets
  2. To see current RAGStack version:

    • Command

    • Result

    pip3 show ragstack-ai
    Name: ragstack-ai
    Version: 0.1.1
    Summary: RAGStack
  3. To upgrade the RAGStack package:

    • Command

    • Result

    pip3 install ragstack-ai --upgrade
    Requirement already satisfied: ragstack-ai in ./ragstack-venv/lib/python3.11/site-packages (0.1.0)
    Collecting ragstack-ai
      Downloading ragstack_ai-0.1.1-py3-none-any.whl.metadata (2.4 kB)
    …installing packages…
    Successfully installed astrapy-0.6.1 backoff-2.2.1 chardet-5.2.0 emoji-2.8.0 filetype-1.2.0 h2-4.1.0 hpack-4.0.0 httpcore-1.0.2 httpx-0.25.1 hyperframe-6.0.1 joblib-1.3.2 langdetect-1.0.9 lxml-4.9.3 nltk-3.8.1 python-iso639-2023.6.15 python-magic-0.4.27 ragstack-ai-0.1.1 rapidfuzz-3.5.2 tabulate-0.9.0 unstructured-0.10.30
  4. If you don’t have a vector database, create one at https://astra.datastax.com/.

    The Astra DB Serverless application token is associated automatically with the Database Administrator permission. An auth token example: AstraCS:WSnyFUhRxsrg…​).

    The Astra API endpoint is available in the Astra Portal. Its format is https://<ASTRA_DB_ID>-<ASTRA_DB_REGION>.apps.astra.datastax.com.

    Create an OpenAI key at https://platform.openai.com/.

  5. Create a .env file in the root of your program with the values from your Astra Connect tab.

    ASTRA_DB_APPLICATION_TOKEN="<AstraCS:...>"
    ASTRA_DB_API_ENDPOINT="<Astra DB API endpoint>"
    OPENAI_API_KEY="sk-..."
    ASTRA_DB_COLLECTION="test"

RAG workflow

With your environment set up, you’re ready to create a RAG workflow in Python.

  1. Import the necessary dependencies:

    import os
    from dotenv import load_dotenv
    from datasets import load_dataset
    from langchain_astradb import AstraDBVectorStore
    from langchain_openai import OpenAIEmbeddings
    from langchain_core.documents import Document
    from langchain.prompts import ChatPromptTemplate
    from langchain_openai import ChatOpenAI
    from langchain_core.output_parsers import StrOutputParser
    from langchain_core.runnables import RunnablePassthrough
  2. Load the environment variables:

    load_dotenv()
    ASTRA_DB_APPLICATION_TOKEN = os.environ.get("ASTRA_DB_APPLICATION_TOKEN")
    ASTRA_DB_API_ENDPOINT = os.environ.get("ASTRA_DB_API_ENDPOINT")
    OPEN_AI_API_KEY = os.environ.get("OPENAI_API_KEY")
    ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")

Construct information base

  1. Declare the embeddings model, create your vector database, and configure their required parameters.

    embedding = OpenAIEmbeddings()
    vstore = AstraDBVectorStore(
        embedding=embedding,
        collection_name=ASTRA_DB_COLLECTION,
        token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
        api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
    )
  2. Load a small dataset of quotes with the Python dataset module.

    philo_dataset = load_dataset("datastax/philosopher-quotes")["train"]
    print("An example entry:")
    print(philo_dataset[16])
  3. Process metadata and convert to a Document object:

    docs = []
    for entry in philo_dataset:
        metadata = {"author": entry["author"]}
        if entry["tags"]:
            # Add metadata tags to the metadata dictionary
            for tag in entry["tags"].split(";"):
                metadata[tag] = "y"
        # Add a LangChain document with the quote and metadata tags
        doc = Document(page_content=entry["quote"], metadata=metadata)
        docs.append(doc)
  4. Compute embeddings:

    inserted_ids = vstore.add_documents(docs)
    print(f"\nInserted {len(inserted_ids)} documents.")

Basic retrieval

Confirm your vector store is populated by printing the vectors in your collection:

print(vstore.astra_db.collection(ASTRA_DB_COLLECTION).find())

Generation with augmented context

  1. Retrieve context from your vector database, pass it to OpenAI with a prompt question, and print the response.

    retriever = vstore.as_retriever(search_kwargs={'k': 3})
    
    prompt_template = """
    Answer the question based only on the supplied context. If you don't know the answer, say you don't know the answer.
    Context: {context}
    Question: {question}
    Your answer:
    """
    prompt = ChatPromptTemplate.from_template(prompt_template)
    model = ChatOpenAI(openai_api_key=OPEN_AI_API_KEY)
    
    chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | model
        | StrOutputParser()
    )
    
    response = chain.invoke("In the given context, what subject are philosophers most concerned with?")
    print(response)
  2. You should get a response like this:

    An example entry:
    {'author': 'aristotle', 'quote': 'Love well, be loved and do something of value.', 'tags': 'love;ethics'}
    
    Inserted 450 documents.
    The subject that philosophers are most concerned with in the given context is truth.

Complete code example

  • Python

  • Result

import os
from dotenv import load_dotenv
from datasets import load_dataset
from langchain_astradb import AstraDBVectorStore
from langchain_openai import OpenAIEmbeddings
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

load_dotenv()

ASTRA_DB_APPLICATION_TOKEN = os.environ.get("ASTRA_DB_APPLICATION_TOKEN")
ASTRA_DB_API_ENDPOINT = os.environ.get("ASTRA_DB_API_ENDPOINT")
OPEN_AI_API_KEY = os.environ.get("OPENAI_API_KEY")
ASTRA_DB_COLLECTION = os.environ.get("ASTRA_DB_COLLECTION")

embedding = OpenAIEmbeddings()
vstore = AstraDBVectorStore(
    embedding=embedding,
    collection_name="test",
    token=os.environ["ASTRA_DB_APPLICATION_TOKEN"],
    api_endpoint=os.environ["ASTRA_DB_API_ENDPOINT"],
)

philo_dataset = load_dataset("datastax/philosopher-quotes")["train"]
print("An example entry:")
print(philo_dataset[16])

docs = []
for entry in philo_dataset:
    metadata = {"author": entry["author"]}
    if entry["tags"]:
        # Add metadata tags to the metadata dictionary
        for tag in entry["tags"].split(";"):
            metadata[tag] = "y"
    # Add a LangChain document with the quote and metadata tags
    doc = Document(page_content=entry["quote"], metadata=metadata)
    docs.append(doc)

inserted_ids = vstore.add_documents(docs)
print(f"\nInserted {len(inserted_ids)} documents.")

print(vstore.astra_db.collection(ASTRA_DB_COLLECTION).find())

retriever = vstore.as_retriever(search_kwargs={'k': 3})

prompt_template = """
Answer the question based only on the supplied context. If you don't know the answer, say you don't know the answer.
Context: {context}
Question: {question}
Your answer:
"""
prompt = ChatPromptTemplate.from_template(prompt_template)
model = ChatOpenAI(openai_api_key=OPEN_AI_API_KEY)

chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | model
    | StrOutputParser()
)

response = chain.invoke("In the given context, what subject are philosophers most concerned with?")
print(response)
An example entry:
{'author': 'aristotle', 'quote': 'Love well, be loved and do something of value.', 'tags': 'love;ethics'}

Inserted 450 documents.
The subject that philosophers are most concerned with in the given context is truth.

Cleanup

To clear data from your vector database but keep the collection, use the vstore.clear() method.

To delete the collection from your vector database, use the vstore.delete_collection() method. Alternatively, you can use the CLI to delete the collection:

curl -v -s --location \
--request POST https://${ASTRA_DB_ID}-${ASTRA_DB_REGION}.apps.astra.datastax.com/api/json/v1/default_keyspace \
--header "X-Cassandra-Token: $ASTRA_DB_APPLICATION_TOKEN" \
--header "Content-Type: application/json" \
--header "Accept: application/json" \
--data '{
  "deleteCollection": {
    "name": "test"
  }
}'

What’s next?

Was this helpful?

Give Feedback

How can we improve the documentation?

© 2024 DataStax | Privacy policy | Terms of use

Apache, Apache Cassandra, Cassandra, Apache Tomcat, Tomcat, Apache Lucene, Apache Solr, Apache Hadoop, Hadoop, Apache Pulsar, Pulsar, Apache Spark, Spark, Apache TinkerPop, TinkerPop, Apache Kafka and Kafka are either registered trademarks or trademarks of the Apache Software Foundation or its subsidiaries in Canada, the United States and/or other countries. Kubernetes is the registered trademark of the Linux Foundation.

General Inquiries: +1 (650) 389-6000, info@datastax.com