Vector store support for Cassandra

docs/extras/modules/data_connection/vectorstores/integrations/cassandra.ipynb

@@ -0,0 +1,269 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "683953b3",
+   "metadata": {},
+   "source": [
+    "# Cassandra\n",
+    "\n",
+    ">[Apache Cassandra®](https://cassandra.apache.org) is a NoSQL, row-oriented, highly scalable and highly available database.\n",
+    "\n",
+    "Newest Cassandra releases natively [support](https://cwiki.apache.org/confluence/display/CASSANDRA/CEP-30%3A+Approximate+Nearest+Neighbor(ANN)+Vector+Search+via+Storage-Attached+Indexes) Vector Similarity Search.\n",
+    "\n",
+    "To run this notebook you need either a running Cassandra cluster equipped with Vector Search capabilities (in pre-release at the time of writing) or a DataStax Astra DB instance running in the cloud (you can get one for free at [datastax.com](https://astra.datastax.com)). Check [cassio.org](https://cassio.org/start_here/) for more information."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4c41cad-08ef-4f72-a545-2151e4598efe",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "!pip install \"cassio>=0.0.5\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b7e46bb0",
+   "metadata": {},
+   "source": [
+    "### Please provide database connection parameters and secrets:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "36128a32",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import getpass\n",
+    "\n",
+    "database_mode = (input('\\n(L)ocal Cassandra or (A)stra DB? ')).upper()\n",
+    "\n",
+    "keyspace_name = input('\\nKeyspace name? ')\n",
+    "\n",
+    "if database_mode == 'A':\n",
+    "    ASTRA_DB_APPLICATION_TOKEN = getpass.getpass('\\nAstra DB Token (\"AstraCS:...\") ')\n",
+    "    #\n",
+    "    ASTRA_DB_SECURE_BUNDLE_PATH = input('Full path to your Secure Connect Bundle? ')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4f22aac2",
+   "metadata": {},
+   "source": [
+    "#### depending on whether local or cloud-based Astra DB, create the corresponding database connection \"Session\" object"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "677f8576",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from cassandra.cluster import Cluster\n",
+    "from cassandra.auth import PlainTextAuthProvider\n",
+    "\n",
+    "if database_mode == 'L':\n",
+    "    cluster = Cluster()\n",
+    "    session = cluster.connect()\n",
+    "elif database_mode == 'A':\n",
+    "    ASTRA_DB_CLIENT_ID = \"token\"\n",
+    "    cluster = Cluster(\n",
+    "        cloud={\n",
+    "            \"secure_connect_bundle\": ASTRA_DB_SECURE_BUNDLE_PATH,\n",
+    "        },\n",
+    "        auth_provider=PlainTextAuthProvider(\n",
+    "            ASTRA_DB_CLIENT_ID,\n",
+    "            ASTRA_DB_APPLICATION_TOKEN,\n",
+    "        ),\n",
+    "    )\n",
+    "    session = cluster.connect()\n",
+    "else:\n",
+    "    raise NotImplementedError"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "320af802-9271-46ee-948f-d2453933d44b",
+   "metadata": {},
+   "source": [
+    "### Please provide OpenAI access key\n",
+    "\n",
+    "We want to use `OpenAIEmbeddings` so we have to get the OpenAI API Key."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ffea66e4-bc23-46a9-9580-b348dfe7b7a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e98a139b",
+   "metadata": {},
+   "source": [
+    "### Creation and usage of the Vector Store"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aac9563e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
+    "from langchain.text_splitter import CharacterTextSplitter\n",
+    "from langchain.vectorstores import Cassandra\n",
+    "from langchain.document_loaders import TextLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3c3999a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.document_loaders import TextLoader\n",
+    "loader = TextLoader('../../../state_of_the_union.txt')\n",
+    "documents = loader.load()\n",
+    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
+    "docs = text_splitter.split_documents(documents)\n",
+    "\n",
+    "embedding_function = OpenAIEmbeddings()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6e104aee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "table_name = 'my_vector_db_table'\n",
+    "\n",
+    "docsearch = Cassandra.from_documents(\n",
+    "    documents=docs,\n",
+    "    embedding=embedding_function,\n",
+    "    session=session,\n",
+    "    keyspace=keyspace_name,\n",
+    "    table_name=table_name,\n",
+    ")\n",
+    "\n",
+    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
+    "docs = docsearch.similarity_search(query)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f509ee02",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## if you already have an index, you can load it and use it like this:\n",
+    "\n",
+    "# docsearch_preexisting = Cassandra(\n",
+    "#     embedding=embedding_function,\n",
+    "#     session=session,\n",
+    "#     keyspace=keyspace_name,\n",
+    "#     table_name=table_name,\n",
+    "# )\n",
+    "\n",
+    "# docsearch_preexisting.similarity_search(query, k=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c608226",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(docs[0].page_content)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d46d1452",
+   "metadata": {},
+   "source": [
+    "### Maximal Marginal Relevance Searches\n",
+    "\n",
+    "In addition to using similarity search in the retriever object, you can also use `mmr` as retriever.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a359ed74",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "retriever = docsearch.as_retriever(search_type=\"mmr\")\n",
+    "matched_docs = retriever.get_relevant_documents(query)\n",
+    "for i, d in enumerate(matched_docs):\n",
+    "    print(f\"\\n## Document {i}\\n\")\n",
+    "    print(d.page_content)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7c477287",
+   "metadata": {},
+   "source": [
+    "Or use `max_marginal_relevance_search` directly:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9ca82740",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "found_docs = docsearch.max_marginal_relevance_search(query, k=2, fetch_k=10)\n",
+    "for i, doc in enumerate(found_docs):\n",
+    "    print(f\"{i + 1}.\", doc.page_content, \"\\n\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

langchain/vectorstores/__init__.py

@@ -5,6 +5,7 @@
 from langchain.vectorstores.awadb import AwaDB
 from langchain.vectorstores.azuresearch import AzureSearch
 from langchain.vectorstores.base import VectorStore
+from langchain.vectorstores.cassandra import Cassandra
 from langchain.vectorstores.chroma import Chroma
 from langchain.vectorstores.clickhouse import Clickhouse, ClickhouseSettings
 from langchain.vectorstores.deeplake import DeepLake
@@ -37,6 +38,7 @@
     "AtlasDB",
     "AwaDB",
     "AzureSearch",
+    "Cassandra",
     "Chroma",
     "Clickhouse",
     "ClickhouseSettings",