1 files changed, 106 insertions, 0 deletions
diff --git a/server/continuedev/plugins/recipes/AddTransformRecipe/steps.py b/server/continuedev/plugins/recipes/AddTransformRecipe/steps.py
new file mode 100644
index 00000000..61638374
--- /dev/null
+++ b/server/continuedev/plugins/recipes/AddTransformRecipe/steps.py
@@ -0,0 +1,106 @@
+import os
+from textwrap import dedent
+
+from ....core.main import Step
+from ....core.sdk import ContinueSDK, Models
+from ....core.steps import MessageStep
+from ....libs.util.paths import find_data_file
+
+AI_ASSISTED_STRING = "(✨ AI-Assisted ✨)"
+
+
+class SetUpChessPipelineStep(Step):
+    hide: bool = True
+    name: str = "Setup Chess.com API dlt Pipeline"
+
+    async def describe(self, models: Models):
+        return "This step will create a new dlt pipeline that loads data from the chess.com API."
+
+    async def run(self, sdk: ContinueSDK):
+        # running commands to get started when creating a new dlt pipeline
+        await sdk.run(
+            [
+                "python3 -m venv .env",
+                "source .env/bin/activate",
+                "pip install dlt",
+                "dlt --non-interactive init chess duckdb",
+                "pip install -r requirements.txt",
+                "pip install pandas streamlit",  # Needed for the pipeline show step later
+            ],
+            name="Set up Python environment",
+            description=dedent(
+                """\
+            - Create a Python virtual environment: `python3 -m venv .env`
+            - Activate the virtual environment: `source .env/bin/activate`
+            - Install dlt: `pip install dlt`
+            - Create a new dlt pipeline called "chess" that loads data into a local DuckDB instance: `dlt init chess duckdb`
+            - Install the Python dependencies for the pipeline: `pip install -r requirements.txt`"""
+            ),
+        )
+
+
+class AddTransformStep(Step):
+    hide: bool = True
+
+    # e.g. "Use the `python-chess` library to decode the moves in the game data"
+    transform_description: str
+
+    async def run(self, sdk: ContinueSDK):
+        source_name = "chess"
+        filename = f"{source_name}_pipeline.py"
+        abs_filepath = os.path.join(sdk.ide.workspace_directory, filename)
+
+        # Open the file and highlight the function to be edited
+        await sdk.ide.setFileOpen(abs_filepath)
+
+        await sdk.run_step(
+            MessageStep(
+                message=dedent(
+                    """\
+                This step will customize your resource function with a transform of your choice:
+                - Add a filter or map transformation depending on your request
+                - Load the data into a local DuckDB instance
+                - Open up a Streamlit app for you to view the data"""
+                ),
+                name="Write transformation function",
+            )
+        )
+
+        with open(find_data_file("dlt_transform_docs.md")) as f:
+            dlt_transform_docs = f.read()
+
+        prompt = dedent(
+            f"""\
+            Task: Write a transform function using the description below and then use `add_map` or `add_filter` from the `dlt` library to attach it a resource.
+
+            Description: {self.transform_description}
+
+            Here are some docs pages that will help you better understand how to use `dlt`.
+                          
+            {dlt_transform_docs}"""
+        )
+
+        # edit the pipeline to add a transform function and attach it to a resource
+        await sdk.edit_file(
+            filename=filename,
+            prompt=prompt,
+            name=f"Writing transform function {AI_ASSISTED_STRING}",
+        )
+
+        await sdk.wait_for_user_confirmation(
+            "Press Continue to confirm that the changes are okay before we run the pipeline."
+        )
+
+        # run the pipeline and load the data
+        await sdk.run(
+            f"python3 {filename}",
+            name="Run the pipeline",
+            description=f"Running `python3 {filename}` to load the data into a local DuckDB instance",
+        )
+
+        # run a streamlit app to show the data
+        await sdk.run(
+            f"dlt pipeline {source_name}_pipeline show",
+            name="Show data in a Streamlit app",
+            description=f"Running `dlt pipeline {source_name} show` to show the data in a Streamlit app, where you can view and play with the data.",
+        )