haystack-tutorials

1
import tomllib
2
import argparse
3
import json
4

5

6
def read_index(path):
7
    with open(path, "rb") as f:
8
        return tomllib.load(f)
9

10

11
if __name__ == "__main__":
12
    parser = argparse.ArgumentParser(usage="""python generate_matrix.py --haystack-version v1.18.1""")
13
    parser.add_argument("--index", dest="index", default="index.toml")
14
    parser.add_argument("--notebooks", dest="notebooks", nargs="+", default=[])
15
    parser.add_argument("--haystack-version", dest="version", required=True)
16
    parser.add_argument("--include-main", dest="main", action="store_true")
17

18
    args = parser.parse_args()
19
    index = read_index(args.index)
20

21
    matrix = []
22
    for tutorial in index["tutorial"]:
23
        notebook = tutorial["notebook"]
24

25
        if args.notebooks and notebook not in args.notebooks:
26
            # If the user specified a list of notebooks to run, only run those
27
            # otherwise run all of them
28
            continue
29

30
        if tutorial.get("needs_gpu", False):
31
            # We're not running GPU tutorials on GitHub Actions
32
            # since we don't have a GPUs there
33
            continue
34

35
        if not tutorial.get("colab", True):
36
            # This tutorial doesn't have any runnable Python code
37
            # so there's nothing to test
38
            continue
39

40
        if tutorial.get("haystack_2", False):
41
            # Haystack 2.0 tutorials should be skipped for now
42
            continue
43

44
        version = tutorial.get("haystack_version", args.version)
45
        if version[0] != "v":
46
            version = f"v{version}"
47

48
        matrix.append({"notebook": notebook[:-6], "haystack_version": version})
49

50
        if args.main and "haystack_version" not in tutorial:
51
            # If a tutorial doesn't specify a version, we also test it on main
52
            matrix.append({"notebook": notebook[:-6], "haystack_version": "main"})
53

54
    print(json.dumps(matrix))
55