Integrate Haystack with Astra DB Serverless

query_builder 15 min

Haystack can use Astra DB Serverless to store and retrieve vectors for ML applications.

Prerequisites

The code samples on this page assume the following:

You have an active Astra account.
You have created a Serverless (Vector) database.
You have created an application token with the Database Administrator role.
You have created an OpenAI API key.
You have installed Python 3.8+ and pip 23.0+.

You have installed the required dependencies:

pip install astra-haystack sentence-transformers python-dotenv

Create a document store

Create a .env file in the root of your program.
Populate the file with the Astra token and endpoint values from the Database Details section of your database’s Overview tab, and your OpenAI API key. If you’re running this code in Google Colab, you’ll be prompted for the .env values instead.
.env
```
ASTRA_DB_API_ENDPOINT=API_ENDPOINT
ASTRA_DB_APPLICATION_TOKEN=TOKEN
OPENAI_API_KEY=API_KEY
```
The endpoint format is https://ASTRA_DB_ID-ASTRA_DB_REGION.apps.astra.datastax.com/api/json/v1/ASTRA_DB_KEYSPACE_NAME/ASTRA_DB_COLLECTION_NAME.

Import your dependencies.

haystack-rag.py

from haystack import Document, Pipeline
from haystack.components.builders.answer_builder import AnswerBuilder
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.generators import OpenAIGenerator
from haystack.components.writers import DocumentWriter
from haystack.document_stores.types import DuplicatePolicy
from haystack_integrations.document_stores.astra import AstraDocumentStore
from haystack_integrations.components.retrievers.astra import AstraEmbeddingRetriever
from dotenv import load_dotenv

load_dotenv()

Create a AstraDocumentStore object.
haystack-rag.py
```
document_store = AstraDocumentStore(
    collection_name="haystack",
    duplicates_policy=DuplicatePolicy.SKIP,
    embedding_dimension=384,
)

# ...
```
The collection_name parameter is the name of the collection in your Astra DB Serverless database. This example uses haystack as the collection name.

duplicates_policy is set to SKIP to avoid inserting duplicate documents.

The embedding_dim parameter is the dimension of the vector embedding. This example uses the all-MiniLM-L6-v2 sentence transformer, which has an embedding dimension of 384. This number should match the embedding_dim you declare in the AstraDocumentStore object.

For more on sentence transformers, see pretrained models.

Create an indexing pipeline

One of Haystack’s core features is the ability to construct reusable pipelines of components.

In this example, you’ll create a pipeline that indexes documents in the AstraDocumentStore and prints the number of documents embedded in the store.

Create a list containing three Document objects. You’ll index these documents into the AstraDocumentStore with the pipeline. Print the number of documents in the store to confirm the loading was successful.

haystack-rag.py

# ...

documents = [
    Document(content="There are over 7,000 languages spoken around the world today."),
    Document(
        content="Elephants have been observed to behave in a way that indicates"
            " a high level of self-awareness, such as recognizing themselves in mirrors."
    ),
    Document(
        content="In certain parts of the world, like the Maldives, Puerto Rico, "
            "and San Diego, you can witness the phenomenon of bioluminescent waves."
    ),
]

print(document_store.count_documents())

# ...

Create a Haystack indexing pipeline.

The first component in the pipeline is a SentenceTransformersDocumentEmbedder, which transforms the documents into vectors.
The second component is a DocumentWriter, which writes the list of documents to the AstraDocumentStore.

The index_pipeline.connect() method connects the output of the first component to the input of the second component.

haystack-rag.py

# ...

embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2"

index_pipeline = Pipeline()
index_pipeline.add_component(instance=SentenceTransformersDocumentEmbedder(model=embedding_model_name), name="embedder")
index_pipeline.add_component(instance=DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP), name="writer")
index_pipeline.connect("embedder.documents", "writer.documents")

# ...

When you run the complete code, this step runs the pipeline, embeds your documents, and prints the document count to confirm your AstraDocumentStore is populated.
haystack-rag.py
```
# ...

index_pipeline.run({"embedder": {"documents": documents}})
print(document_store.count_documents())

# ...
```

Create a RAG pipeline

Use the populated AstraDocumentStore in a Haystack RAG pipeline with the AstraEmbeddingRetriever as the document store retriever.

Define a template for the OpenAI prompt.

haystack-rag.py

# ...

prompt_template = """
                Given these documents, answer the question.
                Documents:
                {% for doc in documents %}
                    {{ doc.content }}
                {% endfor %}
                Question: {{question}}
                Answer:
                """

# ...

Create the Haystack RAG pipeline.

This might look like a lot of code, but it works just like the indexing pipeline you created before. Define each component, and wire them into a series.

haystack-rag.py

# ...

rag_pipeline = Pipeline()

