diff --git a/llama_stack/providers/impls/meta_reference/agents/agent_instance.py b/llama_stack/providers/impls/meta_reference/agents/agent_instance.py
index 4199b2e96..7cf86b1f3 100644
--- a/llama_stack/providers/impls/meta_reference/agents/agent_instance.py
+++ b/llama_stack/providers/impls/meta_reference/agents/agent_instance.py
@@ -38,7 +38,7 @@ from .tools.builtin import (
     SearchTool,
     WolframAlphaTool,
 )
-from .tools.safety import with_safety
+from .tools.safety import SafeTool
 
 
 def make_random_string(length: int = 8):
@@ -87,7 +87,7 @@ class ChatAgent(ShieldRunnerMixin):
                 continue
 
             builtin_tools.append(
-                with_safety(
+                SafeTool(
                     tool,
                     safety_api,
                     tool_defn.input_shields,
diff --git a/llama_stack/providers/impls/meta_reference/agents/tools/safety.py b/llama_stack/providers/impls/meta_reference/agents/tools/safety.py
index df0cd4bd3..fb95786d1 100644
--- a/llama_stack/providers/impls/meta_reference/agents/tools/safety.py
+++ b/llama_stack/providers/impls/meta_reference/agents/tools/safety.py
@@ -41,17 +41,3 @@ class SafeTool(BaseTool, ShieldRunnerMixin):
             await self.run_multiple_shields(messages, self.output_shields)
 
         return res
-
-
-def with_safety(
-    tool: BaseTool,
-    safety_api: Safety,
-    input_shields: List[str] = None,
-    output_shields: List[str] = None,
-) -> SafeTool:
-    return SafeTool(
-        tool,
-        safety_api,
-        input_shields=input_shields,
-        output_shields=output_shields,
-    )