chore(rename): move llama_stack.distribution to llama_stack.core (#2975)

We would like to rename the term `template` to `distribution`. To
prepare for that, this is a precursor.

cc @leseb

llama_stack/core/utils/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.

llama_stack/core/utils/config.py

@@ -0,0 +1,30 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from typing import Any
+
+
+def redact_sensitive_fields(data: dict[str, Any]) -> dict[str, Any]:
+    """Redact sensitive information from config before printing."""
+    sensitive_patterns = ["api_key", "api_token", "password", "secret"]
+
+    def _redact_value(v: Any) -> Any:
+        if isinstance(v, dict):
+            return _redact_dict(v)
+        elif isinstance(v, list):
+            return [_redact_value(i) for i in v]
+        return v
+
+    def _redact_dict(d: dict[str, Any]) -> dict[str, Any]:
+        result = {}
+        for k, v in d.items():
+            if any(pattern in k.lower() for pattern in sensitive_patterns):
+                result[k] = "********"
+            else:
+                result[k] = _redact_value(v)
+        return result
+
+    return _redact_dict(data)

llama_stack/core/utils/config_dirs.py

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+from pathlib import Path
+
+LLAMA_STACK_CONFIG_DIR = Path(os.getenv("LLAMA_STACK_CONFIG_DIR", os.path.expanduser("~/.llama/")))
+
+DISTRIBS_BASE_DIR = LLAMA_STACK_CONFIG_DIR / "distributions"
+
+DEFAULT_CHECKPOINT_DIR = LLAMA_STACK_CONFIG_DIR / "checkpoints"
+
+RUNTIME_BASE_DIR = LLAMA_STACK_CONFIG_DIR / "runtime"
+
+EXTERNAL_PROVIDERS_DIR = LLAMA_STACK_CONFIG_DIR / "providers.d"

llama_stack/core/utils/config_resolution.py

@@ -0,0 +1,125 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from enum import StrEnum
+from pathlib import Path
+
+from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="config_resolution")
+
+
+TEMPLATE_DIR = Path(__file__).parent.parent.parent.parent / "llama_stack" / "templates"
+
+
+class Mode(StrEnum):
+    RUN = "run"
+    BUILD = "build"
+
+
+def resolve_config_or_template(
+    config_or_template: str,
+    mode: Mode = Mode.RUN,
+) -> Path:
+    """
+    Resolve a config/template argument to a concrete config file path.
+
+    Args:
+        config_or_template: User input (file path, template name, or built distribution)
+        mode: Mode resolving for ("run", "build", "server")
+
+    Returns:
+        Path to the resolved config file
+
+    Raises:
+        ValueError: If resolution fails
+    """
+
+    # Strategy 1: Try as file path first
+    config_path = Path(config_or_template)
+    if config_path.exists() and config_path.is_file():
+        logger.info(f"Using file path: {config_path}")
+        return config_path.resolve()
+
+    # Strategy 2: Try as template name (if no .yaml extension)
+    if not config_or_template.endswith(".yaml"):
+        template_config = _get_template_config_path(config_or_template, mode)
+        if template_config.exists():
+            logger.info(f"Using template: {template_config}")
+            return template_config
+
+    # Strategy 3: Try as built distribution name
+    distrib_config = DISTRIBS_BASE_DIR / f"llamastack-{config_or_template}" / f"{config_or_template}-{mode}.yaml"
+    if distrib_config.exists():
+        logger.info(f"Using built distribution: {distrib_config}")
+        return distrib_config
+
+    distrib_config = DISTRIBS_BASE_DIR / f"{config_or_template}" / f"{config_or_template}-{mode}.yaml"
+    if distrib_config.exists():
+        logger.info(f"Using built distribution: {distrib_config}")
+        return distrib_config
+
+    # Strategy 4: Failed - provide helpful error
+    raise ValueError(_format_resolution_error(config_or_template, mode))
+
+
+def _get_template_config_path(template_name: str, mode: Mode) -> Path:
+    """Get the config file path for a template."""
+    return TEMPLATE_DIR / template_name / f"{mode}.yaml"
+
+
+def _format_resolution_error(config_or_template: str, mode: Mode) -> str:
+    """Format a helpful error message for resolution failures."""
+    from llama_stack.core.utils.config_dirs import DISTRIBS_BASE_DIR
+
+    template_path = _get_template_config_path(config_or_template, mode)
+    distrib_path = DISTRIBS_BASE_DIR / f"llamastack-{config_or_template}" / f"{config_or_template}-{mode}.yaml"
+    distrib_path2 = DISTRIBS_BASE_DIR / f"{config_or_template}" / f"{config_or_template}-{mode}.yaml"
+
+    available_templates = _get_available_templates()
+    templates_str = ", ".join(available_templates) if available_templates else "none found"
+
+    return f"""Could not resolve config or template '{config_or_template}'.
+
+Tried the following locations:
+  1. As file path: {Path(config_or_template).resolve()}
+  2. As template: {template_path}
+  3. As built distribution: ({distrib_path}, {distrib_path2})
+
+Available templates: {templates_str}
+
+Did you mean one of these templates?
+{_format_template_suggestions(available_templates, config_or_template)}
+"""
+
+
+def _get_available_templates() -> list[str]:
+    """Get list of available template names."""
+    if not TEMPLATE_DIR.exists() and not DISTRIBS_BASE_DIR.exists():
+        return []
+
+    return list(
+        set(
+            [d.name for d in TEMPLATE_DIR.iterdir() if d.is_dir() and not d.name.startswith(".")]
+            + [d.name for d in DISTRIBS_BASE_DIR.iterdir() if d.is_dir() and not d.name.startswith(".")]
+        )
+    )
+
+
+def _format_template_suggestions(templates: list[str], user_input: str) -> str:
+    """Format template suggestions for error messages, showing closest matches first."""
+    if not templates:
+        return "  (no templates found)"
+
+    import difflib
+
+    # Get up to 3 closest matches with similarity threshold of 0.3 (lower = more permissive)
+    close_matches = difflib.get_close_matches(user_input, templates, n=3, cutoff=0.3)
+    display_templates = close_matches if close_matches else templates[:3]
+
+    suggestions = [f"  - {t}" for t in display_templates]
+    return "\n".join(suggestions)