# SentenceTransformersTextEmbedder transforms the question into a vector
rag_pipeline.add_component(instance=SentenceTransformersTextEmbedder(model=embedding_model_name), name="embedder")

# AstraEmbeddingRetriever retrieves the most similar documents from the AstraDocumentStore
rag_pipeline.add_component(instance=AstraEmbeddingRetriever(document_store=document_store), name="retriever")

# PromptBuilder creates a prompt for the OpenAI API from the defined prompt_template
rag_pipeline.add_component(instance=PromptBuilder(template=prompt_template), name="prompt_builder")

# OpenAIGenerator generates an answer from the prompt
rag_pipeline.add_component(instance=OpenAIGenerator(), name="llm")

# AnswerBuilder extracts the answer from the OpenAI response and metadata (`meta`)
rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder")

# The rag_pipeline.connect() method connects each component to the next
rag_pipeline.connect("embedder", "retriever")
rag_pipeline.connect("retriever", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder", "llm")
rag_pipeline.connect("llm.replies", "answer_builder.replies")
rag_pipeline.connect("llm.meta", "answer_builder.meta")
rag_pipeline.connect("retriever", "answer_builder.documents")

When the pipeline runs, your question passes through the pipeline and returns an answer, thanks to the AstraEmbeddingRetriever.

question = "How many languages are there in the world today?"
result = rag_pipeline.run(
    {
        "embedder": {"text": question},
        "retriever": {"top_k": 2},
        "prompt_builder": {"question": question},
        "answer_builder": {"query": question},
    }
)

print(result)

If you get Disabling parallelism to avoid deadlocks… errors when running, set TOKENIZERS_PARALLELISM=false in your environment variables.

Complete code example

haystack-rag.py

from haystack import Document, Pipeline
from haystack.components.builders.answer_builder import AnswerBuilder
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.generators import OpenAIGenerator
from haystack.components.writers import DocumentWriter
from haystack.document_stores.types import DuplicatePolicy
from haystack_integrations.document_stores.astra import AstraDocumentStore
from haystack_integrations.components.retrievers.astra import AstraEmbeddingRetriever
from dotenv import load_dotenv

load_dotenv()

# Create document store
document_store = AstraDocumentStore(
    collection_name="haystack",
    duplicates_policy=DuplicatePolicy.SKIP,
    embedding_dimension=384,
)

# Add documents
documents = [
    Document(content="There are over 7,000 languages spoken around the world today."),
    Document(
        content="Elephants have been observed to behave in a way that indicates"
        " a high level of self-awareness, such as recognizing themselves in mirrors."
    ),
    Document(
        content="In certain parts of the world, like the Maldives, Puerto Rico, "
        "and San Diego, you can witness the phenomenon of bioluminescent waves."
    ),
]
print(document_store.count_documents())

embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2"

# Indexing pipeline
index_pipeline = Pipeline()
index_pipeline.add_component(
    instance=SentenceTransformersDocumentEmbedder(model=embedding_model_name),
    name="embedder",
)
index_pipeline.add_component(instance=DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP), name="writer")
index_pipeline.connect("embedder.documents", "writer.documents")

# Run the indexing pipeline and print document count
index_pipeline.run({"embedder": {"documents": documents}})
print(document_store.count_documents())

# Build rag pipeline
prompt_template = """
                Given these documents, answer the question.
                Documents:
                {% for doc in documents %}
                    {{ doc.content }}
                {% endfor %}
                Question: {{question}}
                Answer:
                """

rag_pipeline = Pipeline()
rag_pipeline.add_component(instance=SentenceTransformersTextEmbedder(model=embedding_model_name), name="embedder")
rag_pipeline.add_component(instance=AstraEmbeddingRetriever(document_store=document_store), name="retriever")
rag_pipeline.add_component(instance=PromptBuilder(template=prompt_template), name="prompt_builder")
rag_pipeline.add_component(instance=OpenAIGenerator(), name="llm")
rag_pipeline.add_component(instance=AnswerBuilder(), name="answer_builder")
rag_pipeline.connect("embedder", "retriever")
rag_pipeline.connect("retriever", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder", "llm")
rag_pipeline.connect("llm.replies", "answer_builder.replies")
rag_pipeline.connect("llm.meta", "answer_builder.meta")
rag_pipeline.connect("retriever", "answer_builder.documents")

# Run the pipeline
question = "How many languages are there in the world today?"
result = rag_pipeline.run(
    {
        "embedder": {"text": question},
        "retriever": {"top_k": 2},
        "prompt_builder": {"question": question},
        "answer_builder": {"query": question},
    }
)

print(result)

Integrate Haystack with Astra DB Serverless

Prerequisites

Create a document store

Create an indexing pipeline

Create a RAG pipeline

Complete code example

Was this helpful?

Give Feedback