Merge branch 'main' into models_api_2

llama_stack/apis/memory/client.py

@@ -13,12 +13,12 @@ from typing import Any, Dict, List, Optional
 
 import fire
 import httpx
-
-from llama_stack.distribution.datatypes import RemoteProviderConfig
 from termcolor import cprint
 
-from .memory import *  # noqa: F403
-from .common.file_utils import data_url_from_file
+from llama_stack.distribution.datatypes import RemoteProviderConfig
+
+from llama_stack.apis.memory import *  # noqa: F403
+from llama_stack.providers.utils.memory.file_utils import data_url_from_file
 
 
 async def get_client_impl(config: RemoteProviderConfig, _deps: Any) -> Memory:

llama_stack/cli/stack/build.py

@@ -212,7 +212,7 @@ class StackBuild(Subcommand):
                 providers_for_api = all_providers[api]
 
                 api_provider = prompt(
-                    "> Enter the API provider for the {} API: (default=meta-reference): ".format(
+                    "> Enter provider for the {} API: (default=meta-reference): ".format(
                         api.value
                     ),
                     validator=Validator.from_callable(

llama_stack/cli/stack/configure.py

@@ -53,46 +53,61 @@ class StackConfigure(Subcommand):
         from termcolor import cprint
 
         docker_image = None
+        build_config_file = Path(args.config)
+
+        if build_config_file.exists():
+            with open(build_config_file, "r") as f:
+                build_config = BuildConfig(**yaml.safe_load(f))
+                self._configure_llama_distribution(build_config, args.output_dir)
+            return
+
+        # if we get here, we need to try to find the conda build config file
+        cprint(
+            f"Could not find {build_config_file}. Trying conda build name instead...",
+            color="green",
+        )
         conda_dir = Path(os.getenv("CONDA_PREFIX")).parent / f"llamastack-{args.config}"
         build_config_file = Path(conda_dir) / f"{args.config}-build.yaml"
 
-        if not build_config_file.exists():
-            cprint(
-                f"Could not find {build_config_file}. Trying docker image name instead...",
-                color="green",
-            )
-            docker_image = args.config
+        if build_config_file.exists():
+            with open(build_config_file, "r") as f:
+                build_config = BuildConfig(**yaml.safe_load(f))
 
-            builds_dir = BUILDS_BASE_DIR / ImageType.docker.value
-            if args.output_dir:
-                builds_dir = Path(output_dir)
-            os.makedirs(builds_dir, exist_ok=True)
+            self._configure_llama_distribution(build_config, args.output_dir)
+            return
 
-            script = pkg_resources.resource_filename(
-                "llama_stack", "distribution/configure_container.sh"
-            )
-            script_args = [script, docker_image, str(builds_dir)]
+        # if we get here, we need to try to find the docker image
+        cprint(
+            f"Could not find {build_config_file}. Trying docker image name instead...",
+            color="green",
+        )
+        docker_image = args.config
+        builds_dir = BUILDS_BASE_DIR / ImageType.docker.value
+        if args.output_dir:
+            builds_dir = Path(output_dir)
+        os.makedirs(builds_dir, exist_ok=True)
 
-            return_code = run_with_pty(script_args)
+        script = pkg_resources.resource_filename(
+            "llama_stack", "distribution/configure_container.sh"
+        )
+        script_args = [script, docker_image, str(builds_dir)]
 
-            # we have regenerated the build config file with script, now check if it exists
-            if return_code != 0:
-                self.parser.error(
-                    f"Can not find {build_config_file}. Please run llama stack build first or check if docker image exists"
-                )
+        return_code = run_with_pty(script_args)
 
-            build_name = docker_image.removeprefix("llamastack-")
-            saved_file = str(builds_dir / f"{build_name}-run.yaml")
-            cprint(
-                f"YAML configuration has been written to {saved_file}. You can now run `llama stack run {saved_file}`",
-                color="green",
+        # we have regenerated the build config file with script, now check if it exists
+        if return_code != 0:
+            self.parser.error(
+                f"Failed to configure container {docker_image} with return code {return_code}. Please run `llama stack build first`. "
             )
             return
 
-        with open(build_config_file, "r") as f:
-            build_config = BuildConfig(**yaml.safe_load(f))
-
-        self._configure_llama_distribution(build_config, args.output_dir)
+        build_name = docker_image.removeprefix("llamastack-")
+        saved_file = str(builds_dir / f"{build_name}-run.yaml")
+        cprint(
+            f"YAML configuration has been written to {saved_file}. You can now run `llama stack run {saved_file}`",
+            color="green",
+        )
+        return
 
     def _configure_llama_distribution(
         self,

llama_stack/distribution/configure.py

@@ -30,12 +30,8 @@ def make_routing_entry_type(config_class: Any):
 def configure_api_providers(
     config: StackRunConfig, spec: DistributionSpec
 ) -> StackRunConfig:
-    cprint("Configuring APIs to serve...", "white", attrs=["bold"])
-    print("Enter comma-separated list of APIs to serve:")
-
     apis = config.apis_to_serve or list(spec.providers.keys())
     config.apis_to_serve = [a for a in apis if a != "telemetry"]
-    print("")
 
     apis = [v.value for v in stack_apis()]
     all_providers = api_providers()

llama_stack/providers/impls/meta_reference/agents/safety.py

@@ -7,15 +7,14 @@
 from typing import List
 
 from llama_models.llama3.api.datatypes import Message, Role, UserMessage
+from termcolor import cprint
 
 from llama_stack.apis.safety import (
     OnViolationAction,
-    RunShieldRequest,
     Safety,
     ShieldDefinition,
     ShieldResponse,
 )
-from termcolor import cprint
 
 
 class SafetyException(Exception):  # noqa: N818
@@ -45,10 +44,8 @@ class ShieldRunnerMixin:
             messages[0] = UserMessage(content=messages[0].content)
 
         res = await self.safety_api.run_shields(
-            RunShieldRequest(
-                messages=messages,
-                shields=shields,
-            )
+            messages=messages,
+            shields=shields,
         )
 
         results = res.responses

llama_stack/providers/impls/meta_reference/inference/config.py

@@ -11,10 +11,10 @@ from llama_models.datatypes import ModelFamily
 from llama_models.schema_utils import json_schema_type
 from llama_models.sku_list import all_registered_models, resolve_model
 
-from llama_stack.apis.inference import QuantizationConfig
-
 from pydantic import BaseModel, Field, field_validator
 
+from llama_stack.apis.inference import QuantizationConfig
+
 
 @json_schema_type
 class MetaReferenceImplConfig(BaseModel):
@@ -24,7 +24,7 @@ class MetaReferenceImplConfig(BaseModel):
     )
     quantization: Optional[QuantizationConfig] = None
     torch_seed: Optional[int] = None
-    max_seq_len: int
+    max_seq_len: int = 4096
     max_batch_size: int = 1
 
     @field_validator("model")