minor change

This commit is contained in:
Swapna Lekkala 2025-08-11 15:40:27 -07:00
parent 3eef4c6d0f
commit a64451ad42
2 changed files with 8 additions and 8 deletions

View file

@ -11,11 +11,7 @@ from string import Template
from typing import Any from typing import Any
from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem
from llama_stack.apis.inference import ( from llama_stack.apis.inference import Inference, Message, UserMessage
Inference,
Message,
UserMessage,
)
from llama_stack.apis.safety import ( from llama_stack.apis.safety import (
RunShieldResponse, RunShieldResponse,
Safety, 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()} SAFETY_CODE_TO_CATEGORIES_MAP = {v: k for k, v in SAFETY_CATEGORIES_TO_CODE_MAP.items()}
DEFAULT_LG_V3_SAFETY_CATEGORIES = [ DEFAULT_LG_V3_SAFETY_CATEGORIES = [
CAT_VIOLENT_CRIMES, CAT_VIOLENT_CRIMES,
CAT_NON_VIOLENT_CRIMES, CAT_NON_VIOLENT_CRIMES,

View file

@ -166,7 +166,8 @@ def test_unsafe_input_with_code_scanner_run_moderation(client_with_models, code_
if not code_safety_enabled: if not code_safety_enabled:
pytest.skip(f"Testing code scanner shields is not supported for model_providers {model_providers}") 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: 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) print("Insecure hashed password:", hashed_password)
``` ```
""" """
] * 2
moderation_object = client_with_models.moderations.create( moderation_object = client_with_models.moderations.create(
input=insecure_code, input=insecure_code,
model=code_scanner_shield_id, 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: if not code_safety_enabled:
pytest.skip(f"Testing code scanner shields is not supported for model_providers {model_providers}") 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: Extract the first 5 characters from a string:
``` ```
text = "Hello World" 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 "" return text[:5] if text else ""
``` ```
""" """
] * 2
moderation_object = client_with_models.moderations.create( moderation_object = client_with_models.moderations.create(
input=secure_code, input=secure_code,
model=code_scanner_shield_id, model=code_scanner_shield_id,
) )
assert moderation_object.results[0].flagged is False, "Code scanner should not have flagged the code as insecure" assert moderation_object.results[0].flagged is False, "Code scanner should not have flagged the code as insecure"