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

@ -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: ...