An end-to-end LLM reference implementation providing a Q&A interface for Airflow and Astronomer
"""Test the validity of all DAGs. **USED BY DEV PARSE COMMAND DO NOT EDIT**"""
from contextlib import contextmanager
import logging
import os

import pytest

from airflow.models import DagBag, Variable, Connection
from airflow.hooks.base import BaseHook


# The following code patches errors caused by missing OS Variables, Airflow Connections, and Airflow Variables

# =========== MONKEYPATCH BaseHook.get_connection() ===========
def basehook_get_connection_monkeypatch(key: str,*args, **kwargs):
    print(f"Attempted to fetch connection during parse returning an empty Connection object for {key}")
    return Connection(key)


BaseHook.get_connection = basehook_get_connection_monkeypatch
# # =========== /MONKEYPATCH BASEHOOK.GET_CONNECTION() ===========

# =========== MONKEYPATCH OS.GETENV() ===========
def os_getenv_monkeypatch(key: str, *args, default=None, **kwargs):
	print(f"Attempted to fetch os environment variable during parse, returning a mocked value for {key}")
	if key == 'JENKINS_HOME' and default is None: # fix https://github.com/astronomer/astro-cli/issues/601
		return None
	if default:
		return default
	return "NON_DEFAULT_OS_ENV_VALUE"


os.getenv = os_getenv_monkeypatch
# # =========== /MONKEYPATCH OS.GETENV() ===========

# =========== MONKEYPATCH VARIABLE.GET() ===========

class magic_dict(dict):
	def __init__(self,*args,**kwargs):
		self.update(*args,**kwargs)

	def __getitem__(self,key):
		return {}.get(key, 'MOCKED_KEY_VALUE')


def variable_get_monkeypatch(key: str, default_var=None, deserialize_json=False):
	print(f"Attempted to get Variable value during parse, returning a mocked value for {key}")

	if default_var:
		return default_var
	if deserialize_json:
		return magic_dict()
	return "NON_DEFAULT_MOCKED_VARIABLE_VALUE"

Variable.get = variable_get_monkeypatch
# # =========== /MONKEYPATCH VARIABLE.GET() ===========


@contextmanager
def suppress_logging(namespace):
	"""
	Suppress logging within a specific namespace to keep tests "clean" during build
	"""
	logger = logging.getLogger(namespace)
	old_value = logger.disabled
	logger.disabled = True
	try:
		yield
	finally:
		logger.disabled = old_value

def get_import_errors():
	"""
	Generate a tuple for import errors in the dag bag
	"""
	with suppress_logging('airflow') :
		dag_bag = DagBag(include_examples=False)

		def strip_path_prefix(path):
			return os.path.relpath(path ,os.environ.get('AIRFLOW_HOME'))


		# we prepend "(None,None)" to ensure that a test object is always created even if its a no op.
		return [(None,None)] +[ ( strip_path_prefix(k) , v.strip() ) for k,v in dag_bag.import_errors.items()]


@pytest.mark.parametrize("rel_path,rv", get_import_errors(), ids=[x[0] for x in get_import_errors()])
def test_file_imports(rel_path,rv):
	""" Test for import errors on a file """
	if rel_path and rv : #Make sure our no op test doesn't raise an error
		raise Exception(f"{rel_path} failed to import with message \n {rv}")
airflow/.astro/test_dag_integrity_default.py
astronomer/ask-astro
"""
This DAG is responsible for pulling requests from Firestore and passing them
through LangSmith's evaluators. If something is correct, useful, and public,
it marks it as an example to be shown to users.
"""

import os
from typing import Any

from datetime import datetime

from airflow.decorators import dag, task
from airflow.operators.empty import EmptyOperator

from langsmith import Client

from airflow.providers.google.firebase.hooks.firestore import CloudFirestoreHook
from google.cloud import firestore
from google.oauth2 import service_account

from langchain.evaluation import load_evaluator, StringEvaluator
from langchain.evaluation.schema import EvaluatorType


def get_firestore_client():
    """
    This function returns a Firestore client.
    """

    return firestore.Client(
        project="astronomer-success",
    )


@task
def get_unprocessed_runs():
    """
    This task pulls runs from Firestore that haven't been processed yet.
    """
    firestore_client = get_firestore_client()

    # get all requests that don't have a `is_example` field
    unprocessed_runs = (
        firestore_client.collection("ask-astro-dev-requests")
        .where("is_processed", "==", False)
        .get()
    )

    return [obj.to_dict() for obj in unprocessed_runs]


@task
def process_run(run: dict[str, Any]):
    """
    This task processes a run by passing it through LangSmith's evaluators.
    """
    print("Processing run", run)

    prompt = run["prompt"]
    run_id = run["langchain_run_id"]
    response = run["response"]
    sources = run["sources"]

    langsmith_client = Client()
    firestore_client = get_firestore_client()

    feedback = {}
    for criteria in [
        "helpfulness",
        {
            "publicness": "Does the prompt or answer contain only public (non-private) info? If yes, return Y. If no, return N."
        },
        {
            "on-topicness": "Is the prompt or answer related to Apache Airflow, Astronomer, or Data Engineering? If yes, return Y. If no, return N."
        },
    ]:
        evaluator = load_evaluator(
            EvaluatorType.CRITERIA,
            criteria=criteria,
        )

        if not isinstance(evaluator, StringEvaluator):
            raise ValueError("Evaluator must be a StringEvaluator")

        result = evaluator.evaluate_strings(
            prediction=response,
            input=prompt,
            reference=sources,
        )

        print("Evaluated for criteria", criteria)
        print(result)

        if "score" in result:
            if isinstance(criteria, str):
                key = criteria
            elif isinstance(criteria, dict):
                key = list(criteria.keys())[0]
            else:
                key = "unknown"

            langsmith_client.create_feedback(
                run_id=run_id,
                key=key,
                score=result["score"],
                comment=result["reasoning"],
            )
            feedback[key] = result["score"]

    # if all the evaluators agree that the response is correct, useful, and public,
    # mark it as an example and processed
    update_dict = {"is_processed": True}
    if all([score > 0.9 for score in feedback.values()]):
        update_dict["is_example"] = True
        print("Marking run as example")

    firestore_client.collection("ask-astro-dev-requests").document(run["uuid"]).set(
        update_dict, merge=True
    )

    return feedback


@dag(
    schedule="@daily",
    start_date=datetime(2023, 1, 1),
    catchup=False,
)
def find_example_runs():
    begin = EmptyOperator(task_id="begin")
    unprocessed_runs = get_unprocessed_runs()
    end = EmptyOperator(task_id="end")

    begin >> unprocessed_runs >> process_run.expand(run=unprocessed_runs) >> end


find_example_runs()

airflow/dags/feedback/find_example_runs.py
astronomer/ask-astro
from datetime import datetime 
import pandas as pd
import html2text
import requests
from bs4 import BeautifulSoup
import lxml

from typing import List

from airflow.decorators import dag, task
from airflow.exceptions import AirflowException
from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from weaviate_provider.hooks.weaviate import WeaviateHook
from weaviate_provider.operators.weaviate import (
    WeaviateCheckSchemaOperator,
    WeaviateRetrieveAllOperator,
)
from weaviate.util import generate_uuid5
from langchain.text_splitter import (
    MarkdownHeaderTextSplitter, 
    RecursiveCharacterTextSplitter
)
from langchain.schema import Document

_WEAVIATE_CONN_ID = 'weaviate_test'
_SLACK_CONN_ID = 'slack_api_default'

default_args = {
    "retries": 3,
    }

# @dag(schedule_interval="0 5 * * *", start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
@dag(schedule_interval=None, start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
def ask_astro_load_blogs():
    """
    This DAG performs incremental load for any data sources that have changed.  Initial load via 
    ask_astro_load_bulk imported data from a point-in-time data capture.

    This DAG checks to make sure the latest schema exists.  If it does not exist a slack message
    is sent to notify admins.
    """

    def remove_existing_objects(loaded_docs_file_path:str, new_df:pd.DataFrame, class_name:str):
        """
        Helper function to check if existing content needs to 
        be deleted before update. 

        Existing objects (based on 'docLink') with differing uuid or sha 
        will be deleted.

        Returned df includes only the objects that need to be (re)imported.
        """

        weaviate_hook = WeaviateHook(_WEAVIATE_CONN_ID)
        weaviate_hook.client = weaviate_hook.get_conn()

        current_objects_df = pd.read_parquet(loaded_docs_file_path)\
                               .drop(['vector'], axis=1)\
                               .groupby('docLink').agg(set)

        update_objects_df = new_df.groupby('docLink').agg(set)\
                              .join(current_objects_df, rsuffix='_old')
        update_objects_df = update_objects_df[update_objects_df['sha'] != update_objects_df['sha_old']]

        #remove existing objects
        update_objects_df['id'].dropna()\
            .apply(lambda x: [weaviate_hook.client.data_object.delete(uuid=uuid, 
                                                                      class_name=class_name) 
                                for uuid in list(x)])

        objects_to_import = new_df.merge(update_objects_df.reset_index()['docLink'], on='docLink', how='right')

        return objects_to_import

    _check_schema = WeaviateCheckSchemaOperator(task_id='check_schema', 
                                                weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                class_object_data='file://include/data/schema.json')

    @task.branch(retries=0)
    def alert_schema_branch(schema_exists:bool) -> str:
        """
        Check if schema is no longer valid
        """
        # WeaviateHook(_WEAVIATE_CONN_ID).get_conn().schema.delete_all()
        if schema_exists:
            return [
                "extract_astro_blogs",
                "fetch_loaded_docs"
            ]
        elif not schema_exists:
            return ["slack_schema_alert"]
        else:
            return None

    _slack_schema_alert = SlackAPIPostOperator(task_id='slack_schema_alert', 
                                               channel='#airflow_notices',
                                               retries=0,
                                               slack_conn_id = _SLACK_CONN_ID,
                                               text='ask_astro_load_http897 DAG error.  Schema mismatch.')

    @task(trigger_rule='none_failed')
    def extract_astro_blogs():

        cutoff_date = datetime.strptime('2023-01-19', '%Y-%m-%d')

        base_url = "https://www.astronomer.io"
        page_url = base_url+"/blog/{page}/#archive"
        headers={}
        links=[]
        dates=[]
        page=1
        
        response = requests.get(page_url.format(page=page), headers=headers)
        while response.ok:

            soup = BeautifulSoup(response.text, 'lxml')
            cards = soup.find_all(class_="post-card__cover")
            links.extend([base_url+card.find('a', href=True)['href'] for card in cards])
            meta = soup.find_all(class_="post-card__meta")
            dates.extend([post.find('time')['datetime'] for post in meta])

            page = page + 1
            response = requests.get(page_url.format(page=page), headers=headers)

        df = pd.DataFrame(zip(links, dates), columns=['docLink','date'])
        df['date'] = pd.to_datetime(df['date']).dt.date
        df = df[df['date'] > cutoff_date.date()]
        df.drop('date', inplace=True, axis=1)
        df.drop_duplicates(inplace=True)

        df['content'] = df['docLink'].apply(lambda x: requests.get(x).text)
        df['title'] = df['content'].apply(lambda x: BeautifulSoup(x, 'lxml')\
                                          .find(class_="post-card__meta")\
                                          .find(class_='title').get_text())
        
        df['content'] = df['content'].apply(lambda x: BeautifulSoup(x, 'lxml').find(class_='prose').get_text())
        df['content'] = df.apply(lambda x: x.title + ' ' + html2text.html2text(html=x.content), axis=1)
        df.drop('title', axis=1, inplace=True)
        df['sha'] = df['content'].apply(generate_uuid5)
        df['docSource'] = 'astro blog'
        df.reset_index(drop=True, inplace=True)

        blog_dfs = df[['docSource', 'sha', 'content', 'docLink']]

        return blog_dfs
    
    @task(trigger_rule='none_failed')
    def split_data(blog_dfs:List[pd.DataFrame]):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        splits markdown content on markdown headers.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', 'docs', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Chunked content in markdown format.

        """

        df = pd.concat([blog_dfs], axis=0, ignore_index=True)

        splitter = RecursiveCharacterTextSplitter()
        df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_documents([Document(page_content=x)]))

        df = df[df['doc_chunks'].apply(lambda x: len(x))>0].reset_index(drop=True)
        df = df.explode('doc_chunks', ignore_index=True)
        df['content'] = df['doc_chunks'].apply(lambda x: html2text.html2text(x.page_content).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))

        df.drop(['doc_chunks'], inplace=True, axis=1)
        df.reset_index(inplace=True, drop=True)

        return df

    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_data(blog_docs:pd.DataFrame, class_name:str, loaded_docs_file_path:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata (the git sha, document url,  
        the document source (ie. astro) and a concatenation of the headers).

        Any existing documents with the same docLink but differing UUID or sha will be 
        deleted prior to import.

        Vectorization includes the headers for bm25 search.
        """
        df = pd.concat([blog_docs], ignore_index=True)

        df['uuid'] = df.apply(lambda x: generate_uuid5(x.to_dict()), axis=1)

        df = remove_existing_objects(loaded_docs_file_path=loaded_docs_file_path, 
                                     new_df=df, 
                                     class_name=class_name)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "error_threshold": 10}
    
    _alert_schema_branch = alert_schema_branch(_check_schema.output)

    @task
    def fail_schema():
        raise AirflowException('Failing DAG for schema failure.')
    
    _fail_schema = fail_schema()

    blog_docs = extract_astro_blogs()
    
    split_md_docs = split_data(blog_dfs=blog_docs)

    _loaded_docs = WeaviateRetrieveAllOperator(task_id='fetch_loaded_docs', 
                                               weaviate_conn_id=_WEAVIATE_CONN_ID,
                                               trigger_rule='none_failed',
                                               class_name='Docs', 
                                               replace_existing=True,
                                               output_file='file://include/data/loaded_docs.parquet')
    
    _unimported_md = import_data(blog_docs=split_md_docs,
                                 class_name='Docs', 
                                 loaded_docs_file_path=_loaded_docs.output)

    _check_schema >> \
        _alert_schema_branch >> \
            [_slack_schema_alert, blog_docs]
    
    _loaded_docs >> _unimported_md
    _slack_schema_alert >> _fail_schema

ask_astro_load_blogs()
airflow/dags/ingestion/ask-astro-load-blogs.py
astronomer/ask-astro
from datetime import datetime 
from pathlib import Path
import pandas as pd
import pypandoc
import html2text
import re

from typing import List

from airflow.decorators import dag, task
from airflow.exceptions import AirflowException
from airflow.providers.github.hooks.github import GithubHook
from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from weaviate_provider.hooks.weaviate import WeaviateHook
from weaviate_provider.operators.weaviate import (
    WeaviateCheckSchemaOperator,
    WeaviateRetrieveAllOperator,
)
from weaviate.util import generate_uuid5
from langchain.text_splitter import (
    MarkdownHeaderTextSplitter, 
    RecursiveCharacterTextSplitter
)
from langchain.schema import Document

_WEAVIATE_CONN_ID = 'weaviate_test'
_GITHUB_CONN_ID = 'github_default'
_SLACK_CONN_ID = 'slack_api_default'

markdown_docs_sources = [
    {'doc_dir': 'learn', 'repo_base': 'astronomer/docs'}, 
    {'doc_dir': 'astro', 'repo_base': 'astronomer/docs'}
    ]
rst_docs_sources = [
    {'doc_dir': 'docs', 'repo_base': 'apache/airflow'},
    ]
code_samples_sources = [
    {'doc_dir': 'code-samples', 'repo_base': 'astronomer/docs'},
    ]
issues_docs_sources = [
    {'doc_dir': 'issues', 'repo_base': 'apache/airflow'}
]

rst_exclude_docs = ['changelog.rst', 'commits.rst']

default_args = {
    "retries": 3,
    }

