Rename Distribution -> DistributionSpec, simplify RemoteProviders

llama_toolchain/cli/distribution/configure.py

@@ -8,7 +8,7 @@ import argparse
 import json
 import shlex
 
-from pathlib import Path
+from typing import Any, Dict
 
 import yaml
 from termcolor import cprint
@@ -32,83 +32,81 @@ class DistributionConfigure(Subcommand):
         self.parser.set_defaults(func=self._run_distribution_configure_cmd)
 
     def _add_arguments(self):
-        from llama_toolchain.distribution.registry import available_distributions
-
         self.parser.add_argument(
             "--name",
             type=str,
             help="Name of the distribution to configure",
-            default="local-source",
-            choices=[d.name for d in available_distributions()],
+            required=True,
         )
 
     def _run_distribution_configure_cmd(self, args: argparse.Namespace) -> None:
-        from llama_toolchain.distribution.registry import resolve_distribution
+        from llama_toolchain.distribution.registry import resolve_distribution_spec
 
-        dist = resolve_distribution(args.name)
-        if dist is None:
-            self.parser.error(f"Could not find distribution {args.name}")
+        config_file = DISTRIBS_BASE_DIR / args.name / "config.yaml"
+        if not config_file.exists():
+            self.parser.error(
+                f"Could not find {config_file}. Please run `llama distribution install` first"
+            )
             return
 
-        env_file = DISTRIBS_BASE_DIR / dist.name / "conda.env"
-        # read this file to get the conda env name
-        assert env_file.exists(), f"Could not find conda env file {env_file}"
-        with open(env_file, "r") as f:
-            conda_env = f.read().strip()
+        # we need to find the spec from the name
+        with open(config_file, "r") as f:
+            config = yaml.safe_load(f)
 
-        configure_llama_distribution(dist, conda_env)
+        dist = resolve_distribution_spec(config["spec"])
+        if dist is None:
+            raise ValueError(f"Could not find any registered spec `{config['spec']}`")
+
+        configure_llama_distribution(dist, config)
 
 
-def configure_llama_distribution(dist: "Distribution", conda_env: str):
+def configure_llama_distribution(dist: "Distribution", config: Dict[str, Any]):
     from llama_toolchain.common.exec import run_command
     from llama_toolchain.common.prompt_for_config import prompt_for_config
     from llama_toolchain.common.serialize import EnumEncoder
-    from llama_toolchain.distribution.datatypes import RemoteProviderSpec
     from llama_toolchain.distribution.dynamic import instantiate_class_type
 
     python_exe = run_command(shlex.split("which python"))
     # simple check
+    conda_env = config["conda_env"]
     if conda_env not in python_exe:
         raise ValueError(
             f"Please re-run configure by activating the `{conda_env}` conda environment"
         )
 
-    existing_config = None
-    config_path = Path(DISTRIBS_BASE_DIR) / dist.name / "config.yaml"
-    if config_path.exists():
+    existing_config = config
+    if "providers" in existing_config:
         cprint(
-            f"Configuration already exists for {dist.name}. Will overwrite...",
+            f"Configuration already exists for {config['name']}. Will overwrite...",
             "yellow",
             attrs=["bold"],
         )
-        with open(config_path, "r") as fp:
-            existing_config = yaml.safe_load(fp)
 
     provider_configs = {}
     for api, provider_spec in dist.provider_specs.items():
-        if isinstance(provider_spec, RemoteProviderSpec):
-            provider_configs[api.value] = provider_spec.dict()
-        else:
-            cprint(f"Configuring API surface: {api.value}", "white", attrs=["bold"])
-            config_type = instantiate_class_type(provider_spec.config_class)
-            config = prompt_for_config(
-                config_type,
-                (
-                    config_type(**existing_config["providers"][api.value])
-                    if existing_config and api.value in existing_config["providers"]
-                    else None
-                ),
-            )
-            provider_configs[api.value] = {
-                "provider_id": provider_spec.provider_id,
-                **config.dict(),
-            }
+        cprint(f"Configuring API surface: {api.value}", "white", attrs=["bold"])
+        config_type = instantiate_class_type(provider_spec.config_class)
+        config = prompt_for_config(
+            config_type,
+            (
+                config_type(**existing_config["providers"][api.value])
+                if existing_config
+                and "providers" in existing_config
+                and api.value in existing_config["providers"]
+                else None
+            ),
+        )
+        provider_configs[api.value] = {
+            "provider_id": provider_spec.provider_id,
+            **config.dict(),
+        }
 
     dist_config = {
         "providers": provider_configs,
-        "conda_env": conda_env,
+        **existing_config,
     }
 
+    config_path = DISTRIBS_BASE_DIR / existing_config["name"] / "config.yaml"
     with open(config_path, "w") as fp:
         dist_config = json.loads(json.dumps(dist_config, cls=EnumEncoder))
         fp.write(yaml.dump(dist_config, sort_keys=False))