No RunShieldRequest

This commit is contained in:
Ashwin Bharambe 2024-09-18 20:38:21 -07:00
parent f3f5873e9e
commit 8cdc2f0cfb

View file

@ -7,15 +7,14 @@
from typing import List from typing import List
from llama_models.llama3.api.datatypes import Message, Role, UserMessage from llama_models.llama3.api.datatypes import Message, Role, UserMessage
from termcolor import cprint
from llama_stack.apis.safety import ( from llama_stack.apis.safety import (
OnViolationAction, OnViolationAction,
RunShieldRequest,
Safety, Safety,
ShieldDefinition, ShieldDefinition,
ShieldResponse, ShieldResponse,
) )
from termcolor import cprint
class SafetyException(Exception): # noqa: N818 class SafetyException(Exception): # noqa: N818
@ -45,10 +44,8 @@ class ShieldRunnerMixin:
messages[0] = UserMessage(content=messages[0].content) messages[0] = UserMessage(content=messages[0].content)
res = await self.safety_api.run_shields( res = await self.safety_api.run_shields(
RunShieldRequest( messages=messages,
messages=messages, shields=shields,
shields=shields,
)
) )
results = res.responses results = res.responses