# @dag(schedule_interval="0 5 * * *", start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
@dag(schedule_interval=None, start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
def ask_astro_load_github():
    """
    This DAG performs incremental load for any data sources that have changed.  Initial load via 
    ask_astro_load_bulk imported data from a point-in-time data capture.

    This DAG checks to make sure the latest schema exists.  If it does not exist a slack message
    is sent to notify admins.
    """

    def remove_existing_objects(loaded_docs_file_path:str, new_df:pd.DataFrame, class_name:str):
        """
        Helper function to check if existing content needs to 
        be deleted before update. 

        Existing objects (based on 'docLink') with differing uuid or sha 
        will be deleted.

        Returned df includes only the objects that need to be (re)imported.
        """

        weaviate_hook = WeaviateHook(_WEAVIATE_CONN_ID)
        weaviate_hook.client = weaviate_hook.get_conn()

        current_objects_df = pd.read_parquet(loaded_docs_file_path)\
                               .drop(['vector'], axis=1)\
                               .groupby('docLink').agg(set)

        update_objects_df = new_df.groupby('docLink').agg(set)\
                              .join(current_objects_df, rsuffix='_old')
        update_objects_df = update_objects_df[update_objects_df['sha'] != update_objects_df['sha_old']]

        #remove existing objects
        update_objects_df['id'].dropna()\
            .apply(lambda x: [weaviate_hook.client.data_object.delete(uuid=uuid, 
                                                                      class_name=class_name) 
                                for uuid in list(x)])

        objects_to_import = new_df.merge(update_objects_df.reset_index()['docLink'], on='docLink', how='right')

        return objects_to_import

    _check_schema = WeaviateCheckSchemaOperator(task_id='check_schema', 
                                                weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                class_object_data='file://include/data/schema.json')

    @task.branch(retries=0)
    def alert_schema_branch(schema_exists:bool) -> str:
        """
        Check if schema is no longer valid
        """
        # WeaviateHook(_WEAVIATE_CONN_ID).get_conn().schema.delete_all()
        if schema_exists:
            return [
                "extract_github_markdown", 
                "extract_github_rst", 
                "extract_github_python", 
                "extract_github_issues",
            ]
        elif not schema_exists:
            return ["slack_schema_alert"]
        else:
            return None

    _slack_schema_alert = SlackAPIPostOperator(task_id='slack_schema_alert', 
                                               channel='#airflow_notices',
                                               retries=0,
                                               slack_conn_id = _SLACK_CONN_ID,
                                               text='ask_astro_load_incremental DAG error.  Schema mismatch.')
       
    @task(trigger_rule='none_failed')
    def extract_github_markdown(source:dict):
        """
        This task downloads github content as markdown documents in a 
        pandas dataframe.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Entire document content in markdown format.
        """

        downloaded_docs = []
        
        gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        repo = gh_hook.client.get_repo(source['repo_base'])
        contents = repo.get_contents(source['doc_dir'])

        while contents:

            file_content = contents.pop(0)
            if file_content.type == "dir":
                contents.extend(repo.get_contents(file_content.path))

            elif Path(file_content.name).suffix == '.md':

                print(file_content.name)
                
                row = {
                    "docLink": file_content.html_url, 
                    "sha": file_content.sha,
                    "content": file_content.decoded_content.decode(),
                    "docSource": source['doc_dir'], 
                }

                downloaded_docs.append(row)
                
        df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task(trigger_rule='none_failed')
    def extract_github_rst(source:dict):
        """
        This task downloads github content as rst documents 
        in a pandas dataframe.

        The 'content' field is converted from RST to Markdown (via pypandoc).  After 
        removing the preamble (apache license), any empty lines and 'include' footers 
        any empty docs are removed.  Document links and references are not included 
        in the content.

        Dataframe fields are:
        'docSource': ie. 'docs'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Entire document in markdown format.
        """

        downloaded_docs = []
        
        gh_hook = GithubHook(_GITHUB_CONN_ID)

        repo = gh_hook.client.get_repo(source['repo_base'])
        contents = repo.get_contents(source['doc_dir'])

        apache_license_text = Path('include/data/apache_license.rst').read_text()

        while contents:

            file_content = contents.pop(0)
            if file_content.type == "dir":
                contents.extend(repo.get_contents(file_content.path))

            elif Path(file_content.name).suffix == '.rst' and file_content.name not in rst_exclude_docs:

                print(file_content.name)

                row = {
                    "docLink": file_content.html_url, 
                    "sha": file_content.sha,
                    "content": file_content.decoded_content.decode(),
                    "docSource": source['doc_dir'], 
                }

                downloaded_docs.append(row)
                
        df = pd.DataFrame(downloaded_docs)

        df['content'] = df['content'].apply(lambda x: x.replace(apache_license_text, ''))
        df['content'] = df['content'].apply(lambda x: re.sub(r".*include.*", "", x))
        df['content'] = df['content'].apply(lambda x: re.sub(r'^\s*$', "", x))
        df = df[df['content']!='']
        df['content'] = df['content'].apply(lambda x: pypandoc.convert_text(source=x, to='md', 
                                                                            format='rst',
                                                                            extra_args=['--atx-headers']))

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task(trigger_rule='none_failed')
    def extract_github_python(source:dict):
        """
        This task downloads github content as python code in a pandas dataframe.

        The 'content' field of the dataframe is currently not split as the context 
        window is large enough. Code for splitting is provided but commented out.

        Dataframe fields are:
        'docSource': ie. 'code-samples'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': The python code
        'header': a placeholder of 'python' for bm25 search
        """
    
        downloaded_docs = []

        gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        repo = gh_hook.client.get_repo(source['repo_base'])
        contents = repo.get_contents(source['doc_dir'])

        while contents:
            file_content = contents.pop(0)

            if file_content.type == "dir":
                contents.extend(repo.get_contents(file_content.path))

            elif Path(file_content.name).suffix == '.py':
                print(file_content.name)
                                
                row = {
                    "docLink": file_content.html_url, 
                    "sha": file_content.sha,
                    "content": file_content.decoded_content.decode(),
                    "docSource": source['doc_dir'], 
                    "header": 'python', 
                }
                
                downloaded_docs.append(row)

        df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df
    
    @task(trigger_rule='none_failed')
    def extract_github_issues(source:dict):
        """
        This task downloads github issues as markdown documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': repo name + 'issues'
        'docLink': URL for the specific question/answer.
        'content': The base64 encoded content of the question/answer in markdown format.
        'header': document type. (ie. 'question' or 'answer')
        """
        
        gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        repo = gh_hook.client.get_repo(source['repo_base'])
        issues = repo.get_issues()

        issue_autoresponse_text = 'Thanks for opening your first issue here!'
        pr_autoresponse_text = 'Congratulations on your first Pull Request and welcome to the Apache Airflow community!'
        drop_content = [issue_autoresponse_text, pr_autoresponse_text]

        issue_markdown_template = "## ISSUE TITLE: {title}\nDATE: {date}\nBY: {user}\nSTATE: {state}\n{body}\n{comments}"
        comment_markdown_template = "#### COMMENT: {user} on {date}\n{body}\n"

        downloaded_docs = []
        page_num = 0
        
        page = issues.get_page(page_num)
        
        while page:

            for issue in page:
                print(issue.number)
                comments=[]
                for comment in issue.get_comments():
                    #TODO: this is very slow.  Look for vectorized approach.
                    if not any(substring in comment.body for substring in drop_content):
                        comments.append(comment_markdown_template.format(user=comment.user.login, 
                                                                         date=issue.created_at.strftime("%m-%d-%Y"), 
                                                                         body=comment.body))
                downloaded_docs.append({
                    "docLink": issue.html_url, 
                    "sha": '',
                    "content": issue_markdown_template.format(title=issue.title, 
                                                              date=issue.created_at.strftime("%m-%d-%Y"), 
                                                              user=issue.user.login,
                                                              state=issue.state, 
                                                              body=issue.body,
                                                              comments='\n'.join(comments)),
                    "docSource": f"{source['repo_base']} {source['doc_dir']}", 
                    "header": f"{source['repo_base']} issue", 
                })
            page_num=page_num+1
            page = issues.get_page(page_num)
                
        df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df['sha'] = df.apply(generate_uuid5, axis=1)

        return df
    
    @task(trigger_rule='none_failed')
    def split_data(md_dfs:List[pd.DataFrame], 
                   rst_dfs:List[pd.DataFrame], 
                   code_dfs:List[pd.DataFrame], 
                   issues_dfs:List[pd.DataFrame]):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        splits markdown content on markdown headers.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', 'docs', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Chunked content in markdown format.

        """

        md_df = pd.concat(md_dfs, axis=0, ignore_index=True)
        rst_df = pd.concat(rst_dfs, axis=0, ignore_index=True)
        issues_df = pd.concat(issues_dfs, axis=0, ignore_index=True)
        code_df = pd.concat(code_dfs, axis=0, ignore_index=True)
        df = pd.concat([md_df, rst_df, issues_df, code_df], axis=0, ignore_index=True)

        # headers_to_split_on = [
        #     ("#", "Header 1"),
        #     ("##", "Header 2"),
        #     # ("###", "Header 3"),
        # ]

        # splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
        # df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_text(x))
                
        splitter = RecursiveCharacterTextSplitter()
        df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_documents([Document(page_content=x)]))

        df = df[df['doc_chunks'].apply(lambda x: len(x))>0].reset_index(drop=True)
        # _ = df['doc_chunks'].apply(lambda x: x[0].metadata.update({'Header 1': 'Summary'}) if x[0].metadata == {} else x[0] )
        df = df.explode('doc_chunks', ignore_index=True)
        df['content'] = df['doc_chunks'].apply(lambda x: html2text.html2text(x.page_content).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))
        # df['header'] = df['doc_chunks'].apply(lambda x: '. '.join(list(x.metadata.values())))

        df.drop(['doc_chunks'], inplace=True, axis=1)
        df.reset_index(inplace=True, drop=True)

        return df

    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_data(md_df:pd.DataFrame, class_name:str, loaded_docs_file_path:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata (the git sha, document url,  
        the document source (ie. astro) and a concatenation of the headers).

        Any existing documents with the same docLink but differing UUID or sha will be 
        deleted prior to import.

        Vectorization includes the headers for bm25 search.
        """

        df = pd.concat([md_df], ignore_index=True)

        df['uuid'] = df.apply(lambda x: generate_uuid5(x.to_dict()), axis=1)

        df = remove_existing_objects(loaded_docs_file_path=loaded_docs_file_path, new_df=df, class_name=class_name)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "error_threshold": 10}
        
    _alert_schema_branch = alert_schema_branch(_check_schema.output)

    @task
    def fail_schema():
        raise AirflowException('Failing DAG for schema failure.')
    
    _fail_schema = fail_schema()

    md_docs = extract_github_markdown.partial().expand(source=markdown_docs_sources)
    rst_docs = extract_github_rst.partial().expand(source=rst_docs_sources)
    issues_docs = extract_github_issues.partial().expand(source=issues_docs_sources)
    code_samples = extract_github_python.partial().expand(source=code_samples_sources)
    
    split_md_docs = split_data(md_dfs=md_docs, 
                               rst_dfs=rst_docs, 
                               code_dfs=code_samples,
                               issues_dfs=issues_docs)

    _loaded_docs = WeaviateRetrieveAllOperator(task_id='fetch_loaded_docs', 
                                               weaviate_conn_id=_WEAVIATE_CONN_ID,
                                               trigger_rule='none_failed',
                                               class_name='Docs', 
                                               replace_existing=True,
                                               output_file='file://include/data/loaded_docs.parquet')
    
    _unimported_data = import_data(md_df=split_md_docs, 
                                 class_name='Docs', 
                                 loaded_docs_file_path=_loaded_docs.output)
    
    _check_schema >> \
        _alert_schema_branch >> \
            [_slack_schema_alert, md_docs, rst_docs, code_samples, issues_docs]
    
    _loaded_docs >> _unimported_data
    _slack_schema_alert >> _fail_schema

ask_astro_load_github()
airflow/dags/ingestion/ask-astro-load-github.py
astronomer/ask-astro
from datetime import datetime 
from stackapi import StackAPI
from pathlib import Path
import pandas as pd
import pypandoc
import html2text
import re
import requests

from typing import List

from airflow.decorators import dag, task
from airflow.exceptions import AirflowException
from airflow.providers.github.hooks.github import GithubHook
from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from airflow.providers.slack.hooks.slack import SlackHook
from weaviate_provider.hooks.weaviate import WeaviateHook
from weaviate_provider.operators.weaviate import (
    WeaviateCheckSchemaOperator,
    WeaviateRetrieveAllOperator,
)
from weaviate.util import generate_uuid5
from langchain.text_splitter import (
    MarkdownHeaderTextSplitter, 
    RecursiveCharacterTextSplitter
)

_WEAVIATE_CONN_ID = 'weaviate_default' #'weaviate_wcs' #
_GITHUB_CONN_ID = 'github_default'
_SLACK_CONN_ID = 'slack_api_default'

slack_archive_hostname = 'http://apache-airflow.slack-archives.org/'

markdown_docs_sources = [
    {'doc_dir': 'learn', 'repo_base': 'astronomer/docs'}, 
    {'doc_dir': 'astro', 'repo_base': 'astronomer/docs'}
    ]
rst_docs_sources = [
    {'doc_dir': 'docs', 'repo_base': 'apache/airflow'},
    ]
code_samples_sources = [
    {'doc_dir': 'code-samples', 'repo_base': 'astronomer/docs'},
    ]
issues_docs_sources = [
    {'doc_dir': 'issues', 'repo_base': 'apache/airflow'}
]
slack_channel_sources = [
    {'channel_name': 'troubleshooting', 
      'channel_id': 'CCQ7EGB1P', 
      'team_id': 'TCQ18L22Z', 
      'team_name' : 'Airflow Slack Community',
      'slack_api_conn_id' : _SLACK_CONN_ID},
]
http_json_sources = [
    {'name': 'registry_cell_types',
     'base_url': 'https://api.astronomer.io/registryV2/v1alpha1/organizations/public/modules?limit=1000',
     'headers': {},
     'count_field': 'totalCount'}
]

rst_exclude_docs = ['changelog.rst', 'commits.rst']

stackoverflow_cutoff_date = '2023-06-04'
stackoverflow_tags = [
    'airflow',
]

default_args = {
    "retries": 3,
    }

