No RunShieldRequest

togethercomputer · Sep 19, 2024 · 8cdc2f0 · 8cdc2f0
1 parent f3f5873
commit 8cdc2f0
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/llama_stack/providers/impls/meta_reference/agents/safety.py b/llama_stack/providers/impls/meta_reference/agents/safety.py
@@ -7,15 +7,14 @@
 from typing import List
 
 from llama_models.llama3.api.datatypes import Message, Role, UserMessage
+from termcolor import cprint
 
 from llama_stack.apis.safety import (
     OnViolationAction,
-    RunShieldRequest,
     Safety,
     ShieldDefinition,
     ShieldResponse,
 )
-from termcolor import cprint
 
 
 class SafetyException(Exception):  # noqa: N818
@@ -45,10 +44,8 @@ async def run_shields(
             messages[0] = UserMessage(content=messages[0].content)
 
         res = await self.safety_api.run_shields(
-            RunShieldRequest(
-                messages=messages,
-                shields=shields,
-            )
+            messages=messages,
+            shields=shields,
         )
 
         results = res.responses