llama_stack/core/utils/context.py

@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from collections.abc import AsyncGenerator
+from contextvars import ContextVar
+
+
+def preserve_contexts_async_generator[T](
+    gen: AsyncGenerator[T, None], context_vars: list[ContextVar]
+) -> AsyncGenerator[T, None]:
+    """
+    Wraps an async generator to preserve context variables across iterations.
+    This is needed because we start a new asyncio event loop for each streaming request,
+    and we need to preserve the context across the event loop boundary.
+    """
+    # Capture initial context values
+    initial_context_values = {context_var.name: context_var.get() for context_var in context_vars}
+
+    async def wrapper() -> AsyncGenerator[T, None]:
+        while True:
+            try:
+                # Restore context values before any await
+                for context_var in context_vars:
+                    context_var.set(initial_context_values[context_var.name])
+
+                item = await gen.__anext__()
+
+                # Update our tracked values with any changes made during this iteration
+                for context_var in context_vars:
+                    initial_context_values[context_var.name] = context_var.get()
+
+                yield item
+
+            except StopAsyncIteration:
+                break
+
+    return wrapper()

llama_stack/core/utils/dynamic.py

@@ -0,0 +1,13 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import importlib
+
+
+def instantiate_class_type(fully_qualified_name):
+    module_name, class_name = fully_qualified_name.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, class_name)

llama_stack/core/utils/exec.py

