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

improved registration flow

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-08 11:07:41 -08:00
parent 39f0c5f544
commit 0eaca98229
10 changed files with 95 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -7,7 +7,7 @@
import pytest
import pytest_asyncio
from llama_stack.apis.shields import Shield, ShieldType
from llama_stack.apis.shields import ShieldType
from llama_stack.distribution.datatypes import Api, Provider
from llama_stack.providers.inline.safety.meta_reference import (
@ -95,10 +95,10 @@ async def safety_stack(inference_model, safety_model, request):
shields_impl = impls[Api.shields]
# Register the appropriate shield based on provider type
provider_id = safety_fixture.providers[0].provider_id
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
@ -107,12 +107,11 @@ async def safety_stack(inference_model, safety_model, request):
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
# Create shield
shield = Shield(
identifier=identifier,
shield_type=ShieldType.llama_guard,
provider_id=provider_id,
shield = await shields_impl.register_shield(
shield_id=identifier,
shield_type=shield_type,
params=shield_config,
)