Astra DB Vector Store

This page provides a quickstart for using Astra DB as a Vector Store.

DataStax Astra DB is a serverless vector-capable database built on Apache Cassandra® and made conveniently available through an easy-to-use JSON API.

Setup

Use of the integration requires the corresponding Python package:

pip install -qU "langchain-astradb>=0.3.3"

[notice] A new release of pip is available: 24.0 -> 24.1.2
[notice] To update, run: pip install --upgrade pip
Note: you may need to restart the kernel to use updated packages.

Credentials

Head to astra.datastax.com, create an account, and create a new database. Once the database has been initialized, click on the Generate Token button under Application Tokens. Then set the ASTRA_DB_API_ENDPOINT and the ASTRA_DB_APPLICATION_TOKEN in the cell below.

• the API Endpoint looks like https://01234567-89ab-cdef-0123-456789abcdef-us-east1.apps.astra.datastax.com
• the Token looks like AstraCS:6gBhNmsk135....
• you may optionally provide a Namespace such as my_namespace

import getpass

ASTRA_DB_API_ENDPOINT = getpass.getpass("ASTRA_DB_API_ENDPOINT = ")
ASTRA_DB_APPLICATION_TOKEN = getpass.getpass("ASTRA_DB_APPLICATION_TOKEN = ")

desired_namespace = getpass.getpass("ASTRA_DB_KEYSPACE = ")
if desired_namespace:
    ASTRA_DB_KEYSPACE = desired_namespace
else:
    ASTRA_DB_KEYSPACE = None

Instantiation

There are two ways to create an Astra DB vector store, which differ in how the embeddings are computed.

Explicit embeddings. You can separately instantiate a langchain_core.embeddings.Embeddings class and pass it to the AstraDBVectorStore constructor, just like with most other LangChain vector stores.

Integrated embedding computation. Alternatively, you can use the Vectorize feature of Astra DB and simply specify the name of a supported embedding model when creating the store. The embedding computations are entirely handled within the database. (To proceed with this method, you must have enabled the desired embedding integration for your database, as described in the docs.)

Selecting an embedding function

To run the rest of the code in this notebook, you will need to define an embedding function. If you just want to run the code as fast as possible, use the FakeEmbeddings class which requires no setup. If you would like to use an actual embedding function, we are going to show how to use the langchain_ollama embedding function. You can read more about how to set it up at this page.

Fake Embedding

from langchain_core.embeddings import FakeEmbeddings

embedding_function = FakeEmbeddings(size=4096)

Ollama Embedding

from langchain_ollama import OllamaEmbeddings

embedding_function = OllamaEmbeddings(model="llama3")

Explicit Instantiation

Once you have selected an embedding function, we are ready to instantiate our vector store.

from langchain_astradb import AstraDBVectorStore

vector_store = AstraDBVectorStore(
    embedding=embedding_function,
    collection_name="astra_vector_demo",
    api_endpoint=ASTRA_DB_API_ENDPOINT,
    token=ASTRA_DB_APPLICATION_TOKEN,
    namespace=ASTRA_DB_KEYSPACE,
)

API Reference:AstraDBVectorStore

Integrated Instantiation

Here it is assumed that you have

• enabled the OpenAI integration in your Astra DB organization,
• added an API Key named "OPENAI_API_KEY" to the integration, and scoped it to the database you are using.

For more details please consult the documentation.

from astrapy.info import CollectionVectorServiceOptions

openai_vectorize_options = CollectionVectorServiceOptions(
    provider="openai",
    model_name="text-embedding-3-small",
    authentication={
        "providerKey": "OPENAI_API_KEY",
    },
)

vector_store = AstraDBVectorStore(
    collection_name="langchain_example_collection",
    api_endpoint=ASTRA_DB_API_ENDPOINT,
    token=ASTRA_DB_APPLICATION_TOKEN,
    namespace=ASTRA_DB_KEYSPACE,
    collection_vector_service_options=openai_vectorize_options,
)

/Users/isaachershenson/.pyenv/versions/3.11.9/lib/python3.11/site-packages/langchain_astradb/utils/astradb.py:248: UserWarning: Astra DB collection 'langchain_example_collection' is detected as having indexing turned on for all fields (either created manually or by older versions of this plugin). This implies stricter limitations on the amount of text each string in a document can store. Consider reindexing anew on a fresh collection to be able to store longer texts.
  if not self._validate_indexing_policy(

Manage vector store

Add items to vector store

from langchain_core.documents import Document

document_1 = Document(
    page_content="foo",
    metadata={"source": "https://example.com"}
)

document_2 = Document(
    page_content="bar",
    metadata={"source": "https://another-example.com"}
)

document_3 = Document(
    page_content="baz",
    metadata={"source": "https://example.com"}
)

documents = [document_1, document_2, document_3]

vector_store.add_documents(documents=documents,ids=["1","2","3"])

API Reference:Document

['1', '2', '3']

Delete items from vector store

vector_store.delete(ids=["3"])

True

Query vector store

Once your vector store has been created and the relevant documents have been added you will most likely wish to query it during the running of your chain or agent.

Query directly

Performing a simple similarity search can be done as follows:

results = vector_store.similarity_search("thud", k=3, filter={"source": "https://example.com"})
for res in results:
    print(f"* {res.page_content} [{res.metadata}]")

* foo [{'source': 'https://example.com'}]

You can also search with score:

results = vector_store.similarity_search_with_score("thud", k=1, filter={"source": "https://another-example.com"})
for res, score in results:
    print(f"* [SIM={score:3f}] {res.page_content} [{res.metadata}]")

* [SIM=0.627522] bar [{'source': 'https://another-example.com'}]

There are a variety of other search methods that are not covered in this notebook, such as MMR search or searching by vector. For a full list of the search abilites check out the API reference.

Query by turning into retriever

You can also transform the vector store into a retriever for easier usage in your chains.

retriever = vector_store.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={"k": 1, "score_threshold": 0.5},
)
retriever.invoke("thud")

[Document(metadata={'source': 'https://example.com'}, page_content='foo')]

Using retriever in a simple RAG chain (note this requires a functional chat model, in this case we use OpenAI which requires you set OPENAI_API_KEY in your environment):

from langchain_openai import ChatOpenAI
from langchain import hub
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough


llm = ChatOpenAI(model="gpt-3.5-turbo-0125")

prompt = hub.pull("rlm/rag-prompt")

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

rag_chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

rag_chain.invoke("thud")

API Reference:ChatOpenAI | StrOutputParser | RunnablePassthrough

"I'm sorry, I don't know the answer to that question based on the provided context."

For more, check out a complete RAG template using Astra DB here.

Cleanup

If you want to completely delete the collection from your Astra DB instance, run this.

(You will lose the data you stored in it.)

vector_store.delete_collection()

API reference

For detailed documentation of all AstraDBVectorStore features and configurations head to the API reference:https://api.python.langchain.com/en/latest/vectorstores/langchain_astradb.vectorstores.AstraDBVectorStore.html