phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

Resource oriented design for shields (#399)

Browse source 
* init

* working bedrock tests

* bedrock test for inference fixes

* use env vars for bedrock guardrail vars

* add register in meta reference

* use correct shield impl in meta ref

* dont add together fixture

* right naming

* minor updates

* improved registration flow

* address feedback

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>
This commit is contained in:
Dinesh Yeduguru 2024-11-08 12:16:11 -08:00 • committed by GitHub
parent 7ee9f8d8ac
commit d800a16acd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 262 additions and 124 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -7,12 +7,15 @@
import pytest
import pytest_asyncio
from llama_stack.apis.shields import ShieldType
from llama_stack.distribution.datatypes import Api, Provider
from llama_stack.providers.inline.safety.meta_reference import (
LlamaGuardShieldConfig,
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
@ -47,7 +50,20 @@ def safety_meta_reference(safety_model) -> ProviderFixture:
)
SAFETY_FIXTURES = ["meta_reference", "remote"]
@pytest.fixture(scope="session")
def safety_bedrock() -> ProviderFixture:
return ProviderFixture(
providers=[
Provider(
provider_id="bedrock",
provider_type="remote::bedrock",
config=BedrockSafetyConfig().model_dump(),
)
],
)
SAFETY_FIXTURES = ["meta_reference", "bedrock", "remote"]
@pytest_asyncio.fixture(scope="session")
@ -74,4 +90,29 @@ async def safety_stack(inference_model, safety_model, request):
providers,
provider_data,
)
return impls[Api.safety], impls[Api.shields]
safety_impl = impls[Api.safety]
shields_impl = impls[Api.shields]
# Register the appropriate shield based on provider type
provider_type = safety_fixture.providers[0].provider_type
shield_config = {}
shield_type = ShieldType.llama_guard
identifier = "llama_guard"
if provider_type == "meta-reference":
shield_config["model"] = safety_model
elif provider_type == "remote::together":
shield_config["model"] = safety_model
elif provider_type == "remote::bedrock":
identifier = get_env_or_fail("BEDROCK_GUARDRAIL_IDENTIFIER")
shield_config["guardrailVersion"] = get_env_or_fail("BEDROCK_GUARDRAIL_VERSION")
shield_type = ShieldType.generic_content_shield
shield = await shields_impl.register_shield(
shield_id=identifier,
shield_type=shield_type,
params=shield_config,
)
return safety_impl, shields_impl, shield