tests: clean astra between notebook tests (#326)

Author: nicoloboschi

.github/actions/lint/action.yml

@@ -15,3 +15,14 @@ runs:
       run: |
         tox -c ragstack-e2e-tests -e lint
         tox -e lint-yaml
+
+    - name: Notebook
+      shell: bash
+      run: |
+        python scripts/format-example-notebooks.py
+        if [ -n "$(git status --porcelain)" ]; then
+          echo "Notebooks are not formatted"
+          echo "Please run 'python scripts/format-example-notebooks.py' and commit the changes."
+          git status
+          exit 1
+        fi

.github/workflows/ci-e2e-tests.yml

@@ -27,9 +27,13 @@ jobs:
       astradb-prod-region: "us-east-1"
       astradb-prod-cloud: "aws"
       is-scheduled: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
+      # yamllint disable-line rule:line-length
       is-ragstack-dev-cron: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'schedule' && github.event.schedule == '0 0/4 * * *') }}
+      # yamllint disable-line rule:line-length
       is-ragstack-latest-release-cron: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'schedule' && github.event.schedule == '0 1/4 * * *') }}
+      # yamllint disable-line rule:line-length
       is-langchain-dev-cron: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'schedule' && github.event.schedule == '0 2/4 * * *') }}
+      # yamllint disable-line rule:line-length
       is-llamaindex-dev-cron: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'schedule' && github.event.schedule == '0 3/4 * * *') }}
     steps:
       - uses: actions/checkout@v4

examples/notebooks/FLARE.ipynb

@@ -53,7 +53,14 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "nbmake": {
+     "post_cell_execute": [
+      "from conftest import before_notebook",
+      "before_notebook()"
+     ]
+    }
+   },
    "outputs": [],
    "source": [
     "! pip install ragstack-ai"
@@ -699,4 +706,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

examples/notebooks/QA_with_cassio.ipynb

@@ -19,11 +19,11 @@
     "This notebook guides you through setting up [RAGStack](https://www.datastax.com/products/ragstack) using [Astra Vector Search](https://docs.datastax.com/en/astra-serverless/docs/vector-search/overview.html), [OpenAI](https://openai.com/about), and [CassIO](https://cassio.org/) to implement a generative Q&A over your own documentation.\n",
     "\n",
     "## Astra Vector Search\n",
-    "Astra vector search enables developers to search a database by context or meaning rather than keywords or literal values. This is done by using “embeddings”. Embeddings are a type of representation used in machine learning where high-dimensional or complex data is mapped onto vectors in a lower-dimensional space. These vectors capture the semantic properties of the input data, meaning that similar data points have similar embeddings.\n",
+    "Astra vector search enables developers to search a database by context or meaning rather than keywords or literal values. This is done by using \u201cembeddings\u201d. Embeddings are a type of representation used in machine learning where high-dimensional or complex data is mapped onto vectors in a lower-dimensional space. These vectors capture the semantic properties of the input data, meaning that similar data points have similar embeddings.\n",
     "Reference: [Astra Vector Search](https://docs.datastax.com/en/astra-serverless/docs/vector-search/overview.html)\n",
     "\n",
     "## CassIO\n",
-    "CassIO is the ultimate solution for seamlessly integrating Apache Cassandra® with generative artificial intelligence and other machine learning workloads. This powerful Python library simplifies the complicated process of accessing the advanced features of the Cassandra database, including vector search capabilities. With CassIO, developers can fully concentrate on designing and perfecting their AI systems without any concerns regarding the complexities of integration with Cassandra.\n",
+    "CassIO is the ultimate solution for seamlessly integrating Apache Cassandra\u00ae with generative artificial intelligence and other machine learning workloads. This powerful Python library simplifies the complicated process of accessing the advanced features of the Cassandra database, including vector search capabilities. With CassIO, developers can fully concentrate on designing and perfecting their AI systems without any concerns regarding the complexities of integration with Cassandra.\n",
     "Reference: [CassIO](https://cassio.org/)\n",
     "\n",
     "## OpenAI\n",
@@ -72,6 +72,12 @@
     },
     "editable": true,
     "id": "a6d88d66",
+    "nbmake": {
+     "post_cell_execute": [
+      "from conftest import before_notebook",
+      "before_notebook()"
+     ]
+    },
     "outputId": "dc543d17-3fb2-4362-cc4e-0050bd7787ba",
     "slideshow": {
      "slide_type": ""
@@ -543,4 +549,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

examples/notebooks/RAG_with_cassio.ipynb

@@ -67,6 +67,12 @@
      "base_uri": "https://localhost:8080/"
     },
     "id": "2953d95b",
+    "nbmake": {
+     "post_cell_execute": [
+      "from conftest import before_notebook",
+      "before_notebook()"
+     ]
+    },
     "outputId": "f1d1a4dc-9984-405d-bf28-74697815ea12"
    },
    "outputs": [],
@@ -472,4 +478,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

examples/notebooks/__init__.py

@@ -55,7 +55,14 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "nbmake": {
+     "post_cell_execute": [
+      "from conftest import before_notebook",
+      "before_notebook()"
+     ]
+    }
+   },
    "outputs": [],
    "source": [
     "%pip install ragstack-ai"
@@ -1035,4 +1042,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}

examples/notebooks/advancedRAG.ipynb

@@ -13,7 +13,7 @@
    "source": [
     "# Astra DB\n",
     "\n",
-    "This page provides a quickstart for using [Astra DB](https://docs.datastax.com/en/astra/home/astra.html) and [Apache Cassandra®](https://cassandra.apache.org/) as a Vector Store.\n",
+    "This page provides a quickstart for using [Astra DB](https://docs.datastax.com/en/astra/home/astra.html) and [Apache Cassandra\u00ae](https://cassandra.apache.org/) as a Vector Store.\n",
     "\n",
     "_Note: in addition to access to the database, an OpenAI API Key is required to run the full example._"
    ]
@@ -39,7 +39,14 @@
    "cell_type": "code",
    "execution_count": null,
    "id": "8d00fcf4-9798-4289-9214-d9734690adfc",
-   "metadata": {},
+   "metadata": {
+    "nbmake": {
+     "post_cell_execute": [
+      "from conftest import before_notebook",
+      "before_notebook()"
+     ]
+    }
+   },
    "outputs": [],
    "source": [
     "!pip install --quiet datasets pypdf"
@@ -902,4 +909,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

examples/notebooks/astradb.ipynb

@@ -28,7 +28,7 @@ def get_required_env(name) -> str:
 )
 
 
-def try_delete_with_backoff(collection: str, sleep=1, max_tries=5):
+def try_delete_with_backoff(collection: str, sleep=1, max_tries=2):
     try:
         logging.info(f"deleting collection {collection}")
         response = client.delete_collection(
@@ -45,10 +45,8 @@ def try_delete_with_backoff(collection: str, sleep=1, max_tries=5):
         try_delete_with_backoff(collection, sleep * 2, max_tries)
 
 
-@pytest.fixture
-def cleanup_astra():
+def before_notebook():
     collections = client.get_collections().get("status").get("collections")
     logging.info(f"Existing collections: {collections}")
     for collection in collections:
         try_delete_with_backoff(collection)
-    yield