Orchestrate Weaviate operations with Apache Airflow

Weaviate is an open source vector database, which store high-dimensional embeddings of objects like text, images, audio or video. The Weaviate Airflow provider offers modules to easily integrate Weaviate with Airflow.

In this tutorial you'll use Airflow to ingest movie descriptions into Weaviate, use Weaviate's automatic vectorization to create vectors for the descriptions, and query Weaviate for movies that are thematically close to user-provided concepts.

Other ways to learn

There are multiple resources for learning about this topic. See also:

• Webinar: Power your LLMOps with Airflow’s Weaviate Provider.

Why use Airflow with Weaviate?

Weaviate allows you to store objects alongside their vector embeddings and to query these objects based on their similarity. Vector embeddings are key components of many modern machine learning models such as LLMs or ResNet.

Integrating Weaviate with Airflow into one end-to-end machine learning pipeline allows you to:

• Use Airflow's data-driven scheduling to run operations on Weaviate based on upstream events in your data ecosystem, such as when a new model is trained or a new dataset is available.
• Run dynamic queries based on upstream events in your data ecosystem or user input via Airflow params against Weaviate to retrieve objects with similar vectors.
• Add Airflow features like retries and alerts to your Weaviate operations.

Time to complete

This tutorial takes approximately 30 minutes to complete.

Assumed knowledge

To get the most out of this tutorial, make sure you have an understanding of:

• The basics of Weaviate. See Weaviate Introduction.
• Airflow fundamentals, such as writing DAGs and defining tasks. See Get started with Apache Airflow.
• Airflow decorators. Introduction to the TaskFlow API and Airflow decorators.
• Airflow hooks. See Hooks 101.
• Airflow connections. See Managing your Connections in Apache Airflow.

Prerequisites

• The Astro CLI.
• (Optional) An OpenAI API key of at least tier 1 if you want to use OpenAI for vectorization. The tutorial can be completed using local vectorization with text2vec-transformers if you don't have an OpenAI API key.

This tutorial uses a local Weaviate instance created as a Docker container. You do not need to install the Weaviate client locally.

info

The example code from this tutorial is also available on GitHub.

Step 1: Configure your Astro project

1. Create a new Astro project:

   $ mkdir astro-weaviate-tutorial && cd astro-weaviate-tutorial
   $ astro dev init

2. Add the following two packages to your requirements.txt file to install the Weaviate Airflow provider and the Weaviate Python client in your Astro project:

   apache-airflow-providers-weaviate==1.0.0
   weaviate-client==3.25.3 

3. This tutorial uses a local Weaviate instance and a text2vec-transformer model, with each running in a Docker container. To add additional containers to your Astro project, create a new file in your project's root directory called docker-compose.override.yml and add the following:

   version: '3.1'
   services:           
     weaviate:
       image: cr.weaviate.io/semitechnologies/weaviate:latest
       command: "--host 0.0.0.0 --port '8081' --scheme http"
       ports:
         - 8081:8081
       environment:
         QUERY_DEFAULTS_LIMIT: 25
         PERSISTENCE_DATA_PATH: '/var/lib/weaviate'
         DEFAULT_VECTORIZER_MODULE: 'text2vec-transformers'
         ENABLE_MODULES: 'text2vec-transformers, text2vec-openai'
         CLUSTER_HOSTNAME: 'node1'
         AUTHENTICATION_APIKEY_ENABLED: 'true'
         AUTHENTICATION_APIKEY_ALLOWED_KEYS: 'readonlykey,adminkey'
         AUTHENTICATION_APIKEY_USERS: 'jane@doe.com,john@doe.com'
         TRANSFORMERS_INFERENCE_API: 'http://t2v-transformers:8080'
       networks:
         - airflow
     t2v-transformers:
       image: semitechnologies/transformers-inference:sentence-transformers-multi-qa-MiniLM-L6-cos-v1
       environment:
         ENABLE_CUDA: 0 # set to 1 to enable
       ports:
         - 8082:8080
       networks:
         - airflow

4. To create an Airflow connection to the local Weaviate instance, add the following environment variable to your .env file. You only need to provide an X-OpenAI-Api-Key if you plan on using the OpenAI API for vectorization.

   AIRFLOW_CONN_WEAVIATE_DEFAULT='{
       "conn_type": "weaviate",
       "host": "http://weaviate:8081/",
       "extra": {
           "token": "adminkey",
           "additional_headers": {
               "X-OpenAI-Api-Key": "YOUR OPEN API KEY" # optional
           }
       }
   }'

tip

See the Weaviate documentation on environment variables, modules, and client instantiation for more information on configuring a Weaviate instance and connection.

Step 2: Add your data

The DAG in this tutorial runs a query on vectorized movie descriptions from IMDB. If you run the project locally, Astronomer recommends testing the pipeline with a small subset of the data. If you use a remote vectorizer like text2vec-openai, you can use larger parts of the full dataset.

