Add monitoring for APIs

airflow/Dockerfile

@@ -1,5 +1,5 @@
 # syntax=quay.io/astronomer/airflow-extensions:latest
 
-FROM quay.io/astronomer/astro-runtime:9.3.0-base
+FROM quay.io/astronomer/astro-runtime:9.4.0-base
 
 COPY include/airflow_provider_weaviate-0.0.1-py3-none-any.whl /tmp

airflow/dags/ingestion/ask-astro-load-blogs.py

@@ -9,6 +9,7 @@
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
+WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsProd")
 
 blog_cutoff_date = datetime.strptime("2023-01-19", "%Y-%m-%d")
 
@@ -30,7 +31,7 @@ def ask_astro_load_blogs():
         weaviate_conn_id=_WEAVIATE_CONN_ID,
         retries=10,
         retry_delay=30,
-    ).partial(class_name="Docs", primary_key="docLink").expand(dfs=[split_md_docs])
+    ).partial(class_name=WEAVIATE_CLASS, primary_key="docLink").expand(dfs=[split_md_docs])
 
 
 ask_astro_load_blogs()

airflow/dags/ingestion/ask-astro-load-github.py

@@ -10,7 +10,7 @@
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 _GITHUB_CONN_ID = "github_ro"
-
+WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsProd")
 markdown_docs_sources = [
     {"doc_dir": "learn", "repo_base": "astronomer/docs"},
     {"doc_dir": "astro", "repo_base": "astronomer/docs"},
@@ -71,7 +71,7 @@ def ask_astro_load_github():
         weaviate_conn_id=_WEAVIATE_CONN_ID,
         retries=10,
         retry_delay=30,
-    ).partial(class_name="Docs", primary_key="docLink").expand(dfs=[split_md_docs, split_code_docs])
+    ).partial(class_name=WEAVIATE_CLASS, primary_key="docLink").expand(dfs=[split_md_docs, split_code_docs])
 
     issues_docs >> md_docs >> rst_docs >> code_samples
 

airflow/dags/ingestion/ask-astro-load-registry.py

@@ -9,6 +9,7 @@
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
+WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsProd")
 
 
 @dag(schedule_interval="0 5 * * *", start_date=datetime(2023, 9, 27), catchup=False, is_paused_upon_creation=True)
@@ -32,7 +33,7 @@ def ask_astro_load_registry():
         weaviate_conn_id=_WEAVIATE_CONN_ID,
         retries=10,
         retry_delay=30,
-    ).partial(class_name="Docs", primary_key="docLink").expand(dfs=[split_md_docs, split_code_docs])
+    ).partial(class_name=WEAVIATE_CLASS, primary_key="docLink").expand(dfs=[split_md_docs, split_code_docs])
 
 
 ask_astro_load_registry()

airflow/dags/ingestion/ask-astro-load-slack.py

@@ -9,7 +9,7 @@
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
-
+WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsProd")
 slack_channel_sources = [
     {
         "channel_name": "troubleshooting",
@@ -38,7 +38,7 @@ def ask_astro_load_slack():
         weaviate_conn_id=_WEAVIATE_CONN_ID,
         retries=10,
         retry_delay=30,
-    ).partial(class_name="Docs", primary_key="docLink").expand(dfs=[split_md_docs])
+    ).partial(class_name=WEAVIATE_CLASS, primary_key="docLink").expand(dfs=[split_md_docs])
 
 
 ask_astro_load_slack()

airflow/dags/ingestion/ask-astro-load-stackoverflow.py

@@ -9,7 +9,7 @@
 ask_astro_env = os.environ.get("ASK_ASTRO_ENV", "")
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
-
+WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsProd")
 stackoverflow_cutoff_date = "2021-09-01"
 
 stackoverflow_tags = [
@@ -38,7 +38,7 @@ def ask_astro_load_stackoverflow():
         weaviate_conn_id=_WEAVIATE_CONN_ID,
         retries=10,
         retry_delay=30,
-    ).partial(class_name="Docs", primary_key="docLink").expand(dfs=[split_md_docs])
+    ).partial(class_name=WEAVIATE_CLASS, primary_key="docLink").expand(dfs=[split_md_docs])
 
 
 ask_astro_load_stackoverflow()