@dag(schedule_interval="0 5 * * *", start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
def ask_astro_load_incremental():
    """
    This DAG performs incremental load for any data sources that have changed.  Initial load via 
    ask_astro_load_bulk imported data from a point-in-time data capture specified by 
    'stackoverflow_cutoff_date'.

    The DAG checks to make sure the latest schema exists.  If it does not exist a slack message
    is sent to notify admins.
    """

    def remove_existing_objects(loaded_docs_file_path:str, new_df:pd.DataFrame, class_name:str):
        """
        Helper function to check if existing content needs to 
        be deleted before update. 

        Existing objects (based on 'docLink') with differing uuid or sha 
        will be deleted.

        Returned df includes only the objects that need to be (re)imported.
        """

        weaviate_hook = WeaviateHook(_WEAVIATE_CONN_ID)
        weaviate_hook.client = weaviate_hook.get_conn()

        current_objects_df = pd.read_parquet(loaded_docs_file_path)\
                               .drop(['vector'], axis=1)\
                               .groupby('docLink').agg(set)

        update_objects_df = new_df.groupby('docLink').agg(set)\
                              .join(current_objects_df, rsuffix='_old')
        update_objects_df = update_objects_df[update_objects_df['sha'] != update_objects_df['sha_old']]

        #remove existing objects
        update_objects_df['id'].dropna()\
            .apply(lambda x: [weaviate_hook.client.data_object.delete(uuid=uuid, 
                                                                      class_name=class_name) 
                                for uuid in list(x)])

        objects_to_import = new_df.merge(update_objects_df.reset_index()['docLink'], on='docLink', how='right')

        return objects_to_import

    _check_schema = WeaviateCheckSchemaOperator(task_id='check_schema', 
                                                weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                class_object_data='file://include/data/schema.json')

    @task.branch(retries=0)
    def alert_schema_branch(schema_exists:bool) -> str:
        """
        Check if schema is no longer valid
        """
        # WeaviateHook(_WEAVIATE_CONN_ID).get_conn().schema.delete_all()
        if schema_exists:
            return [
                "extract_github_markdown", 
                "extract_github_rst", 
                "extract_github_python", 
                "extract_stack_overflow",
                "extract_slack",
                "extract_github_issues",
                "extract_registry",
            ]
        elif not schema_exists:
            return ["slack_schema_alert"]
        else:
            return None

    _slack_schema_alert = SlackAPIPostOperator(task_id='slack_schema_alert', 
                                               channel='#airflow_notices',
                                               retries=0,
                                               slack_conn_id = _SLACK_CONN_ID,
                                               text='ask_astro_load_incremental DAG error.  Schema mismatch.')
       
    @task(trigger_rule='none_failed')
    def extract_github_markdown(source:dict):
        """
        This task downloads github content as markdown documents in a 
        pandas dataframe.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Entire document content in markdown format.
        """

        downloaded_docs = []
        
        gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        repo = gh_hook.client.get_repo(source['repo_base'])
        contents = repo.get_contents(source['doc_dir'])

        while contents:

            file_content = contents.pop(0)
            if file_content.type == "dir":
                contents.extend(repo.get_contents(file_content.path))

            elif Path(file_content.name).suffix == '.md':

                print(file_content.name)
                
                row = {
                    "docLink": file_content.html_url, 
                    "sha": file_content.sha,
                    "content": file_content.decoded_content.decode(),
                    "docSource": source['doc_dir'], 
                }

                downloaded_docs.append(row)
                
        df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task(trigger_rule='none_failed')
    def extract_github_rst(source:dict):
        """
        This task downloads github content as rst documents 
        in a pandas dataframe.

        The 'content' field is converted from RST to Markdown (via pypandoc).  After 
        removing the preamble (apache license), any empty lines and 'include' footers 
        any empty docs are removed.  Document links and references are not included 
        in the content.

        Dataframe fields are:
        'docSource': ie. 'docs'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Entire document in markdown format.
        """

        downloaded_docs = []
        
        gh_hook = GithubHook(_GITHUB_CONN_ID)

        repo = gh_hook.client.get_repo(source['repo_base'])
        contents = repo.get_contents(source['doc_dir'])

        apache_license_text = Path('include/data/apache_license.rst').read_text()

        while contents:

            file_content = contents.pop(0)
            if file_content.type == "dir":
                contents.extend(repo.get_contents(file_content.path))

            elif Path(file_content.name).suffix == '.rst' and file_content.name not in rst_exclude_docs:

                print(file_content.name)

                row = {
                    "docLink": file_content.html_url, 
                    "sha": file_content.sha,
                    "content": file_content.decoded_content.decode(),
                    "docSource": source['doc_dir'], 
                }

                downloaded_docs.append(row)
                
        df = pd.DataFrame(downloaded_docs)

        df['content'] = df['content'].apply(lambda x: x.replace(apache_license_text, ''))
        df['content'] = df['content'].apply(lambda x: re.sub(r".*include.*", "", x))
        df['content'] = df['content'].apply(lambda x: re.sub(r'^\s*$', "", x))
        df = df[df['content']!='']
        df['content'] = df['content'].apply(lambda x: pypandoc.convert_text(source=x, to='md', 
                                                                            format='rst',
                                                                            extra_args=['--atx-headers']))

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task(trigger_rule='none_failed')
    def extract_github_python(source:dict):
        """
        This task downloads github content as python code in a pandas dataframe.

        The 'content' field of the dataframe is currently not split as the context 
        window is large enough. Code for splitting is provided but commented out.

        Dataframe fields are:
        'docSource': ie. 'code-samples'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': The python code
        'header': a placeholder of 'python' for bm25 search
        """
    
        downloaded_docs = []

        gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        repo = gh_hook.client.get_repo(source['repo_base'])
        contents = repo.get_contents(source['doc_dir'])

        while contents:
            file_content = contents.pop(0)

            if file_content.type == "dir":
                contents.extend(repo.get_contents(file_content.path))

            elif Path(file_content.name).suffix == '.py':
                print(file_content.name)
                                
                row = {
                    "docLink": file_content.html_url, 
                    "sha": file_content.sha,
                    "content": file_content.decoded_content.decode(),
                    "docSource": source['doc_dir'], 
                    "header": 'python', 
                }
                
                downloaded_docs.append(row)

        df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df
    
    @task(trigger_rule='none_failed')
    def extract_stack_overflow(tag:dict, stackoverflow_cutoff_date:str):
        """
        This task generates stack overflow questions and answers as markdown 
        documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': 'stackoverflow' plus the tag name (ie. 'airflow')
        'docLink': URL for the specific question/answer.
        'content': The base64 encoded content of the question/answer in markdown format.
        'header': document type. (ie. 'question' or 'answer')
        """

        question_template = "TITLE: {title}\nDATE: {date}\nBY: {user}\nSCORE: {score}\n{body}{question_comments}"
        answer_template = "DATE: {date}\nBY: {user}\nSCORE: {score}\n{body}{answer_comments}"
        comment_template = "{user} on {date} [Score: {score}]: {body}\n"
        
        SITE = StackAPI(name='stackoverflow', max_pagesize=100, max_pages=10000000)
        
        fromdate = datetime.strptime(stackoverflow_cutoff_date, "%Y-%m-%d")
        filter = '!-(5KXGCFLp3w9.-7QsAKFqaf5yFPl**9q*_hsHzYGjJGQ6BxnCMvDYijFE'    

        questions_dict = SITE.fetch(endpoint='questions', tagged=tag, fromdate=fromdate, filter=filter)
        items = questions_dict.pop('items')

        #TODO: add backoff logic.  For now just fail the task if we can't fetch all results due to api rate limits.
        assert not questions_dict['has_more']

        questions_df = pd.DataFrame(items)
        questions_df = questions_df[questions_df['answer_count']>=1]
        questions_df = questions_df[questions_df['score']>=1]
        questions_df.reset_index(inplace=True, drop=True)

        answers_df = questions_df.explode('answers')
        
        #consolidate comments for the original question
        questions_df['comments'] = questions_df['comments'].fillna('')
        questions_df['comments'] = questions_df['comments'].apply(lambda x: [comment_template.format(
            user=comment['owner']['user_id'],
            date=datetime.fromtimestamp(comment['creation_date']).strftime("%Y-%m-%d"),
            score=comment['score'],
            body=comment['body_markdown']) for comment in x])
        questions_df['question_comments'] = questions_df['comments'].apply(lambda x: ''.join(x))

        #format question content
        questions_df['question_text'] = questions_df.apply(lambda x: question_template.format(
            title=x.title,
            user=x.owner['user_id'],
            date=datetime.fromtimestamp(x.creation_date).strftime("%Y-%m-%d"),
            score=x.score,
            body=x.body_markdown,
            question_comments=x.question_comments), axis=1)
        questions_df = questions_df[['link', 'question_id', 'question_text']].set_index('question_id')

        questions_df['docSource'] = f'stackoverflow {tag}'
        questions_df = questions_df[['docSource', 'link', 'question_text']]
        questions_df.columns = ['docSource', 'docLink', 'content']
        questions_df['header'] = 'question'


        # consolidate comments for each answer
        answers_df['answer_comments'] = answers_df['answers'].apply(lambda x: [comment_template.format(
            user=comment['owner']['user_id'],
            date=datetime.fromtimestamp(comment['creation_date']).strftime("%Y-%m-%d"),
            score=comment['score'],
            body=comment['body_markdown']) for comment in x.get('comments', '')])
        answers_df['answer_comments'] = answers_df['answer_comments'].apply(lambda x: ''.join(x))

        answers_df['answer_text'] = answers_df[['answers', 'answer_comments']].apply(lambda x: answer_template.format(
            score=x.answers['score'],
            date=datetime.fromtimestamp(x.answers['creation_date']).strftime("%Y-%m-%d"),
            user=x.answers['owner']['user_id'],
            body=x.answers['body_markdown'],
            answer_comments=x.answer_comments), axis=1)
        answers_df = answers_df.groupby('question_id')['answer_text'].apply(lambda x: ''.join(x))

        answers_df = questions_df.join(answers_df).apply(lambda x: pd.Series([
            f'stackoverflow {tag}',
            x.docLink, 
            x.answer_text]), axis=1)
        answers_df.columns=['docSource', 'docLink','content']
        answers_df['header'] = 'answer'

        df = pd.concat([questions_df, answers_df], axis=0).reset_index(drop=True)

        df['content'] = df['content'].apply(lambda x: html2text.html2text(x).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))

        # df.to_parquet('include/data/stackoverflow_incremental.parquet')
        # df = pd.read_parquet('include/data/stackoverflow_incremental.parquet')
        df['sha'] = df.apply(generate_uuid5, axis=1)

        return df

    @task(trigger_rule='none_failed')
    def extract_slack(source:dict):
        """
        This task downloads archived slack messages as documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': slack team and channel names
        'docLink': URL for the specific message/reply
        'content': The message/reply content in markdown format.
        'header': document type. (ie. 'question' or 'answer')

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """
        #### THIS IS A PLACEHOLDER AS WE WAIT FOR SLACK BOT ACCESS TO AIRFLOW ####

        # slack_client = SlackHook(slack_conn_id = _SLACK_CONN_ID).client
        
        # channel_info = slack_client.conversations_info(channel=source['channel_id']).data['channel']
        # assert channel_info['is_member'] or not channel_info['is_private']

        # history = slack_client.conversations_history(channel=source['channel_id'])
        # df = pd.DataFrame(history.data['messages'])

        # #if channel has no replies yet they will not be in the messages
        # if 'thread_ts' not in df:
        #     df['thread_ts'] = ''

        # df = df[['user', 'text', 'ts', 'thread_ts', 'client_msg_id', 'type']].drop_duplicates()
        # df.rename({'client_msg_id': 'sha'}, axis=1, inplace=True)
        # df['docLink'] = df['ts'].apply(lambda x: f"https://app.slack.com/client/{source['team_id']}/{source['channel_id']}/p{str(x).replace('.','')}")
        # df['docSource'] = source['channel_name']
        # df['type'] = df.apply(lambda x: 'question' if x.ts == x.thread_ts else 'answer', axis=1)
        # df['ts'] = df['ts'].apply(lambda x: datetime.fromtimestamp(float(x)))
        # df['content'] = df.apply(lambda x: f"# slack {source['team_name']} {source['channel_name']}\n## {x.type} \n### [{x.ts}] <@{x.user}>\n\n{x.text}", axis=1)
        
        # return df[['docSource', 'sha', 'content', 'docLink']]
        return pd.DataFrame()

    @task(trigger_rule='none_failed')
    def extract_github_issues(source:dict):
        """
        This task downloads github issues as markdown documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': repo name + 'issues'
        'docLink': URL for the specific question/answer.
        'content': The base64 encoded content of the question/answer in markdown format.
        'header': document type. (ie. 'question' or 'answer')
        """
        
        gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        repo = gh_hook.client.get_repo(source['repo_base'])
        issues = repo.get_issues()

        issue_autoresponse_text = 'Thanks for opening your first issue here!'
        pr_autoresponse_text = 'Congratulations on your first Pull Request and welcome to the Apache Airflow community!'
        drop_content = [issue_autoresponse_text, pr_autoresponse_text]

        issue_markdown_template = "## ISSUE TITLE: {title}\nDATE: {date}\nBY: {user}\nSTATE: {state}\n{body}\n{comments}"
        comment_markdown_template = "#### COMMENT: {user} on {date}\n{body}\n"

        downloaded_docs = []
        page_num = 0
        
        page = issues.get_page(page_num)
        
        while page:

            for issue in page:
                print(issue.number)
                comments=[]
                for comment in issue.get_comments():
                    #TODO: this is very slow.  Look for vectorized approach.
                    if not any(substring in comment.body for substring in drop_content):
                        comments.append(comment_markdown_template.format(user=comment.user.login, 
                                                                         date=issue.created_at.strftime("%m-%d-%Y"), 
                                                                         body=comment.body))
                downloaded_docs.append({
                    "docLink": issue.html_url, 
                    "sha": '',
                    "content": issue_markdown_template.format(title=issue.title, 
                                                              date=issue.created_at.strftime("%m-%d-%Y"), 
                                                              user=issue.user.login,
                                                              state=issue.state, 
                                                              body=issue.body,
                                                              comments='\n'.join(comments)),
                    "docSource": f"{source['repo_base']} {source['doc_dir']}", 
                    "header": f"{source['repo_base']} issue", 
                })
            page_num=page_num+1
            page = issues.get_page(page_num)
                
        df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df['sha'] = df.apply(generate_uuid5, axis=1)

        return df

    @task(trigger_rule='none_failed')
    def extract_registry(source:dict):

        data_class = source['base_url'].split('/')[-1].split('?')[0]

        response = requests.get(source['base_url'], headers=source['headers']).json()
        total_count = response[source['count_field']]
        data = response.get(data_class, [])
        while len(data) < total_count-1:
            response = requests.get(f"{source['base_url']}&offset={len(data)+1}").json()
            data.extend(response.get(data_class, []))

        df = pd.DataFrame(data)
        df.rename({'githubUrl': 'docLink', 'searchId': 'sha'}, axis=1, inplace=True)
        df['docSource'] = source['name']
        df['description'] = df['description'].apply(lambda x: html2text.html2text(x) if x else 'No Description')
        df['content'] = df.apply(lambda x: f"# Registry\n## {x.providerName}__{x.version}__{x['name']}\n\n{x.description})", axis=1)

        df = df[['docSource', 'sha', 'content', 'docLink']]

        return df
    
    @task(trigger_rule='none_failed')
    def split_markdown(md_dfs:List[pd.DataFrame], rst_dfs:List[pd.DataFrame], slack_dfs:List[pd.DataFrame], reg_dfs:List[pd.DataFrame]):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        splits markdown content on markdown headers.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', 'docs', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Chunked content in markdown format.

        """

        md_df = pd.concat(md_dfs, axis=0, ignore_index=True)
        rst_df = pd.concat(rst_dfs, axis=0, ignore_index=True)
        slack_df = pd.concat(slack_dfs, axis=0, ignore_index=True)
        reg_df = pd.concat(reg_dfs, axis=0, ignore_index=True)
        df = pd.concat([md_df, rst_df, slack_df, reg_df], axis=0, ignore_index=True)

        headers_to_split_on = [
            ("#", "Header 1"),
            ("##", "Header 2"),
            # ("###", "Header 3"),
        ]

        # splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
        splitter = RecursiveCharacterTextSplitter()
        
        df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_text(x))
        df = df[df['doc_chunks'].apply(lambda x: len(x))>0].reset_index(drop=True)
        _ = df['doc_chunks'].apply(lambda x: x[0].metadata.update({'Header 1': 'Summary'}) if x[0].metadata == {} else x[0] )
        df = df.explode('doc_chunks', ignore_index=True)
        df['content'] = df['doc_chunks'].apply(lambda x: html2text.html2text(x.page_content).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))
        df['header'] = df['doc_chunks'].apply(lambda x: '. '.join(list(x.metadata.values())))

        df.drop(['doc_chunks'], inplace=True, axis=1)
        df.reset_index(inplace=True, drop=True)

        return df

    @task()
    def split_python(python_dfs:List[pd.DataFrame]):
        """
        This concatenates multiple dataframes from upstream dynamic tasks and 
        splits python code in a pandas dataframe.

        This task is a concatenation and passthru. The 'content' field of the dataframe 
        is currently not split as the context window is large enough. Code for splitting
        is provided but commented out.

        Dataframe fields are:
        'docSource': ie. 'code-samples'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': The base64 encoded python code
        'header': a placeholder of 'python' for bm25 search
        """
    
        df = pd.concat(python_dfs, axis=0, ignore_index=True)

        #chunk code and use each chunk as a separate doc for vectorization
        # python_splitter = RecursiveCharacterTextSplitter.from_language(
        #         language=Language.PYTHON, chunk_size=50, chunk_overlap=0
        #     )
        # df['doc_chunks'] = df['content'].apply(lambda x: python_splitter.create_documents([base64.b64decode(x).decode()]))
        # df = df.explode('doc_chunks', ignore_index=True)
        # df['content'] = df['doc_chunks'].apply(lambda x: x.page_content.replace('\n',' '))
        # df.drop(['doc_chunks'], inplace=True, axis=1)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        # df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_md(md_docs:pd.DataFrame, 
                  stack_docs:pd.DataFrame, 
                  issues_docs:pd.DataFrame, 
                  class_name:str, 
                  loaded_docs_file_path:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata (the git sha, document url,  
        the document source (ie. astro) and a concatenation of the headers).

        Any existing documents with the same docLink but differing UUID or sha will be 
        deleted prior to import.

        Vectorization includes the headers for bm25 search.
        """

        stack_df = pd.concat(stack_docs, axis=0, ignore_index=True)
        issues_df = pd.concat(issues_docs, axis=0, ignore_index=True)
        df = pd.concat([stack_df, issues_df, md_docs], ignore_index=True)

        df['uuid'] = df.apply(lambda x: generate_uuid5(x.to_dict()), axis=1)

        df = remove_existing_objects(loaded_docs_file_path=loaded_docs_file_path, 
                                     new_df=df, 
                                     class_name=class_name)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "error_threshold": 10}
    
    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_python_code(code_dfs: List[pd.DataFrame], class_name:str, loaded_docs_file_path:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata which is comprised of
        the git sha, document url,  the document source (ie. code-samples) and a simulated 
        header (ie. python).

        Any existing documents with the same docLink but differing UUID or sha will be 
        deleted prior to import.
        """

        df = pd.concat(code_dfs, axis=0, ignore_index=True)

        df['uuid'] = df.apply(generate_uuid5, axis=1)

        df = remove_existing_objects(loaded_docs_file_path=loaded_docs_file_path, 
                                     new_df=df, 
                                     class_name=class_name)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "error_threshold": 0}
        
    _alert_schema_branch = alert_schema_branch(_check_schema.output)

    @task
    def fail_schema():
        raise AirflowException('Failing DAG for schema failure.')
    _fail_schema = fail_schema()

    md_docs = extract_github_markdown.partial().expand(source=markdown_docs_sources)
    rst_docs = extract_github_rst.partial().expand(source=rst_docs_sources)
    issues_md = extract_github_issues.partial().expand(source=issues_docs_sources)
    code_samples = extract_github_python.partial().expand(source=code_samples_sources)
    stackoverflow_md = extract_stack_overflow.partial(stackoverflow_cutoff_date=stackoverflow_cutoff_date).expand(tag=stackoverflow_tags)
    slack_md = extract_slack.partial().expand(source=slack_channel_sources)
    registry_md = extract_registry.partial().expand(source=http_json_sources)
    
    split_md_docs = split_markdown(md_dfs=md_docs, rst_dfs=rst_docs, slack_dfs=slack_md, reg_dfs=registry_md)
    split_python_code = split_python(code_samples)

    _loaded_docs = WeaviateRetrieveAllOperator(task_id='fetch_loaded_docs', 
                                               weaviate_conn_id=_WEAVIATE_CONN_ID,
                                               trigger_rule='none_failed',
                                               class_name='Docs', 
                                               replace_existing=True,
                                               output_file='file://include/data/loaded_docs.parquet')
    
    _unimported_md = import_md(md_docs=split_md_docs, 
                               stack_docs=stackoverflow_md, 
                               issues_docs=issues_md,
                               class_name='Docs', 
                               loaded_docs_file_path=_loaded_docs.output)
    
    _unimported_code = import_python_code(code_dfs=[split_python_code],
                                          class_name='Docs', 
                                          loaded_docs_file_path=_loaded_docs.output)

    _check_schema >> \
        _alert_schema_branch >> \
            [_slack_schema_alert, md_docs, rst_docs, code_samples, stackoverflow_md, issues_md, slack_md, registry_md]
    
    _loaded_docs >> [_unimported_md, _unimported_code]
    _slack_schema_alert >> _fail_schema

ask_astro_load_incremental()


# _backup = WeaviateBackupOperator(task_id='backup_to_fs', 
#                                  weaviate_conn_id='weaviate_default',
#                                  backend='filesystem', 
#                                  id="backup_base",
#                                 #  id="backup_fs_"+"{{ ts_nodash }}",
#                                  include=['Docs'],
#                                  )
#Restore only works with local weaviate
# _restore = WeaviateRestoreOperator(task_id='restore_from_fs',
#                                     weaviate_conn_id='weaviate_default',
#                                     backend='filesystem', 
#                                     id='backup_base',
#                                     include=['Docs'],
#                                     replace_existing=True)
airflow/dags/ingestion/ask-astro-load-incremental.py
astronomer/ask-astro
from datetime import datetime 
import pandas as pd
import html2text
import requests

from typing import List

from airflow.decorators import dag, task
from airflow.exceptions import AirflowException
from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from weaviate_provider.hooks.weaviate import WeaviateHook
from weaviate_provider.operators.weaviate import (
    WeaviateCheckSchemaOperator,
    WeaviateRetrieveAllOperator,
)
from weaviate.util import generate_uuid5
from langchain.text_splitter import (
    MarkdownHeaderTextSplitter, 
    RecursiveCharacterTextSplitter
)
from langchain.schema import Document

_WEAVIATE_CONN_ID = 'weaviate_test'
_SLACK_CONN_ID = 'slack_api_default'

default_args = {
    "retries": 3,
    }

# @dag(schedule_interval="0 5 * * *", start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
@dag(schedule_interval=None, start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
def ask_astro_load_registry():
    """
    This DAG performs incremental load for any data sources that have changed.  Initial load via 
    ask_astro_load_bulk imported data from a point-in-time data capture.

    This DAG checks to make sure the latest schema exists.  If it does not exist a slack message
    is sent to notify admins.
    """

    def remove_existing_objects(loaded_docs_file_path:str, new_df:pd.DataFrame, class_name:str):
        """
        Helper function to check if existing content needs to 
        be deleted before update. 

        Existing objects (based on 'docLink') with differing uuid or sha 
        will be deleted.

        Returned df includes only the objects that need to be (re)imported.
        """

        weaviate_hook = WeaviateHook(_WEAVIATE_CONN_ID)
        weaviate_hook.client = weaviate_hook.get_conn()

        current_objects_df = pd.read_parquet(loaded_docs_file_path)\
                               .drop(['vector'], axis=1)\
                               .groupby('docLink').agg(set)

        update_objects_df = new_df.groupby('docLink').agg(set)\
                              .join(current_objects_df, rsuffix='_old')
        update_objects_df = update_objects_df[update_objects_df['sha'] != update_objects_df['sha_old']]

        #remove existing objects
        update_objects_df['id'].dropna()\
            .apply(lambda x: [weaviate_hook.client.data_object.delete(uuid=uuid, 
                                                                      class_name=class_name) 
                                for uuid in list(x)])

        objects_to_import = new_df.merge(update_objects_df.reset_index()['docLink'], on='docLink', how='right')

        return objects_to_import

    _check_schema = WeaviateCheckSchemaOperator(task_id='check_schema', 
                                                weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                class_object_data='file://include/data/schema.json')

    @task.branch(retries=0)
    def alert_schema_branch(schema_exists:bool) -> str:
        """
        Check if schema is no longer valid
        """
        # WeaviateHook(_WEAVIATE_CONN_ID).get_conn().schema.delete_all()
        if schema_exists:
            return [
                "extract_astro_registry_cell_types",
                "fetch_loaded_docs"
            ]
        elif not schema_exists:
            return ["slack_schema_alert"]
        else:
            return None

    _slack_schema_alert = SlackAPIPostOperator(task_id='slack_schema_alert', 
                                               channel='#airflow_notices',
                                               retries=0,
                                               slack_conn_id = _SLACK_CONN_ID,
                                               text='ask_astro_load_http897 DAG error.  Schema mismatch.')

    @task(trigger_rule='none_failed')
    def extract_astro_registry_cell_types():

        base_url='https://api.astronomer.io/registryV2/v1alpha1/organizations/public/modules?limit=1000'
        headers={}

        json_data_class = 'modules'
        response = requests.get(base_url, headers=headers).json()
        total_count = response['totalCount']
        data = response.get(json_data_class, [])

        while len(data) < total_count-1:
            response = requests.get(f"{base_url}&offset={len(data)+1}").json()
            data.extend(response.get(json_data_class, []))

        md_template = "# Registry\n## {providerName}__{version}__{name}\n\n{description})"
        
        df = pd.DataFrame(data)
        df.rename({'githubUrl': 'docLink', 'searchId': 'sha'}, axis=1, inplace=True)
        df['docSource'] = 'registry_cell_types'
        df['description'] = df['description'].apply(lambda x: html2text.html2text(x) if x else 'No Description')
        df['content'] = df.apply(lambda x: md_template.format(providerName=x.providerName, 
                                                              version=x.version, 
                                                              name=x.name,
                                                              description=x.description), axis=1)

        reg_dfs = df[['docSource', 'sha', 'content', 'docLink']]

        return reg_dfs
    
    @task(trigger_rule='none_failed')
    def split_data(reg_dfs:List[pd.DataFrame]):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        splits markdown content on markdown headers.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', 'docs', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Chunked content in markdown format.

        """

        df = pd.concat([reg_dfs], axis=0, ignore_index=True)

        splitter = RecursiveCharacterTextSplitter()
        df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_documents([Document(page_content=x)]))

        df = df[df['doc_chunks'].apply(lambda x: len(x))>0].reset_index(drop=True)
        df = df.explode('doc_chunks', ignore_index=True)
        df['content'] = df['doc_chunks'].apply(lambda x: html2text.html2text(x.page_content).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))

        df.drop(['doc_chunks'], inplace=True, axis=1)
        df.reset_index(inplace=True, drop=True)

        return df

    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_data(md_docs:pd.DataFrame, class_name:str, loaded_docs_file_path:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata (the git sha, document url,  
        the document source (ie. astro) and a concatenation of the headers).

        Any existing documents with the same docLink but differing UUID or sha will be 
        deleted prior to import.

        Vectorization includes the headers for bm25 search.
        """
        df = pd.concat([md_docs], ignore_index=True)

        df['uuid'] = df.apply(lambda x: generate_uuid5(x.to_dict()), axis=1)

        df = remove_existing_objects(loaded_docs_file_path=loaded_docs_file_path, 
                                     new_df=df, 
                                     class_name=class_name)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "error_threshold": 10}
    
    _alert_schema_branch = alert_schema_branch(_check_schema.output)

    @task
    def fail_schema():
        raise AirflowException('Failing DAG for schema failure.')
    
    _fail_schema = fail_schema()

    registry_md = extract_astro_registry_cell_types()
    
    split_md_docs = split_data(reg_dfs=registry_md)

    _loaded_docs = WeaviateRetrieveAllOperator(task_id='fetch_loaded_docs', 
                                               weaviate_conn_id=_WEAVIATE_CONN_ID,
                                               trigger_rule='none_failed',
                                               class_name='Docs', 
                                               replace_existing=True,
                                               output_file='file://include/data/loaded_docs.parquet')
    
    _unimported_md = import_data(md_docs=split_md_docs,
                               class_name='Docs', 
                               loaded_docs_file_path=_loaded_docs.output)

    _check_schema >> \
        _alert_schema_branch >> \
            [_slack_schema_alert, registry_md]
    
    _loaded_docs >> _unimported_md
    _slack_schema_alert >> _fail_schema

ask_astro_load_registry()

def test():
    from weaviate_provider.operators.weaviate import WeaviateImportDataOperator

    self = WeaviateImportDataOperator(task_id='test', 
                                      weaviate_conn_id=_WEAVIATE_CONN_ID,
                                      data=df, class_name='Docs', 
                                      uuid_column='uuid')
airflow/dags/ingestion/ask-astro-load-registry.py
astronomer/ask-astro
from datetime import datetime 
from stackapi import StackAPI
import pandas as pd
import numpy as np
import html2text

from typing import List

from airflow.decorators import dag, task
from airflow.exceptions import AirflowException
from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from airflow.providers.slack.hooks.slack import SlackHook
from weaviate_provider.hooks.weaviate import WeaviateHook
from weaviate_provider.operators.weaviate import (
    WeaviateCheckSchemaOperator,
    WeaviateRetrieveAllOperator,
)
from weaviate.util import generate_uuid5
from langchain.text_splitter import (
    MarkdownHeaderTextSplitter, 
    RecursiveCharacterTextSplitter
)
from langchain.schema import Document


_WEAVIATE_CONN_ID = 'weaviate_default' #'weaviate_wcs' #
_GITHUB_CONN_ID = 'github_default'
_SLACK_CONN_ID = 'slack_api_default'

slack_archive_hostname = 'http://apache-airflow.slack-archives.org/'

slack_channel_sources = [
    {'channel_name': 'troubleshooting', 
      'channel_id': 'CCQ7EGB1P', 
      'team_id': 'TCQ18L22Z', 
      'team_name' : 'Airflow Slack Community',
      'slack_api_conn_id' : _SLACK_CONN_ID},
    {'channel_name': 'dev', 
      'channel_id': 'C05LWJWSV9P', 
      'team_id': 'T04GW479DRB', 
      'team_name' : 'dev',
      'slack_api_conn_id' : _SLACK_CONN_ID},
]

default_args = {
    "retries": 3,
    }

@dag(schedule_interval="0 5 * * *", start_date=datetime(2023, 9, 11), catchup=False, default_args=default_args)
def ask_astro_load_slack():
    """
    This DAG performs incremental load for any data sources that have changed.  Initial load via 
    ask_astro_load_bulk imported data from a point-in-time data capture.

    This DAG checks to make sure the latest schema exists.  If it does not exist a slack message
    is sent to notify admins.
    """

    def remove_existing_objects(loaded_docs_file_path:str, new_df:pd.DataFrame, class_name:str):
        """
        Helper function to check if existing content needs to 
        be deleted before update. 

        Existing objects (based on 'docLink') with differing uuid or sha 
        will be deleted.

        Returned df includes only the objects that need to be (re)imported.
        """

        weaviate_hook = WeaviateHook(_WEAVIATE_CONN_ID)
        weaviate_hook.client = weaviate_hook.get_conn()

        current_objects_df = pd.read_parquet(loaded_docs_file_path)\
                               .drop(['vector'], axis=1)\
                               .groupby('docLink').agg(set)

        update_objects_df = new_df.groupby('docLink').agg(set)\
                              .join(current_objects_df, rsuffix='_old')
        update_objects_df = update_objects_df[update_objects_df['sha'] != update_objects_df['sha_old']]

        #remove existing objects
        update_objects_df['id'].dropna()\
            .apply(lambda x: [weaviate_hook.client.data_object.delete(uuid=uuid, 
                                                                      class_name=class_name) 
                                for uuid in list(x)])

        objects_to_import = new_df.merge(update_objects_df.reset_index()['docLink'], on='docLink', how='right')

        return objects_to_import

    _check_schema = WeaviateCheckSchemaOperator(task_id='check_schema', 
                                                weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                class_object_data='file://include/data/schema.json')

    @task.branch(retries=0)
    def alert_schema_branch(schema_exists:bool) -> str:
        """
        Check if schema is no longer valid
        """
        # WeaviateHook(_WEAVIATE_CONN_ID).get_conn().schema.delete_all()
        if schema_exists:
            return [
                "extract_slack",
            ]
        elif not schema_exists:
            return ["slack_schema_alert"]
        else:
            return None

    _slack_schema_alert = SlackAPIPostOperator(task_id='slack_schema_alert', 
                                               channel='#airflow_notices',
                                               retries=0,
                                               slack_conn_id = _SLACK_CONN_ID,
                                               text='ask_astro_load_incremental DAG error.  Schema mismatch.')
       
    @task(trigger_rule='none_failed')
    def extract_slack(source:dict):
        """
        This task downloads archived slack messages as documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': slack team and channel names
        'docLink': URL for the specific message/reply
        'content': The message/reply content in markdown format.
        'header': document type. (ie. 'question' or 'answer')

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """
        #### THIS IS A PLACEHOLDER AS WE WAIT FOR SLACK BOT ACCESS TO AIRFLOW ####

        message_md_format = "# slack: {team_name}\n\n## {channel_name}\n\n{content}"
        reply_md_format = "### [{ts}] <@{user}>\n\n{text}"
        link_format = "https://app.slack.com/client/{team_id}/{channel_id}/p{ts}"
        
        slack_client = SlackHook(slack_conn_id = source['slack_api_conn_id']).client
        
        channel_info = slack_client.conversations_info(channel=source['channel_id']).data['channel']
        assert channel_info['is_member'] or not channel_info['is_private']

        history = slack_client.conversations_history(channel=source['channel_id'])
        replies = []
        df = pd.DataFrame(history.data['messages'])

        #if channel has no replies yet thread_ts will not be present
        if 'thread_ts' not in df:
            df['thread_ts'] = np.nan
        else:
            for ts in df[df['thread_ts'].notna()]['thread_ts']:
                reply=slack_client.conversations_replies(channel=source['channel_id'], ts=ts)
                replies.extend(reply.data['messages'])
            df = pd.concat([df, pd.DataFrame(replies)], axis=0)

        df = df[['user', 'text', 'ts', 'thread_ts', 'client_msg_id', 'type']]\
                .drop_duplicates()\
                .reset_index(drop=True)
        
        df['thread_ts'] = df['thread_ts'].astype(float)
        df['ts'] = df['ts'].astype(float)
        df['thread_ts'].fillna(value=df.ts, inplace=True)
        
        df['content'] = df.apply(lambda x: reply_md_format.format(ts=datetime.fromtimestamp(x.ts),
                                                                  user=x.user,
                                                                  text=x.text), axis=1)

        df = df.sort_values('ts').groupby('thread_ts').agg({'content': '\n'.join}).reset_index()

        df['content'] = df['content'].apply(lambda x: message_md_format.format(team_name=source['team_name'],
                                                                     channel_name=source['channel_name'],
                                                                     content=x))
       
        df['docLink'] = df['thread_ts'].apply(lambda x: link_format.format(team_id=source['team_id'],
                                                                    channel_id=source['channel_id'],
                                                                    ts=str(x).replace('.','')))
        df['docSource'] = source['channel_name']

        df['sha'] = df['content'].apply(generate_uuid5)
        
        return df[['docSource', 'sha', 'content', 'docLink']]
        # return pd.DataFrame()

    @task(trigger_rule='none_failed')
    def split_data(slack_dfs:List[pd.DataFrame]):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        splits markdown content on markdown headers.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', 'docs', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Chunked content in markdown format.

        """

        df = pd.concat(slack_dfs, axis=0, ignore_index=True)

        # headers_to_split_on = [
        #     ("#", "Header 1"),
        #     ("##", "Header 2"),
        #     # ("###", "Header 3"),
        # ]

        # splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
        # df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_text(x))
                
        splitter = RecursiveCharacterTextSplitter()
        df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_documents([Document(page_content=x)]))

        df = df[df['doc_chunks'].apply(lambda x: len(x))>0].reset_index(drop=True)
        # _ = df['doc_chunks'].apply(lambda x: x[0].metadata.update({'Header 1': 'Summary'}) if x[0].metadata == {} else x[0] )
        df = df.explode('doc_chunks', ignore_index=True)
        df['content'] = df['doc_chunks'].apply(lambda x: html2text.html2text(x.page_content).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))
        # df['header'] = df['doc_chunks'].apply(lambda x: '. '.join(list(x.metadata.values())))

        df.drop(['doc_chunks'], inplace=True, axis=1)
        df.reset_index(inplace=True, drop=True)

        return df

    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_data(md_docs:pd.DataFrame, class_name:str, loaded_docs_file_path:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata (the git sha, document url,  
        the document source (ie. astro) and a concatenation of the headers).

        Any existing documents with the same docLink but differing UUID or sha will be 
        deleted prior to import.

        Vectorization includes the headers for bm25 search.
        """

        df = pd.concat([md_docs], ignore_index=True)

        df['uuid'] = df.apply(lambda x: generate_uuid5(x.to_dict()), axis=1)

        df = remove_existing_objects(loaded_docs_file_path=loaded_docs_file_path, 
                                     new_df=df, 
                                     class_name=class_name)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "error_threshold": 10}
        
    _alert_schema_branch = alert_schema_branch(_check_schema.output)

    @task
    def fail_schema():
        raise AirflowException('Failing DAG for schema failure.')
    _fail_schema = fail_schema()

    slack_md = extract_slack.partial().expand(source=slack_channel_sources)
    
    split_md_docs = split_data(slack_dfs=slack_md)

    _loaded_docs = WeaviateRetrieveAllOperator(task_id='fetch_loaded_docs', 
                                               weaviate_conn_id=_WEAVIATE_CONN_ID,
                                               trigger_rule='none_failed',
                                               class_name='Docs', 
                                               replace_existing=True,
                                               output_file='file://include/data/loaded_docs.parquet')
    
    _unimported_md = import_data(md_docs=split_md_docs, 
                                 class_name='Docs', 
                                 loaded_docs_file_path=_loaded_docs.output)
    
    _check_schema >> \
        _alert_schema_branch >> \
            [_slack_schema_alert, slack_md]
    
    _loaded_docs >> _unimported_md
    _slack_schema_alert >> _fail_schema

ask_astro_load_slack()
airflow/dags/ingestion/ask-astro-load-slack.py
astronomer/ask-astro
from datetime import datetime 
from stackapi import StackAPI
from pathlib import Path
import pandas as pd
import pypandoc
import html2text
import re
import requests

from typing import List

from airflow.decorators import dag, task, task_group
from airflow.exceptions import AirflowException
from airflow.providers.github.hooks.github import GithubHook
from airflow.providers.slack.operators.slack import SlackAPIPostOperator
from airflow.providers.slack.hooks.slack import SlackHook
from weaviate_provider.hooks.weaviate import WeaviateHook
from weaviate_provider.operators.weaviate import (
    WeaviateCreateSchemaOperator,
    WeaviateCheckSchemaOperator,
    WeaviateImportDataOperator,
    )
from weaviate.util import generate_uuid5
from langchain.text_splitter import (
    MarkdownHeaderTextSplitter, 
    RecursiveCharacterTextSplitter,
)
from langchain.schema import Document

_WEAVIATE_CONN_ID = 'weaviate_test'
_GITHUB_CONN_ID = 'github_default'
_SLACK_CONN_ID = 'slack_api_default'

# doc_dir:baseurl pairs for dynamic tasks
markdown_docs_sources = [
    {'doc_dir': 'learn', 'repo_base': 'astronomer/docs'}, 
    {'doc_dir': 'astro', 'repo_base': 'astronomer/docs'}
    ]
rst_docs_sources = [
    {'doc_dir': 'docs', 'repo_base': 'apache/airflow'},
    ]
code_samples_sources = [
    {'doc_dir': 'code-samples', 'repo_base': 'astronomer/docs'},
    ]
issues_docs_sources = [
    {'doc_dir': 'issues', 'repo_base': 'apache/airflow'}
]
slack_channel_sources = [
    {'channel_name': 'troubleshooting', 
      'channel_id': 'CCQ7EGB1P', 
      'team_id': 'TCQ18L22Z', 
      'team_name' : 'Airflow Slack Community',
      'slack_api_conn_id' : 'TBD'}
]
http_json_sources = [
    {'name': 'registry_cell_types',
     'base_url': 'https://api.astronomer.io/registryV2/v1alpha1/organizations/public/modules?limit=1000',
     'headers': {},
     'count_field': 'totalCount'}
]

rst_exclude_docs = ['changelog.rst', 'commits.rst']

stackoverflow_cutoff_date = '2021-09-01'
stackoverflow_tags = [
    'airflow',
]

weaviate_doc_count = {
    'Docs': 7307,
}

default_args = {
    "retries": 3,
    }

@dag(schedule_interval=None, start_date=datetime(2023, 8, 1), catchup=False, default_args=default_args)
def ask_astro_load_bulk():
    """
    This DAG performs the initial load of data from sources.  While the code to generate these datasets
    is included for each function, the data is frozen as a parquet file for simple ingest and the steps
    to create are commented out.
    """
    _check_schema = WeaviateCheckSchemaOperator(task_id='check_schema', 
                                                weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                class_object_data='file://include/data/schema.json')

    @task.branch
    def recreate_schema_branch(schema_exists:bool) -> str:
        # WeaviateHook(_WEAVIATE_CONN_ID).get_conn().schema.delete_all()
        if schema_exists:
            return ["check_object_count"]
        elif not schema_exists:
            return ["create_schema"]
        else:
            return None

    @task.branch
    def check_object_count(weaviate_doc_count:dict, class_name:str) -> str:
        try:
            weaviate_hook = WeaviateHook(_WEAVIATE_CONN_ID)
            weaviate_hook.client = weaviate_hook.get_conn()
            response = weaviate_hook.run(f'{{Aggregate {{ {class_name} {{ meta {{ count }} }} }} }}')
        except Exception as e:
            if e.status_code == 422 and 'no graphql provider present' in e.message:
                response = None

        if response and \
            response['data']['Aggregate'][class_name][0]['meta']['count'] >= weaviate_doc_count[class_name]:
            print('Initial Upload complete. Skipping')
            return None
        else:
            return ["extract_github_markdown", 
                    "extract_github_rst", 
                    "extract_github_python", 
                    "extract_stack_overflow",
                    "extract_slack",
                    "extract_registry",
                    "extract_github_issues"]
       
    _create_schema = WeaviateCreateSchemaOperator(task_id='create_schema', 
                                                  weaviate_conn_id=_WEAVIATE_CONN_ID,
                                                  class_object_data='file://include/data/schema.json',
                                                  existing='fail')
    
    @task(trigger_rule='none_failed')
    def extract_github_markdown(source:dict):
        """
        This task downloads github content as markdown documents in a 
        pandas dataframe.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Entire document content in markdown format.

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """

        # downloaded_docs = []
        
        # gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        # repo = gh_hook.client.get_repo(source['repo_base'])
        # contents = repo.get_contents(source['doc_dir'])

        # while contents:

        #     file_content = contents.pop(0)
        #     if file_content.type == "dir":
        #         contents.extend(repo.get_contents(file_content.path))

        #     elif Path(file_content.name).suffix == '.md':

        #         print(file_content.name)
                
        #         row = {
        #             "docLink": file_content.html_url, 
        #             "sha": file_content.sha,
        #             "content": file_content.decoded_content.decode(),
        #             "docSource": source['doc_dir'], 
        #         }

        #         downloaded_docs.append(row)
                
        # df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task(trigger_rule='none_failed')
    def extract_github_rst(source:dict):
        """
        This task downloads github content as rst documents 
        in a pandas dataframe.

        The 'content' field is converted from RST to Markdown (via pypandoc).  After 
        removing the preamble (apache license), any empty lines and 'include' footers 
        any empty docs are removed.  Document links and references are not included 
        in the content.

        Dataframe fields are:
        'docSource': ie. 'docs'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Entire document in markdown format.

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """

        # downloaded_docs = []
        
        # gh_hook = GithubHook(_GITHUB_CONN_ID)

        # repo = gh_hook.client.get_repo(source['repo_base'])
        # contents = repo.get_contents(source['doc_dir'])

        # apache_license_text = Path('include/data/apache_license.rst').read_text()

        # while contents:

        #     file_content = contents.pop(0)
        #     if file_content.type == "dir":
        #         contents.extend(repo.get_contents(file_content.path))

        #     elif Path(file_content.name).suffix == '.rst' and file_content.name not in rst_exclude_docs:

        #         print(file_content.name)

        #         row = {
        #             "docLink": file_content.html_url, 
        #             "sha": file_content.sha,
        #             "content": file_content.decoded_content.decode(),
        #             "docSource": source['doc_dir'], 
        #         }

        #         downloaded_docs.append(row)
                
        # df = pd.DataFrame(downloaded_docs)

        # df['content'] = df['content'].apply(lambda x: x.replace(apache_license_text, ''))
        # df['content'] = df['content'].apply(lambda x: re.sub(r".*include.*", "", x))
        # df['content'] = df['content'].apply(lambda x: re.sub(r'^\s*$', "", x))
        # df = df[df['content']!='']
        # df['content'] = df['content'].apply(lambda x: pypandoc.convert_text(source=x, to='md', 
        #                                                                     format='rst',
        #                                                                     extra_args=['--atx-headers']))

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df

    @task(trigger_rule='none_failed')
    def extract_github_python(source:dict):
        """
        This task downloads github content as python code in a pandas dataframe.

        The 'content' field of the dataframe is currently not split as the context 
        window is large enough. Code for splitting is provided but commented out.

        Dataframe fields are:
        'docSource': ie. 'code-samples'
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': The python code
        'header': a placeholder of 'python' for bm25 search

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """
    
        # downloaded_docs = []

        # gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        # repo = gh_hook.client.get_repo(source['repo_base'])
        # contents = repo.get_contents(source['doc_dir'])

        # while contents:
        #     file_content = contents.pop(0)

        #     if file_content.type == "dir":
        #         contents.extend(repo.get_contents(file_content.path))

        #     elif Path(file_content.name).suffix == '.py':
        #         print(file_content.name)
                                
        #         row = {
        #             "docLink": file_content.html_url, 
        #             "sha": file_content.sha,
        #             "content": file_content.decoded_content.decode(),
        #             "docSource": source['doc_dir'], 
        #             "header": 'python', 
        #         }
                
        #         downloaded_docs.append(row)

        # df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")

        return df
    
    @task(trigger_rule='none_failed')
    def extract_stack_overflow(tag:dict, stackoverflow_cutoff_date:str):
        """
        This task generates stack overflow questions and answers as markdown 
        documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': 'stackoverflow' plus the tag name (ie. 'airflow')
        'docLink': URL for the specific question/answer.
        'content': The base64 encoded content of the question/answer in markdown format.
        'header': document type. (ie. 'question' or 'answer')

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """

        # question_template = "TITLE: {title}\nDATE: {date}\nBY: {user}\nSCORE: {score}\n{body}{question_comments}"
        # answer_template = "DATE: {date}\nBY: {user}\nSCORE: {score}\n{body}{answer_comments}"
        # comment_template = "{user} on {date} [Score: {score}]: {body}\n"

        # post_types = {
        #     '1': 'Question',
        #     '2': 'Answer',
        #     '3': 'Wiki',
        #     '4': 'TagWikiExcerpt',
        #     '5': 'TagWiki',
        #     '6': 'ModeratorNomination',
        #     '7': 'WikiPlaceholder',
        #     '8': 'PrivilegeWiki'}
        # posts_columns = {
        #     '_COL_0':'post_id', 
        #     '_COL_1':'type', 
        #     '_COL_3':'parent_id', 
        #     '_COL_4':'created_on', 
        #     '_COL_6':'score', 
        #     '_COL_8':'body',
        #     '_COL_9':'user_id', 
        #     '_COL_10':'user_name',
        #     '_COL_15':'title',
        #     '_COL_17':'answer_count'}
        # comments_columns = {
        #     '_COL_0':'comment_id', 
        #     '_COL_1':'post_id', 
        #     '_COL_2':'comment_score', 
        #     '_COL_3':'comment_body',
        #     '_COL_4':'comment_created_on',
        #     '_COL_5':'comment_user_name',
        #     '_COL_6':'comment_user_id'}
        
        # posts_df = pd.read_parquet('include/data/StackOverflow/Posts/data_0_0_0.snappy.parquet')[posts_columns.keys()]
        # comments_df = pd.concat([
        #     pd.read_parquet('include/data/StackOverflow/Comments/data_0_0_0.snappy.parquet')[comments_columns.keys()],
        #     pd.read_parquet('include/data/StackOverflow/Comments/data_0_1_0.snappy.parquet')[comments_columns.keys()]
        #     ], ignore_index=True)
        
        # posts_df.rename(posts_columns, axis=1, inplace=True)
        # posts_df['type'] = posts_df['type'].apply(lambda x: post_types[x])
        # posts_df['created_on'] = pd.to_datetime(posts_df['created_on'])
        # posts_df[['post_id', 'parent_id', 'user_id', 'user_name']] = posts_df[['post_id', 'parent_id', 'user_id', 'user_name']].astype(str)
        # posts_df = posts_df[posts_df['created_on'] >= stackoverflow_cutoff_date]
        # posts_df['user_id'] = posts_df.apply(lambda x: x.user_id or x.user_name or 'Unknown User', axis=1)
        # posts_df.reset_index(inplace=True, drop=True)

        # comments_df.rename(comments_columns, axis=1, inplace=True)
        # comments_df['comment_created_on'] = pd.to_datetime(comments_df['comment_created_on'])
        # comments_df['comment_user_id'] = comments_df.apply(lambda x: x.comment_user_id or x.comment_user_name or 'Unknown User', axis=1)
        # comments_df[['post_id', 'comment_user_id', 'comment_user_name']] = comments_df[['post_id', 'comment_user_id', 'comment_user_name']].astype(str)
        # comments_df[['comment_score']] = comments_df[['comment_score']].astype(int)
        # comments_df['comment_text'] = comments_df.apply(lambda x: comment_template.format(user=x.comment_user_id,
        #                                                                                   date=x.comment_created_on,
        #                                                                                   score=x.comment_score,
        #                                                                                   body=x.comment_body), axis=1)
        # comments_df = comments_df[['post_id', 'comment_text']].groupby('post_id').agg(list)
        # comments_df['comment_text'] = comments_df['comment_text'].apply(lambda x: '\n'.join(x))
        # comments_df.reset_index(inplace=True)

        # questions_df = posts_df[posts_df['type']=='Question']
        # questions_df = questions_df.drop('parent_id', axis=1)
        # questions_df.rename({'body':'question_body', 'post_id': 'question_id'}, axis=1, inplace=True)
        # questions_df[['answer_count','score']] = questions_df[['answer_count','score']].astype(int)
        # questions_df = questions_df[questions_df['score']>=1]
        # questions_df = questions_df[questions_df['answer_count']>=1]
        # questions_df.reset_index(inplace=True, drop=True)
        # questions_df = pd.merge(questions_df, comments_df, left_on='question_id', right_on='post_id', how='left')
        # questions_df['comment_text'].fillna('', inplace=True)
        # questions_df.drop('post_id', axis=1, inplace=True)
        # questions_df['link'] = questions_df['question_id'].apply(lambda x: f"https://stackoverflow.com/questions/{x}")
        # questions_df['question_text'] = questions_df.apply(lambda x: question_template.format(title=x.title,
        #                                                                                       user=x.user_id,
        #                                                                                       date=x.created_on,
        #                                                                                       score=x.score,
        #                                                                                       body=x.question_body,
        #                                                                                       question_comments=x.comment_text), axis=1)
        # questions_df = questions_df[['link', 'question_id', 'question_text']].set_index('question_id')
        # questions_df['docSource'] = f'stackoverflow {tag}'
        # questions_df = questions_df[['docSource', 'link', 'question_text']]
        # questions_df.columns = ['docSource', 'docLink', 'content']
        # questions_df['header'] = 'question'

        # answers_df = posts_df[posts_df['type']=='Answer'][['created_on', 'score', 'user_id', 'post_id', 'parent_id', 'body']]
        # answers_df.rename({'body':'answer_body', 'post_id': 'answer_id', 'parent_id': 'question_id'}, axis=1, inplace=True)
        # answers_df.reset_index(inplace=True, drop=True)
        # answers_df = pd.merge(answers_df, comments_df, left_on='answer_id', right_on='post_id', how='left')
        # answers_df['comment_text'].fillna('', inplace=True)
        # answers_df.drop('post_id', axis=1, inplace=True)
        # answers_df['link'] = answers_df['question_id'].apply(lambda x: f"https://stackoverflow.com/questions/{x}")
        # answers_df['answer_text'] = answers_df.apply(lambda x: answer_template.format(user=x.user_id,
        #                                                                               date=x.created_on,
        #                                                                               score=x.score,
        #                                                                               body=x.answer_body,
        #                                                                               answer_comments=x.comment_text), axis=1)
        # answers_df = answers_df.groupby('question_id')['answer_text'].apply(lambda x: ''.join(x))

        # answers_df = questions_df.join(answers_df).apply(lambda x: pd.Series([
        #     f'stackoverflow {tag}',
        #     x.docLink, 
        #     x.answer_text]), axis=1)
        # answers_df.columns=['docSource', 'docLink','content']
        # answers_df['header'] = 'answer'

        # df = pd.concat([questions_df, answers_df], axis=0).reset_index(drop=True)

        # df.to_parquet('include/data/stackoverflow_base.parquet')
        df = pd.read_parquet('include/data/stackoverflow_base.parquet')
        df['sha'] = df.apply(generate_uuid5, axis=1)

        return df

    @task(trigger_rule='none_failed')
    def extract_slack(source:dict):
        """
        This task downloads archived slack messages as documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': slack team and channel names
        'docLink': URL for the specific message/reply
        'content': The message/reply content in markdown format.
        'header': document type. (ie. 'question' or 'answer')

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """

        df = pd.read_parquet('include/data/slack/troubleshooting.parquet')

        message_md_format = "# slack: {team_name}\n\n## {channel_name}\n\n{content}"
        reply_md_format = "### [{ts}] <@{user}>\n\n{text}"
        link_format = "https://app.slack.com/client/{team_id}/{channel_id}/p{ts}"

        df = df[['user', 'text', 'ts', 'thread_ts', 'client_msg_id', 'type']]\
                .drop_duplicates()\
                .reset_index(drop=True)
        
        df['thread_ts'] = df['thread_ts'].astype(float)
        df['ts'] = df['ts'].astype(float)

        df['thread_ts'].fillna(value=df.ts, inplace=True)
        
        df['content'] = df.apply(lambda x: reply_md_format.format(ts=datetime.fromtimestamp(x.ts),
                                                                  user=x.user,
                                                                  text=x.text), axis=1)

        df = df.sort_values('ts').groupby('thread_ts').agg({'content': '\n'.join}).reset_index()

        df['content'] = df['content'].apply(lambda x: message_md_format.format(team_name=source['team_name'],
                                                                     channel_name=source['channel_name'],
                                                                     content=x))
       
        df['docLink'] = df['thread_ts'].apply(lambda x: link_format.format(team_id=source['team_id'],
                                                                    channel_id=source['channel_id'],
                                                                    ts=str(x).replace('.','')))
        df['docSource'] = source['channel_name']

        df['sha'] = df['content'].apply(generate_uuid5)

        df = df[['docSource', 'sha', 'content', 'docLink']]

        return df

    @task(trigger_rule='none_failed')
    def extract_github_issues(source:dict):
        """
        This task downloads github issues as markdown documents in a pandas dataframe.

        Dataframe fields are:
        'docSource': repo name + 'issues'
        'docLink': URL for the specific question/answer.
        'content': The base64 encoded content of the question/answer in markdown format.
        'header': document type. (ie. 'airflow issue')

        Code is provided for the processing of questions and answers but is 
        commented out as the historical data is provided as a parquet file.
        """
        # gh_hook = GithubHook(_GITHUB_CONN_ID)
        
        # repo = gh_hook.client.get_repo(source['repo_base'])
        # issues = repo.get_issues()

        # issue_autoresponse_text = 'Thanks for opening your first issue here!'
        # pr_autoresponse_text = 'Congratulations on your first Pull Request and welcome to the Apache Airflow community!'
        # drop_content = [issue_autoresponse_text, pr_autoresponse_text]

        # issue_markdown_template = "## ISSUE TITLE: {title}\nDATE: {date}\nBY: {user}\nSTATE: {state}\n{body}\n{comments}"
        # comment_markdown_template = "#### COMMENT: {user} on {date}\n{body}\n"

        # downloaded_docs = []
        # page_num = 0
        
        # page = issues.get_page(page_num)
        
        # while page:

        #     for issue in page:
        #         print(issue.number)
        #         comments=[]
        #         for comment in issue.get_comments():
        #             #TODO: this is very slow.  Look for vectorized approach.
        #             if not any(substring in comment.body for substring in drop_content):
        #                 comments.append(comment_markdown_template.format(user=comment.user.login, 
        #                                                                  date=issue.created_at.strftime("%m-%d-%Y"), 
        #                                                                  body=comment.body))
        #         downloaded_docs.append({
        #             "docLink": issue.html_url, 
        #             "sha": '',
        #             "content": issue_markdown_template.format(title=issue.title, 
        #                                                       date=issue.created_at.strftime("%m-%d-%Y"), 
        #                                                       user=issue.user.login,
        #                                                       state=issue.state, 
        #                                                       body=issue.body,
        #                                                       comments='\n'.join(comments)),
        #             "docSource": f"{source['repo_base']} {source['doc_dir']}", 
        #             "header": f"{source['repo_base']} issue", 
        #         })
        #     page_num=page_num+1
        #     page = issues.get_page(page_num)
                
        # df = pd.DataFrame(downloaded_docs)

        # df.to_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df = pd.read_parquet(f"include/data/{source['repo_base']}/{source['doc_dir']}.parquet")
        df['sha'] = df.apply(generate_uuid5, axis=1)

        return df

    @task(trigger_rule='none_failed')
    def extract_registry(source:dict):

        # data_class = source['base_url'].split('/')[-1].split('?')[0]

        # response = requests.get(source['base_url'], headers=source['headers']).json()
        # total_count = response[source['count_field']]
        # data = response.get(data_class, [])
        # while len(data) < total_count-1:
        #     response = requests.get(f"{source['base_url']}&offset={len(data)+1}").json()
        #     data.extend(response.get(data_class, []))

        # df = pd.DataFrame(data)
        # df.rename({'githubUrl': 'docLink', 'searchId': 'sha'}, axis=1, inplace=True)
        # df['docSource'] = source['name']
        # df['description'] = df['description'].apply(lambda x: html2text.html2text(x) if x else 'No Description')
        # df['content'] = df.apply(lambda x: md_template.format(providerName=x.providerName, 
        #                                                       version=x.version, 
        #                                                       name=x.name,
        #                                                       description=x.description), axis=1)

        # df = df[['docSource', 'sha', 'content', 'docLink']]

        # df.to_parquet('include/data/registry.parquet')
        
        df = pd.read_parquet('include/data/registry.parquet')

        return df
    
    @task()
    def split_data(md_dfs:List[pd.DataFrame], 
                   rst_dfs:List[pd.DataFrame], 
                   slack_dfs:List[pd.DataFrame], 
                   stackoverflow_dfs:List[pd.DataFrame], 
                   code_dfs:List[pd.DataFrame], 
                   issues_dfs:List[pd.DataFrame], 
                   reg_dfs:List[pd.DataFrame]):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        splits markdown content on markdown headers.

        Dataframe fields are:
        'docSource': ie. 'astro', 'learn', 'docs', etc.
        'sha': the github sha for the document
        'docLink': URL for the specific document in github.
        'content': Chunked content in markdown format.

        """

        md_df = pd.concat(md_dfs, axis=0, ignore_index=True)
        rst_df = pd.concat(rst_dfs, axis=0, ignore_index=True)
        slack_df = pd.concat(slack_dfs, axis=0, ignore_index=True)
        stackoverflow_df = pd.concat(stackoverflow_dfs, axis=0, ignore_index=True)
        code_df = pd.concat(code_dfs, axis=0, ignore_index=True)
        issues_df = pd.concat(issues_dfs, axis=0, ignore_index=True)
        reg_df = pd.concat(reg_dfs, axis=0, ignore_index=True)

        df = pd.concat([md_df, 
                        rst_df, 
                        slack_df, 
                        reg_df, 
                        code_df, 
                        stackoverflow_df,
                        issues_df,
                    ], axis=0, ignore_index=True)

        # headers_to_split_on = [
        #     ("#", "Header 1"),
        #     ("##", "Header 2"),
        #     # ("###", "Header 3"),
        # ]

        # splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
        splitter = RecursiveCharacterTextSplitter()
        df['doc_chunks'] = df['content'].apply(lambda x: splitter.split_documents([Document(page_content=x)]))

        df = df[df['doc_chunks'].apply(lambda x: len(x))>0].reset_index(drop=True)
        # _ = df['doc_chunks'].apply(lambda x: x[0].metadata.update({'Header 1': 'Summary'}) if x[0].metadata == {} else x[0] )
        df = df.explode('doc_chunks', ignore_index=True)
        df['content'] = df['doc_chunks'].apply(lambda x: html2text.html2text(x.page_content).replace('\n',' '))
        df['content'] = df['content'].apply(lambda x: x.replace('\\',''))
        # df['header'] = df['doc_chunks'].apply(lambda x: '. '.join(list(x.metadata.values())))

        df.drop(['doc_chunks'], inplace=True, axis=1)
        df.drop(['header'], inplace=True, axis=1)
        df.reset_index(inplace=True, drop=True)

        return df

    @task.weaviate_import(weaviate_conn_id=_WEAVIATE_CONN_ID)
    def import_data(md_docs:pd.DataFrame, class_name:str):
        """
        This task concatenates multiple dataframes from upstream dynamic tasks and 
        vectorizes with import to weaviate.

        A 'uuid' is generated based on the content and metadata (the git sha, document url,  
        the document source (ie. astro) and a concatenation of the headers).

        Vectorization includes the headers for bm25 search.
        """
        
        df = pd.concat([md_docs], ignore_index=True)

        df['uuid'] = df.apply(lambda x: generate_uuid5(x.to_dict()), axis=1)

        print(f"Passing {len(df)} objects for import.")

        return {"data": df, "class_name": class_name, "uuid_column": "uuid", "batch_size": 1000, "error_threshold": 12}
    
    _recreate_schema_branch = recreate_schema_branch(_check_schema.output)
    _check_object_count = check_object_count(weaviate_doc_count, 'Docs')

    md_docs = extract_github_markdown.partial().expand(source=markdown_docs_sources)
    rst_docs = extract_github_rst.partial().expand(source=rst_docs_sources)
    issues_md = extract_github_issues.partial().expand(source=issues_docs_sources)
    code_samples = extract_github_python.partial().expand(source=code_samples_sources)
    stackoverflow_md = extract_stack_overflow.partial(stackoverflow_cutoff_date=stackoverflow_cutoff_date).expand(tag=stackoverflow_tags)
    slack_md = extract_slack.partial().expand(source=slack_channel_sources)
    registry_md = extract_registry.partial().expand(source=http_json_sources)
    
    split_md_docs = split_data(md_dfs=md_docs, 
                               rst_dfs=rst_docs, 
                               stackoverflow_dfs=stackoverflow_md,
                               code_dfs=code_samples,
                               slack_dfs=slack_md,
                               issues_dfs=issues_md,
                               reg_dfs=registry_md)

    _unimported_md = import_data(md_docs=split_md_docs, class_name='Docs')
    
    _check_schema >> _recreate_schema_branch >> [_create_schema, _check_object_count] 
    _check_object_count >> [md_docs, rst_docs, code_samples, stackoverflow_md, issues_md, slack_md, registry_md]
    _create_schema >> [md_docs, rst_docs, code_samples, stackoverflow_md, issues_md, slack_md, registry_md]

ask_astro_load_bulk()

def test():
    from weaviate_provider.hooks.weaviate import WeaviateHook
    weaviate_client = WeaviateHook(_WEAVIATE_CONN_ID).get_conn()
    search = weaviate_client.query\
        .get(properties = ['content'], class_name='Docs')\
        .with_limit(2000)\
        .with_where({"path": ["docSource"], 
                     "operator": "Equal", 
                     "valueText": "registry_cell_types"}).do()
    len(search['data']['Get']['Docs'])

    self = WeaviateImportDataOperator(task_id='test',
                                      data=df,
                                      class_name='Docs',
                                      uuid_column='uuid')


airflow/dags/ingestion/ask-astro-load.py
astronomer/ask-astro
"""
Tasks for interacting with Weaviate.
"""

import time
import weaviate

import pandas as pd
from airflow.decorators import task
from logging import getLogger

from itertools import chain
from weaviate.util import generate_uuid5
from weaviate_provider.hooks.weaviate import WeaviateHook

from include.config import Connections

logger = getLogger(__name__)


@task()
def weaviate_upsert(outputs: list[list[dict[str, str]]], weaviate_class: str):
    """
    This task takes in a list of documents and upserts them into Weaviate. Note that
    Weaviate doesn't officially support upserting, so this first deletes any existing
    documents with the same docLink, then inserts the new documents.
    """
    docs = chain(*outputs)

    weaviate_hook = WeaviateHook(Connections.weaviate)
    client = weaviate_hook.get_conn()

    # the weaviate_import decorator expects a dataframe, so we need to convert
    df = pd.DataFrame(docs)
    df["uuid"] = df.apply(generate_uuid5, axis=1)

    def upsert(uuid: str, doc: dict[str, str], delay: int = 0):
        if delay > 0:
            logger.info("Sleeping for %s seconds", delay)
            time.sleep(delay)
        if delay > 60:
            raise Exception("Rate limit reached")

        try:
            if client.data_object.exists(uuid=uuid, class_name=weaviate_class):
                logger.info("Object with uuid %s exists, replacing", uuid)

                client.data_object.replace(
                    uuid=uuid, class_name=weaviate_class, data_object=doc
                )
            else:
                logger.info("Object with uuid %s does not exist, inserting", uuid)

                client.data_object.create(
                    class_name=weaviate_class, uuid=uuid, data_object=doc
                )
        except weaviate.UnexpectedStatusCodeException as e:
            if "Rate limit reached" in str(e):
                new_delay = delay * 2 if delay > 0 else 1
                logger.info("Rate limit reached, sleeping for %s seconds", new_delay)
                upsert(uuid, doc, new_delay)
            else:
                logger.error("Failed to insert document %s: %s", uuid, e)
        except Exception as e:
            logger.error("Failed to insert document %s: %s", uuid, e)

    for _, row in df.iterrows():
        doc = row.to_dict()
        uuid = doc.pop("uuid")

        upsert(uuid, doc)

airflow/include/tasks/weaviate_upsert.py
astronomer/ask-astro
"""
Initialize the Sanic app and route requests to the Slack app.
"""
import os
import logging
from logging import getLogger

from sanic import Sanic, Request

from ask_astro.slack.app import slack_app, app_handler
from ask_astro.slack.controllers import register_controllers
from ask_astro.rest.controllers import register_routes

# set the logging level based on an env var
logging.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))