Create a new file called movie_data.txt in the include directory, then copy and paste the following information:

1 ::: Arrival (2016) ::: sci-fi ::: A linguist works with the military to communicate with alien lifeforms after twelve mysterious spacecraft appear around the world.
2 ::: Don't Look Up (2021) ::: drama ::: Two low-level astronomers must go on a giant media tour to warn humankind of an approaching comet that will destroy planet Earth.
3 ::: Primer (2004) ::: sci-fi ::: Four friends/fledgling entrepreneurs, knowing that there's something bigger and more innovative than the different error-checking devices they've built, wrestle over their new invention.
4 ::: Serenity (2005) ::: sci-fi ::: The crew of the ship Serenity try to evade an assassin sent to recapture telepath River.
5 ::: Upstream Colour (2013) ::: romance ::: A man and woman are drawn together, entangled in the life cycle of an ageless organism. Identity becomes an illusion as they struggle to assemble the loose fragments of wrecked lives.
6 ::: The Matrix (1999) ::: sci-fi ::: When a beautiful stranger leads computer hacker Neo to a forbidding underworld, he discovers the shocking truth--the life he knows is the elaborate deception of an evil cyber-intelligence.
7 ::: Inception (2010) ::: sci-fi ::: A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O., but his tragic past may doom the project and his team to disaster.

Step 3: Create your DAG

1. In your dags folder, create a file called query_movie_vectors.py.

