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

init

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-07 12:09:14 -08:00
parent 7ee9f8d8ac
commit d960f9b60f
16 changed files with 140 additions and 105 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -40,33 +40,12 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
async def shutdown(self) -> None:
pass
async def register_shield(self, shield: ShieldDef) -> None:
async def register_shield(self, shield: Shield) -> None:
raise ValueError("Registering dynamic shields is not supported")
async def list_shields(self) -> List[ShieldDef]:
response = self.bedrock_client.list_guardrails()
shields = []
for guardrail in response["guardrails"]:
# populate the shield def with the guardrail id and version
shield_def = ShieldDef(
identifier=guardrail["id"],
shield_type=ShieldType.generic_content_shield.value,
params={
"guardrailIdentifier": guardrail["id"],
"guardrailVersion": guardrail["version"],
},
)
self.registered_shields.append(shield_def)
shields.append(shield_def)
return shields
async def run_shield(
self, identifier: str, messages: List[Message], params: Dict[str, Any] = None
self, shield: Shield, messages: List[Message], params: Dict[str, Any] = None
) -> RunShieldResponse:
shield_def = await self.shield_store.get_shield(identifier)
if not shield_def:
raise ValueError(f"Unknown shield {identifier}")
"""This is the implementation for the bedrock guardrails. The input to the guardrails is to be of this format
```content = [
{
@ -81,7 +60,7 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
They contain content, role . For now we will extract the content and default the "qualifiers": ["query"]
"""
shield_params = shield_def.params
shield_params = shield.params
logger.debug(f"run_shield::{shield_params}::messages={messages}")
# - convert the messages into format Bedrock expects
@ -93,7 +72,7 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
)
response = self.bedrock_runtime_client.apply_guardrail(
guardrailIdentifier=shield_params["guardrailIdentifier"],
guardrailIdentifier=shield.identifier,
guardrailVersion=shield_params["guardrailVersion"],
source="OUTPUT", # or 'INPUT' depending on your use case
content=content_messages,