From cd4880126b158d98a9036f0e13f818751ce6bcb3 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Sat, 21 Sep 2024 23:02:47 -0700
Subject: [PATCH] kill with_safety(), not needed

---
 .../impls/meta_reference/agents/agent_instance.py  |  4 ++--
 .../impls/meta_reference/agents/tools/safety.py    | 14 --------------
 2 files changed, 2 insertions(+), 16 deletions(-)

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,
-    )