2. Copy the following code into the file. If you want to use text2vec-openai for vectorization, change the VECTORIZER variable to text2vec-openai and make sure you provide an OpenAI API key in the AIRFLOW_CONN_WEAVIATE_DEFAULT in your .env file.

   """
   ## Use the Airflow Weaviate Provider to generate and query vectors for movie descriptions
   
   This DAG runs a simple MLOps pipeline that uses the Weaviate Provider to import 
   movie descriptions, generate vectors for them, and query the vectors for movies based on
   concept descriptions.
   """
   
   from airflow.decorators import dag, task
   from airflow.models.param import Param
   from airflow.operators.empty import EmptyOperator
   from airflow.models.baseoperator import chain
   from airflow.providers.weaviate.hooks.weaviate import WeaviateHook
   from airflow.providers.weaviate.operators.weaviate import WeaviateIngestOperator
   from weaviate.util import generate_uuid5
   from pendulum import datetime
   import re
   
   
   WEAVIATE_USER_CONN_ID = "weaviate_default"
   TEXT_FILE_PATH = "include/movie_data.txt"
   # the base class name is used to create a unique class name for the vectorizer
   # note that it is best practice to capitalize the first letter of the class name
   CLASS_NAME_BASE = "Movie"
   
   # set the vectorizer to text2vec-openai if you want to use the openai model
   # note that using the OpenAI vectorizer requires a valid API key in the
   # AIRFLOW_CONN_WEAVIATE_DEFAULT connection.
   # If you want to use a different vectorizer
   # (https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules)
   # make sure to also add it to the weaviate configuration's `ENABLE_MODULES` list
   # for example in the docker-compose.override.yml file
   VECTORIZER = "text2vec-transformers"
   
   # the class name is a combination of the base class name and the vectorizer
   CLASS_NAME = CLASS_NAME_BASE + "_" + VECTORIZER.replace("-", "_")
   
   
   @dag(
       start_date=datetime(2023, 9, 1),
       schedule=None,
       catchup=False,
       tags=["weaviate"],
       params={
           "movie_concepts": Param(
               ["innovation", "friends"],
               type="array",
               description=(
                   "What kind of movie do you want to watch today?"
                   + " Add one concept per line."
               ),
           ),
       },
   )
   def query_movie_vectors():
       @task.branch
       def check_for_class(conn_id: str, class_name: str) -> bool:
           "Check if the provided class already exists and decide on the next step."
           hook = WeaviateHook(conn_id)
           client = hook.get_client()
   
           if not client.schema.get()["classes"]:
               print("No classes found in this weaviate instance.")
               return "create_class"
   
           existing_classes_names_with_vectorizer = [
               x["class"] for x in client.schema.get()["classes"]
           ]
   
           if class_name in existing_classes_names_with_vectorizer:
               print(f"Schema for class {class_name} exists.")
               return "class_exists"
           else:
               print(f"Class {class_name} does not exist yet.")
               return "create_class"
   
       @task
       def create_class(conn_id: str, class_name: str, vectorizer: str):
           "Create a class with the provided name and vectorizer."
           weaviate_hook = WeaviateHook(conn_id)
   
           class_obj = {
               "class": class_name,
               "vectorizer": vectorizer,
           }
           weaviate_hook.create_class(class_obj)
   
       class_exists = EmptyOperator(task_id="class_exists")
   
       def import_data_func(text_file_path: str, class_name: str):
           "Read the text file and create a list of dicts for ingestion to Weaviate."
           with open(text_file_path, "r") as f:
               lines = f.readlines()
   
               num_skipped_lines = 0
               data = []
               for line in lines:
                   parts = line.split(":::")
                   title_year = parts[1].strip()
                   match = re.match(r"(.+) \((\d{4})\)", title_year)
                   try:
                       title, year = match.groups()
                       year = int(year)
                   # skip malformed lines
                   except:
                       num_skipped_lines += 1
                       continue
   
                   genre = parts[2].strip()
                   description = parts[3].strip()
   
                   data.append(
                       {
                           "movie_id": generate_uuid5(
                               identifier=[title, year, genre, description],
                               namespace=class_name,
                           ),
                           "title": title,
                           "year": year,
                           "genre": genre,
                           "description": description,
                       }
                   )
                   # note that to import pre-computed vectors you would add them to the dict
                   # with the key `Vector`
   
               print(
                   f"Created a list with {len(data)} elements while skipping {num_skipped_lines} lines."
               )
               return data
   
       import_data = WeaviateIngestOperator(
           task_id="import_data",
           conn_id=WEAVIATE_USER_CONN_ID,
           class_name=CLASS_NAME,
           input_json=import_data_func(
               text_file_path=TEXT_FILE_PATH, class_name=CLASS_NAME
           ),
           trigger_rule="none_failed",
       )
   
       @task
       def query_embeddings(weaviate_conn_id: str, class_name: str, **context):
           "Query the Weaviate instance for movies based on the provided concepts."
           hook = WeaviateHook(weaviate_conn_id)
           movie_concepts = context["params"]["movie_concepts"]
   
           movie = hook.query_without_vector(
               " ".join(movie_concepts),
               class_name,
               [
                   "title",
                   "year",
                   "genre",
                   "description",
               ],
               limit=1,
           )
   
           if movie["data"]["Get"][class_name] == []:
               print("No movies found for this query. Try again.")
               return
   
           movie_title = movie["data"]["Get"][class_name][0]["title"]
           movie_year = movie["data"]["Get"][class_name][0]["year"]
           movie_description = movie["data"]["Get"][class_name][0]["description"]
   
           print(
               "Your movie query was for the following concepts: "
               + " ".join(movie_concepts)
           )
           print(
               "You should watch the following movie(s): "
               + str(movie_title)
               + " ("
               + str(movie_year)
               + ")"
           )
           print("Description: " + str(movie_description))
   
           return movie
   
       chain(
           check_for_class(conn_id=WEAVIATE_USER_CONN_ID, class_name=CLASS_NAME),
           [
               create_class(
                   conn_id=WEAVIATE_USER_CONN_ID,
                   class_name=CLASS_NAME,
                   vectorizer=VECTORIZER,
               ),
               class_exists,
           ],
           import_data,
           query_embeddings(weaviate_conn_id=WEAVIATE_USER_CONN_ID, class_name=CLASS_NAME),
       )
   
   
   query_movie_vectors()

   This DAG consists of five tasks to make a simple ML orchestration pipeline.

   • The check_for_class task uses the WeaviateHook to check if a class of the name CLASS_NAME already exists in your Weaviate instance. The task is defined using the @task.branch decorator and returns the the id of the task to run next based on whether the class of interest exists. If the class exists, the DAG runs the empty class_exists task. If the class does not exist, the DAG runs the create_class task.
   • The create_class task uses the WeaviateHook to create a class with the CLASS_NAME and specified VECTORIZER in your Weaviate instance.
   • The import_data task is defined using the WeaviateIngestOperator and ingests the data into Weaviate. You can run any Python code on the data before ingesting it into Weaviate by providing a callable to the input_json parameter. This makes it possible to create your own embeddings or complete other transformations before ingesting the data. In this example we use automatic schema inference and vector creation by Weaviate.
   • The query_embeddings task uses the WeaviateHook to connect to the Weaviate instance and run a query. The query returns the most similar movie to the concepts provided by the user when running the DAG in the next step.

Step 4: Run your DAG

1. Run astro dev start in your Astro project to start Airflow and open the Airflow UI at localhost:8080.

2. In the Airflow UI, run the query_movie_vectors DAG by clicking the play button. Then, provide Airflow params for movie_concepts.

   Note that if you are running the project locally on a larger dataset, the import_data task might take a longer time to complete because Weaviate generates the vector embeddings in this task.

   

3. View your movie suggestion in the task logs of the query_embeddings task:

   [2023-11-13, 13:29:56 UTC] {logging_mixin.py:154} INFO - Your movie query was for the following concepts: innovation friends
   [2023-11-13, 13:29:56 UTC] {logging_mixin.py:154} INFO - You should watch the following movie(s): Primer (2004)
   [2023-11-13, 13:29:56 UTC] {logging_mixin.py:154} INFO - Description: Four friends/fledgling entrepreneurs, knowing that there's something bigger and more innovative than the different error-checking devices they've built, wrestle over their new invention.

Conclusion

Congratulations! You used Airflow and Weaviate to get your next movie suggestion!