Orchestrate MongoDB operations with Apache Airflow

1
"""
2
## Tutorial DAG: Load and query video game descriptions with MongoDB and OpenAI
3
"""
4
5
import logging
6
import os
7
8
from airflow.decorators import dag, task
9
from airflow.models.baseoperator import chain
10
from airflow.models.param import Param
11
from airflow.operators.empty import EmptyOperator
12
from airflow.providers.mongo.hooks.mongo import MongoHook
13
from pendulum import datetime
14
15
t_log = logging.getLogger("airflow.task")
16
17
_MONGO_DB_CONN = os.getenv("MONGO_DB_CONN", "mongodb_default")
18
_MONGO_DB_DATABASE_NAME = os.getenv("MONGO_DB_DATABASE_NAME", "games")
19
_MONGO_DB_COLLECTION_NAME = os.getenv("MONGO_DB_COLLECTION_NAME", "games_nostalgia")
20
_MONGO_DB_SEARCH_INDEX_NAME = os.getenv("MONGO_DB_SEARCH_INDEX_NAME", "find_me_a_game")
21
_MONGO_DB_VECTOR_COLUMN_NAME = os.getenv("MONGO_DB_VECTOR_COLUMN_NAME", "vector")
22
23
_OPENAI_EMBEDDING_MODEL = os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
24
_OPENAI_EMBEDDING_MODEL_DIMENSIONS = os.getenv(
25
    "OPENAI_EMBEDDING_MODEL_DIMENSIONS", 1536
26
)
27
28
_DATA_TEXT_FILE_PATH = os.getenv("DATA_TEXT_FILE_PATH", "include/games.txt")
29
30
_COLLECTION_EXISTS_TASK_ID = "collection_already_exists"
31
_CREATE_COLLECTION_TASK_ID = "create_collection"
32
_CREATE_INDEX_TASK_ID = "create_search_index"
33
_INDEX_EXISTS_TASK_ID = "search_index_already_exists"
34
35
36
def _get_mongodb_database(
37
    mongo_db_conn_id: str = _MONGO_DB_CONN,
38
    mongo_db_database_name: str = _MONGO_DB_DATABASE_NAME,
39
):
40
    """
41
    Get the MongoDB database.
42
    Args:
43
        mongo_db_conn_id (str): The connection ID for the MongoDB connection.
44
        mongo_db_database_name (str): The name of the database.
45
    Returns:
46
        The MongoDB database.
47
    """
48
    hook = MongoHook(mongo_conn_id=mongo_db_conn_id)
49
    client = hook.get_conn()
50
    return client[mongo_db_database_name]
51
52
53
def _create_openai_embeddings(text: str, model: str):
54
    """
55
    Create embeddings for a text with the OpenAI API.
56
    Args:
57
        text (str): The text to create embeddings for.
58
        model (str): The OpenAI model to use.
59
    Returns:
60
        The embeddings for the text.
61
    """
62
    from openai import OpenAI
63
64
    client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
65
    response = client.embeddings.create(input=text, model=model)
66
    embeddings = response.data[0].embedding
67
68
    return embeddings
69
70
71
@dag(
72
    start_date=datetime(2024, 10, 1),
73
    schedule=None,
74
    catchup=False,
75
    max_consecutive_failed_dag_runs=5,
76
    tags=["mongodb"],
77
    doc_md=__doc__,
78
    params={
79
        "game_concepts": Param(
80
            ["fantasy", "quests"],
81
            type="array",
82
            description=(
83
                "What kind of game do you want to play today?"
84
                + " Add one concept per line."
85
            ),
86
        ),
87
    },
88
)
89
def query_game_vectors():
90
91
    @task.branch
92
    def check_for_collection() -> str:
93
        "Check if the provided collection already exists and decide on the next step."
94
        database = _get_mongodb_database()
95
        collection_list = database.list_collection_names()
96
        if _MONGO_DB_COLLECTION_NAME in collection_list:
97
            return _COLLECTION_EXISTS_TASK_ID
98
        else:
99
            return _CREATE_COLLECTION_TASK_ID
100
101
    @task(task_id=_CREATE_COLLECTION_TASK_ID)
102
    def create_collection():
103
        "Create a new collection in the database."
104
        database = _get_mongodb_database()
105
        database.create_collection(_MONGO_DB_COLLECTION_NAME)
