llama-stack-mirror/llama_toolchain/safety/api/config.py
2024-07-21 23:48:38 -07:00

19 lines
466 B
Python

from typing import List, Optional
from pydantic import BaseModel
class LlamaGuardShieldConfig(BaseModel):
model_dir: str
excluded_categories: List[str]
disable_input_check: bool = False
disable_output_check: bool = False
class PromptGuardShieldConfig(BaseModel):
model_dir: str
class SafetyConfig(BaseModel):
llama_guard_shield: Optional[LlamaGuardShieldConfig] = None
prompt_guard_shield: Optional[PromptGuardShieldConfig] = None