Begin adding a /safety/run_shield API

This commit is contained in:
Ashwin Bharambe 2024-07-31 21:57:10 -07:00
parent 1bc81eae7b
commit c253c1c9ad
2 changed files with 28 additions and 7 deletions

View file

@ -44,13 +44,6 @@ class ShieldDefinition(BaseModel):
execution_config: Optional[RestAPIExecutionConfig] = None execution_config: Optional[RestAPIExecutionConfig] = None
@json_schema_type
class ShieldCall(BaseModel):
call_id: str
shield_type: ShieldType
arguments: Dict[str, str]
@json_schema_type @json_schema_type
class ShieldResponse(BaseModel): class ShieldResponse(BaseModel):
shield_type: ShieldType shield_type: ShieldType

View file

@ -0,0 +1,28 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from .datatypes import * # noqa: F403
from typing import Protocol
from llama_models.llama3_1.api.datatypes import Message
# this dependency is annoying and we need a forked up version anyway
from pyopenapi import webmethod
@json_schema_type
class RunShieldRequest(BaseModel):
shield_type: ShieldType
messages: List[Message]
class SafetyCheck(Protocol):
@webmethod(route="/safety/run_shield")
async def run_shield(
self,
request: RunShieldRequest,
) -> ShieldResponse: ...