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,24 @@ 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)
+      - name: output-ws-folder
+        description: Output folder path (Databricks Workspace path starting with /Workspace/)
+      - name: source-dialect
+        description: Source dialect name (e.g., 'snowflake', 'teradata')
+      - name: catalog-name
+        description: Databricks Catalog name (already existing with relevant permission)
+      - name: schema-name
+        description: Databricks Schema name (already existing with relevant permission)
+      - name: volume
+        description: Databricks UC Volume name (already existing with relevant permission)
+      - name: foundational-model
+        description: Databricks Model Serving Endpoint
+
   - name: reconcile
     description: Reconcile source and target data residing on Databricks
 
@@ -63,6 +81,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,9 +36,11 @@
 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
+from databricks.labs.switch.lsp import get_switch_dialects
 
 
 # Subclass to allow controlled access to protected methods.
@@ -92,7 +94,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 +242,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.
@@ -730,6 +740,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 +749,15 @@ 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")
+        # Decision was made not to prompt when include_llm_transpiler is set, and we expect users to use llm-transpile and pass all the arguments
+        is_interactive = False
     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)
 
 
@@ -818,6 +835,133 @@ def analyze(
         logger.debug(f"User: {ctx.current_user}")
 
 
+def _validate_llm_transpile_args(
+    input_source: str | None,
+    output_ws_folder: str | None,
+    source_dialect: str | None,
+    prompts: Prompts,
+) -> tuple[str, str, str]:
+
+    _switch_dialects = get_switch_dialects()
+
+    # Validate presence after attempting to source from config
+    if not input_source:
+        input_source = prompts.question("Enter input SQL path")
+    if not output_ws_folder:
+        output_ws_folder = prompts.question("Enter output workspace folder must start with /Workspace/")
+    if not source_dialect:
+        source_dialect = prompts.choice("Select the source dialect", sorted(_switch_dialects))
+
+    # Validate input_source path exists (local path)
+    if not Path(input_source).exists():
+        raise_validation_exception(f"Invalid path for '--input-source': Path '{input_source}' does not exist.")
+
+    # Validate output_ws_folder is a workspace path
+    if not str(output_ws_folder).startswith("/Workspace/"):
+        raise_validation_exception(
+            f"Invalid value for '--output-ws-folder': workspace output path must start with /Workspace/. Got: {output_ws_folder!r}"
+        )
+
+    if source_dialect not in _switch_dialects:
+        raise_validation_exception(
+            f"Invalid value for '--source-dialect': {source_dialect!r} must be one of: {', '.join(sorted(_switch_dialects))}"
+        )
+
+    return input_source, output_ws_folder, source_dialect
+
+
+@lakebridge.command
+def llm_transpile(
+    *,
+    w: WorkspaceClient,
+    input_source: str | None = None,
+    output_ws_folder: str | None = None,
+    source_dialect: str | None = None,
+    catalog_name: str | None = None,
+    schema_name: str | None = None,
+    volume: str | None = None,
+    foundational_model: str | None = None,
+    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", "llm-transpile")
+    user = ctx.current_user
+    logger.debug(f"User: {user}")
+
+    logger.info(
+        """Please read and accept the following comments before proceeding:\n
+        This Feature leverages a large language model (LLM) to analyse and convert your provided content, code and data.\n
+        You consent to your content being transmitted to, processed by, and returned from the LLM hosted by Databricks foundational models or other external models you may configure during the runtime.\n
+        The outputs of the LLM are generated automatically without human review, and may contain inaccuracies or errors. \n
+        You are responsible for reviewing and validating all outputs before relying on them for any critical or production use. \n
+        By running this feature you accept these conditions.
+    """
+    )
+
+    prompts = ctx.prompts
+    resource_configurator = ctx.resource_configurator
+
+    # If CLI args are missing, try to read them from config.yml
+    input_source, output_ws_folder, source_dialect = _validate_llm_transpile_args(
+        input_source,
+        output_ws_folder,
+        source_dialect,
+        prompts,
+    )
+
+    if catalog_name is None:
+        catalog_name = resource_configurator.prompt_for_catalog_setup(default_catalog_name="lakebridge")
+
+    if schema_name is None:
+        schema_name = resource_configurator.prompt_for_schema_setup(catalog=catalog_name, default_schema_name="switch")
+
+    if volume is None:
+        volume = resource_configurator.prompt_for_volume_setup(
+            catalog=catalog_name, schema=schema_name, default_volume_name="switch_volume"
+        )
+
+    resource_configurator.has_necessary_access(catalog_name, schema_name, volume)
+
+    if foundational_model is None:
+        foundational_model = resource_configurator.prompt_for_foundation_model_choice()
+
+    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"])
+
+    try:
+        ctx.add_user_agent_extra("transpiler_source_dialect", source_dialect)
+        job_runner = SwitchRunner(ctx.workspace_client, ctx.installation)
+        volume_input_path = job_runner.upload_to_volume(
+            local_path=Path(input_source),
+            catalog=catalog_name,
+            schema=schema_name,
+            volume=volume,
+        )
+
+        response = job_runner.run(
+            volume_input_path=volume_input_path,
+            output_ws_folder=output_ws_folder,
+            source_tech=source_dialect,
+            catalog=catalog_name,
+            schema=schema_name,
+            volume=volume,
+            foundational_model=foundational_model,
+            job_id=job_id,
+        )
+        json.dump(response, sys.stdout, indent=2)
+    except Exception as ex:
+        raise RuntimeError(ex) from ex
+
+
 if __name__ == "__main__":
     app = lakebridge
     logger = app.get_logger()

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"[Recommended] {default_choice}", *other_models]
+        selected = self._prompts.choice("Select a Foundation Model serving endpoint:", choices, sort=True)
+
+        if selected.startswith("[Recommended]"):
+            selected = default_choice
+        return selected
+
     def has_necessary_catalog_access(
         self, catalog_name: str, user_name: str, privilege_sets: tuple[set[Privilege], ...]
     ):