Remove the "ShieldType" concept (#430)

# What does this PR do? This PR kills the notion of "ShieldType". The impetus for this is the realization: > Why is keyword llama-guard appearing so many times everywhere, sometimes with hyphens, sometimes with underscores? Now that we have a notion of "provider specific resource identifiers" and "user specific aliases" for those and the fact that this works with models ("Llama3.1-8B-Instruct" <> "fireworks/llama-3pv1-..."), we can follow the same rules for Shields. So each Safety provider can make up a notion of identifiers it has registered. This already happens with Bedrock correctly. We just generalize it for Llama Guard, Prompt Guard, etc. For Llama Guard, we further simplify by just adopting the underlying model name itself as the identifier! No confusion necessary. While doing this, I noticed a bug in our DistributionRegistry where we weren't scoping identifiers by type. Fixed. ## Feature/Issue validation/testing/test plan Ran (inference, safety, memory, agents) tests with ollama and fireworks providers.
2025-06-27 18:50:41 +00:00 · 2024-11-12 12:37:24 -08:00 · 2024-11-12 12:37:24 -08:00 · 983d6ce2df
commit 983d6ce2df
parent 09269e2a44
26 changed files with 150 additions and 209 deletions
--- a/docs/resources/llama-stack-spec.html
+++ b/docs/resources/llama-stack-spec.html
@ -21,7 +21,7 @@
    "info": {
        "title": "[DRAFT] Llama Stack Specification",
        "version": "0.0.1",
-        "description": "This is the specification of the llama stack that provides\n                a set of endpoints and their corresponding interfaces that are tailored to\n                best leverage Llama Models. The specification is still in draft and subject to change.\n                Generated at 2024-11-12 11:16:58.657871"
+        "description": "This is the specification of the llama stack that provides\n                a set of endpoints and their corresponding interfaces that are tailored to\n                best leverage Llama Models. The specification is still in draft and subject to change.\n                Generated at 2024-11-12 11:39:48.665782"
    },
    "servers": [
        {
@ -5743,9 +5743,6 @@
                        "const": "shield",
                        "default": "shield"
                    },
-                    "shield_type": {
-                        "$ref": "#/components/schemas/ShieldType"
-                    },
                    "params": {
                        "type": "object",
                        "additionalProperties": {
@ -5777,20 +5774,10 @@
                    "identifier",
                    "provider_resource_id",
                    "provider_id",
-                    "type",
-                    "shield_type"
+                    "type"
                ],
                "title": "A safety shield resource that can be used to check content"
            },
-            "ShieldType": {
-                "type": "string",
-                "enum": [
-                    "generic_content_shield",
-                    "llama_guard",
-                    "code_scanner",
-                    "prompt_guard"
-                ]
-            },
            "Trace": {
                "type": "object",
                "properties": {
@ -7262,9 +7249,6 @@
                    "shield_id": {
                        "type": "string"
                    },
-                    "shield_type": {
-                        "$ref": "#/components/schemas/ShieldType"
-                    },
                    "provider_shield_id": {
                        "type": "string"
                    },
@ -7299,8 +7283,7 @@
                },
                "additionalProperties": false,
                "required": [
-                    "shield_id",
-                    "shield_type"
+                    "shield_id"
                ]
            },
            "RunEvalRequest": {
@ -7854,13 +7837,19 @@
    ],
    "tags": [
        {
-            "name": "Inference"
+            "name": "MemoryBanks"
+        },
+        {
+            "name": "BatchInference"
        },
        {
            "name": "Agents"
        },
        {
-            "name": "Telemetry"
+            "name": "Inference"
+        },
+        {
+            "name": "DatasetIO"
        },
        {
            "name": "Eval"
@ -7869,43 +7858,37 @@
            "name": "Models"
        },
        {
-            "name": "Inspect"
-        },
-        {
-            "name": "EvalTasks"
+            "name": "PostTraining"
        },
        {
            "name": "ScoringFunctions"
        },
        {
-            "name": "Memory"
-        },
-        {
-            "name": "Safety"
-        },
-        {
-            "name": "DatasetIO"
-        },
-        {
-            "name": "MemoryBanks"
+            "name": "Datasets"
        },
        {
            "name": "Shields"
        },
        {
-            "name": "PostTraining"
+            "name": "Telemetry"
        },
        {
-            "name": "Datasets"
+            "name": "Inspect"
        },
        {
-            "name": "Scoring"
+            "name": "Safety"
        },
        {
            "name": "SyntheticDataGeneration"
        },
        {
-            "name": "BatchInference"
+            "name": "Memory"
+        },
+        {
+            "name": "Scoring"
+        },
+        {
+            "name": "EvalTasks"
        },
        {
            "name": "BuiltinTool",
@ -8255,10 +8238,6 @@
            "name": "Shield",
            "description": "A safety shield resource that can be used to check content\n\n<SchemaDefinition schemaRef=\"#/components/schemas/Shield\" />"
        },
-        {
-            "name": "ShieldType",
-            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/ShieldType\" />"
-        },
        {
            "name": "Trace",
            "description": "<SchemaDefinition schemaRef=\"#/components/schemas/Trace\" />"
@ -8614,7 +8593,6 @@
                "Session",
                "Shield",
                "ShieldCallStep",
-                "ShieldType",
                "SpanEndPayload",
                "SpanStartPayload",
                "SpanStatus",
--- a/docs/resources/llama-stack-spec.yaml
+++ b/docs/resources/llama-stack-spec.yaml
@ -2227,11 +2227,8 @@ components:
          type: string
        shield_id:
          type: string
-        shield_type:
-          $ref: '#/components/schemas/ShieldType'
      required:
      - shield_id
-      - shield_type
      type: object
    RestAPIExecutionConfig:
      additionalProperties: false
@ -2698,8 +2695,6 @@ components:
          type: string
        provider_resource_id:
          type: string
-        shield_type:
-          $ref: '#/components/schemas/ShieldType'
        type:
          const: shield
          default: shield
@ -2709,7 +2704,6 @@ components:
      - provider_resource_id
      - provider_id
      - type
-      - shield_type
      title: A safety shield resource that can be used to check content
      type: object
    ShieldCallStep:
@ -2736,13 +2730,6 @@ components:
      - step_id
      - step_type
      type: object
-    ShieldType:
-      enum:
-      - generic_content_shield
-      - llama_guard
-      - code_scanner
-      - prompt_guard
-      type: string
    SpanEndPayload:
      additionalProperties: false
      properties:
@ -3397,7 +3384,7 @@ info:
  description: "This is the specification of the llama stack that provides\n     \
    \           a set of endpoints and their corresponding interfaces that are tailored\
    \ to\n                best leverage Llama Models. The specification is still in\
-    \ draft and subject to change.\n                Generated at 2024-11-12 11:16:58.657871"
+    \ draft and subject to change.\n                Generated at 2024-11-12 11:39:48.665782"
  title: '[DRAFT] Llama Stack Specification'
  version: 0.0.1
 jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema
@ -4761,24 +4748,24 @@ security:
 servers:
 - url: http://any-hosted-llama-stack.com
 tags:
- name: Inference
+- name: MemoryBanks
+- name: BatchInference
 - name: Agents
- name: Telemetry
+- name: Inference
+- name: DatasetIO
 - name: Eval
 - name: Models
- name: Inspect
- name: EvalTasks
- name: ScoringFunctions
- name: Memory
- name: Safety
- name: DatasetIO
- name: MemoryBanks
- name: Shields
 - name: PostTraining
+- name: ScoringFunctions
 - name: Datasets
- name: Scoring
+- name: Shields
+- name: Telemetry
+- name: Inspect
+- name: Safety
 - name: SyntheticDataGeneration
- name: BatchInference
+- name: Memory
+- name: Scoring
+- name: EvalTasks
 - description: <SchemaDefinition schemaRef="#/components/schemas/BuiltinTool" />
  name: BuiltinTool
 - description: <SchemaDefinition schemaRef="#/components/schemas/CompletionMessage"
@ -5046,8 +5033,6 @@ tags:

    <SchemaDefinition schemaRef="#/components/schemas/Shield" />'
  name: Shield
- description: <SchemaDefinition schemaRef="#/components/schemas/ShieldType" />
-  name: ShieldType
 - description: <SchemaDefinition schemaRef="#/components/schemas/Trace" />
  name: Trace
 - description: 'Checkpoint created during training runs
@ -5343,7 +5328,6 @@ x-tagGroups:
  - Session
  - Shield
  - ShieldCallStep
-  - ShieldType
  - SpanEndPayload
  - SpanStartPayload
  - SpanStatus