logger = getLogger(__name__)

api = Sanic(name="ask_astro")


# route slack requests to the slack app
@api.get("/slack/oauth_redirect", name="oauth_redirect")
@api.get("/slack/install", name="install")
@api.post("/slack/events", name="events")
async def endpoint(req: Request):
    "Forward requests to the Slack bolt hander"
    return await app_handler.handle(req)


server_port = int(os.environ.get("PORT", 8080))

register_controllers(slack_app)
register_routes(api)

if __name__ == "__main__":
    api.run(host="0.0.0.0", port=server_port)

api/ask_astro/app.py
astronomer/ask-astro
from ask_astro.clients.weaviate_ import docsearch
from ask_astro.config import AzureOpenAIParams
from langchain import LLMChain
from langchain.chains import ConversationalRetrievalChain
from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT
from langchain.chains.question_answering import load_qa_chain
from langchain.chat_models import AzureChatOpenAI
from langchain.prompts import (
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    ChatPromptTemplate,
    MessagesPlaceholder,
)
from langchain.retrievers import MultiQueryRetriever

with open("ask_astro/templates/combine_docs_chat_prompt.txt", "r") as system_prompt_fd:
    messages = [
        SystemMessagePromptTemplate.from_template(system_prompt_fd.read()),
        MessagesPlaceholder(variable_name="messages"),
        HumanMessagePromptTemplate.from_template("{question}"),
    ]