@@ -0,0 +1,143 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import logging
+import os
+import signal
+import subprocess
+import sys
+
+from termcolor import cprint
+
+log = logging.getLogger(__name__)
+
+import importlib
+import json
+from pathlib import Path
+
+from llama_stack.core.utils.image_types import LlamaStackImageType
+
+
+def formulate_run_args(image_type: str, image_name: str) -> list[str]:
+    env_name = ""
+
+    if image_type == LlamaStackImageType.CONDA.value:
+        current_conda_env = os.environ.get("CONDA_DEFAULT_ENV")
+        env_name = image_name or current_conda_env
+        if not env_name:
+            cprint(
+                "No current conda environment detected, please specify a conda environment name with --image-name",
+                color="red",
+                file=sys.stderr,
+            )
+            return
+
+        def get_conda_prefix(env_name):
+            # Conda "base" environment does not end with "base" in the
+            # prefix, so should be handled separately.
+            if env_name == "base":
+                return os.environ.get("CONDA_PREFIX")
+            # Get conda environments info
+            conda_env_info = json.loads(subprocess.check_output(["conda", "info", "--envs", "--json"]).decode())
+            envs = conda_env_info["envs"]
+            for envpath in envs:
+                if os.path.basename(envpath) == env_name:
+                    return envpath
+            return None
+
+        cprint(f"Using conda environment: {env_name}", color="green", file=sys.stderr)
+        conda_prefix = get_conda_prefix(env_name)
+        if not conda_prefix:
+            cprint(
+                f"Conda environment {env_name} does not exist.",
+                color="red",
+                file=sys.stderr,
+            )
+            return
+
+        build_file = Path(conda_prefix) / "llamastack-build.yaml"
+        if not build_file.exists():
+            cprint(
+                f"Build file {build_file} does not exist.\n\nPlease run `llama stack build` or specify the correct conda environment name with --image-name",
+                color="red",
+                file=sys.stderr,
+            )
+            return
+    else:
+        # else must be venv since that is the only valid option left.
+        current_venv = os.environ.get("VIRTUAL_ENV")
+        env_name = image_name or current_venv
+        if not env_name:
+            cprint(
+                "No current virtual environment detected, please specify a virtual environment name with --image-name",
+                color="red",
+                file=sys.stderr,
+            )
+            return
+        cprint(f"Using virtual environment: {env_name}", file=sys.stderr)
+
+    script = importlib.resources.files("llama_stack") / "core/start_stack.sh"
+    run_args = [
+        script,
+        image_type,
+        env_name,
+    ]
+
+    return run_args
+
+
+def in_notebook():
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True
+
+
+def run_command(command: list[str]) -> int:
+    """
+    Run a command with interrupt handling and output capture.
+    Uses subprocess.run with direct stream piping for better performance.
+
+    Args:
+        command (list): The command to run.
+
+    Returns:
+        int: The return code of the command.
+    """
+    original_sigint = signal.getsignal(signal.SIGINT)
+    ctrl_c_pressed = False
+
+    def sigint_handler(signum, frame):
+        nonlocal ctrl_c_pressed
+        ctrl_c_pressed = True
+        log.info("\nCtrl-C detected. Aborting...")
+
+    try:
+        # Set up the signal handler
+        signal.signal(signal.SIGINT, sigint_handler)
+
+        # Run the command with stdout/stderr piped directly to system streams
+        result = subprocess.run(
+            command,
+            text=True,
+            check=False,
+        )
+        return result.returncode
+    except subprocess.SubprocessError as e:
+        log.error(f"Subprocess error: {e}")
+        return 1
+    except Exception as e:
+        log.exception(f"Unexpected error: {e}")
+        return 1
+    finally:
+        # Restore the original signal handler
+        signal.signal(signal.SIGINT, original_sigint)

llama_stack/core/utils/image_types.py

@@ -0,0 +1,13 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import enum
+
+
+class LlamaStackImageType(enum.Enum):
+    CONTAINER = "container"
+    CONDA = "conda"
+    VENV = "venv"

llama_stack/core/utils/model_utils.py

@@ -0,0 +1,13 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from pathlib import Path
+
+from .config_dirs import DEFAULT_CHECKPOINT_DIR
+
+
+def model_local_dir(descriptor: str) -> str:
+    return str(Path(DEFAULT_CHECKPOINT_DIR) / (descriptor.replace(":", "-")))

llama_stack/core/utils/prompt_for_config.py

