AstraDB_langchain_quickstart_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6d60a143-5cc8-43b9-8f2b-d5353aaa7f3c",
   "metadata": {},
   "source": [
    "# Integrate LangChain with Astra DB Serverless\n",
    "\n",
    "For more information, visit the DataStax [Astra DB docs page](https://docs.datastax.com/en/astra-db-serverless/integrations/langchain.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9764e550-225f-42b9-83ec-ae3c84d796fd",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "- An active Astra account.\n",
    "- An active [Serverless (Vector) database](https://docs.datastax.com/en/astra-db-serverless/get-started/quickstart.html#create-a-database-and-store-your-credentials).\n",
    "- An Open AI account and an [OpenAI API key](https://platform.openai.com/).\n",
    "\n",
    "_This guide uses OpenAI to generate embeddings. You can get embeddings directly from OpenAI, or you can use Astra DB’s built-in OpenAI embedding provider integration (also known as a \"vectorize integration\")._\n",
    "\n",
    "_If you want to use the built-in OpenAI integration, you must [configure the OpenAI embedding provider integration](https://docs.datastax.com/en/astra-db-serverless/integrations/embedding-providers/openai.html) before you begin. In the integration settings, note the **API key name**, and make sure that your database is in the key’s scope._\n",
    "\n",
    "- The following Python dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943defcb-ace5-4d78-82e3-37b0a24dbc1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet \\\n",
    "    \"langchain>=0.3,<0.4\" \\\n",
    "    \"langchain-astradb>=0.6,<0.7\" \\\n",
    "    \"langchain-openai>=0.3,<0.4\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9847b646-fe12-4b12-a887-667f2d44cb65",
   "metadata": {},
   "source": [
    "## Connect to the Serverless (Vector) database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a751f02e-5897-4dad-9c8e-3a4dfa4b2811",
   "metadata": {},
   "source": [
    "### Import dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c9a6912-7f94-4bab-b0fa-4c61b10251e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "from getpass import getpass\n",
    "\n",
    "from astrapy.info import VectorServiceOptions\n",
    "from langchain_astradb import AstraDBVectorStore\n",
    "\n",
    "from langchain_core.documents import Document\n",
    "from langchain_openai import OpenAIEmbeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d863d3c-8778-4109-a1ae-f6b81fee6fc9",
   "metadata": {},
   "source": [
    "### Set secrets and connection parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1df88a14-9358-4cff-90cd-3d523ae0cd75",
   "metadata": {},
   "source": [
    "Get an application token and Data API endpoint for your database:\n",
    "\n",
    "- In the [Astra Portal](https://astra.datastax.com/) navigation menu, click Databases, and then click the name of your Serverless (Vector) database.\n",
    "- On the Overview tab, find the Database Details section.\n",
    "- In API Endpoint, click Copy to get your database’s Data API endpoint in the form of `https://ASTRA_DB_ID-ASTRA_DB_REGION.apps.astra.datastax.com`.\n",
    "- Click Generate Token to create an [application token](https://docs.datastax.com/en/astra-db-serverless/administration/manage-application-tokens.html) scoped to your database."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f2fae75-020e-4c9f-aef6-02e820c96c90",
   "metadata": {},
   "source": [
    "#### Astra DB parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47802c94-ec4a-460b-ba01-cbeff5c9c060",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"ASTRA_DB_API_ENDPOINT\"] = input(\"ASTRA_DB_API_ENDPOINT =\")\n",
    "os.environ[\"ASTRA_DB_APPLICATION_TOKEN\"] = getpass(\"ASTRA_DB_APPLICATION_TOKEN =\")\n",
    "\n",
    "if _keyspace := input(\"ASTRA_DB_KEYSPACE (optional) =\"):\n",
    "    os.environ[\"ASTRA_DB_KEYSPACE\"] = _keyspace\n",
    "\n",
    "os.environ[\"ASTRA_DB_API_KEY_NAME\"] = input(\"ASTRA_DB_API_KEY_NAME (required for 'vectorize') =\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81819613-86a9-432e-b38b-c1f8ee4ea811",
   "metadata": {},
   "source": [
    "#### OpenAI parameter (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0201947e-b795-44c7-80bd-a7cdc96e604c",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_KEY\"] = getpass(\"OPENAI_API_KEY (required for explicit embeddings) =\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb0accf-a225-4333-9164-746f3a8df124",
   "metadata": {},
   "source": [
    "### Load environment variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bfaaff0-f2f1-4245-82fb-92f3a4b44b31",
   "metadata": {},
   "outputs": [],
   "source": [
    "ASTRA_DB_APPLICATION_TOKEN = os.environ[\"ASTRA_DB_APPLICATION_TOKEN\"]\n",
    "ASTRA_DB_API_ENDPOINT = os.environ[\"ASTRA_DB_API_ENDPOINT\"]\n",
    "ASTRA_DB_KEYSPACE = os.environ.get(\"ASTRA_DB_KEYSPACE\")\n",
    "ASTRA_DB_API_KEY_NAME = os.environ.get(\"ASTRA_DB_API_KEY_NAME\") or None\n",
    "\n",
    "OPENAI_API_KEY = os.environ.get(\"OPENAI_API_KEY\") or None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78db8321-0ba0-4620-b552-e7e924804405",
   "metadata": {},
   "source": [
    "## Create embeddings from text\n",
    "\n",
    "### Create a vector store\n",
    "\n",
    "> **Choose** between [server-side embedding computation](https://docs.datastax.com/en/astra-db-serverless/databases/embedding-generation.html) (\"vectorize\") or explicit embeddings by editing the following as desired. Then run the cell. _(If using \"vectorize\", you must have configured an embedding provider in your Astra DB. Conversely, if opting for explicit embeddings, the OpenAI API Key must have been set in the notebook.)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a76d54f4-eb26-48fb-8548-3fc62800cffd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Edit if necessary, then run the cell\n",
    "\n",
    "USE_VECTORIZE = True  # server-side embeddings\n",
    "# USE_VECTORIZE = False  # explicit embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d1c373-5ce3-4b4b-b685-1fa50bee7e73",
   "metadata": {},
   "source": [
    "Depending on the choice of embedding computation, the parameters are slightly different.\n",
    "\n",
    "When creating the LangChain vector store, you specify the database and a collection name. The collection is created automatically if it does not exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f327321-bf8b-4db4-ab2b-71f63ab72829",
   "metadata": {},
   "outputs": [],
   "source": [
    "if USE_VECTORIZE:\n",
    "    vectorize_options = VectorServiceOptions(\n",
    "        provider=\"openai\",  # Change these if using another embedding provider/model\n",
    "        model_name=\"text-embedding-3-small\",\n",
    "        authentication={\"providerKey\": ASTRA_DB_API_KEY_NAME},\n",
    "    )\n",
    "    vector_store = AstraDBVectorStore(\n",
    "        collection_name=\"langchain_integration_demo_vectorize\",\n",
    "        token=ASTRA_DB_APPLICATION_TOKEN,\n",
    "        api_endpoint=ASTRA_DB_API_ENDPOINT,\n",
    "        namespace=ASTRA_DB_KEYSPACE,\n",
    "        collection_vector_service_options=vectorize_options,\n",
    "    )\n",
    "\n",
    "if not USE_VECTORIZE:\n",
    "    embedding = OpenAIEmbeddings()\n",
    "    vector_store = AstraDBVectorStore(\n",
    "        collection_name=\"langchain_integration_demo\",\n",
    "        embedding=embedding,\n",
    "        token=ASTRA_DB_APPLICATION_TOKEN,\n",
    "        api_endpoint=ASTRA_DB_API_ENDPOINT,\n",
    "        namespace=ASTRA_DB_KEYSPACE,\n",
    "    )\n",
    "\n",
    "\n",
    "## If you already have a populated vector collection, try this instead\n",
    "## (and then skip the load+process+insert phases if you are so inclined):\n",
    "\n",
    "# vector_store = AstraDBVectorStore(\n",
    "#     collection_name=\"INSERT_YOUR_COLLECTION_NAME\",\n",
    "#     embedding=EMBEDDING,  # omit for vectorize; else, must be the same used for the data on DB\n",
    "#     token=ASTRA_DB_APPLICATION_TOKEN,\n",
    "#     api_endpoint=ASTRA_DB_API_ENDPOINT,\n",
    "#     namespace=ASTRA_DB_KEYSPACE,\n",
    "#     autodetect_collection=True,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d001b63a-7ff2-49e7-865a-a955504912de",
   "metadata": {},
   "source": [
    "### Load data\n",
    "\n",
    "Load a small dataset of philosophical quotes from this repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9859326a-92b5-4a31-87d8-4f32f9cd4c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "philo_dataset = requests.get(\n",
    "    \"https://raw.githubusercontent.com/\"\n",
    "    \"datastaxdevs/mini-demo-astradb-langchain/\"\n",
    "    \"refs/heads/main/data/philosopher-quotes.json\"\n",
    ").json()\n",
    "\n",
    "print(\"An example entry:\")\n",
    "print(philo_dataset[16])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d116494e-e3e0-4345-a3cc-268da354fe2e",
   "metadata": {},
   "source": [
    "### Process dataset\n",
    "\n",
    "Transform the dataset into ready-to-insert LangChain `Document` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b43669a4-0f13-4abc-8197-3b4336f7586c",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents_to_insert = []\n",
    "\n",
    "for entry_idx, entry in enumerate(philo_dataset):\n",
    "    metadata = {\n",
    "        \"author\": entry[\"author\"],\n",
    "        **entry[\"metadata\"],\n",
    "    }\n",
    "    # Construct the Document, with the quote and metadata tags\n",
    "    new_document = Document(\n",
    "        id=entry[\"_id\"],\n",
    "        page_content=entry[\"quote\"],\n",
    "        metadata=metadata,\n",
    "    )\n",
    "    documents_to_insert.append(new_document)\n",
    "\n",
    "print(f\"Ready to insert {len(documents_to_insert)} documents.\")\n",
    "print(f\"Example document: {documents_to_insert[16]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "438a4548-988c-4c32-83c7-0502f3b1624a",
   "metadata": {},
   "source": [
    "### Insert documents\n",
    "\n",
    "This step will compute vector embedding and save all entries in the vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80126487-5cfc-42d5-9643-0706f4f1662b",
   "metadata": {},
   "outputs": [],
   "source": [
    "inserted_ids = vector_store.add_documents(documents_to_insert)\n",
    "\n",
    "print(f\"\\nInserted {len(inserted_ids)} documents: {', '.join(inserted_ids[:3])} ...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ded8ac62-73dd-4258-8cdd-bb4d754cc851",
   "metadata": {},
   "source": [
    "## Verify the integration\n",
    "\n",
    "Find quotes semantically similar to a given input query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f22c8276-8d87-4b15-a6d1-5aadf649bf87",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = vector_store.similarity_search(\"Our life is what we make of it\", k=3)\n",
    "\n",
    "for res in results:\n",
    "    print(f\"* {res.page_content} [{res.metadata}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd961154-e33d-4eda-96ce-c7dd887354b0",
   "metadata": {},
   "source": [
    "## Further usage patterns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f2f7cc7-75f3-43d9-a3d6-f7f59fed3950",
   "metadata": {},
   "source": [
    "### Use `add_texts`\n",
    "\n",
    "You can store documents through `add_texts` and supply three parallel lists for the texts, the metadata and the IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0277d8-f3e5-4526-84e1-6e365aaef493",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = [\n",
    "    \"I think, therefore I am.\",\n",
    "    \"To the things themselves!\",\n",
    "]\n",
    "metadatas = [\n",
    "    {\"author\": \"descartes\", \"knowledge\": \"y\"},\n",
    "    {\"author\": \"husserl\", \"knowledge\": \"y\"},\n",
    "]\n",
    "ids = [\n",
    "    \"desc_999\",\n",
    "    \"huss_888\",\n",
    "]\n",
    "inserted_ids_2 = vector_store.add_texts(texts=texts, metadatas=metadatas, ids=ids)\n",
    "print(f\"\\nInserted {len(inserted_ids_2)} documents.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f69ceadb-68be-494a-abac-4af9cd62493e",
   "metadata": {},
   "source": [
    "### Return similarity scores from a search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d79ef95c-2a1e-4e5d-85a1-73ad902c3ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = vector_store.similarity_search_with_score(\"Our life is what we make of it\", k=3)\n",
    "for res, score in results:\n",
    "    print(f\"* [{score:.3f}] {res.page_content} [{res.metadata}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b86a4db1-67ef-48d2-9939-cd24c0273fd3",
   "metadata": {},
   "source": [
    "### Similarity search with metadata filtering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fcacba8-a865-4c90-8617-a6e379d71b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = vector_store.similarity_search(\n",
    "    \"Our life is what we make of it\",\n",
    "    k=3,\n",
    "    filter={\"author\": \"aristotle\"},\n",
    ")\n",
    "for res in results:\n",
    "    print(f\"* {res.page_content} [{res.metadata}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d22175b-4a7f-4ebb-936a-72db24551f56",
   "metadata": {},
   "source": [
    "### MMR (maximal marginal relevance) similarity search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "523a1b18-8b81-484a-b79e-c8fa74ac5d68",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = vector_store.max_marginal_relevance_search(\n",
    "    \"Our life is what we make of it\",\n",
    "    k=3,\n",
    "    filter={\"author\": \"aristotle\"},\n",
    ")\n",
    "for res in results:\n",
    "    print(f\"* {res.page_content} [{res.metadata}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14cb448d-91d1-4edc-8047-adcfa87b5afc",
   "metadata": {},
   "source": [
    "### Delete documents from the store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8d031f8-4d5e-4ab7-a6a6-020f14a480af",
   "metadata": {},
   "source": [
    "#### Delete by document ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71ab0207-9535-4cc0-8557-ef030c08e73c",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_1 = vector_store.delete(inserted_ids[:3])\n",
    "print(f\"delete result = {delete_1}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee19e34-6826-49ec-93cd-692e1af3894e",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_2 = vector_store.delete(inserted_ids[2:5])\n",
    "print(f\"delete result = {delete_2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdc96eda-0047-485b-962c-60fe329ab1b3",
   "metadata": {},
   "source": [
    "#### Retrieve and then delete\n",
    "\n",
    "Sometimes you do not have the IDs, ... but you might want to run a search and then delete the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5bc1cb3-76f7-4d2d-9e00-9919b813da5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ids_to_delete = []\n",
    "for res_doc, res_score, res_id in vector_store.similarity_search_with_score_id(\n",
    "    \"Philosophy has no goals\",\n",
    "    k=2,\n",
    "):\n",
    "    print(f\"* [SIM={res_score:.3f}] {res_doc.page_content} [{res_doc.metadata}]\")\n",
    "    ids_to_delete.append(res_id)\n",
    "\n",
    "print(f\"\\nDeleting IDs = {ids_to_delete} ...\")\n",
    "success = vector_store.delete(ids_to_delete)\n",
    "print(f\"Deletion succeeded = {success}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44ea752-826f-4d3b-8ae5-4a23eef38ede",
   "metadata": {},
   "source": [
    "Now try again the same search:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e24396-49e6-4fd5-90c4-badd13274b3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for res_doc, res_score, res_id in vector_store.similarity_search_with_score_id(\n",
    "    \"Philosophy has no goals\",\n",
    "    k=2,\n",
    "):\n",
    "    print(f\"* [SIM={res_score:.3f}] {res_doc.page_content} [{res_doc.metadata}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8280918-25ea-40f4-82c1-78d0b9b27278",
   "metadata": {},
   "source": [
    "#### Delete the **whole** stored data\n",
    "\n",
    "> _Warning: use with caution. Data loss!_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7414400-3e1b-45e4-bac7-9b49fcd4cba4",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store.clear()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533e5923-a1a1-40f0-976d-4b1af2cb03da",
   "metadata": {},
   "source": [
    "## Cleanup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c08d41e-65f8-4096-8b9d-9f12ed91c485",
   "metadata": {},
   "source": [
    "Completely delete the collection, thereby freeing the associated resources on Astra DB:\n",
    "\n",
    "> _Warning: use with caution. Data loss!_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fe88518-0494-4f08-b0eb-35305ce24ac4",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store.delete_collection()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41653b13-c903-4e7f-9806-afed3f4d726c",
   "metadata": {},
   "source": [
    "## Next steps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d6c352-cdcc-46a5-9683-376f55d16b45",
   "metadata": {},
   "source": [
    "- [This quickstart on DataStax documentation](https://docs.datastax.com/en/astra-db-serverless/integrations/langchain.html)\n",
    "- [`AstraDBVectorStore` in LangChain docs](https://python.langchain.com/docs/integrations/providers/astradb/#vector-store)\n",
    "- [`AstraDBVectorStore`, API Reference](https://python.langchain.com/api_reference/astradb/vectorstores/langchain_astradb.vectorstores.AstraDBVectorStore.html#langchain_astradb.vectorstores.AstraDBVectorStore)"
   ]
  }
 ],
 "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}