retriever = MultiQueryRetriever.from_llm(
    llm=AzureChatOpenAI(
        **AzureOpenAIParams.us_east,
        deployment_name="gpt-35-turbo",
        temperature=0,
    ),
    retriever=docsearch.as_retriever(),
)

answer_question_chain = ConversationalRetrievalChain(
    retriever=retriever,
    return_source_documents=True,
    question_generator=LLMChain(
        llm=AzureChatOpenAI(
            **AzureOpenAIParams.us_east,
            deployment_name="gpt-35-turbo-16k",
            temperature=0.3,
        ),
        prompt=CONDENSE_QUESTION_PROMPT,
    ),
    combine_docs_chain=load_qa_chain(
        AzureChatOpenAI(
            **AzureOpenAIParams.us_east2,
            deployment_name="gpt-4-32k",
            temperature=0.5,
        ),
        chain_type="stuff",
        prompt=ChatPromptTemplate.from_messages(messages),
    ),
)

api/ask_astro/chains/answer_question.py
astronomer/ask-astro
"Firestore client to handle database operations"

from google.cloud import firestore

# auth is handled implicitly by the environment
firestore_client = firestore.AsyncClient()

__all__ = ["firestore_client"]

api/ask_astro/clients/firestore.py
astronomer/ask-astro
from langsmith import Client

