Add nbtest CI job (#179)

Author: miguelgrinberg

.github/workflows/tests.yml

@@ -0,0 +1,44 @@
+name: tests
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - notebooks/**
+  pull_request:
+    branches:
+      - main
+    paths:
+      - notebooks/**
+jobs:
+  notebook-tests:
+    strategy:
+      matrix:
+        es_stack:
+          - 8.11.4
+          - 8.12.0
+    runs-on: ubuntu-latest
+    services:
+      elasticsearch:
+        image: docker.elastic.co/elasticsearch/elasticsearch:${{ matrix.es_stack }}
+        env:
+          discovery.type: single-node
+          xpack.security.enabled: false
+          xpack.security.http.ssl.enabled: false
+          xpack.license.self_generated.type: trial
+        ports:
+          - 9200:9200
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Setup python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+      - name: Setup nbtest
+        run: make nbtest
+      - name: Warm up
+        continue-on-error: true
+        run: sleep 30 && PATCH_ES=1 ELASTIC_CLOUD_ID=foo ELASTIC_API_KEY=bar bin/nbtest notebooks/search/00-quick-start.ipynb
+      - name: Run tests
+        run: PATCH_ES=1 FORCE_COLOR=1 make -s test

Makefile

@@ -1,28 +1,20 @@
+# this is the list of notebooks that are integrated with the testing framework
+NOTEBOOKS = $(shell bin/find-notebooks-to-test.sh)
+
 .PHONY: install pre-commit nbtest test notebooks
 
 test: nbtest notebooks
 
-notebooks: search document-chunking model-upgrades langchain
-
-search:
-	$(MAKE) -C notebooks/search
-
-document-chunking:
-	$(MAKE) -C notebooks/document-chunking
-
-model-upgrades:
-	$(MAKE) -C notebooks/model-upgrades
-
-langchain:
-	$(MAKE) -C notebooks/langchain
+notebooks:
+	bin/nbtest $(NOTEBOOKS)
 
 install: pre-commit nbtest
 
 pre-commit:
 	python -m venv .venv
-	.venv/bin/pip install -r requirements-dev.txt
+	.venv/bin/pip install -qqq -r requirements-dev.txt
 	.venv/bin/pre-commit install
 
 nbtest:
 	python3 -m venv .venv
-	.venv/bin/pip install elastic-nbtest
+	.venv/bin/pip install -qqq elastic-nbtest

bin/find-notebooks-to-test.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# add any notebooks that are currently not testable to the exempt list
+EXEMPT_NOTEBOOKS=(
+    "notebooks/search/07-inference.ipynb"
+    "notebooks/search/08-learning-to-rank.ipynb"
+    "notebooks/langchain/langchain-vector-store.ipynb"
+    "notebooks/langchain/self-query-retriever-examples/chatbot-example.ipynb"
+    "notebooks/langchain/self-query-retriever-examples/chatbot-with-bm25-only-example.ipynb"
+    "notebooks/langchain/self-query-retriever-examples/langchain-self-query-retriever.ipynb"
+    "notebooks/langchain/multi-query-retriever-examples/chatbot-with-multi-query-retriever.ipynb"
+    "notebooks/langchain/multi-query-retriever-examples/langchain-multi-query-retriever.ipynb"
+    "notebooks/generative-ai/question-answering.ipynb"
+    "notebooks/generative-ai/chatbot.ipynb"
+    "notebooks/integrations/amazon-bedrock/langchain-qa-example.ipynb"
+    "notebooks/integrations/llama-index/intro.ipynb"
+    "notebooks/integrations/gemini/vector-search-gemini-elastic.ipynb"
+    "notebooks/integrations/gemini/qa-langchain-gemini-elasticsearch.ipynb"
+    "notebooks/integrations/openai/openai-KNN-RAG.ipynb"
+)
+
+ALL_NOTEBOOKS=$(find notebooks -name "*.ipynb" | grep -v "_nbtest" | grep -v ".ipynb_checkpoints" | sort)
+for notebook in $ALL_NOTEBOOKS; do
+    if [[ ! "${EXEMPT_NOTEBOOKS[@]}" =~ $notebook ]]; then
+        echo $notebook
+    fi
+done

bin/mocks/elasticsearch.py

@@ -0,0 +1,41 @@
+import os
+import sys
+
+
+def patch_elasticsearch():
+    # preserve the original import path
+    saved_path = sys.path.copy()
+
+    # remove the path entry that refers to this directory
+    for path in sys.path:
+        if not path.startswith('/'):
+            path = os.path.join(os.getcwd(), path)
+        if __file__ == os.path.join(path, 'elasticsearch.py'):
+            sys.path.remove(path)
+            break
+
+    # remove this module, and import the real one instead
+    del sys.modules['elasticsearch']
+    import elasticsearch
+
+    # restore the import path
+    sys.path = saved_path
+
+    # preserve the original Elasticsearch.__init__ method   
+    orig_es_init = elasticsearch.Elasticsearch.__init__
+
+    # patched version of Elasticsearch.__init__ that connects to self-hosted
+    # regardless of connection arguments given
+    def patched_es_init(self, *args, **kwargs):
+        if 'cloud_id' in kwargs:
+            assert kwargs['cloud_id'] == 'foo'
+        if 'api_key' in kwargs:
+            assert kwargs['api_key'] == 'bar'
+        return orig_es_init(self, 'http://localhost:9200')
+
+    # patch Elasticsearch.__init__
+    elasticsearch.Elasticsearch.__init__ = patched_es_init
+
+
+patch_elasticsearch()
+del patch_elasticsearch

bin/nbtest

@@ -4,5 +4,18 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 if [[ ! -f $SCRIPT_DIR/../.venv/bin/nbtest ]]; then
     make nbtest
 fi
+
+if [[ "$PATCH_ES" != "" ]]; then
+    # here we do some Python dark magic to patch the elasticsearch package to
+    # connect to a locally hosted instance in spite of connection arguments
+    # given
+    export ELASTIC_CLOUD_ID=foo
+    export ELASTIC_API_KEY=bar
+    export PYTHONPATH=$SCRIPT_DIR/mocks
+
+    # ensure elasticsearch is installed so that it can be patched
+    $SCRIPT_DIR/../.venv/bin/pip install -qqq elasticsearch
+fi
+
 source $SCRIPT_DIR/../.venv/bin/activate
 $SCRIPT_DIR/../.venv/bin/nbtest $*

notebooks/document-chunking/Makefile

@@ -0,0 +1,2 @@
+masks:
+- 'Score: [0-9]+\.[0-9][0-9]*'

notebooks/integrations/hugging-face/.nbtest.yml

@@ -47,7 +47,7 @@
    },
    "outputs": [],
    "source": [
-    "!python3 -m pip -qU install sentence-transformers eland elasticsearch transformers"
+    "!python3 -m pip install sentence-transformers eland elasticsearch transformers"
    ]
   },
   {
@@ -60,7 +60,8 @@
     "from elasticsearch import Elasticsearch\n",
     "from getpass import getpass\n",
     "from urllib.request import urlopen\n",
-    "import json"
+    "import json\n",
+    "from time import sleep"
    ]
   },
   {
@@ -111,7 +112,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!eland_import_hub_model --cloud-id $ELASTIC_CLOUD_ID --hub-model-id sentence-transformers/all-MiniLM-L6-v2 --task-type text_embedding --es-api-key $ELASTIC_API_KEY --start"
+    "!eland_import_hub_model --cloud-id $ELASTIC_CLOUD_ID --hub-model-id sentence-transformers/all-MiniLM-L6-v2 --task-type text_embedding --es-api-key $ELASTIC_API_KEY --start --clear-previous"
    ]
   },
   {
@@ -301,7 +302,8 @@
     "for title in titles:\n",
     "    actions.append({\"index\": {\"_index\": \"blogs\"}})\n",
     "    actions.append(title)\n",
-    "es.bulk(index=\"blogs\", operations=actions)"
+    "es.bulk(index=\"blogs\", operations=actions)\n",
+    "sleep(5)"
    ]
   },
   {
@@ -423,7 +425,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.10.13"
   },
   "vscode": {
    "interpreter": {