Resource oriented design for shields (#399)

* init

* working bedrock tests

* bedrock test for inference fixes

* use env vars for bedrock guardrail vars

* add register in meta reference

* use correct shield impl in meta ref

* dont add together fixture

* right naming

* minor updates

* improved registration flow

* address feedback

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>
This commit is contained in:
Dinesh Yeduguru 2024-11-08 12:16:11 -08:00 committed by GitHub
parent 7ee9f8d8ac
commit d800a16acd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 262 additions and 124 deletions

View file

@ -41,13 +41,13 @@ class SafetyClient(Safety):
pass
async def run_shield(
self, shield_type: str, messages: List[Message]
self, shield_id: str, messages: List[Message]
) -> RunShieldResponse:
async with httpx.AsyncClient() as client:
response = await client.post(
f"{self.base_url}/safety/run_shield",
json=dict(
shield_type=shield_type,
shield_id=shield_id,
messages=[encodable_dict(m) for m in messages],
),
headers={
@ -80,7 +80,7 @@ async def run_main(host: str, port: int, image_path: str = None):
)
cprint(f"User>{message.content}", "green")
response = await client.run_shield(
shield_type="llama_guard",
shield_id="llama_guard",
messages=[message],
)
print(response)
@ -91,7 +91,7 @@ async def run_main(host: str, port: int, image_path: str = None):
]:
cprint(f"User>{message.content}", "green")
response = await client.run_shield(
shield_type="llama_guard",
shield_id="llama_guard",
messages=[message],
)
print(response)