Add llm-transpile command with Switch runner and integration tests

.gitignore

@@ -22,3 +22,4 @@ remorph_transpile/
 /linter/src/main/antlr4/library/gen/
 .databricks-login.json
 .mypy_cache
+.env

labs.yml

@@ -46,6 +46,19 @@ commands:
       {{range .}}{{.total_files_processed}}\t{{.total_queries_processed}}\t{{.analysis_error_count}}\t{{.parsing_error_count}}\t{{.validation_error_count}}\t{{.generation_error_count}}\t{{.error_log_file}}
       {{end}}
 
+  - name: llm-transpile
+    description: Transpile source code to Databricks using LLM Transpiler (Switch)
+    flags:
+      - name: input-source
+        description: Input Script Folder or File (local path)
+        default: null
+      - name: output-ws-folder
+        description: Output folder path (Databricks Workspace path starting with /Workspace/)
+        default: null
+      - name: source-dialect
+        description: Source dialect name (e.g., 'snowflake', 'teradata')
+        default: null
+
   - name: reconcile
     description: Reconcile source and target data residing on Databricks
 
@@ -63,6 +76,9 @@ commands:
       - name: interactive
         description: (Optional) Whether installing in interactive mode (`true|false|auto`); configuration settings are prompted for when interactive
         default: auto
+      - name: include-llm-transpiler
+        description: (Optional) Whether to include LLM-based transpiler in installation (`true|false`)
+        default: "false"
 
   - name: describe-transpile
     description: Describe installed transpilers

pyproject.toml

@@ -40,7 +40,9 @@ dependencies = [
   "SQLAlchemy~=2.0.40",
   "pygls~=2.0.0a2",
   "duckdb~=1.2.2",
+  "databricks-switch-plugin~=0.1.4",  # Temporary, until Switch is migrated to be a transpiler (LSP) plugin.
   "requests>=2.28.1,<3"  # Matches databricks-sdk (and 'types-requests' below), to avoid conflicts.
+
 ]
 
 [project.urls]
@@ -447,7 +449,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

@@ -36,6 +36,7 @@
 from databricks.labs.lakebridge.transpiler.lsp.lsp_engine import LSPEngine
 from databricks.labs.lakebridge.transpiler.repository import TranspilerRepository
 from databricks.labs.lakebridge.transpiler.sqlglot.sqlglot_engine import SqlglotEngine
+from databricks.labs.lakebridge.transpiler.switch_runner import SwitchRunner
 from databricks.labs.lakebridge.transpiler.transpile_engine import TranspileEngine
 
 from databricks.labs.lakebridge.transpiler.transpile_status import ErrorSeverity
@@ -92,7 +93,7 @@ def _remove_warehouse(ws: WorkspaceClient, warehouse_id: str):
 
 
 @lakebridge.command
