From 86834ee6c2e556559aef993fb9cb27812b4291e5 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Mon, 30 Sep 2024 16:36:53 -0700
Subject: [PATCH] Update configure.py to use multiple routing keys for safety

---
 llama_stack/distribution/configure.py         | 16 +++++-----
 .../docker/llamastack-local-cpu/run.yaml      | 17 +---------
 .../docker/llamastack-local-gpu/run.yaml      | 17 +---------
 tests/examples/local-run.yaml                 | 32 +------------------
 4 files changed, 11 insertions(+), 71 deletions(-)

diff --git a/llama_stack/distribution/configure.py b/llama_stack/distribution/configure.py
index 879738c00..d3b807d4a 100644
--- a/llama_stack/distribution/configure.py
+++ b/llama_stack/distribution/configure.py
@@ -117,18 +117,18 @@ def configure_api_providers(
             if api_str == "safety":
                 # TODO: add support for other safety providers, and simplify safety provider config
                 if p == "meta-reference":
-                    for shield_type in MetaReferenceShieldType:
-                        routing_entries.append(
-                            RoutableProviderConfig(
-                                routing_key=shield_type.value,
-                                provider_id=p,
-                                config=cfg.dict(),
-                            )
+                    routing_entries.append(
+                        RoutableProviderConfig(
+                            routing_key=[s.value for s in MetaReferenceShieldType],
+                            provider_id=p,
+                            config=cfg.dict(),
                         )
+                    )
                 else:
                     cprint(
-                        f"[WARN] Interactive configuration of safety provider {p} is not supported, please manually configure safety shields types in routing_table of run.yaml",
+                        f"[WARN] Interactive configuration of safety provider {p} is not supported. Please look for `{routing_key}` in run.yaml and replace it appropriately.",
                         "yellow",
+                        attrs=["bold"],
                     )
                     routing_entries.append(
                         RoutableProviderConfig(
diff --git a/llama_stack/distribution/templates/docker/llamastack-local-cpu/run.yaml b/llama_stack/distribution/templates/docker/llamastack-local-cpu/run.yaml
index 6a4b2e464..0a845582c 100644
--- a/llama_stack/distribution/templates/docker/llamastack-local-cpu/run.yaml
+++ b/llama_stack/distribution/templates/docker/llamastack-local-cpu/run.yaml
@@ -42,22 +42,7 @@ routing_table:
     config:
       llama_guard_shield: null
       prompt_guard_shield: null
-    routing_key: llama_guard
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: code_scanner_guard
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: injection_shield
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: jailbreak_shield
+    routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
   memory:
   - provider_id: meta-reference
     config: {}
diff --git a/llama_stack/distribution/templates/docker/llamastack-local-gpu/run.yaml b/llama_stack/distribution/templates/docker/llamastack-local-gpu/run.yaml
index 2969479dc..66f6cfcef 100644
--- a/llama_stack/distribution/templates/docker/llamastack-local-gpu/run.yaml
+++ b/llama_stack/distribution/templates/docker/llamastack-local-gpu/run.yaml
@@ -45,22 +45,7 @@ routing_table:
     config:
       llama_guard_shield: null
       prompt_guard_shield: null
-    routing_key: llama_guard
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: code_scanner_guard
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: injection_shield
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield: null
-      prompt_guard_shield: null
-    routing_key: jailbreak_shield
+    routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
   memory:
   - provider_id: meta-reference
     config: {}
diff --git a/tests/examples/local-run.yaml b/tests/examples/local-run.yaml
index 2ae975cdc..cbe36193c 100644
--- a/tests/examples/local-run.yaml
+++ b/tests/examples/local-run.yaml
@@ -50,37 +50,7 @@ routing_table:
         disable_output_check: false
       prompt_guard_shield:
         model: Prompt-Guard-86M
-    routing_key: llama_guard
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield:
-        model: Llama-Guard-3-8B
-        excluded_categories: []
-        disable_input_check: false
-        disable_output_check: false
-      prompt_guard_shield:
-        model: Prompt-Guard-86M
-    routing_key: code_scanner_guard
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield:
-        model: Llama-Guard-3-8B
-        excluded_categories: []
-        disable_input_check: false
-        disable_output_check: false
-      prompt_guard_shield:
-        model: Prompt-Guard-86M
-    routing_key: injection_shield
-  - provider_id: meta-reference
-    config:
-      llama_guard_shield:
-        model: Llama-Guard-3-8B
-        excluded_categories: []
-        disable_input_check: false
-        disable_output_check: false
-      prompt_guard_shield:
-        model: Prompt-Guard-86M
-    routing_key: jailbreak_shield
+    routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
   memory:
   - provider_id: meta-reference
     config: {}