langsmith_client = Client()

api/ask_astro/clients/langsmith_.py
astronomer/ask-astro
import weaviate
from weaviate import Client as WeaviateClient
from ask_astro.config import AzureOpenAIParams, WeaviateConfig
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Weaviate

embeddings = OpenAIEmbeddings(
    **AzureOpenAIParams.us_east,
    deployment="text-embedding-ada-002",
    model="text-embedding-ada-002",
)

client = WeaviateClient(
    url=WeaviateConfig.url,
    auth_client_secret=weaviate.AuthApiKey(api_key=WeaviateConfig.api_key),
    additional_headers={
        "X-Openai-Api-Key": WeaviateConfig.OpenAIApiKey,
    },
)
docsearch = Weaviate(
    client=client,
    index_name=WeaviateConfig.index_name,
    text_key=WeaviateConfig.text_key,
    attributes=WeaviateConfig.attributes,
)

api/ask_astro/clients/weaviate_.py
astronomer/ask-astro
"""Exports config variables that are used throughout the code."""

# TODO: Refactor with https://sanic.dev/en/guide/deployment/configuration.html

import os
import json


class FirestoreCollections:
    "Contains the names of the collections in the Firestore database."
    installation_store = os.environ["FIRESTORE_INSTALLATION_STORE_COLLECTION"]
    state_store = os.environ["FIRESTORE_STATE_STORE_COLLECTION"]
    messages = os.environ["FIRESTORE_MESSAGES_COLLECTION"]
    mentions = os.environ["FIRESTORE_MENTIONS_COLLECTION"]
    actions = os.environ["FIRESTORE_ACTIONS_COLLECTION"]
    responses = os.environ["FIRESTORE_RESPONSES_COLLECTION"]
    reactions = os.environ["FIRESTORE_REACTIONS_COLLECTION"]
    shortcuts = os.environ["FIRESTORE_SHORTCUTS_COLLECTION"]
    teams = os.environ["FIRESTORE_TEAMS_COLLECTION"]
    requests = os.environ["FIRESTORE_REQUESTS_COLLECTION"]


