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

Kill the notion of shield_type

Browse source
This commit is contained in:
Ashwin Bharambe 2024-11-12 11:41:23 -08:00
parent 09269e2a44
commit b1c3a95485
20 changed files with 87 additions and 161 deletions
Download patch file Download diff file Expand all files Collapse all files

1
llama_stack/apis/datasets/datasets.py
View file

@ -38,7 +38,6 @@ class Dataset(CommonDatasetFields, Resource):
return self.provider_resource_id
@json_schema_type
class DatasetInput(CommonDatasetFields, BaseModel):
dataset_id: str
provider_id: Optional[str] = None

1
llama_stack/apis/eval_tasks/eval_tasks.py
View file

@ -34,7 +34,6 @@ class EvalTask(CommonEvalTaskFields, Resource):
return self.provider_resource_id
@json_schema_type
class EvalTaskInput(CommonEvalTaskFields, BaseModel):
eval_task_id: str
provider_id: Optional[str] = None

1
llama_stack/apis/memory_banks/memory_banks.py
View file

@ -122,7 +122,6 @@ MemoryBank = Annotated[
]
@json_schema_type
class MemoryBankInput(BaseModel):
memory_bank_id: str
params: BankParams

1
llama_stack/apis/models/models.py
View file

@ -32,7 +32,6 @@ class Model(CommonModelFields, Resource):
return self.provider_resource_id
@json_schema_type
class ModelInput(CommonModelFields):
model_id: str
provider_id: Optional[str] = None

1
llama_stack/apis/scoring_functions/scoring_functions.py
View file

@ -96,7 +96,6 @@ class ScoringFn(CommonScoringFnFields, Resource):
return self.provider_resource_id
@json_schema_type
class ScoringFnInput(CommonScoringFnFields, BaseModel):
scoring_fn_id: str
provider_id: Optional[str] = None

6
llama_stack/apis/shields/client.py
View file

