Add Switch transpiler with --include-llm-transpiler flag

labs.yml

@@ -64,6 +64,9 @@ commands:
       - name: interactive
         description: "Whether installing in interactive mode, which may prompt for configuration settings."
         default: auto
+      - name: include-llm-transpiler
+        description: "Include LLM-based transpiler in installation"
+        default: false
   - name: describe-transpile
     description: "Describe installed transpilers."
     table_template: |-

pyproject.toml

@@ -447,7 +447,7 @@ bad-functions = ["map", "input"]
 # ignored-parents =
 
 # Maximum number of arguments for function / method.
-max-args = 12
+max-args = 13
 
 # Maximum number of attributes for a class (see R0902).
 max-attributes = 13

src/databricks/labs/lakebridge/cli.py

@@ -625,6 +625,7 @@ def install_transpile(
     w: WorkspaceClient,
     artifact: str | None = None,
     interactive: str | None = None,
+    include_llm_transpiler: bool = False,
     transpiler_repository: TranspilerRepository = TranspilerRepository.user_home(),
 ) -> None:
     """Install or upgrade the Lakebridge transpilers."""
@@ -633,9 +634,13 @@ def install_transpile(
     ctx.add_user_agent_extra("cmd", "install-transpile")
     if artifact:
         ctx.add_user_agent_extra("artifact-overload", Path(artifact).name)
+    if include_llm_transpiler:
+        ctx.add_user_agent_extra("include-llm-transpiler", "true")
     user = w.current_user
     logger.debug(f"User: {user}")
-    transpile_installer = installer(w, transpiler_repository, is_interactive=is_interactive)
+    transpile_installer = installer(
+        w, transpiler_repository, is_interactive=is_interactive, include_llm_transpiler=include_llm_transpiler
+    )
     transpile_installer.run(module="transpile", artifact=artifact)
 
 

src/databricks/labs/lakebridge/deployment/installation.py

@@ -13,6 +13,8 @@
 
 from databricks.labs.lakebridge.config import LakebridgeConfiguration
 from databricks.labs.lakebridge.deployment.recon import ReconDeployment
+from databricks.labs.lakebridge.transpiler.repository import TranspilerRepository
+from databricks.labs.lakebridge.transpiler.installers import SwitchInstaller
 
 logger = logging.getLogger("databricks.labs.lakebridge.install")
 
@@ -116,9 +118,25 @@ def uninstall(self, config: LakebridgeConfiguration):
                 f"Won't remove transpile validation schema `{config.transpile.schema_name}` "
                 f"from catalog `{config.transpile.catalog_name}`. Please remove it manually."
             )
+            self._uninstall_switch_job()
 
         if config.reconcile:
             self._recon_deployment.uninstall(config.reconcile)
 
         self._installation.remove()
         logger.info("Uninstallation completed successfully.")
+
+    def _uninstall_switch_job(self) -> None:
+        """Remove Switch transpiler job if exists."""
+        repository = TranspilerRepository.user_home()
+        switch_installer = SwitchInstaller(repository, self._ws, self._installation)
+        # Get configured resources before uninstalling for logging purpose
+        resources = switch_installer.get_configured_resources()
+        switch_installer.uninstall()
+
+        if resources:
+            logger.info(
+                f"Won't remove Switch resources: catalog=`{resources['catalog']}`, "
+                f"schema=`{resources['schema']}`, volume=`{resources['volume']}`. "
+                "Please remove them manually if needed."
+            )

src/databricks/labs/lakebridge/install.py

@@ -2,7 +2,7 @@
 import logging
 import os
 import webbrowser
-from collections.abc import Set, Callable, Sequence
+from collections.abc import Set, Sequence
 from pathlib import Path
 from typing import Any, cast
 
@@ -28,6 +28,7 @@
 from databricks.labs.lakebridge.transpiler.installers import (
     BladebridgeInstaller,
     MorpheusInstaller,
+    SwitchInstaller,
     TranspilerInstaller,
 )
 from databricks.labs.lakebridge.transpiler.repository import TranspilerRepository
@@ -50,10 +51,12 @@ def __init__(
         environ: dict[str, str] | None = None,
         *,
         is_interactive: bool = True,
+        include_llm_transpiler: bool = False,
         transpiler_repository: TranspilerRepository = TranspilerRepository.user_home(),
-        transpiler_installers: Sequence[Callable[[TranspilerRepository], TranspilerInstaller]] = (
+        transpiler_installers: Sequence[type[TranspilerInstaller]] = (
             BladebridgeInstaller,
             MorpheusInstaller,
+            SwitchInstaller,
         ),
     ):
         self._ws = ws
@@ -65,6 +68,7 @@ def __init__(
         self._ws_installation = workspace_installation
         # TODO: Refactor the 'prompts' property in preference to using this flag, which should be redundant.
         self._is_interactive = is_interactive
+        self._include_llm_transpiler = include_llm_transpiler
         self._transpiler_repository = transpiler_repository
         self._transpiler_installer_factories = transpiler_installers
 
@@ -77,7 +81,12 @@ def __init__(
 
     @property
     def _transpiler_installers(self) -> Set[TranspilerInstaller]:
-        return frozenset(factory(self._transpiler_repository) for factory in self._transpiler_installer_factories)
+        factories = self._transpiler_installer_factories
+        if not self._include_llm_transpiler:
+            if SwitchInstaller in factories:
+                logger.info("Skipping Switch installation (LLM transpiler not requested)")
+            factories = tuple(f for f in factories if f != SwitchInstaller)
+        return frozenset(factory(self._transpiler_repository, self._ws, self._installation) for factory in factories)
 
     def run(
         self,
@@ -386,6 +395,7 @@ def installer(
     transpiler_repository: TranspilerRepository,
     *,
     is_interactive: bool,
+    include_llm_transpiler: bool = False,
 ) -> WorkspaceInstaller:
     app_context = ApplicationContext(_verify_workspace_client(ws))
     return WorkspaceInstaller(
@@ -398,6 +408,7 @@ def installer(
         app_context.workspace_installation,
         transpiler_repository=transpiler_repository,
         is_interactive=is_interactive,
+        include_llm_transpiler=include_llm_transpiler,
     )