106
107
    collection_already_exists = EmptyOperator(task_id=_COLLECTION_EXISTS_TASK_ID)
108
    collection_ready = EmptyOperator(
109
        task_id="collection_ready", trigger_rule="none_failed"
110
    )
111
112
    @task
113
    def extract() -> list:
114
        """
115
        Extract the games from the text file.
116
        Returns:
117
            list: A list with the games.
118
        """
119
        import re
120
121
        with open(_DATA_TEXT_FILE_PATH, "r") as f:
122
            games = f.readlines()
123
124
        games_list = []
125
126
        for game in games:
127
128
            parts = game.split(":::")
129
            title_year = parts[1].strip()
130
            match = re.match(r"(.+) \((\d{4})\)", title_year)
131
132
            title, year = match.groups()
133
            year = int(year)
134
135
            genre = parts[2].strip()
136
            description = parts[3].strip()
137
138
            game_data = {
139
                "title": title,
140
                "year": year,
141
                "genre": genre,
142
                "description": description,
143
            }
144
145
            games_list.append(game_data)
146
147
        return games_list
148
149
    @task(map_index_template="{{ game_str }}")
150
    def transform_create_embeddings(game: dict) -> dict:
151
        """
152
        Create embeddings for the game description.
153
        Args:
154
            game (dict): A dictionary with the game's data.
155
        Returns:
156
            dict: The game's data with the embeddings.
157
        """
158
        embeddings = _create_openai_embeddings(
159
            text=game.get("description"), model=_OPENAI_EMBEDDING_MODEL
160
        )
161
        game[_MONGO_DB_VECTOR_COLUMN_NAME] = embeddings
162
163
        # optional: setting the custom map index
164
        from airflow.operators.python import get_current_context
165
166
        context = get_current_context()
167
        context["game_str"] = f"{game['title']} ({game['year']}) - {game['genre']}"
168
169
        return game
170
171
    @task(trigger_rule="none_failed", map_index_template="{{ game_str }}")
172
    def load_data_to_mongo_db(game_data: dict) -> None:
173
        """
174
        Load the game data to the MongoDB collection.
175
        Args:
176
            game_data (dict): A dictionary with the game's data.
177
        """
178
179
        database = _get_mongodb_database()
180
        collection = database[_MONGO_DB_COLLECTION_NAME]
181
182
        filter_query = {
183
            "title": game_data["title"],
184
            "year": game_data["year"],
185
            "genre": game_data["genre"],
186
        }
187
188
        game_str = f"{game_data['title']} ({game_data['year']}) - {game_data['genre']}"
189
190
        existing_document = collection.find_one(filter_query)
191
192
        if existing_document:
193
            if existing_document.get("description") != game_data["description"]:
194
                collection.update_one(
195
                    filter_query, {"$set": {"description": game_data["description"]}}
196
                )
197
                t_log.info(f"Updated description for record: {game_str}")
198
            else:
199
                t_log.info(f"Skipped duplicate record: {game_str}")
200
        else:
201
            collection.update_one(
202
                filter_query, {"$setOnInsert": game_data}, upsert=True
203
            )
204
            t_log.info(f"Inserted record: {game_str}")
205
206
        # optional: setting the custom map index
207
        from airflow.operators.python import get_current_context
208
209
        context = get_current_context()
210
        context["game_str"] = game_str
211
212
    @task.branch
213
    def check_for_search_index() -> str:
214
        "Check if the provided index already exists and decide on the next step."
215
        database = _get_mongodb_database()
216
        collection = database[_MONGO_DB_COLLECTION_NAME]
217
        index_list = collection.list_search_indexes().to_list()
218
        index_name_list = [index.get("name") for index in index_list]
219
        if _MONGO_DB_SEARCH_INDEX_NAME in index_name_list:
220
            return _INDEX_EXISTS_TASK_ID
221
        else:
222
            return _CREATE_INDEX_TASK_ID
223
224
    @task(task_id=_CREATE_INDEX_TASK_ID)
225
    def create_search_index():
226
        """
227
        Create a search index model for the MongoDB collection.
228
        """
229
        from pymongo.operations import SearchIndexModel
230
231
        database = _get_mongodb_database()
232
        collection = database[_MONGO_DB_COLLECTION_NAME]
