phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-30 23:51:00 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

use env vars for bedrock guardrail vars

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-07 20:15:42 -08:00
parent 98c09323a9
commit 932b524449
2 changed files with 3 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

12
llama_stack/providers/tests/safety/conftest.py
View file

@ -71,18 +71,6 @@ def pytest_addoption(parser):
default=None,
help="Specify the safety model to use for testing",
)
parser.addoption(
"--bedrock-guardrail-id",
action="store",
default=None,
help="Specify the guard rail ID to use for testing bedrock",
)
parser.addoption(
"--bedrock-guardrail-version",
action="store",
default=None,
help="Specify the guard rail version to use for testing bedrock",
)
SAFETY_MODEL_PARAMS = [

7
llama_stack/providers/tests/safety/fixtures.py
View file

@ -15,6 +15,7 @@ from llama_stack.providers.inline.safety.meta_reference import (
SafetyConfig,
)
from llama_stack.providers.remote.safety.bedrock import BedrockSafetyConfig
from llama_stack.providers.tests.env import get_env_or_fail
from llama_stack.providers.tests.resolver import resolve_impls_for_test_v2
from ..conftest import ProviderFixture, remote_stack_fixture
@ -120,10 +121,8 @@ async def safety_stack(inference_model, safety_model, request):
elif provider_type == "remote::together":
shield_config["model"] = safety_model
elif provider_type == "remote::bedrock":
identifier = request.config.getoption("--bedrock-guardrail-id", None)
shield_config["guardrailVersion"] = request.config.getoption(
"--bedrock-guardrail-version", None
)
identifier = get_env_or_fail("BEDROCK_GUARDRAIL_IDENTIFIER")
shield_config["guardrailVersion"] = get_env_or_fail("BEDROCK_GUARDRAIL_VERSION")
# Create shield
shield = Shield(