diff --git a/llama_toolchain/stack.py b/llama_toolchain/stack.py
index 875bc5802..1e2976ab3 100644
--- a/llama_toolchain/stack.py
+++ b/llama_toolchain/stack.py
@@ -15,6 +15,7 @@ from llama_toolchain.telemetry.api import * # noqa: F403
from llama_toolchain.post_training.api import * # noqa: F403
from llama_toolchain.reward_scoring.api import * # noqa: F403
from llama_toolchain.synthetic_data_generation.api import * # noqa: F403
+from llama_toolchain.safety.api import * # noqa: F403
class LlamaStack(
@@ -22,6 +23,7 @@ class LlamaStack(
BatchInference,
AgenticSystem,
RewardScoring,
+ Safety,
SyntheticDataGeneration,
Datasets,
Telemetry,
diff --git a/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.html b/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.html
index f5f5fa154..211290ce1 100644
--- a/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.html
+++ b/rfcs/RFC-0001-llama-stack-assets/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-09-10 01:13:08.531639"
+ "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-10 16:42:15.870336"
},
"servers": [
{
@@ -1433,6 +1433,36 @@
}
}
},
+ "/safety/run_shields": {
+ "post": {
+ "responses": {
+ "200": {
+ "description": "OK",
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/RunShieldResponse"
+ }
+ }
+ }
+ }
+ },
+ "tags": [
+ "Safety"
+ ],
+ "parameters": [],
+ "requestBody": {
+ "content": {
+ "application/json": {
+ "schema": {
+ "$ref": "#/components/schemas/RunShieldRequestWrapper"
+ }
+ }
+ },
+ "required": true
+ }
+ }
+ },
"/post_training/supervised_fine_tune": {
"post": {
"responses": {
@@ -2409,7 +2439,7 @@
"items": {
"oneOf": [
{
- "$ref": "#/components/schemas/BraveSearchToolDefinition"
+ "$ref": "#/components/schemas/SearchToolDefinition"
},
{
"$ref": "#/components/schemas/WolframAlphaToolDefinition"
@@ -2624,34 +2654,6 @@
"instructions"
]
},
- "BraveSearchToolDefinition": {
- "type": "object",
- "properties": {
- "input_shields": {
- "type": "array",
- "items": {
- "$ref": "#/components/schemas/ShieldDefinition"
- }
- },
- "output_shields": {
- "type": "array",
- "items": {
- "$ref": "#/components/schemas/ShieldDefinition"
- }
- },
- "type": {
- "type": "string",
- "const": "brave_search"
- },
- "remote_execution": {
- "$ref": "#/components/schemas/RestAPIExecutionConfig"
- }
- },
- "additionalProperties": false,
- "required": [
- "type"
- ]
- },
"BuiltinShield": {
"type": "string",
"enum": [
@@ -2785,19 +2787,76 @@
"params": {
"type": "object",
"additionalProperties": {
- "type": "string"
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
}
},
"headers": {
"type": "object",
"additionalProperties": {
- "type": "string"
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
}
},
"body": {
"type": "object",
"additionalProperties": {
- "type": "string"
+ "oneOf": [
+ {
+ "type": "null"
+ },
+ {
+ "type": "boolean"
+ },
+ {
+ "type": "number"
+ },
+ {
+ "type": "string"
+ },
+ {
+ "type": "array"
+ },
+ {
+ "type": "object"
+ }
+ ]
}
}
},
@@ -2816,6 +2875,42 @@
"DELETE"
]
},
+ "SearchToolDefinition": {
+ "type": "object",
+ "properties": {
+ "input_shields": {
+ "type": "array",
+ "items": {
+ "$ref": "#/components/schemas/ShieldDefinition"
+ }
+ },
+ "output_shields": {
+ "type": "array",
+ "items": {
+ "$ref": "#/components/schemas/ShieldDefinition"
+ }
+ },
+ "type": {
+ "type": "string",
+ "const": "brave_search"
+ },
+ "engine": {
+ "type": "string",
+ "enum": [
+ "bing",
+ "brave"
+ ]
+ },
+ "remote_execution": {
+ "$ref": "#/components/schemas/RestAPIExecutionConfig"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "type",
+ "engine"
+ ]
+ },
"ShieldDefinition": {
"type": "object",
"properties": {
@@ -2959,7 +3054,7 @@
"items": {
"oneOf": [
{
- "$ref": "#/components/schemas/BraveSearchToolDefinition"
+ "$ref": "#/components/schemas/SearchToolDefinition"
},
{
"$ref": "#/components/schemas/WolframAlphaToolDefinition"
@@ -5549,6 +5644,68 @@
"score"
]
},
+ "RunShieldRequest": {
+ "type": "object",
+ "properties": {
+ "messages": {
+ "type": "array",
+ "items": {
+ "oneOf": [
+ {
+ "$ref": "#/components/schemas/UserMessage"
+ },
+ {
+ "$ref": "#/components/schemas/SystemMessage"
+ },
+ {
+ "$ref": "#/components/schemas/ToolResponseMessage"
+ },
+ {
+ "$ref": "#/components/schemas/CompletionMessage"
+ }
+ ]
+ }
+ },
+ "shields": {
+ "type": "array",
+ "items": {
+ "$ref": "#/components/schemas/ShieldDefinition"
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "messages",
+ "shields"
+ ]
+ },
+ "RunShieldRequestWrapper": {
+ "type": "object",
+ "properties": {
+ "request": {
+ "$ref": "#/components/schemas/RunShieldRequest"
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "request"
+ ]
+ },
+ "RunShieldResponse": {
+ "type": "object",
+ "properties": {
+ "responses": {
+ "type": "array",
+ "items": {
+ "$ref": "#/components/schemas/ShieldResponse"
+ }
+ }
+ },
+ "additionalProperties": false,
+ "required": [
+ "responses"
+ ]
+ },
"DoraFinetuningConfig": {
"type": "object",
"properties": {
@@ -6062,36 +6219,39 @@
}
],
"tags": [
- {
- "name": "SyntheticDataGeneration"
- },
- {
- "name": "RewardScoring"
- },
- {
- "name": "Datasets"
- },
- {
- "name": "Memory"
- },
- {
- "name": "AgenticSystem"
- },
{
"name": "BatchInference"
},
{
- "name": "PostTraining"
- },
- {
- "name": "Evaluations"
+ "name": "Safety"
},
{
"name": "Telemetry"
},
+ {
+ "name": "RewardScoring"
+ },
+ {
+ "name": "Memory"
+ },
+ {
+ "name": "SyntheticDataGeneration"
+ },
{
"name": "Inference"
},
+ {
+ "name": "Evaluations"
+ },
+ {
+ "name": "PostTraining"
+ },
+ {
+ "name": "Datasets"
+ },
+ {
+ "name": "AgenticSystem"
+ },
{
"name": "BatchChatCompletionRequest",
"description": ""
@@ -6224,10 +6384,6 @@
"name": "AgentConfig",
"description": ""
},
- {
- "name": "BraveSearchToolDefinition",
- "description": ""
- },
{
"name": "BuiltinShield",
"description": ""
@@ -6256,6 +6412,10 @@
"name": "RestAPIMethod",
"description": ""
},
+ {
+ "name": "SearchToolDefinition",
+ "description": ""
+ },
{
"name": "ShieldDefinition",
"description": ""
@@ -6620,6 +6780,18 @@
"name": "ScoredMessage",
"description": ""
},
+ {
+ "name": "RunShieldRequest",
+ "description": ""
+ },
+ {
+ "name": "RunShieldRequestWrapper",
+ "description": ""
+ },
+ {
+ "name": "RunShieldResponse",
+ "description": ""
+ },
{
"name": "DoraFinetuningConfig",
"description": ""
@@ -6697,6 +6869,7 @@
"Memory",
"PostTraining",
"RewardScoring",
+ "Safety",
"SyntheticDataGeneration",
"Telemetry"
]
@@ -6726,7 +6899,6 @@
"BatchCompletionRequest",
"BatchCompletionRequestWrapper",
"BatchCompletionResponse",
- "BraveSearchToolDefinition",
"BuiltinShield",
"BuiltinTool",
"CancelEvaluationJobRequest",
@@ -6814,10 +6986,14 @@
"RewardScoringRequestWrapper",
"RewardScoringResponse",
"Run",
+ "RunShieldRequest",
+ "RunShieldRequestWrapper",
+ "RunShieldResponse",
"SamplingParams",
"SamplingStrategy",
"ScoredDialogGenerations",
"ScoredMessage",
+ "SearchToolDefinition",
"Session",
"ShieldCallStep",
"ShieldDefinition",
diff --git a/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.yaml b/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.yaml
index 3c3475fff..322645813 100644
--- a/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.yaml
+++ b/rfcs/RFC-0001-llama-stack-assets/llama-stack-spec.yaml
@@ -25,7 +25,7 @@ components:
tools:
items:
oneOf:
- - $ref: '#/components/schemas/BraveSearchToolDefinition'
+ - $ref: '#/components/schemas/SearchToolDefinition'
- $ref: '#/components/schemas/WolframAlphaToolDefinition'
- $ref: '#/components/schemas/PhotogenToolDefinition'
- $ref: '#/components/schemas/CodeInterpreterToolDefinition'
@@ -218,7 +218,7 @@ components:
tools:
items:
oneOf:
- - $ref: '#/components/schemas/BraveSearchToolDefinition'
+ - $ref: '#/components/schemas/SearchToolDefinition'
- $ref: '#/components/schemas/WolframAlphaToolDefinition'
- $ref: '#/components/schemas/PhotogenToolDefinition'
- $ref: '#/components/schemas/CodeInterpreterToolDefinition'
@@ -635,25 +635,6 @@ components:
required:
- completion_message_batch
type: object
- BraveSearchToolDefinition:
- additionalProperties: false
- properties:
- input_shields:
- items:
- $ref: '#/components/schemas/ShieldDefinition'
- type: array
- output_shields:
- items:
- $ref: '#/components/schemas/ShieldDefinition'
- type: array
- remote_execution:
- $ref: '#/components/schemas/RestAPIExecutionConfig'
- type:
- const: brave_search
- type: string
- required:
- - type
- type: object
BuiltinShield:
enum:
- llama_guard
@@ -2011,17 +1992,35 @@ components:
properties:
body:
additionalProperties:
- type: string
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
type: object
headers:
additionalProperties:
- type: string
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
type: object
method:
$ref: '#/components/schemas/RestAPIMethod'
params:
additionalProperties:
- type: string
+ oneOf:
+ - type: 'null'
+ - type: boolean
+ - type: number
+ - type: string
+ - type: array
+ - type: object
type: object
url:
$ref: '#/components/schemas/URL'
@@ -2103,6 +2102,43 @@ components:
- started_at
- metadata
type: object
+ RunShieldRequest:
+ additionalProperties: false
+ properties:
+ messages:
+ items:
+ oneOf:
+ - $ref: '#/components/schemas/UserMessage'
+ - $ref: '#/components/schemas/SystemMessage'
+ - $ref: '#/components/schemas/ToolResponseMessage'
+ - $ref: '#/components/schemas/CompletionMessage'
+ type: array
+ shields:
+ items:
+ $ref: '#/components/schemas/ShieldDefinition'
+ type: array
+ required:
+ - messages
+ - shields
+ type: object
+ RunShieldRequestWrapper:
+ additionalProperties: false
+ properties:
+ request:
+ $ref: '#/components/schemas/RunShieldRequest'
+ required:
+ - request
+ type: object
+ RunShieldResponse:
+ additionalProperties: false
+ properties:
+ responses:
+ items:
+ $ref: '#/components/schemas/ShieldResponse'
+ type: array
+ required:
+ - responses
+ type: object
SamplingParams:
additionalProperties: false
properties:
@@ -2161,6 +2197,31 @@ components:
- message
- score
type: object
+ SearchToolDefinition:
+ additionalProperties: false
+ properties:
+ engine:
+ enum:
+ - bing
+ - brave
+ type: string
+ input_shields:
+ items:
+ $ref: '#/components/schemas/ShieldDefinition'
+ type: array
+ output_shields:
+ items:
+ $ref: '#/components/schemas/ShieldDefinition'
+ type: array
+ remote_execution:
+ $ref: '#/components/schemas/RestAPIExecutionConfig'
+ type:
+ const: brave_search
+ type: string
+ required:
+ - type
+ - engine
+ type: object
Session:
additionalProperties: false
properties:
@@ -2777,7 +2838,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-09-10 01:13:08.531639"
+ \ draft and subject to change.\n Generated at 2024-09-10 16:42:15.870336"
title: '[DRAFT] Llama Stack Specification'
version: 0.0.1
jsonSchemaDialect: https://json-schema.org/draft/2020-12/schema
@@ -3706,6 +3767,24 @@ paths:
description: OK
tags:
- Telemetry
+ /safety/run_shields:
+ post:
+ parameters: []
+ requestBody:
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/RunShieldRequestWrapper'
+ required: true
+ responses:
+ '200':
+ content:
+ application/json:
+ schema:
+ $ref: '#/components/schemas/RunShieldResponse'
+ description: OK
+ tags:
+ - Safety
/synthetic_data_generation/generate:
post:
parameters: []
@@ -3729,16 +3808,17 @@ security:
servers:
- url: http://any-hosted-llama-stack.com
tags:
-- name: SyntheticDataGeneration
-- name: RewardScoring
-- name: Datasets
-- name: Memory
-- name: AgenticSystem
- name: BatchInference
-- name: PostTraining
-- name: Evaluations
+- name: Safety
- name: Telemetry
+- name: RewardScoring
+- name: Memory
+- name: SyntheticDataGeneration
- name: Inference
+- name: Evaluations
+- name: PostTraining
+- name: Datasets
+- name: AgenticSystem
- description:
name: BatchChatCompletionRequest
@@ -3844,9 +3924,6 @@ tags:
name: CompletionResponseStreamChunk
- description:
name: AgentConfig
-- description:
- name: BraveSearchToolDefinition
- description:
name: BuiltinShield
- description:
name: RestAPIMethod
+- description:
+ name: SearchToolDefinition
- description:
name: ShieldDefinition
@@ -4163,6 +4243,15 @@ tags:
name: ScoredDialogGenerations
- description:
name: ScoredMessage
+- description:
+ name: RunShieldRequest
+- description:
+ name: RunShieldRequestWrapper
+- description:
+ name: RunShieldResponse
- description:
name: DoraFinetuningConfig
@@ -4232,6 +4321,7 @@ x-tagGroups:
- Memory
- PostTraining
- RewardScoring
+ - Safety
- SyntheticDataGeneration
- Telemetry
- name: Types
@@ -4258,7 +4348,6 @@ x-tagGroups:
- BatchCompletionRequest
- BatchCompletionRequestWrapper
- BatchCompletionResponse
- - BraveSearchToolDefinition
- BuiltinShield
- BuiltinTool
- CancelEvaluationJobRequest
@@ -4346,10 +4435,14 @@ x-tagGroups:
- RewardScoringRequestWrapper
- RewardScoringResponse
- Run
+ - RunShieldRequest
+ - RunShieldRequestWrapper
+ - RunShieldResponse
- SamplingParams
- SamplingStrategy
- ScoredDialogGenerations
- ScoredMessage
+ - SearchToolDefinition
- Session
- ShieldCallStep
- ShieldDefinition