airflow/dags/ingestion/ask-astro-load.py

@@ -15,6 +15,7 @@
 
 _WEAVIATE_CONN_ID = f"weaviate_{ask_astro_env}"
 _GITHUB_CONN_ID = "github_ro"
+WEAVIATE_CLASS = os.environ.get("WEAVIATE_CLASS", "DocsProd")
 
 markdown_docs_sources = [
     {"doc_dir": "learn", "repo_base": "astronomer/docs"},
@@ -83,7 +84,7 @@ def ask_astro_load_bulk():
         task_id="create_schema",
         weaviate_conn_id=_WEAVIATE_CONN_ID,
         class_object_data="file://include/data/schema.json",
-        existing="fail",
+        existing="ignore",
     )
 
     @task.branch(trigger_rule="none_failed")
@@ -238,12 +239,12 @@ def extract_astro_blogs():
     split_code_docs = task(split.split_python).expand(dfs=python_code_tasks)
 
     task.weaviate_import(ingest.import_data, weaviate_conn_id=_WEAVIATE_CONN_ID, retries=10, retry_delay=30).partial(
-        class_name="Docs"
+        class_name=WEAVIATE_CLASS
     ).expand(dfs=[split_md_docs, split_code_docs])
 
     _import_baseline = task.weaviate_import(
         ingest.import_baseline, trigger_rule="none_failed", weaviate_conn_id=_WEAVIATE_CONN_ID
-    )(class_name="Docs", seed_baseline_url=seed_baseline_url)
+    )(class_name=WEAVIATE_CLASS, seed_baseline_url=seed_baseline_url)
 
     _check_schema >> [_check_seed_baseline, _create_schema]
 

airflow/dags/monitor/monitor.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import json
+import logging
 import os
 import tempfile
 from datetime import datetime
@@ -11,6 +14,9 @@
 from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator
 from airflow.utils.trigger_rule import TriggerRule
 
+# Set up the logger
+logger = logging.getLogger("airflow.task")
+
 monitoring_interval = os.environ.get("MONITORING_INTERVAL", "@daily")
 
 weaviate_conn_id = os.environ.get("WEAVIATE_CONN_ID", WeaviateHook.default_conn_name)
@@ -22,6 +28,47 @@
 
 google_service_account_json_value = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_VALUE", None)
 
+ASK_ASTRO_API_BASE_URL = os.environ.get("ASK_ASTRO_API_BASE_URL", None)
+
+# request ids to monitor from Ask-Astro database.
+ASK_ASTRO_REQUEST_ID_1 = os.environ.get("ASK_ASTRO_REQUEST_ID_1", "05d8882e-56ac-11ee-a818-4200a9fe0102")
+ASK_ASTRO_REQUEST_ID_2 = os.environ.get("ASK_ASTRO_REQUEST_ID_2", "f2d6524c-56ab-11ee-a818-4200a9fe0102")
+
+
+class APIMonitoring:
+    """
+    A class to test API endpoints using GET or POST requests and ensure that they return HTTP 200 status codes.
+    """
+
+    def __init__(self, base_url: str) -> None:
+        self.base_url = base_url
+
+    def test_endpoint(
+        self, endpoint: str, method: str = "GET", data: dict | None = None, headers: dict | None = None
+    ) -> int:
+        """
+        Test an endpoint with the specified method, data, and headers.
+
+        :param endpoint: The endpoint to test.
+        :param method: The HTTP method to use. Defaults to 'GET'.
+        :param data: The data to send in the request. Defaults to None.
+        :param headers: The headers to send in the request. Defaults to None.
+        """
+        try:
+            url = f"{self.base_url}{endpoint}"
+            if method.upper() == "GET":
+                response = requests.get(url, headers=headers)
+            elif method.upper() == "POST":
+                response = requests.post(url, json=data, headers=headers)
+            else:
+                logger.error(f"Unsupported method: {method}")
+                raise Exception(f"Unsupported method: {method}")
+            response.raise_for_status()
+            logger.info(f"Success: {method} {endpoint} returned status code {response.status_code}")
+
+        except requests.exceptions.HTTPError as err:
+            logger.error(f"Failed: {method} {endpoint} returned status code {err.response.status_code}")
+
 
 @task(trigger_rule=TriggerRule.ALL_DONE)
 def slack_status(**context):