class AzureOpenAIParams:
    "Contains the parameters for the Azure OpenAI API."
    us_east = json.loads(os.environ["AZURE_OPENAI_USEAST_PARAMS"])
    us_east2 = json.loads(os.environ["AZURE_OPENAI_USEAST2_PARAMS"])


class ZendeskConfig:
    "Contains the config variables for the Zendesk API."
    credentials = os.environ.get("ZENDESK_CREDENTIALS")
    assignee_group_id = os.environ["ZENDESK_ASSIGNEE_GROUP_ID"]


class SlackAppConfig:
    "Contains the config variables for the Slack app."
    client_id = os.environ["SLACK_CLIENT_ID"]
    client_secret = os.environ["SLACK_CLIENT_SECRET"]
    signing_secret = os.environ["SLACK_SIGNING_SECRET"]


class LangSmithConfig:
    "Contains the config variables for the Langsmith API."
    project_name = os.environ["LANGCHAIN_PROJECT"]


class WeaviateConfig:
    OpenAIApiKey = os.environ["OPENAI_API_KEY"]
    url = os.environ["WEAVIATE_URL"]
    api_key = os.environ["WEAVIATE_API_KEY"]
    index_name = os.environ["WEAVIATE_INDEX_NAME"]
    text_key = os.environ["WEAVIATE_TEXT_KEY"]
    attributes = os.environ.get("WEAVIATE_ATTRIBUTES", "").split(",")

api/ask_astro/config.py
astronomer/ask-astro
from datetime import datetime

from uuid import UUID
from typing import Any

from pydantic.v1 import BaseModel, Field

from langchain.schema import BaseMessage, AIMessage, HumanMessage


class Source(BaseModel):
    "Represents a source for a request."
    name: str = Field(..., description="The name of the source")
    snippet: str = Field(..., description="The snippet of the source")


class AskAstroRequest(BaseModel):
    "Represents a request to ask-astro."
    uuid: UUID = Field(..., description="The UUID of the request")
    prompt: str = Field(..., description="The prompt for the request")
    messages: list[BaseMessage] = Field(
        default_factory=list,
        description="The messages in the request",
    )
    sources: list[Source] = Field(
        default_factory=list,
        description="The sources for the request",
    )
    response: str | None = Field(
        None,
        description="The response to the request",
    )
    langchain_run_id: UUID | None = Field(
        None,
        description="The ID of the langchain run for the request",
    )
    score: int | None = Field(
        None,
        description="The score of the request",
    )
    status: str = Field(..., description="The status of the request")
    response_received_at: int | None = Field(
        None,
        description="The timestamp of when the response was received",
    )
    sent_at: int = Field(
        default_factory=lambda: int(datetime.now().timestamp()),
        description="The timestamp of the request",
    )
    is_processed: bool = Field(
        False,
        description="Whether the request has been processed",
    )
    is_example: bool = Field(
        False,
        description="Whether the request is an example",
    )

    def to_firestore(self) -> dict[str, Any]:
        """
        Returns the Request as a dict for Firestore.
        """
        return {
            "uuid": str(self.uuid),
            "prompt": self.prompt,
            "messages": [
                {
                    **message.dict(),
                    "type": message.type,
                }
                for message in self.messages
            ],
            "sources": [source.dict() for source in self.sources],
            "response": self.response,
            "status": self.status,
            "langchain_run_id": str(self.langchain_run_id),
            "score": self.score,
            "sent_at": self.sent_at,
            "response_received_at": self.response_received_at,
            "is_processed": self.is_processed,
            "is_example": self.is_example,
        }

    @classmethod
    def from_dict(cls, dict: dict[str, Any]) -> "AskAstroRequest":
        """
        Returns the Request from a dict.
        """
        return cls(
            uuid=UUID(dict["uuid"]),
            prompt=dict["prompt"],
            messages=[
                (HumanMessage if msg.get("type") == "human" else AIMessage)(
                    content=msg["content"],
                    additional_kwargs=msg.get("additional_kwargs", {}),
                )
                for msg in dict["messages"]
            ],
            sources=[Source(**source) for source in dict["sources"]],
            response=dict["response"],
            status=dict["status"],
            langchain_run_id=UUID(dict["langchain_run_id"]),
            score=dict["score"],
            sent_at=dict["sent_at"],
            response_received_at=dict.get("response_received_at"),
            is_processed=dict.get("is_processed", False),
            is_example=dict.get("is_example", False),
        )

api/ask_astro/models/request.py
astronomer/ask-astro
"Contains a function to register all controllers with the app."

from sanic import Sanic

from ask_astro.rest.controllers.list_recent_requests import on_list_recent_requests
from ask_astro.rest.controllers.get_request import on_get_request
from ask_astro.rest.controllers.post_request import on_post_request
from ask_astro.rest.controllers.submit_feedback import on_submit_feedback

from logging import getLogger

logger = getLogger(__name__)


def register_routes(api: Sanic):
    """
    Registers all controllers with the app.
    """

    api.add_route(
        on_list_recent_requests,
        "/requests",
        methods=["GET"],
        name="list_recent_requests",
    )
    logger.info("Registered GET /requests controller")

    api.add_route(
        on_get_request,
        "/requests/<request_id:uuid>",
        methods=["GET"],
        name="get_request",
    )
    logger.info("Registered GET /requests/<request_id> controller")

    api.add_route(
        on_post_request,
        "/requests",
        methods=["POST"],
        name="post_request",
    )
    logger.info("Registered POST /requests controller")

    api.add_route(
        on_submit_feedback,
        "/requests/<request_id:uuid>/feedback",
        methods=["POST"],
        name="submit_feedback",
    )
    logger.info("Registered POST /requests/<request_id>/feedback controller")

api/ask_astro/rest/controllers/__init__.py
astronomer/ask-astro
"""
Handles GET requests to the /ask/{question_id} endpoint.
"""

from uuid import UUID
from sanic import json, Request
from sanic_ext import openapi

from ask_astro.config import FirestoreCollections
from ask_astro.clients.firestore import firestore_client
from ask_astro.models.request import AskAstroRequest

from logging import getLogger

logger = getLogger(__name__)


@openapi.definition(response=AskAstroRequest.schema_json())
async def on_get_request(_: Request, request_id: UUID):
    """
    Handles GET requests to the /requests/{request_id} endpoint.
    """
    logger.info("Received GET request for request %s", request_id)
    request = await (
        firestore_client.collection(FirestoreCollections.requests)
        .document(str(request_id))
        .get()
    )
    logger.info("Request %s exists: %s", request_id, request.exists)

    if not request.exists:
        return json({"error": "Question not found"}, status=404)

    return json(request.to_dict(), status=200)

api/ask_astro/rest/controllers/get_request.py
astronomer/ask-astro
"""
Handles GET requests to the /ask/{question_id} endpoint.
"""

from sanic import json, Request
from sanic_ext import openapi

from pydantic.v1 import BaseModel, Field

from ask_astro.config import FirestoreCollections
from ask_astro.clients.firestore import firestore_client
from ask_astro.models.request import AskAstroRequest


class RecentRequestsResponse(BaseModel):
    requests: list[AskAstroRequest] = Field(
        default_factory=list,
        description="The requests",
    )

    def to_dict(self):
        return {
            "requests": [request.to_firestore() for request in self.requests],
        }


@openapi.definition(
    response=RecentRequestsResponse.schema_json(),
)
async def on_list_recent_requests(_: Request):
    """
    Handles GET requests to the /requests endpoint.
    """
    # list the most recent 5 requests
    requests = await (
        firestore_client.collection(FirestoreCollections.requests)
        .order_by("sent_at", direction="DESCENDING")
        .where("status", "==", "complete")
        .where("is_example", "==", True)
        .limit(12)
        .get()
    )

    return json(
        RecentRequestsResponse(
            requests=[
                AskAstroRequest.from_dict(request.to_dict()) for request in requests
            ]
        ).to_dict(),
        status=200,
    )

api/ask_astro/rest/controllers/list_recent_requests.py
astronomer/ask-astro
"""
Handles GET requests to the /ask/{question_id} endpoint.
"""

import uuid

from logging import getLogger

from sanic import json, Request
from sanic_ext import openapi

from pydantic.v1 import BaseModel, Field
from langchain.schema import AIMessage, HumanMessage

from ask_astro.config import FirestoreCollections
from ask_astro.models.request import AskAstroRequest
from ask_astro.services.questions import answer_question
from ask_astro.clients.firestore import firestore_client

logger = getLogger(__name__)


class PostRequestResponse(BaseModel):
    request_uuid: str = Field(..., description="The UUID of the request")


class PostRequestBody(BaseModel):
    prompt: str = Field(..., description="The prompt for the request")
    from_request_uuid: str | None = Field(
        None,
        description="The UUID of the request to continue",
    )


@openapi.definition(
    response=PostRequestResponse.schema(),
    body=PostRequestBody.schema(),
)
async def on_post_request(request: Request):
    """
    Handles POST requests to the /requests endpoint.
    """
    if "prompt" not in request.json:
        return json({"error": "prompt is required"}, status=400)

    messages = []
    if "from_request_uuid" in request.json:
        from_request_uuid = request.json["from_request_uuid"]
        logger.info("Received request to continue %s", from_request_uuid)

        from_request = await (
            firestore_client.collection(FirestoreCollections.requests)
            .document(from_request_uuid)
            .get()
        )
        if not from_request.exists:
            return json(
                {
                    "error": "from_request_uuid not found",
                    "from_request_uuid": from_request_uuid,
                },
                status=404,
            )

        from_request = AskAstroRequest.from_dict(from_request.to_dict())

        messages = from_request.messages if from_request else []
        messages.append(
            HumanMessage(
                content=from_request.prompt,
                additional_kwargs={
                    "ts": from_request.sent_at,
                },
            )
        )
        messages.append(
            AIMessage(
                content=from_request.response,
                additional_kwargs={
                    "ts": from_request.response_received_at,
                    "ask_astro_request_uuid": from_request_uuid,
                },
            )
        )

    req = AskAstroRequest(
        uuid=uuid.uuid1(),
        prompt=request.json["prompt"],
        status="in_progress",
        messages=messages,
    )

    request.app.add_task(lambda: answer_question(req))

    return json(
        PostRequestResponse(
            request_uuid=str(req.uuid),
        ).dict(),
        status=200,
    )

api/ask_astro/rest/controllers/post_request.py
astronomer/ask-astro
"""
Handles POST requests to the /requests/{question_id}/feedback endpoint.
"""

from uuid import UUID

from sanic import Request, HTTPResponse
from sanic_ext import openapi

from pydantic.v1 import BaseModel, Field

from ask_astro.services.feedback import submit_feedback


class PostRequestBody(BaseModel):
    positive: bool = Field(..., description="Whether the feedback is positive")


@openapi.definition(
    body=PostRequestBody.schema(),
)
async def on_submit_feedback(request: Request, request_id: UUID):
    """
    Handles POST requests to the /requests/{request_id}/feedback endpoint.
    """
    positive = request.json["positive"]

    await submit_feedback(str(request_id), positive, {"source": "api"})

    return HTTPResponse(status=200)

api/ask_astro/rest/controllers/submit_feedback.py
astronomer/ask-astro
"Handles app mention events from Slack"
import asyncio

from typing import Any

from ask_astro.config import FirestoreCollections
from ask_astro.clients.firestore import firestore_client
from ask_astro.clients.langsmith_ import langsmith_client

from logging import getLogger

logger = getLogger(__name__)


async def submit_feedback(
    request_id: str, correct: bool, source_info: dict[str, Any] | None
):
    """
    Submits feedback for a request. Writes to firestore and langsmith.
    """
    logger.info("Submitting feedback for request %s: %s", request_id, correct)
    # first, get the request from the database
    request = await (
        firestore_client.collection(FirestoreCollections.requests)
        .document(request_id)
        .get()
    )

    if not request.exists:
        raise ValueError(f"Request {request_id} does not exist")

    langchain_run_id = request.to_dict().get("langchain_run_id")
    if not langchain_run_id:
        raise ValueError(f"Request {request_id} does not have a langchain run id")

    # update the db and langsmith
    async with asyncio.TaskGroup() as tg:
        # update just the score field
        tg.create_task(
            firestore_client.collection(FirestoreCollections.requests)
            .document(request_id)
            .update({"score": 1 if correct else 0})
        )

        tg.create_task(
            asyncio.to_thread(
                lambda: langsmith_client.create_feedback(
                    key="correctness",
                    run_id=langchain_run_id,
                    score=1 if correct else 0,
                    source_info=source_info,
                )
            )
        )

api/ask_astro/services/feedback.py
astronomer/ask-astro
"Handles app mention events from Slack"
import asyncio
import time

from langchain import callbacks

from ask_astro.config import FirestoreCollections
from ask_astro.clients.firestore import firestore_client
from ask_astro.models.request import AskAstroRequest, Source
from ask_astro.chains.answer_question import answer_question_chain

from logging import getLogger

logger = getLogger(__name__)


async def answer_question(request: AskAstroRequest):
    """
    Performs the actual question answering logic. Writes to the request object.
    """
    try:
        # first, mark the request as in_progress and add it to the database
        request.status = "in_progress"
        await firestore_client.collection(FirestoreCollections.requests).document(
            str(request.uuid)
        ).set(request.to_firestore())

        # then, run the question answering chain
        with callbacks.collect_runs() as cb:
            result = await asyncio.to_thread(
                lambda: answer_question_chain(
                    {
                        "question": request.prompt,
                        "chat_history": [],
                        "messages": request.messages,
                    },
                    metadata={"request_id": str(request.uuid)},
                )
            )
            request.langchain_run_id = cb.traced_runs[0].id

        logger.info("Question answering chain finished with result %s", result)

        # update the request in the database
        request.status = "complete"
        request.response = result["answer"]
        request.response_received_at = int(time.time())
        request.sources = [
            Source(
                name=doc.metadata.get("docLink"),
                snippet=doc.page_content,
            )
            for doc in result.get("source_documents", [])
            if doc.metadata.get("docLink", "").startswith("https://")
        ]

        await firestore_client.collection(FirestoreCollections.requests).document(
            str(request.uuid)
        ).set(request.to_firestore())

    except Exception as e:
        # if there's an error, mark the request as errored and add it to the database
        request.status = "error"
        request.response = "Sorry, something went wrong. Please try again later."
        await firestore_client.collection(FirestoreCollections.requests).document(
            str(request.uuid)
        ).set(request.to_firestore())

        # then propogate the error
        raise e

api/ask_astro/services/questions.py
astronomer/ask-astro
"Generates the slack app and the slack app handler."

from slack_bolt.oauth.async_oauth_settings import AsyncOAuthSettings
from slack_bolt.adapter.sanic import AsyncSlackRequestHandler
from slack_bolt.app.async_app import AsyncApp

from ask_astro.config import SlackAppConfig, FirestoreCollections
from ask_astro.stores import (
    AsyncFirestoreInstallationStore,
    AsyncFirestoreOAuthStateStore,
)


oauth_settings = AsyncOAuthSettings(
    client_id=SlackAppConfig.client_id,
    client_secret=SlackAppConfig.client_secret,
    scopes=[
        "commands",
        "app_mentions:read",
        "channels:read",
        "channels:history",
        "groups:read",
        "groups:history",
        "chat:write",
        "reactions:read",
        "reactions:write",
        "users:read",
        "users:read.email",
        "team:read",
        "im:history",
        "mpim:history",
        "files:read",
    ],
    installation_store=AsyncFirestoreInstallationStore(
        collection=FirestoreCollections.installation_store,
    ),
    state_store=AsyncFirestoreOAuthStateStore(
        expiration_seconds=600,
        collection=FirestoreCollections.state_store,
    ),
)


slack_app = AsyncApp(
    signing_secret=SlackAppConfig.signing_secret,
    oauth_settings=oauth_settings,
)
app_handler = AsyncSlackRequestHandler(slack_app)