233
234
        search_index_model = SearchIndexModel(
235
            definition={
236
                "mappings": {
237
                    "dynamic": True,
238
                    "fields": {
239
                        _MONGO_DB_VECTOR_COLUMN_NAME: {
240
                            "type": "knnVector",
241
                            "dimensions": _OPENAI_EMBEDDING_MODEL_DIMENSIONS,
242
                            "similarity": "cosine",
243
                        }
244
                    },
245
                },
246
            },
247
            name=_MONGO_DB_SEARCH_INDEX_NAME,
248
        )
249
250
        collection.create_search_index(model=search_index_model)
251
252
    search_index_already_exists = EmptyOperator(task_id=_INDEX_EXISTS_TASK_ID)
253
254
    @task.sensor(
255
        poke_interval=10, timeout=3600, mode="poke", trigger_rule="none_failed"
256
    )
257
    def wait_for_full_indexing():
258
        """
259
        Wait for the search index to be fully built.
260
        """
261
        from airflow.sensors.base import PokeReturnValue
262
263
        database = _get_mongodb_database()
264
        collection = database[_MONGO_DB_COLLECTION_NAME]
265
266
        index_list = collection.list_search_indexes().to_list()
267
        index = next(
268
            (
269
                index
270
                for index in index_list
271
                if index.get("name") == _MONGO_DB_SEARCH_INDEX_NAME
272
            ),
273
            None,
274
        )
275
276
        if index:
277
            status = index.get("status")
278
            if status == "READY":
279
                t_log.info(f"Search index is {status}. Ready to query.")
280
                condition_met = True
281
            elif status == "FAILED":
282
                raise ValueError("Search index failed to build.")
283
            else:
284
                t_log.info(
285
                    f"Search index is {status}. Waiting for indexing to complete."
286
                )
287
                condition_met = False
288
        else:
289
            raise ValueError("Search index not found.")
290
291
        return PokeReturnValue(is_done=condition_met)
292
293
    @task
294
    def embed_concepts(**context):
295
        """
296
        Create embeddings for the provided concepts.
297
        """
298
        from openai import OpenAI
299
300
        client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
301
        game_concepts = context["params"]["game_concepts"]
302
        game_concepts_str = " ".join(game_concepts)
303
304
        embeddings = client.embeddings.create(
305
            input=game_concepts_str, model=_OPENAI_EMBEDDING_MODEL
306
        )
307
308
        return embeddings.to_dict()
309
310
    @task
311
    def query(query_vector: list):
312
        """
313
        Query the MongoDB collection for games based on the provided concepts.
314
        """
315
316
        db = _get_mongodb_database()
317
        collection = db[_MONGO_DB_COLLECTION_NAME]
318
319
        results = collection.aggregate(
320
            [
321
                {
322
                    "$vectorSearch": {
323
                        "exact": True,
324
                        "index": _MONGO_DB_SEARCH_INDEX_NAME,
325
                        "limit": 1,
326
                        "path": _MONGO_DB_VECTOR_COLUMN_NAME,
327
                        "queryVector": query_vector["data"][0]["embedding"],
328
                    }
329
                }
330
            ]
331
        )
332
333
        results_list = []
334
335
        for result in results:
336
337
            game_id = str(result["_id"])
338
            title = result["title"]
339
            year = result["year"]
340
            genre = result["genre"]
341
            description = result["description"]
342
343
            t_log.info(f"You should play {title}!")
344
            t_log.info(f"It was released in {year} and belongs to the {genre} genre.")
345
            t_log.info(f"Description: {description}")
346
347
            results_list.append(
348
                {
349
                    "game_id": game_id,
350
                    "title": title,
351
                    "year": year,
352
                    "genre": genre,
353
                    "description": description,
354
                }
355
            )
356
357
        return results_list
358
359
    _extract = extract()
360
    _transform_create_embeddings = transform_create_embeddings.expand(game=_extract)
361
    _load_data_to_mongo_db = load_data_to_mongo_db.expand(
362
        game_data=_transform_create_embeddings
363
    )
364
365
    _query = query(embed_concepts())
366
367
    chain(
368
        check_for_collection(),
369
        [create_collection(), collection_already_exists],
370
        collection_ready,
371
    )
372
373
    chain(
374
        collection_ready,
375
        check_for_search_index(),
376
        [create_search_index(), search_index_already_exists],
377
        wait_for_full_indexing(),
378
        _query,
379
    )
380
381
    chain(collection_ready, _load_data_to_mongo_db, _query)
382
383
384
query_game_vectors()