@@ -71,7 +118,8 @@ def check_ui_status():
     endpoint = "https://ask.astronomer.io"
     response = requests.get(endpoint)
     if response.status_code != 200:
-        raise response
+        raise Exception(f"UI check failed with status code: {response.status_code}")
+    logger.info(f"UI check passed with status code: {response.status_code}")
 
 
 @task(trigger_rule=TriggerRule.ALL_DONE)
@@ -88,9 +136,9 @@ def check_weaviate_status():
             count = metadata.get("count")
             break
     if count == 0:
-        print(f"Weavaite class {weaviate_class} is empty!")
+        logger.error(f"Weaviate class {weaviate_class} is empty!")
     else:
-        print(f"{count} record found in Weavaite class {weaviate_class}")
+        logger.info(f"{count} record found in Weaviate class {weaviate_class}")
 
 
 @task(trigger_rule=TriggerRule.ALL_DONE)
@@ -105,14 +153,53 @@ def check_firestore_status():
         firebase_admin.initialize_app()
         app = firebase_admin.get_app(name=firestore_app_name)
 
-        print(f"{app.name} found!")
+        logger.info(f"{app.name} found!")
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def monitor_apis():
+    """
+    Monitor a set of predefined API endpoints using the APIMonitoring class and report on their HTTP status codes.
+    This task will test each API endpoint defined in the endpoints list and will print the status code for each.
+    If any endpoint does not return a 200 status, an exception will be raised.
+    """
+    if not ASK_ASTRO_API_BASE_URL:
+        raise ValueError("base_url cannot be empty.")
+
+    api_monitor = APIMonitoring(base_url=ASK_ASTRO_API_BASE_URL)
+
+    headers = {"accept": "*/*", "Content-Type": "application/json"}
+    request_body = {"prompt": "Example prompt"}
+    feedback_body = {"positive": True}
+
+    endpoints_to_monitor = [
+        ("/slack/install", "GET", headers),
+        ("/requests", "POST", headers, request_body),
+        ("/requests", "GET", headers),
+    ]
+
+    request_uuids = [ASK_ASTRO_REQUEST_ID_1, ASK_ASTRO_REQUEST_ID_2]
+    for request_id in request_uuids:
+        endpoints_to_monitor.append((f"/requests/{request_id}", "GET", headers))
+        endpoints_to_monitor.append((f"/requests/{request_id}/feedback", "POST", headers, feedback_body))
+
+    # Monitor each API endpoint
+    for endpoint_info in endpoints_to_monitor:
+        try:
+            # Unpack the tuple with a default value for 'body'
+            endpoint, method, headers, data = (*endpoint_info, None)[:4]
+            status_code = api_monitor.test_endpoint(endpoint=endpoint, method=method, headers=headers, data=data)
+            logger.info(f"Endpoint {endpoint} returned status code {status_code}")
+        except requests.HTTPError as e:
+            logger.error(f"Endpoint {endpoint_info[0]} failed with status code: {e.response.status_code}")
+            raise Exception(f"Endpoint {endpoint_info[0]} failed with status code: {e.response.status_code}")
 
 
 @dag(
     schedule_interval=monitoring_interval, start_date=datetime(2023, 9, 27), catchup=False, is_paused_upon_creation=True
 )
 def monitoring_dag():
-    [check_ui_status(), check_weaviate_status(), check_firestore_status()] >> slack_status()
+    [check_ui_status(), check_weaviate_status(), check_firestore_status(), monitor_apis()] >> slack_status()
 
 
 monitoring_dag()