Introduce model_store, shield_store, memory_bank_store

This commit is contained in:
Ashwin Bharambe 2024-10-06 16:29:33 -07:00 committed by Ashwin Bharambe
parent e45a417543
commit 91e0063593
19 changed files with 172 additions and 297 deletions

View file

@ -38,7 +38,13 @@ class RunShieldResponse(BaseModel):
violation: Optional[SafetyViolation] = None
class ShieldStore(Protocol):
def get_shield(self, identifier: str) -> ShieldDef: ...
class Safety(Protocol):
shield_store: ShieldStore
@webmethod(route="/safety/run_shield")
async def run_shield(
self, shield_type: str, messages: List[Message], params: Dict[str, Any] = None
@ -46,9 +52,3 @@ class Safety(Protocol):
@webmethod(route="/safety/register_shield")
async def register_shield(self, shield: ShieldDef) -> None: ...
@webmethod(route="/safety/list_shields")
async def list_shields(self) -> List[ShieldDef]: ...
@webmethod(route="/safety/get_shield")
async def get_shield(self, identifier: str) -> Optional[ShieldDef]: ...