fix safety shield -> shield_type

This commit is contained in:
Xi Yan 2024-09-23 10:51:39 -07:00
parent 75357df02d
commit 98da002b4c
4 changed files with 43 additions and 44 deletions

View file

@ -21,7 +21,7 @@
"info": { "info": {
"title": "[DRAFT] Llama Stack Specification", "title": "[DRAFT] Llama Stack Specification",
"version": "0.0.1", "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-09-23 10:24:44.422779" "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-09-23 10:51:15.364842"
}, },
"servers": [ "servers": [
{ {
@ -5588,7 +5588,7 @@
"RunShieldRequest": { "RunShieldRequest": {
"type": "object", "type": "object",
"properties": { "properties": {
"shield": { "shield_type": {
"type": "string" "type": "string"
}, },
"messages": { "messages": {
@ -5638,7 +5638,7 @@
}, },
"additionalProperties": false, "additionalProperties": false,
"required": [ "required": [
"shield", "shield_type",
"messages", "messages",
"params" "params"
] ]
@ -5969,48 +5969,48 @@
} }
], ],
"tags": [ "tags": [
{
"name": "Inference"
},
{
"name": "Agents"
},
{
"name": "PostTraining"
},
{ {
"name": "SyntheticDataGeneration" "name": "SyntheticDataGeneration"
}, },
{ {
"name": "Models" "name": "Models"
}, },
{
"name": "Memory"
},
{
"name": "Datasets"
},
{ {
"name": "BatchInference" "name": "BatchInference"
}, },
{
"name": "PostTraining"
},
{
"name": "Shields"
},
{
"name": "Inference"
},
{
"name": "RewardScoring"
},
{
"name": "Telemetry"
},
{
"name": "Agents"
},
{
"name": "MemoryBanks"
},
{ {
"name": "Safety" "name": "Safety"
}, },
{ {
"name": "Evaluations" "name": "Evaluations"
}, },
{
"name": "Shields"
},
{
"name": "Telemetry"
},
{
"name": "RewardScoring"
},
{
"name": "MemoryBanks"
},
{
"name": "Datasets"
},
{
"name": "Memory"
},
{ {
"name": "BuiltinTool", "name": "BuiltinTool",
"description": "<SchemaDefinition schemaRef=\"#/components/schemas/BuiltinTool\" />" "description": "<SchemaDefinition schemaRef=\"#/components/schemas/BuiltinTool\" />"

View file

@ -1642,10 +1642,10 @@ components:
- type: array - type: array
- type: object - type: object
type: object type: object
shield: shield_type:
type: string type: string
required: required:
- shield - shield_type
- messages - messages
- params - params
type: object type: object
@ -2443,7 +2443,7 @@ info:
description: "This is the specification of the llama stack that provides\n \ description: "This is the specification of the llama stack that provides\n \
\ a set of endpoints and their corresponding interfaces that are tailored\ \ a set of endpoints and their corresponding interfaces that are tailored\
\ to\n best leverage Llama Models. The specification is still in\ \ to\n best leverage Llama Models. The specification is still in\
\ draft and subject to change.\n Generated at 2024-09-23 10:24:44.422779" \ draft and subject to change.\n Generated at 2024-09-23 10:51:15.364842"
title: '[DRAFT] Llama Stack Specification' title: '[DRAFT] Llama Stack Specification'
version: 0.0.1 version: 0.0.1
jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema
@ -3681,20 +3681,20 @@ security:
servers: servers:
- url: http://any-hosted-llama-stack.com - url: http://any-hosted-llama-stack.com
tags: tags:
- name: Inference
- name: Agents
- name: PostTraining
- name: SyntheticDataGeneration - name: SyntheticDataGeneration
- name: Models - name: Models
- name: Memory
- name: Datasets
- name: BatchInference - name: BatchInference
- name: PostTraining
- name: Shields
- name: Inference
- name: RewardScoring
- name: Telemetry
- name: Agents
- name: MemoryBanks
- name: Safety - name: Safety
- name: Evaluations - name: Evaluations
- name: Shields
- name: Telemetry
- name: RewardScoring
- name: MemoryBanks
- name: Datasets
- name: Memory
- description: <SchemaDefinition schemaRef="#/components/schemas/BuiltinTool" /> - description: <SchemaDefinition schemaRef="#/components/schemas/BuiltinTool" />
name: BuiltinTool name: BuiltinTool
- description: <SchemaDefinition schemaRef="#/components/schemas/CompletionMessage" - description: <SchemaDefinition schemaRef="#/components/schemas/CompletionMessage"

View file

@ -40,5 +40,5 @@ class RunShieldResponse(BaseModel):
class Safety(Protocol): class Safety(Protocol):
@webmethod(route="/safety/run_shield") @webmethod(route="/safety/run_shield")
async def run_shield( async def run_shield(
self, shield: str, messages: List[Message], params: Dict[str, Any] = None self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None
) -> RunShieldResponse: ... ) -> RunShieldResponse: ...

View file

@ -162,7 +162,6 @@ class SafetyRouter(Safety):
messages: List[Message], messages: List[Message],
params: Dict[str, Any] = None, params: Dict[str, Any] = None,
) -> RunShieldResponse: ) -> RunShieldResponse:
print(f"Running shield {shield_type}")
return await self.routing_table.get_provider_impl(shield_type).run_shield( return await self.routing_table.get_provider_impl(shield_type).run_shield(
shield_type=shield_type, shield_type=shield_type,
messages=messages, messages=messages,