-def transpile(  # pylint: disable=too-many-arguments
+def transpile(
     *,
     w: WorkspaceClient,
     transpiler_config_path: str | None = None,
@@ -240,6 +241,14 @@ def use_transpiler_config_path(self, transpiler_config_path: str | None) -> None
             )
             self._config = dataclasses.replace(self._config, transpiler_config_path=transpiler_config_path)
 
+        # Switch is installed inside "/Users/<>/.lakebridge/transpilers/Switch/lsp/config.yml
+        if (
+            self._config.transpiler_config_path is not None
+            and Path(self._config.transpiler_config_path).parent.parent.name == "Switch"
+        ):
+            msg = "Switch transpiler is not supported through `transpile` run `llm-transpile` instead."
+            raise RuntimeError(msg)
+
     def use_source_dialect(self, source_dialect: str | None) -> None:
         if source_dialect is not None:
             # Defer validation: depends on the transpiler config path, we'll deal with this later.
@@ -637,6 +646,77 @@ def _override_workspace_client_config(ctx: ApplicationContext, overrides: dict[s
         ctx.connect_config.cluster_id = cluster_id
 
 
+@lakebridge.command
+def llm_transpile(
+    *,
+    w: WorkspaceClient,
+    input_source: str,
+    output_ws_folder: str,
+    source_dialect: str,
+    ctx: ApplicationContext | None = None,
+) -> None:
+    """Transpile source code to Databricks using LLM Transpiler (Switch)"""
+    if ctx is None:
+        ctx = ApplicationContext(w)
+    del w
+    ctx.add_user_agent_extra("cmd", "transpile-switch")
+    user = ctx.current_user
+    logger.debug(f"User: {user}")
+
+    job_list = ctx.install_state.jobs
+    if "Switch" not in job_list:
+        raise RuntimeError(
+            "Switch Job ID not found. "
+            "Please run 'databricks labs lakebridge install-transpile --include-llm-transpiler true' first."
+        )
+    logger.debug("Switch job ID found in InstallState")
+    job_id = int(job_list["Switch"])
+
+    # Upload File to Volume
+    transpile_config = ctx.transpile_config
+    if transpile_config is None or transpile_config.transpiler_options is None:
+        raise RuntimeError(
+            "Transpile configuration config.yml not found in workspace."
+            "Please run 'databricks labs lakebridge install-transpile --include-llm-transpiler true' first."
+        )
+
+    transpile_options = transpile_config.transpiler_options
+    logger.debug(f"Transpiler options: {transpile_options}")
+    if not isinstance(transpile_options, Mapping):
+        raise RuntimeError(
+            "Invalid transpile configuration: transpiler_options must be a mapping. "
+            "Please run 'databricks labs lakebridge install-transpile --include-llm-transpiler true' first."
+        )
+    catalog = transpile_options.get("catalog", None)
+    schema = transpile_options.get("schema", None)
+    volume = transpile_options.get("volume", None)
+
+    if catalog is None or schema is None or volume is None:
+        raise RuntimeError(
+            "Invalid transpile configuration: catalog, schema or volume is missing. "
+            "Please run 'databricks labs lakebridge install-transpile --include-llm-transpiler true' first."
+        )
+    assert isinstance(catalog, str)
+    assert isinstance(schema, str)
+    assert isinstance(volume, str)
+
+    try:
+        job_runner = SwitchRunner(ctx.workspace_client, ctx.installation)
+        volume_input_path = job_runner.upload_to_volume(
+            local_path=Path(input_source), catalog=catalog, schema=schema, volume=volume
+        )
+
+        response = job_runner.run(
+            volume_input_path=volume_input_path,
+            output_ws_folder=output_ws_folder,
+            source_tech=source_dialect,
+            job_id=job_id,
+        )
+        json.dump(response, sys.stdout, indent=2)
+    except Exception as ex:
+        raise RuntimeError(ex) from ex
+
+
 @lakebridge.command
 def reconcile(*, w: WorkspaceClient) -> None:
     """[EXPERIMENTAL] Reconciles source to Databricks datasets"""
@@ -730,6 +810,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."""
@@ -738,9 +819,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=include_llm_transpiler
+    )
     transpile_installer.run(module="transpile", artifact=artifact)
 
 

src/databricks/labs/lakebridge/config.py

@@ -166,6 +166,7 @@ class TranspileConfig:
     error_file_path: str | None = None
     sdk_config: dict[str, str] | None = None
     skip_validation: bool = False
+    include_llm: bool = False
     catalog_name: str = "remorph"
     schema_name: str = "transpiler"
     transpiler_options: JsonValue = None
@@ -274,5 +275,7 @@ class ReconcileConfig:
 
 @dataclass
 class LakebridgeConfiguration:
-    transpile: TranspileConfig | None = None
-    reconcile: ReconcileConfig | None = None
+    transpile: TranspileConfig | None
+    reconcile: ReconcileConfig | None
+    # Temporary flag, indicating whether to include the LLM-based Switch transpiler.
+    include_switch: bool = False

src/databricks/labs/lakebridge/contexts/application.py

@@ -18,6 +18,7 @@
 from databricks.labs.lakebridge.deployment.dashboard import DashboardDeployment
 from databricks.labs.lakebridge.deployment.installation import WorkspaceInstallation
 from databricks.labs.lakebridge.deployment.recon import TableDeployment, JobDeployment, ReconDeployment
+from databricks.labs.lakebridge.deployment.switch import SwitchDeployment
 from databricks.labs.lakebridge.helpers.metastore import CatalogOperations
 
 logger = logging.getLogger(__name__)
@@ -119,13 +120,24 @@ def recon_deployment(self) -> ReconDeployment:
             self.dashboard_deployment,
         )
 
+    @cached_property
+    def switch_deployment(self) -> SwitchDeployment:
+        return SwitchDeployment(
+            self.workspace_client,
+            self.installation,
+            self.install_state,
+            self.product_info,
+            self.job_deployment,
+        )
+
     @cached_property
     def workspace_installation(self) -> WorkspaceInstallation:
         return WorkspaceInstallation(
             self.workspace_client,
             self.prompts,
             self.installation,
             self.recon_deployment,
+            self.switch_deployment,
             self.product_info,
             self.upgrades,
         )

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

@@ -1,14 +1,18 @@
 import logging
 import time
 
+from collections.abc import Iterator
+
 from databricks.labs.blueprint.tui import Prompts
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import DatabricksError
 from databricks.sdk.service.catalog import Privilege, SecurableType
 from databricks.sdk.service.sql import (
     CreateWarehouseRequestWarehouseType,
     EndpointInfoWarehouseType,
     SpotInstancePolicy,
 )
+from databricks.sdk.service.serving import ServingEndpoint
 
 from databricks.labs.lakebridge.helpers.metastore import CatalogOperations
 
@@ -29,8 +33,9 @@ def __init__(self, ws: WorkspaceClient, prompts: Prompts, catalog_ops: CatalogOp
 
     def prompt_for_catalog_setup(
         self,
+        default_catalog_name: str = "remorph",
     ) -> str:
-        catalog_name = self._prompts.question("Enter catalog name", default="remorph")
+        catalog_name = self._prompts.question("Enter catalog name", default=default_catalog_name)
         catalog = self._catalog_ops.get_catalog(catalog_name)
         if catalog:
             logger.info(f"Found existing catalog `{catalog_name}`")
@@ -103,6 +108,35 @@ def warehouse_type(_):
             raise SystemExit("Cannot continue installation, without a valid warehouse. Aborting the installation.")
         return warehouse_id
 
+    def prompt_for_foundation_model_choice(self, default_choice: str = "databricks-claude-sonnet-4-5") -> str:
+        """
+        List Serving Endpoints that expose a foundation model and prompt the user to pick one.
+        Returns the selected endpoint name
+        """
+        endpoints: Iterator[ServingEndpoint] = self._ws.serving_endpoints.list()
+
+        model_endpoints = [
+            ep
+            for ep in endpoints
+            if ep.name
+            and ep.config
+            and ep.config.served_entities
+            and any(getattr(se, "foundation_model", None) is not None for se in ep.config.served_entities)
+        ]
+
+        foundational_model_names = [ep.name for ep in model_endpoints if ep.name]
+
+        if foundational_model_names is None:
+            raise DatabricksError("No Foundation Model serving endpoints found. Aborting the installation.")
+        # This logic is implemented to make the default choice always to appear first in the list
+        other_models = sorted(set(foundational_model_names) - {default_choice})
+        choices = [f"[DEFAULT] {default_choice}", *other_models]
+        selected = self._prompts.choice("Select a Foundation Model serving endpoint:", choices, sort=True)
+
+        if selected.startswith("[DEFAULT]"):
+            selected = default_choice
+        return selected
+
     def has_necessary_catalog_access(
         self, catalog_name: str, user_name: str, privilege_sets: tuple[set[Privilege], ...]
     ):

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

@@ -13,6 +13,7 @@
 
 from databricks.labs.lakebridge.config import LakebridgeConfiguration
 from databricks.labs.lakebridge.deployment.recon import ReconDeployment
+from databricks.labs.lakebridge.deployment.switch import SwitchDeployment
 
 logger = logging.getLogger("databricks.labs.lakebridge.install")
 
@@ -24,13 +25,15 @@ def __init__(
         prompts: Prompts,
         installation: Installation,
         recon_deployment: ReconDeployment,
+        switch_deployment: SwitchDeployment,
         product_info: ProductInfo,
         upgrades: Upgrades,
     ):
         self._ws = ws
         self._prompts = prompts
         self._installation = installation
         self._recon_deployment = recon_deployment
+        self._switch_deployment = switch_deployment
         self._product_info = product_info
         self._upgrades = upgrades
 
@@ -96,6 +99,9 @@ def install(self, config: LakebridgeConfiguration):
         if config.reconcile:
             logger.info("Installing Lakebridge reconcile Metadata components.")
             self._recon_deployment.install(config.reconcile, wheel_path)
+        if config.include_switch:
+            logger.info("Installing Switch transpiler to workspace.")
+            self._switch_deployment.install()
 
     def uninstall(self, config: LakebridgeConfiguration):
         # This will remove all the Lakebridge modules
@@ -116,9 +122,14 @@ 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."""
+        self._switch_deployment.uninstall()