api/ask_astro/slack/app.py
astronomer/ask-astro
"Contains a function to register all controllers with the app."

from slack_bolt.async_app import AsyncApp
from ask_astro.slack.controllers.mention import on_mention
from ask_astro.slack.controllers.feedback.good import handle_feedback_good
from ask_astro.slack.controllers.feedback.bad import handle_feedback_bad

from logging import getLogger

logger = getLogger(__name__)


def register_controllers(app: AsyncApp):
    """
    Registers all controllers with the app.
    """

    app.event("app_mention")(on_mention)
    logger.info("Registered event:app_mention controller")

    app.action("feedback_good")(handle_feedback_good)
    logger.info("Registered action:feedback_good controller")

    app.action("feedback_bad")(handle_feedback_bad)
    logger.info("Registered action:feedback_bad controller")

api/ask_astro/slack/controllers/__init__.py
astronomer/ask-astro
import asyncio

from logging import getLogger

from typing import Any

from slack_bolt.async_app import AsyncAck, AsyncRespond
from slack_sdk.web.async_client import AsyncWebClient
from slack_sdk.errors import SlackApiError
from asyncio import TaskGroup

from ask_astro.services.feedback import submit_feedback

logger = getLogger(__name__)


async def handle_feedback_bad(
    body: dict[str, Any], ack: AsyncAck, respond: AsyncRespond, client: AsyncWebClient
):
    await ack()

    try:
        user = body["user"]["id"]
        channel = body["channel"]["id"]
        thread_ts = body["message"]["thread_ts"]
        message_ts = body["message"]["ts"]
        value = body["actions"][0]["value"]
    except KeyError as e:
        logger.error(f"Missing key: {e}")
        return

    request_id = value.split(":")[0]

    await submit_feedback(
        request_id, False, source_info={"type": "slack", "user": user}
    )

    async with TaskGroup() as tg:
        tg.create_task(
            respond(
                f"😥 Thank you for your feedback, <@{user}>!",
                thread_ts=thread_ts,
                replace_original=False,
                response_type="ephemeral",
            )
        )

        async def update_reaction():
            try:
                await client.reactions_add(
                    name="thumbsdown",
                    channel=channel,
                    timestamp=message_ts,
                )
            except SlackApiError as e:
                # ignore the error if the reaction already exists
                if e.response["error"] != "already_reacted":
                    raise e

            try:
                await client.reactions_remove(
                    name="thumbsup",
                    channel=channel,
                    timestamp=message_ts,
                )
            except:
                pass

        tg.create_task(
            update_reaction(),
        )

api/ask_astro/slack/controllers/feedback/bad.py
astronomer/ask-astro
import asyncio

from logging import getLogger

from typing import Any

from slack_bolt.async_app import AsyncAck, AsyncRespond
from slack_sdk.web.async_client import AsyncWebClient
from slack_sdk.errors import SlackApiError
from asyncio import TaskGroup

from ask_astro.services.feedback import submit_feedback

logger = getLogger(__name__)


async def handle_feedback_good(
    body: dict[str, Any], ack: AsyncAck, respond: AsyncRespond, client: AsyncWebClient
):
    await ack()

    try:
        user = body["user"]["id"]
        channel = body["channel"]["id"]
        thread_ts = body["message"]["thread_ts"]
        message_ts = body["message"]["ts"]
        value = body["actions"][0]["value"]
    except KeyError as e:
        logger.error(f"Missing key: {e}")
        return

    request_id = value.split(":")[0]

    await submit_feedback(request_id, True, source_info={"type": "slack", "user": user})

    async with TaskGroup() as tg:
        tg.create_task(
            respond(
                f"☺️ Thank you for your feedback, <@{user}>!",
                thread_ts=thread_ts,
                replace_original=False,
                response_type="ephemeral",
            )
        )

        async def update_reaction():
            try:
                await client.reactions_add(
                    name="thumbsup",
                    channel=channel,
                    timestamp=message_ts,
                )
            except SlackApiError as e:
                # ignore the error if the reaction already exists
                if e.response["error"] != "already_reacted":
                    raise e

            try:
                await client.reactions_remove(
                    name="thumbsdown",
                    channel=channel,
                    timestamp=message_ts,
                )
            except:
                pass

        tg.create_task(
            update_reaction(),
        )

api/ask_astro/slack/controllers/feedback/good.py
astronomer/ask-astro
"Handles app mention events from Slack"

from asyncio import TaskGroup
import asyncio
import uuid
from typing import Any

from slack_bolt.async_app import AsyncAck, AsyncSay
from slack_sdk.web.async_client import AsyncWebClient
from slack_sdk.errors import SlackApiError
from langchain.schema import HumanMessage, AIMessage

from ask_astro.services.questions import answer_question
from ask_astro.slack.utils import get_blocks, markdown_to_slack
from ask_astro.models.request import AskAstroRequest

from logging import getLogger

logger = getLogger(__name__)


async def on_mention(
    body: dict[str, Any],
    ack: AsyncAck,
    say: AsyncSay,
    client: AsyncWebClient,
):
    await ack()

    # pull out the relevant fields from the body
    try:
        text = body["event"]["text"]
        bot_id = body["authorizations"][0]["user_id"]
        channel_id = body["event"]["channel"]
        ts = body["event"]["ts"]
    except KeyError as e:
        logger.error(f"Missing key: {e}")
        return

    logger.info("Received question (%s): %s", ts, text)
    thread_ts = body["event"].get("thread_ts")
    logger.info("In thread %s", thread_ts)

    # Preprocess text to remove leading bot mentions
    text = text.strip().removeprefix(f"<@{bot_id}>").strip()

    # create the request state object and add it to the database
    request = AskAstroRequest(
        uuid=uuid.uuid1(),
        prompt=text,
        status="in_progress",
    )

    async with TaskGroup() as tg:
        # add a thinking reaction for immediate feedback
        async def try_add_thinking_reaction():
            try:
                await client.reactions_add(
                    name="thought_balloon",
                    channel=channel_id,
                    timestamp=ts,
                )
            except SlackApiError as exc:
                # this is a non-fatal error, so we can just log it and move on
                logger.warning("Failed to add thought_balloon reaction", exc_info=exc)

        tg.create_task(try_add_thinking_reaction())

    logger.info("Acknowledged question")

    try:
        # get messages from the thread, if there is a thread
        if thread_ts:
            # TODO: what happens if there are more than 100 messages?
            slack_messages = await client.conversations_replies(
                channel=channel_id,
                ts=thread_ts,
                limit=100,
            )

            logger.debug("Messages in thread: %s", slack_messages)

            # turn the slack messages into a list of langchain messages
            request.messages = [
                (AIMessage if msg.get("bot_id") == bot_id else HumanMessage)(
                    content=msg["text"], additional_kwargs=msg
                )
                for msg in slack_messages["messages"]
                if msg.get("type") == "message" and msg.get("ts") != ts
            ]

        await answer_question(request)
        response = markdown_to_slack(request.response)

        source_links = ", ".join(
            (
                f"<{link}|[{n_}]>"
                for n_, link in enumerate(
                    [*set(s.name for s in request.sources if s.name)],
                    start=1,
                )
            )
        )

        resp = get_blocks(
            "message.jinja2",
            message=f"{response[:2999]}…" if len(response) > 3000 else response,
            sources=f"📚 Related: {source_links}" if source_links else None,
            feedback_value=f"{request.uuid}:{request.langchain_run_id}",
        )

        async with TaskGroup() as tg:
            # remove the reaction from the message
            tg.create_task(
                client.reactions_remove(
                    name="thought_balloon",
                    channel=channel_id,
                    timestamp=ts,
                )
            )

            # write the response
            tg.create_task(
                say(
                    text=response,
                    blocks=resp,
                    thread_ts=ts,
                    unfurl_links=False,
                    unfurl_media=False,
                )
            )

    except Exception as e:
        # remove the thinking reaction from the message
        await client.reactions_remove(
            name="thought_balloon",
            channel=channel_id,
            timestamp=ts,
        )

        # add a failure reaction
        await client.reactions_add(
            name="x",
            channel=channel_id,
            timestamp=ts,
        )

        # write the failure message
        await say(
            text="Sorry, I couldn't answer your question. Please try again later.",
            thread_ts=ts,
        )

        raise e

api/ask_astro/slack/controllers/mention.py
astronomer/ask-astro
import json
import re
import jinja2


def markdown_to_slack(message: str):
    regexp_replacements = (
        (re.compile("^- ", flags=re.M), "• "),
        (re.compile("^  - ", flags=re.M), "  ◦ "),
        (re.compile("^    - ", flags=re.M), "    ⬩ "),  # ◆
        (re.compile("^      - ", flags=re.M), "    ◽ "),
        (re.compile("^#+ (.+)$", flags=re.M), r"*\1*"),
        (re.compile("\*\*"), "*"),
        (re.compile(r"\[(.+)]\((.+)\)"), r"<\2|\1>"),
        (re.compile("```\\S+\\n"), r"```\n"),
    )

    for regex, replacement in regexp_replacements:
        message = regex.sub(replacement, message)

    return message


def get_blocks(block: str, **kwargs):
    env = jinja2.Environment(
        loader=jinja2.FileSystemLoader("ask_astro/templates"), autoescape=True
    )
    return json.loads(env.get_template(block).render(kwargs))["blocks"]

api/ask_astro/slack/utils.py
astronomer/ask-astro
"Re-exports stores for easier importing."

from .installation_store import AsyncFirestoreInstallationStore
from .oauth_state_store import AsyncFirestoreOAuthStateStore

api/ask_astro/stores/__init__.py
astronomer/ask-astro
import logging
from logging import Logger
from typing import Optional

import google.cloud.firestore

from slack_sdk.oauth.installation_store.async_installation_store import (
    AsyncInstallationStore,
)
from slack_sdk.oauth.installation_store.models.bot import Bot
from slack_sdk.oauth.installation_store.models.installation import Installation


class AsyncFirestoreInstallationStore(AsyncInstallationStore):
    def __init__(
        self,
        *,
        collection: str,
        historical_data_enabled: bool = True,
        client_id: Optional[str] = None,
        logger: Logger = logging.getLogger(__name__),
    ):
        firestore_client = google.cloud.firestore.AsyncClient()
        self.fp = firestore_client.field_path
        self.collection = firestore_client.collection(collection)
        self.historical_data_enabled = historical_data_enabled
        self.client_id = client_id
        self._logger = logger

    @property
    def logger(self) -> Logger:
        if self._logger is None:
            self._logger = logging.getLogger(__name__)
        return self._logger

    async def async_save(self, installation: Installation):
        none = "none"
        e_id = installation.enterprise_id or none
        t_id = installation.team_id or none

        doc_ref = self.collection.document(f"{e_id}-{t_id}")

        try:
            await doc_ref.create({})
        except google.cloud.exceptions.Conflict:
            pass

        await self.async_save_bot(installation.to_bot())

        if self.historical_data_enabled:
            history_version: str = str(installation.installed_at)

            # per workspace
            entity = installation.__dict__
            await doc_ref.update(
                {
                    self.fp("installer", "latest"): entity,
                    self.fp("installer", history_version): entity,
                },
            )

            # per workspace per user
            u_id = installation.user_id or none
            entity = installation.__dict__
            await doc_ref.update(
                {
                    self.fp("installer", u_id, "latest"): entity,
                    self.fp("installer", u_id, history_version): entity,
                },
            )
        else:
            u_id = installation.user_id or none
            entity = installation.__dict__
            await doc_ref.update(
                {self.fp("installer", u_id, "latest"): entity},
            )

    async def async_save_bot(self, bot: Bot):
        none = "none"
        e_id = bot.enterprise_id or none
        t_id = bot.team_id or none
        doc_ref = self.collection.document(f"{e_id}-{t_id}")

        if self.historical_data_enabled:
            history_version: str = str(bot.installed_at)

            entity = bot.__dict__
            await doc_ref.update(
                {
                    self.fp("bot", "latest"): entity,
                    self.fp(f"bot", history_version): entity,
                },
            )
        else:
            entity = bot.__dict__
            await doc_ref.update(
                {self.fp("bot", "latest"): entity},
            )

    async def async_find_bot(
        self,
        *,
        enterprise_id: Optional[str],
        team_id: Optional[str],
        is_enterprise_install: Optional[bool] = False,
    ) -> Optional[Bot]:
        none = "none"
        e_id = enterprise_id or none
        t_id = team_id or none
        if is_enterprise_install:
            t_id = none

        doc_ref = self.collection.document(f"{e_id}-{t_id}")

        if data := (await doc_ref.get([self.fp("bot", "latest")])).to_dict():
            return Bot(**data["bot"]["latest"])
        else:
            message = f"Installation data missing for enterprise: {e_id}, team: {t_id}"
            self.logger.debug(message)
            return None

    async def async_find_installation(
        self,
        *,
        enterprise_id: Optional[str],
        team_id: Optional[str],
        user_id: Optional[str] = None,
        is_enterprise_install: Optional[bool] = False,
    ) -> Optional[Installation]:
        none = "none"
        e_id = enterprise_id or none
        t_id = team_id or none
        if is_enterprise_install:
            t_id = none

        doc_ref = self.collection.document(f"{e_id}-{t_id}")

        if user_id:
            data = (
                (await doc_ref.get([self.fp("installer", user_id, "latest")]))
                .to_dict()
                .get("installer", {})
                .get(user_id, {})
                .get("latest")
            )
        else:
            data = (
                (await doc_ref.get([self.fp("installer", "latest")]))
                .to_dict()
                .get("installer", {})
                .get("latest")
            )

        if data:
            installation = Installation(**data) if data else None

            if installation is not None and user_id is not None:
                # Retrieve the latest bot token, just in case
                # See also: https://github.com/slackapi/bolt-python/issues/664
                latest_bot_installation = await self.async_find_installation(
                    enterprise_id=enterprise_id,
                    team_id=team_id,
                    is_enterprise_install=is_enterprise_install,
                )
                if (
                    latest_bot_installation is not None
                    and installation.bot_token != latest_bot_installation.bot_token
                ):
                    # NOTE: this logic is based on the assumption that every single installation has bot scopes
                    # If you need to installation patterns without bot scopes in the same S3 bucket,
                    # please fork this code and implement your own logic.
                    installation.bot_id = latest_bot_installation.bot_id
                    installation.bot_user_id = latest_bot_installation.bot_user_id
                    installation.bot_token = latest_bot_installation.bot_token
                    installation.bot_scopes = latest_bot_installation.bot_scopes
                    installation.bot_refresh_token = (
                        latest_bot_installation.bot_refresh_token
                    )
                    installation.bot_token_expires_at = (
                        latest_bot_installation.bot_token_expires_at
                    )

            return installation
        else:
            message = f"Installation data missing for enterprise: {e_id}, team: {t_id}"
            self.logger.debug(message)
            return None

    async def async_delete_bot(
        self, *, enterprise_id: Optional[str], team_id: Optional[str]
    ) -> None:
        none = "none"
        e_id = enterprise_id or none
        t_id = team_id or none
        doc_ref = self.collection.document(f"{e_id}-{t_id}")
        await doc_ref.update({"bot": google.cloud.firestore.DELETE_FIELD})

    async def async_delete_installation(
        self,
        *,
        enterprise_id: Optional[str],
        team_id: Optional[str],
        user_id: Optional[str] = None,
    ) -> None:
        none = "none"
        e_id = enterprise_id or none
        t_id = team_id or none
        doc_ref = self.collection.document(f"{e_id}-{t_id}")
        await doc_ref.update({"installer": google.cloud.firestore.DELETE_FIELD})

api/ask_astro/stores/installation_store.py
astronomer/ask-astro
import datetime
import logging
from datetime import timedelta
from logging import Logger
from typing import Optional

import google.cloud.firestore
from slack_sdk.oauth.state_store.async_state_store import AsyncOAuthStateStore


class AsyncFirestoreOAuthStateStore(AsyncOAuthStateStore):
    def __init__(
        self,
        *,
        collection: str,
        expiration_seconds: int,
        client_id: Optional[str] = None,
        logger: Logger = logging.getLogger(__name__),
    ):
        firestore_client = google.cloud.firestore.AsyncClient()
        self.collection = firestore_client.collection(collection)
        self.expiration_seconds = expiration_seconds

        self.client_id = client_id
        self._logger = logger

    @property
    def logger(self) -> Logger:
        if self._logger is None:
            self._logger = logging.getLogger(__name__)
        return self._logger

    async def async_issue(self, *args, **kwargs) -> str:
        doc_ref = self.collection.document()
        await doc_ref.set({"timestamp": datetime.datetime.now().astimezone()})
        return doc_ref.id

    async def async_consume(self, state: str) -> bool:
        doc_ref = self.collection.document(state)
        created = (await doc_ref.get()).get("timestamp")

        if created:
            expiration = created + timedelta(seconds=self.expiration_seconds)
            still_valid: bool = datetime.datetime.now().astimezone() < expiration
            return still_valid

        return False

api/ask_astro/stores/oauth_state_store.py
astronomer/ask-astro