@ -37,7 +37,6 @@ class ShieldsClient(Shields):
async def register_shield(
self,
shield_id: str,
shield_type: ShieldType,
provider_shield_id: Optional[str],
provider_id: Optional[str],
params: Optional[Dict[str, Any]],
@ -47,7 +46,6 @@ class ShieldsClient(Shields):
f"{self.base_url}/shields/register",
json={
"shield_id": shield_id,
"shield_type": shield_type,
"provider_shield_id": provider_shield_id,
"provider_id": provider_id,
"params": params,
@ -56,12 +54,12 @@ class ShieldsClient(Shields):
)
response.raise_for_status()
async def get_shield(self, shield_type: str) -> Optional[Shield]:
async def get_shield(self, shield_id: str) -> Optional[Shield]:
async with httpx.AsyncClient() as client:
response = await client.get(
f"{self.base_url}/shields/get",
params={
"shield_type": shield_type,
"shield_id": shield_id,
},
headers={"Content-Type": "application/json"},
)

11
llama_stack/apis/shields/shields.py
View file

@ -4,7 +4,6 @@
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
from llama_models.schema_utils import json_schema_type, webmethod
@ -13,16 +12,7 @@ from pydantic import BaseModel
from llama_stack.apis.resource import Resource, ResourceType
@json_schema_type
class ShieldType(Enum):
generic_content_shield = "generic_content_shield"
llama_guard = "llama_guard"
code_scanner = "code_scanner"
prompt_guard = "prompt_guard"
class CommonShieldFields(BaseModel):
shield_type: ShieldType
params: Optional[Dict[str, Any]] = None
@ -59,7 +49,6 @@ class Shields(Protocol):
async def register_shield(
self,
shield_id: str,
shield_type: ShieldType,
provider_shield_id: Optional[str] = None,
provider_id: Optional[str] = None,
params: Optional[Dict[str, Any]] = None,

3
llama_stack/distribution/routers/routers.py
View file

@ -172,13 +172,12 @@ class SafetyRouter(Safety):
async def register_shield(
self,
shield_id: str,
shield_type: ShieldType,
provider_shield_id: Optional[str] = None,
provider_id: Optional[str] = None,
params: Optional[Dict[str, Any]] = None,
) -> Shield:
return await self.routing_table.register_shield(
shield_id, shield_type, provider_shield_id, provider_id, params
shield_id, provider_shield_id, provider_id, params
)
async def run_shield(

2
llama_stack/distribution/routers/routing_tables.py
View file

@ -241,7 +241,6 @@ class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
async def register_shield(
self,
shield_id: str,
shield_type: ShieldType,
provider_shield_id: Optional[str] = None,
provider_id: Optional[str] = None,
params: Optional[Dict[str, Any]] = None,
@ -260,7 +259,6 @@ class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):
params = {}
shield = Shield(
identifier=shield_id,
shield_type=shield_type,
provider_resource_id=provider_shield_id,
provider_id=provider_id,
params=params,

12
llama_stack/providers/inline/safety/code_scanner/code_scanner.py
View file

@ -14,6 +14,12 @@ from .config import CodeScannerConfig
from llama_stack.apis.safety import * # noqa: F403
ALLOWED_CODE_SCANNER_MODEL_IDS = [
"CodeScanner",
"CodeShield",
]
class MetaReferenceCodeScannerSafetyImpl(Safety):
def __init__(self, config: CodeScannerConfig, deps) -> None:
self.config = config
@ -25,8 +31,10 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
pass
async def register_shield(self, shield: Shield) -> None:
if shield.shield_type != ShieldType.code_scanner:
raise ValueError(f"Unsupported safety shield type: {shield.shield_type}")
if shield.provider_resource_id not in ALLOWED_CODE_SCANNER_MODEL_IDS:
raise ValueError(
f"Unsupported Code Scanner ID: {shield.provider_resource_id}. Allowed IDs: {ALLOWED_CODE_SCANNER_MODEL_IDS}"
)
async def run_shield(
self,

26
llama_stack/providers/inline/safety/llama_guard/config.py
View file

@ -6,32 +6,8 @@
from typing import List
from llama_models.sku_list import CoreModelId, safety_models
from pydantic import BaseModel, field_validator
from pydantic import BaseModel
class LlamaGuardConfig(BaseModel):
model: str = "Llama-Guard-3-1B"
excluded_categories: List[str] = []
@field_validator("model")
@classmethod
def validate_model(cls, model: str) -> str:
permitted_models = [
m.descriptor()
for m in safety_models()
if (
m.core_model_id
in {
CoreModelId.llama_guard_3_8b,
CoreModelId.llama_guard_3_1b,
CoreModelId.llama_guard_3_11b_vision,
}
)
]
if model not in permitted_models:
raise ValueError(
f"Invalid model: {model}. Must be one of {permitted_models}"
)
return model

25
llama_stack/providers/inline/safety/llama_guard/llama_guard.py
View file

@ -73,6 +73,11 @@ DEFAULT_LG_V3_SAFETY_CATEGORIES = [
CAT_ELECTIONS,
]
LLAMA_GUARD_MODEL_IDS = [
CoreModelId.llama_guard_3_8b.value,
CoreModelId.llama_guard_3_1b.value,
CoreModelId.llama_guard_3_11b_vision.value,
]
MODEL_TO_SAFETY_CATEGORIES_MAP = {
CoreModelId.llama_guard_3_8b.value: (
@ -118,18 +123,16 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
self.inference_api = deps[Api.inference]
async def initialize(self) -> None:
self.shield = LlamaGuardShield(
model=self.config.model,
inference_api=self.inference_api,
excluded_categories=self.config.excluded_categories,
)
pass
async def shutdown(self) -> None:
pass
async def register_shield(self, shield: Shield) -> None:
if shield.shield_type != ShieldType.llama_guard:
raise ValueError(f"Unsupported shield type: {shield.shield_type}")
if shield.provider_resource_id not in LLAMA_GUARD_MODEL_IDS:
raise ValueError(
f"Unsupported Llama Guard type: {shield.provider_resource_id}. Allowed types: {LLAMA_GUARD_MODEL_IDS}"
)
async def run_shield(
self,
@ -147,7 +150,13 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
if len(messages) > 0 and messages[0].role != Role.user.value:
messages[0] = UserMessage(content=messages[0].content)
return await self.shield.run(messages)
impl = LlamaGuardShield(
model=shield.provider_resource_id,
inference_api=self.inference_api,
excluded_categories=self.config.excluded_categories,
)
return await impl.run(messages)
class LlamaGuardShield:

6
llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py
View file

@ -36,8 +36,10 @@ class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
pass
async def register_shield(self, shield: Shield) -> None:
if shield.shield_type != ShieldType.prompt_guard:
raise ValueError(f"Unsupported shield type: {shield.shield_type}")
if shield.provider_resource_id != PROMPT_GUARD_MODEL:
raise ValueError(
f"Only {PROMPT_GUARD_MODEL} is supported for Prompt Guard. "
)
async def run_shield(
self,

5
llama_stack/providers/remote/safety/bedrock/bedrock.py
View file

@ -20,11 +20,6 @@ from .config import BedrockSafetyConfig
logger = logging.getLogger(__name__)
BEDROCK_SUPPORTED_SHIELDS = [
ShieldType.generic_content_shield,
]
class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
def __init__(self, config: BedrockSafetyConfig) -> None:
self.config = config

8
llama_stack/providers/tests/agents/test_agents.py
View file

@ -81,15 +81,17 @@ async def create_agent_session(agents_impl, agent_config):
class TestAgents:
@pytest.mark.asyncio
async def test_agent_turns_with_safety(self, agents_stack, common_params):
async def test_agent_turns_with_safety(
self, safety_model, agents_stack, common_params
):
agents_impl, _ = agents_stack
agent_id, session_id = await create_agent_session(
agents_impl,
AgentConfig(
**{
**common_params,
"input_shields": ["llama_guard"],
"output_shields": ["llama_guard"],
"input_shields": [safety_model],
"output_shields": [safety_model],
}
),
)

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

@ -9,7 +9,7 @@ import pytest_asyncio
from llama_stack.apis.models import ModelInput
from llama_stack.apis.shields import ShieldInput, ShieldType
from llama_stack.apis.shields import ShieldInput
from llama_stack.distribution.datatypes import Api, Provider
from llama_stack.providers.inline.safety.llama_guard import LlamaGuardConfig
@ -41,7 +41,7 @@ def safety_llama_guard(safety_model) -> ProviderFixture:
Provider(
provider_id="inline::llama-guard",
provider_type="inline::llama-guard",
config=LlamaGuardConfig(model=safety_model).model_dump(),
config=LlamaGuardConfig().model_dump(),
)
],
)
@ -114,20 +114,14 @@ async def safety_stack(inference_model, safety_model, request):
def get_shield_to_register(provider_type: str, safety_model: str) -> ShieldInput:
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":
if 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
params = {"guardrailVersion": get_env_or_fail("BEDROCK_GUARDRAIL_VERSION")}
else:
params = {}
identifier = safety_model
return ShieldInput(
shield_id=identifier,
shield_type=shield_type,
params=shield_config,
params=params,
)

1
llama_stack/providers/tests/safety/test_safety.py
View file

@ -34,7 +34,6 @@ class TestSafety:
for shield in response:
assert isinstance(shield, Shield)
assert shield.shield_type in [v for v in ShieldType]
@pytest.mark.asyncio
async def test_run_shield(self, safety_stack):