mdbt 0.4.27__py3-none-any.whl

mdbt/ai_core.py

@@ -0,0 +1,116 @@
+import os
+import re
+import subprocess
+from typing import Dict
+from typing import List
+
+import openai
+from snowflake.connector import DatabaseError
+
+from mdbt.core import Core
+from mdbt.prompts import Prompts
+
+# Have to load env before import openai package.
+# flake8: noqa: E402
+
+
+class AiCore(Core):
+
+    def __init__(self, model: str = "gpt-4.1", test_mode: bool = False):
+        super().__init__(test_mode=test_mode)
+        self.model = model
+        # Make sure you have OPENAI_API_KEY set in your environment variables.
+        self.client = openai.OpenAI()
+
+        self.prompts = Prompts()
+
+    def send_message(self, _messages: List[Dict[str, str]]) -> object:
+        print("Sending to API")
+        completion = self.client.chat.completions.create(
+            model=self.model, messages=_messages
+        )
+        return completion.choices[0].message.content
+
+    @staticmethod
+    def read_file(path: str) -> str:
+        with open(path, "r") as file:
+            return file.read()
+
+    @staticmethod
+    def is_file_committed(file_path):
+        try:
+            # Check the Git status of the file
+            subprocess.run(
+                ["git", "ls-files", "--error-unmatch", file_path],
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            # If the file is tracked, check if it has any modifications
+            status_result = subprocess.run(
+                ["git", "status", "--porcelain", file_path], stdout=subprocess.PIPE
+            )
+            status_output = status_result.stdout.decode().strip()
+            # If the output is empty, file is committed and has no modifications
+            return len(status_output) == 0
+        except subprocess.CalledProcessError:
+            # The file is either untracked or does not exist
+            return False
+
+    def _get_sample_data_from_snowflake(self, model_names: List[str]) -> Dict[str, str]:
+        """
+        Compiles the target model to SQL, then breaks out each sub query and CTE into a separate SQL strings, executing
+        each to get a sample of the data.
+        Args:
+            model_names: A list of target model names to pull sample data from.
+
+        Returns:
+            A dictionary of model names and their sample data in CSV format.
+        """
+        sample_results = {}
+        for model_name in model_names:
+            print(f"Getting sample data for {model_name}")
+            args = ["--select", model_name]
+            cmd = "compile"
+            results = self.execute_dbt_command_capture(cmd, args)
+            extracted_sql = self.extract_sql(results)
+            sample_sql = self.build_sample_sql(extracted_sql)
+            try:
+                self._cur.execute(sample_sql)
+            except DatabaseError as e:
+                print(f"Error executing sample SQL for {model_name}")
+                print(e)
+                print("\n\n" + sample_sql + "\n\n")
+                raise e
+            tmp_df = self._cur.fetch_pandas_all()
+            sample_results[model_name] = tmp_df.to_csv(index=False)
+        print(f"Sample results: {sample_results}")
+        return sample_results
+
+    @staticmethod
+    def build_sample_sql(sql: str) -> str:
+        sql = f"""
+            with tgt_table as (
+                {sql}
+            )
+            select *
+            from tgt_table
+            sample (10 rows)
+            """
+        return sql
+
+    @staticmethod
+    def extract_sql(log):
+        sql_lines = [line for line in log.splitlines() if not re.match(r"--\s.*", line)]
+
+        keyword_line_index = 0
+        for i, line in enumerate(sql_lines):
+            if "Compiled node" in line:
+                keyword_line_index = i + 1
+                break
+
+        sql_lines = sql_lines[keyword_line_index:]
+
+        # Join the remaining lines and remove escape sequences
+        sql = "\n".join(sql_lines).replace("\x1b[0m", "").strip()
+        return sql

mdbt/build_dbt_docs_ai.py

@@ -0,0 +1,147 @@
+import subprocess
+
+import pyperclip
+from dotenv import find_dotenv
+from dotenv import load_dotenv
+
+from mdbt.ai_core import AiCore
+from mdbt.prompts import Prompts
+
+load_dotenv(find_dotenv("../.env"))
+load_dotenv(find_dotenv(".env"))
+
+
+class BuildDBTDocs(AiCore):
+    """
+    # Make sure you have OPENAI_API_KEY set in your environment variables.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def main(self, model_name, sys_context, is_new=False):
+        if model_name.endswith(".sql"):
+            model_name = model_name[:-4]
+        if not is_new:
+            print(
+                """
+            1) Build new DBT documentation.
+            2) Check existing DBT documentation against model for missing definitions.
+            """
+            )
+            mode = int(input())
+        else:
+            mode = 1
+        print("Getting file.")
+        sql_file_path = self.get_file_path(model_name)
+
+        if "l4" in sql_file_path.lower() or "l3" in sql_file_path.lower():
+            system_instructions = Prompts().dbt_docs_gte_l3_prompt
+        else:
+            system_instructions = Prompts().dbt_docs_lte_l2_prompt
+
+        if sys_context:
+            system_instructions += f"\nContext about system docs are generated for: \n{sys_context}\n"
+
+        sample_data = self._get_sample_data_from_snowflake([model_name])
+
+        system_instructions = system_instructions + sample_data[model_name]
+
+        # Might bring this back in the future.
+        extra_info = ""
+
+        if mode == 1:
+            # Build new documentation
+            user_input = self.build_user_msg_mode_1(sql_file_path, extra_info)
+            yml_file_path = sql_file_path.replace(".sql", ".yml")
+        elif mode == 2:
+            # Check existing documentation
+            yml_file_path = sql_file_path[:-4] + ".yml"
+            user_input = self.build_user_msg_mode_2(
+                sql_file_path, yml_file_path, extra_info
+            )
+        else:
+            print(mode)
+            raise ValueError("Invalid mode")
+
+        messages = [
+            {"role": "user", "content": system_instructions + "\n" + user_input}
+        ]
+
+        assistant_responses = []
+        result = self.send_message(messages)
+        assistant_responses.append(result)
+
+        messages.append({"role": "assistant", "content": assistant_responses[0]})
+        print(assistant_responses[0])
+        output = assistant_responses[0]
+        # Check for ``` at end of output (str) and remove
+        # Remove trailing markdown code fences if present
+        lines = output.split('\n')
+        if lines and '```' in lines[-1].strip():
+            lines = lines[:-1]
+        elif len(lines) > 1 and '```' in lines[-2].strip():
+            # Remove the second-to-last line if it's a code fence
+            lines.pop(-2)
+        output = '\n'.join(lines)
+        if not is_new:
+            clip_or_file = input(
+                f"1. to copy to clipboard\n2, to write to file ({yml_file_path}\n:"
+            )
+        else:
+            clip_or_file = "2"
+
+        if clip_or_file == "1":
+            print("Output copied to clipboard")
+            pyperclip.copy(output)
+        elif clip_or_file == "2":
+            if mode == 2:
+                # Make a backup of the current YML file.
+                self.backup_existing_yml_file(yml_file_path)
+            output = assistant_responses[0].split("\n")
+            # output = output[1:-1]
+            output = output[1:]
+            output = "\n".join(output)
+            with open(yml_file_path, "w") as file:
+                file.write(output)
+            if not self.is_file_committed(yml_file_path):
+                if not is_new:
+                    commit_file = input("Press 1 to add to git, any other key to byapss: ")
+                else:
+                    commit_file = "1"
+
+                if commit_file == "1":
+                    subprocess.run(["git", "add", yml_file_path])
+
+    @staticmethod
+    def backup_existing_yml_file(yml_file_path):
+        with open(yml_file_path, "r") as file:
+            yml_content = file.read()
+        with open(yml_file_path + ".bak", "w") as file:
+            file.write(yml_content)
+
+    def build_user_msg_mode_1(self, _sql_file_path: str, extra_info: str) -> str:
+        self.read_file(_sql_file_path)
+        model_name = _sql_file_path.split("/")[-1].split(".")[0]
+        prompt_str = f"Build new DBT documentation for the following SQL query with model name {model_name}"
+        if len(extra_info):
+            prompt_str += f"\n{extra_info}"
+
+        return prompt_str
+
+    def build_user_msg_mode_2(
+        self, _sql_file_path: str, _yml_file_path: str, extra_info: str
+    ) -> str:
+        self.read_file(_sql_file_path)
+        yml = self.read_file(_yml_file_path)
+        model_name = _sql_file_path.split("/")[-1].split(".")[0]
+        prompt_str = f"Check for missing columns in the following DBT documentation for the following SQL query with model name {model_name}. Identify any columns in the DBT documentation that do not exist in the SQL and comment them out."
+        if len(extra_info):
+            prompt_str += f"\n {extra_info}"
+        prompt_str += f"\nYML File Contents:\n{yml}"
+
+        return prompt_str
+
+
+if __name__ == "__main__":
+    BuildDBTDocs().main("revenue_by_dvm")

mdbt/build_unit_test_data_ai.py

@@ -0,0 +1,129 @@
+import logging
+import os
+import re
+import warnings
+from typing import Dict
+
+import pyperclip
+from dotenv import find_dotenv
+from dotenv import load_dotenv
+
+
+load_dotenv(find_dotenv("../.env"))
+load_dotenv(find_dotenv(".env"))
+# flake8: noqa: E402
+from mdbt.ai_core import AiCore
+
+
+# Have to load env before import openai package.
+warnings.simplefilter(action="ignore", category=FutureWarning)
+logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
+
+
+class BuildUnitTestDataAI(AiCore):
+
+    def __init__(self):
+        super().__init__(model="o3-mini")
+
+    def main(self, model_name: str):
+
+        file_path = self.get_file_path(model_name)
+        # Extract the folder immediately after 'models'. Not sure I need to use this just yet, holding on to it for
+        # later.
+        layer_name = file_path.split("/")[1][:2]
+        sub_folder = file_path.split("/")[2]
+        file_name = os.path.splitext(os.path.basename(file_path))[0]
+
+        test_file_path = (
+            f"tests/unit_tests/{layer_name}/{sub_folder}/test_{file_name}.sql"
+        )
+
+        input_sql_file_name = file_path
+
+        input_sql = self.read_file(input_sql_file_name)
+
+        models_in_model_file = self.extract_model_names(input_sql)
+
+        sample_data = self._get_sample_data_from_snowflake(models_in_model_file)
+
+        prompt = self.build_prompt(
+            self.prompts.build_unit_test_prompt.format(model_name=model_name),
+            model_name,
+            input_sql,
+            sample_data,
+        )
+
+        print(f"##################\n{prompt}\n##################")
+
+        messages = [
+            {
+                "role": "user",
+                "content": "You are helping to build unit tests for DBT (database build tools) models.\n"
+                + prompt,
+            },
+        ]
+
+        response = self.send_message(messages)
+
+        output = self._remove_first_and_last_line_from_string(response)
+        print(output)
+
+        clip_or_file = input(
+            f"1. to copy to clipboard\n2, to write to file ({test_file_path}"
+        )
+
+        if clip_or_file == "1":
+            print("Output copied to clipboard")
+            pyperclip.copy(output)
+        elif clip_or_file == "2":
+            # Check if file exists and ask if it should be overwritten.
+            if os.path.exists(test_file_path):
+                overwrite = input(f"File {test_file_path} exists. Overwrite? (y/n)")
+                if overwrite.lower() == "y":
+                    with open(test_file_path, "w") as file:
+                        file.write(output)
+                    print(f"Output written to {test_file_path}")
+            else:
+                with open(test_file_path, "w") as file:
+                    file.write(output)
+                print(f"Output written to {test_file_path}")
+
+    def _remove_first_and_last_line_from_string(self, s: str) -> str:
+        return "\n".join(s.split("\n")[1:-1])
+
+    @staticmethod
+    def extract_model_names(dbt_script):
+        # Regular expression to find all occurrences of {{ ref('model_name') }}
+        pattern = r"\{\{\s*ref\('([^']+)'\)\s*\}\}"
+        # Find all matches in the script
+        model_names = re.findall(pattern, dbt_script)
+        return model_names
+
+    @staticmethod
+    def build_prompt(
+        prompt_template: str,
+        model_name: str,
+        model_sql,
+        sample_models_and_data: Dict[str, str],
+    ):
+        sample_str = ""
+        for model_name, sample_data in sample_models_and_data.items():
+            sample_str += f"""{model_name}: \n{sample_data}\n"""
+
+        output = f"""
+The model name we are building the test for is {model_name}. In the example, this says "model_name". Put this value in that same place.'
+{prompt_template}
+
+The SQL for the model is:
+{model_sql}
+
+Here is sample data for each input model. This just represents a random sample. Use it to create realistic test data, but try to build the test input data so that it tests the logic found within the model, regardless of the particular combination of sample data. Imagine that certain flags might be true or false, even if that flag is always true or false in the sample data.
+
+{sample_str}
+
+"""
+        return output
+
+
+if __name__ == "__main__":
+    BuildUnitTestDataAI().main("avg_client_rev_per_year")