Update configure.py to use multiple routing keys for safety

This commit is contained in:
Ashwin Bharambe 2024-09-30 16:36:53 -07:00
parent 0996ffb3b3
commit 86834ee6c2
4 changed files with 11 additions and 71 deletions

View file

@ -117,18 +117,18 @@ def configure_api_providers(
if api_str == "safety": if api_str == "safety":
# TODO: add support for other safety providers, and simplify safety provider config # TODO: add support for other safety providers, and simplify safety provider config
if p == "meta-reference": if p == "meta-reference":
for shield_type in MetaReferenceShieldType:
routing_entries.append( routing_entries.append(
RoutableProviderConfig( RoutableProviderConfig(
routing_key=shield_type.value, routing_key=[s.value for s in MetaReferenceShieldType],
provider_id=p, provider_id=p,
config=cfg.dict(), config=cfg.dict(),
) )
) )
else: else:
cprint( cprint(
f"[WARN] Interactive configuration of safety provider {p} is not supported, please manually configure safety shields types in routing_table of run.yaml", f"[WARN] Interactive configuration of safety provider {p} is not supported. Please look for `{routing_key}` in run.yaml and replace it appropriately.",
"yellow", "yellow",
attrs=["bold"],
) )
routing_entries.append( routing_entries.append(
RoutableProviderConfig( RoutableProviderConfig(

View file

@ -42,22 +42,7 @@ routing_table:
config: config:
llama_guard_shield: null llama_guard_shield: null
prompt_guard_shield: null prompt_guard_shield: null
routing_key: llama_guard routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: code_scanner_guard
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: injection_shield
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: jailbreak_shield
memory: memory:
- provider_id: meta-reference - provider_id: meta-reference
config: {} config: {}

View file

@ -45,22 +45,7 @@ routing_table:
config: config:
llama_guard_shield: null llama_guard_shield: null
prompt_guard_shield: null prompt_guard_shield: null
routing_key: llama_guard routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: code_scanner_guard
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: injection_shield
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: jailbreak_shield
memory: memory:
- provider_id: meta-reference - provider_id: meta-reference
config: {} config: {}

View file

@ -50,37 +50,7 @@ routing_table:
disable_output_check: false disable_output_check: false
prompt_guard_shield: prompt_guard_shield:
model: Prompt-Guard-86M model: Prompt-Guard-86M
routing_key: llama_guard routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
- provider_id: meta-reference
config:
llama_guard_shield:
model: Llama-Guard-3-8B
excluded_categories: []
disable_input_check: false
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: code_scanner_guard
- provider_id: meta-reference
config:
llama_guard_shield:
model: Llama-Guard-3-8B
excluded_categories: []
disable_input_check: false
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: injection_shield
- provider_id: meta-reference
config:
llama_guard_shield:
model: Llama-Guard-3-8B
excluded_categories: []
disable_input_check: false
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: jailbreak_shield
memory: memory:
- provider_id: meta-reference - provider_id: meta-reference
config: {} config: {}