From 39f0c5f5442ade707493ff19c8ff7d2467e0bdb7 Mon Sep 17 00:00:00 2001
From: Dinesh Yeduguru <dineshyv@fb.com>
Date: Thu, 7 Nov 2024 22:40:15 -0800
Subject: [PATCH] minor updates

---
 llama_stack/providers/inline/agents/meta_reference/safety.py    | 2 +-
 .../inline/agents/meta_reference/tests/test_chat_agent.py       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/llama_stack/providers/inline/agents/meta_reference/safety.py b/llama_stack/providers/inline/agents/meta_reference/safety.py
index 915ddd303..77525e871 100644
--- a/llama_stack/providers/inline/agents/meta_reference/safety.py
+++ b/llama_stack/providers/inline/agents/meta_reference/safety.py
@@ -37,7 +37,7 @@ class ShieldRunnerMixin:
         responses = await asyncio.gather(
             *[
                 self.safety_api.run_shield(
-                    identifier=identifier,
+                    shield_id=identifier,
                     messages=messages,
                 )
                 for identifier in identifiers
diff --git a/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py b/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py
index bd059c990..6edef0672 100644
--- a/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py
+++ b/llama_stack/providers/inline/agents/meta_reference/tests/test_chat_agent.py
@@ -80,7 +80,7 @@ class MockInferenceAPI:
 
 class MockSafetyAPI:
     async def run_shield(
-        self, shield_identifier: str, messages: List[Message]
+        self, shield_id: str, messages: List[Message]
     ) -> RunShieldResponse:
         return RunShieldResponse(violation=None)