From a64451ad42a01dec440b4ad1619aa73b89715e94 Mon Sep 17 00:00:00 2001 From: Swapna Lekkala Date: Mon, 11 Aug 2025 15:40:27 -0700 Subject: [PATCH] minor change --- .../providers/inline/safety/llama_guard/llama_guard.py | 7 +------ tests/integration/safety/test_safety.py | 9 +++++++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/llama_stack/providers/inline/safety/llama_guard/llama_guard.py b/llama_stack/providers/inline/safety/llama_guard/llama_guard.py index bae744010..a1133f932 100644 --- a/llama_stack/providers/inline/safety/llama_guard/llama_guard.py +++ b/llama_stack/providers/inline/safety/llama_guard/llama_guard.py @@ -11,11 +11,7 @@ from string import Template from typing import Any from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem -from llama_stack.apis.inference import ( - Inference, - Message, - UserMessage, -) +from llama_stack.apis.inference import Inference, Message, UserMessage from llama_stack.apis.safety import ( RunShieldResponse, Safety, @@ -72,7 +68,6 @@ SAFETY_CATEGORIES_TO_CODE_MAP = { } SAFETY_CODE_TO_CATEGORIES_MAP = {v: k for k, v in SAFETY_CATEGORIES_TO_CODE_MAP.items()} - DEFAULT_LG_V3_SAFETY_CATEGORIES = [ CAT_VIOLENT_CRIMES, CAT_NON_VIOLENT_CRIMES, diff --git a/tests/integration/safety/test_safety.py b/tests/integration/safety/test_safety.py index 7a1695f00..c3d9836a8 100644 --- a/tests/integration/safety/test_safety.py +++ b/tests/integration/safety/test_safety.py @@ -166,7 +166,8 @@ def test_unsafe_input_with_code_scanner_run_moderation(client_with_models, code_ if not code_safety_enabled: pytest.skip(f"Testing code scanner shields is not supported for model_providers {model_providers}") - insecure_code = """ + insecure_code = [ + """ Here is an example of code that hashes a password using MD5: ``` @@ -176,6 +177,7 @@ def test_unsafe_input_with_code_scanner_run_moderation(client_with_models, code_ print("Insecure hashed password:", hashed_password) ``` """ + ] * 2 moderation_object = client_with_models.moderations.create( input=insecure_code, model=code_scanner_shield_id, @@ -191,7 +193,8 @@ def test_safe_input_with_code_scanner_run_moderation(client_with_models, code_sc if not code_safety_enabled: pytest.skip(f"Testing code scanner shields is not supported for model_providers {model_providers}") - secure_code = """ + secure_code = [ + """ Extract the first 5 characters from a string: ``` text = "Hello World" @@ -203,10 +206,12 @@ def test_safe_input_with_code_scanner_run_moderation(client_with_models, code_sc return text[:5] if text else "" ``` """ + ] * 2 moderation_object = client_with_models.moderations.create( input=secure_code, model=code_scanner_shield_id, ) + assert moderation_object.results[0].flagged is False, "Code scanner should not have flagged the code as insecure"