@@ -0,0 +1,282 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import inspect
+import json
+import logging
+from enum import Enum
+from typing import Annotated, Any, Literal, Union, get_args, get_origin
+
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+from pydantic_core import PydanticUndefinedType
+
+log = logging.getLogger(__name__)
+
+
+def is_list_of_primitives(field_type):
+    """Check if a field type is a List of primitive types."""
+    origin = get_origin(field_type)
+    if origin is list or origin is list:
+        args = get_args(field_type)
+        if len(args) == 1 and args[0] in (int, float, str, bool):
+            return True
+    return False
+
+
+def is_basemodel_without_fields(typ):
+    return inspect.isclass(typ) and issubclass(typ, BaseModel) and len(typ.__fields__) == 0
+
+
+def can_recurse(typ):
+    return inspect.isclass(typ) and issubclass(typ, BaseModel) and len(typ.__fields__) > 0
+
+
+def get_literal_values(field):
+    """Extract literal values from a field if it's a Literal type."""
+    if get_origin(field.annotation) is Literal:
+        return get_args(field.annotation)
+    return None
+
+
+def is_optional(field_type):
+    """Check if a field type is Optional."""
+    return get_origin(field_type) is Union and type(None) in get_args(field_type)
+
+
+def get_non_none_type(field_type):
+    """Get the non-None type from an Optional type."""
+    return next(arg for arg in get_args(field_type) if arg is not type(None))
+
+
+def manually_validate_field(model: type[BaseModel], field_name: str, value: Any):
+    validators = model.__pydantic_decorators__.field_validators
+    for _name, validator in validators.items():
+        if field_name in validator.info.fields:
+            validator.func(value)
+
+    return value
+
+
+def is_discriminated_union(typ) -> bool:
+    if isinstance(typ, FieldInfo):
+        return typ.discriminator
+    else:
+        if get_origin(typ) is not Annotated:
+            return False
+        args = get_args(typ)
+        return len(args) >= 2 and args[1].discriminator
+
+
+def prompt_for_discriminated_union(
+    field_name,
+    typ,
+    existing_value,
+):
+    if isinstance(typ, FieldInfo):
+        inner_type = typ.annotation
+        discriminator = typ.discriminator
+        default_value = typ.default
+    else:
+        args = get_args(typ)
+        inner_type = args[0]
+        discriminator = args[1].discriminator
+        default_value = args[1].default
+
+    union_types = get_args(inner_type)
+    # Find the discriminator field in each union type
+    type_map = {}
+    for t in union_types:
+        disc_field = t.__fields__[discriminator]
+        literal_values = get_literal_values(disc_field)
+        if literal_values:
+            for value in literal_values:
+                type_map[value] = t
+
+    while True:
+        prompt = f"Enter `{discriminator}` for {field_name} (options: {', '.join(type_map.keys())})"
+        if default_value is not None:
+            prompt += f" (default: {default_value})"
+
+        discriminator_value = input(f"{prompt}: ")
+        if discriminator_value == "" and default_value is not None:
+            discriminator_value = default_value
+
+        if discriminator_value in type_map:
+            chosen_type = type_map[discriminator_value]
+            log.info(f"\nConfiguring {chosen_type.__name__}:")
+
+            if existing_value and (getattr(existing_value, discriminator) != discriminator_value):
+                existing_value = None
+
+            sub_config = prompt_for_config(chosen_type, existing_value)
+            # Set the discriminator field in the sub-config
+            setattr(sub_config, discriminator, discriminator_value)
+            return sub_config
+        else:
+            log.error(f"Invalid {discriminator}. Please try again.")
+
+
+# This is somewhat elaborate, but does not purport to be comprehensive in any way.
+# We should add handling for the most common cases to tide us over.
+#
+# doesn't support List[nested_class] yet or Dicts of any kind. needs a bunch of
+# unit tests for coverage.
+def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel | None = None) -> BaseModel:
+    """
+    Recursively prompt the user for configuration values based on a Pydantic BaseModel.
+
+    Args:
+        config_type: A Pydantic BaseModel class representing the configuration structure.
+
+    Returns:
+        An instance of the config_type with user-provided values.
+    """
+    config_data = {}
+
+    for field_name, field in config_type.__fields__.items():
+        field_type = field.annotation
+        existing_value = getattr(existing_config, field_name) if existing_config else None
+        if existing_value:
+            default_value = existing_value
+        else:
+            default_value = field.default if not isinstance(field.default, PydanticUndefinedType) else None
+        is_required = field.is_required
+
+        # Skip fields with Literal type
+        if get_origin(field_type) is Literal:
+            continue
+
+        # Skip fields with no type annotations
+        if is_basemodel_without_fields(field_type):
+            config_data[field_name] = field_type()
+            continue
+
+        if inspect.isclass(field_type) and issubclass(field_type, Enum):
+            prompt = f"Choose {field_name} (options: {', '.join(e.name for e in field_type)}):"
+            while True:
+                # this branch does not handle existing and default values yet
+                user_input = input(prompt + " ")
+                try:
+                    value = field_type[user_input]
+                    validated_value = manually_validate_field(config_type, field, value)
+                    config_data[field_name] = validated_value
+                    break
+                except KeyError:
+                    log.error(f"Invalid choice. Please choose from: {', '.join(e.name for e in field_type)}")
+            continue
+
+        if is_discriminated_union(field):
+            config_data[field_name] = prompt_for_discriminated_union(field_name, field, existing_value)
+            continue
+
+        if is_optional(field_type) and can_recurse(get_non_none_type(field_type)):
+            prompt = f"Do you want to configure {field_name}? (y/n): "
+            if input(prompt).lower() == "n":
+                config_data[field_name] = None
+                continue
+            nested_type = get_non_none_type(field_type)
+            log.info(f"Entering sub-configuration for {field_name}:")
+            config_data[field_name] = prompt_for_config(nested_type, existing_value)
+        elif is_optional(field_type) and is_discriminated_union(get_non_none_type(field_type)):
+            prompt = f"Do you want to configure {field_name}? (y/n): "
+            if input(prompt).lower() == "n":
+                config_data[field_name] = None
+                continue
+            nested_type = get_non_none_type(field_type)
+            config_data[field_name] = prompt_for_discriminated_union(
+                field_name,
+                nested_type,
+                existing_value,
+            )
+        elif can_recurse(field_type):
+            log.info(f"\nEntering sub-configuration for {field_name}:")
+            config_data[field_name] = prompt_for_config(
+                field_type,
+                existing_value,
+            )
+        else:
+            prompt = f"Enter value for {field_name}"
+            if existing_value is not None:
+                prompt += f" (existing: {existing_value})"
+            elif default_value is not None:
+                prompt += f" (default: {default_value})"
+            if is_optional(field_type):
+                prompt += " (optional)"
+            elif is_required:
+                prompt += " (required)"
+            prompt += ": "
+
+            while True:
+                user_input = input(prompt)
+                if user_input == "":
+                    if default_value is not None:
+                        config_data[field_name] = default_value
+                        break
+                    elif is_optional(field_type) or not is_required:
+                        config_data[field_name] = None
+                        break
+                    else:
+                        log.error("This field is required. Please provide a value.")
+                        continue
+                else:
+                    try:
+                        # Handle Optional types
+                        if is_optional(field_type):
+                            if user_input.lower() == "none":
+                                value = None
+                            else:
+                                field_type = get_non_none_type(field_type)
+                                value = user_input
+
+                        # Handle List of primitives
+                        elif is_list_of_primitives(field_type):
+                            try:
+                                value = json.loads(user_input)
+                                if not isinstance(value, list):
+                                    raise ValueError("Input must be a JSON-encoded list")
+                                element_type = get_args(field_type)[0]
+                                value = [element_type(item) for item in value]
+
+                            except json.JSONDecodeError:
+                                log.error('Invalid JSON. Please enter a valid JSON-encoded list e.g., ["foo","bar"]')
+                                continue
+                            except ValueError as e:
+                                log.error(f"{str(e)}")
+                                continue
+
+                        elif get_origin(field_type) is dict:
+                            try:
+                                value = json.loads(user_input)
+                                if not isinstance(value, dict):
+                                    raise ValueError("Input must be a JSON-encoded dictionary")
+
+                            except json.JSONDecodeError:
+                                log.error("Invalid JSON. Please enter a valid JSON-encoded dict.")
+                                continue
+
+                        # Convert the input to the correct type
+                        elif inspect.isclass(field_type) and issubclass(field_type, BaseModel):
+                            # For nested BaseModels, we assume a dictionary-like string input
+                            import ast
+
+                            value = field_type(**ast.literal_eval(user_input))
+                        else:
+                            value = field_type(user_input)
+
+                    except ValueError:
+                        log.error(f"Invalid input. Expected type: {getattr(field_type, '__name__', str(field_type))}")
+                        continue
+
+                try:
+                    # Validate the field using our manual validation function
+                    validated_value = manually_validate_field(config_type, field_name, value)
+                    config_data[field_name] = validated_value
+                    break
+                except ValueError as e:
+                    log.error(f"Validation error: {str(e)}")
+
+    return config_type(**config_data)

llama_stack/core/utils/serialize.py

@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+from datetime import datetime
+from enum import Enum
+
+
+class EnumEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, Enum):
+            return obj.value
+        elif isinstance(obj, datetime):
+            return obj.isoformat()
+        return super().default(obj)