phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-07-29 15:23:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

nuke safety/list_shields, we don't need it now

Browse source
This commit is contained in:
Ashwin Bharambe 2024-09-22 19:49:08 -07:00
parent bc394882a0
commit 484dc2e4f5
3 changed files with 2 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

24
llama_stack/apis/safety/client.py
View file

@ -13,10 +13,11 @@ import fire
import httpx import httpx
from llama_models.llama3.api.datatypes import * # noqa: F403 from llama_models.llama3.api.datatypes import * # noqa: F403
from llama_stack.distribution.datatypes import RemoteProviderConfig
from pydantic import BaseModel from pydantic import BaseModel
from termcolor import cprint from termcolor import cprint
from llama_stack.distribution.datatypes import RemoteProviderConfig
from llama_stack.apis.safety import * # noqa: F403 from llama_stack.apis.safety import * # noqa: F403
@ -61,24 +62,6 @@ class SafetyClient(Safety):
content = response.json() content = response.json()
return RunShieldResponse(**content) return RunShieldResponse(**content)
async def list_shields(self) -> ListShieldsResponse:
async with httpx.AsyncClient() as client:
response = await client.post(
f"{self.base_url}/safety/list_shields",
json={},
headers={"Content-Type": "application/json"},
timeout=20,
)
if response.status_code != 200:
content = await response.aread()
error = f"Error: HTTP {response.status_code} {content.decode()}"
cprint(error, "red")
raise Exception(error)
content = response.json()
return ListShieldsResponse(**content)
async def run_main(host: str, port: int): async def run_main(host: str, port: int):
client = SafetyClient(f"http://{host}:{port}") client = SafetyClient(f"http://{host}:{port}")
@ -100,9 +83,6 @@ async def run_main(host: str, port: int):
) )
print(response) print(response)
response = await client.list_shields()
print(response)
def main(host: str, port: int): def main(host: str, port: int):
asyncio.run(run_main(host, port)) asyncio.run(run_main(host, port))

8
llama_stack/apis/safety/safety.py
View file

@ -37,16 +37,8 @@ class RunShieldResponse(BaseModel):
violation: Optional[SafetyViolation] = None violation: Optional[SafetyViolation] = None
@json_schema_type
class ListShieldsResponse(BaseModel):
shields: List[str] = None
class Safety(Protocol): class Safety(Protocol):
@webmethod(route="/safety/run_shield") @webmethod(route="/safety/run_shield")
async def run_shield( async def run_shield(
self, shield: str, messages: List[Message], params: Dict[str, Any] = None self, shield: str, messages: List[Message], params: Dict[str, Any] = None
) -> RunShieldResponse: ... ) -> RunShieldResponse: ...
@webmethod(route="/safety/list_shields")
async def list_shields(self) -> ListShieldsResponse: ...

21
llama_stack/providers/impls/meta_reference/safety/safety.py
View file

@ -80,27 +80,6 @@ class MetaReferenceSafetyImpl(Safety):
return RunShieldResponse(violation=violation) return RunShieldResponse(violation=violation)
async def list_shields(self) -> ListShieldsResponse:
supported_sheilds = [
v.value for v in MetaReferenceShieldType if self.is_supported(v)
]
return ListShieldsResponse(shields=supported_sheilds)
def is_supported(self, typ: MetaReferenceShieldType) -> bool:
if typ == MetaReferenceShieldType.llama_guard:
return self.config.llama_guard_shield is not None
if typ == MetaReferenceShieldType.jailbreak_shield:
return self.config.prompt_guard_shield is not None
if typ == MetaReferenceShieldType.injection_shield:
return self.config.prompt_guard_shield is not None
if typ == MetaReferenceShieldType.code_scanner_guard:
return True
return False
def get_shield_impl(self, typ: MetaReferenceShieldType) -> ShieldBase: def get_shield_impl(self, typ: MetaReferenceShieldType) -> ShieldBase:
cfg = self.config cfg = self.config
if typ == MetaReferenceShieldType.llama_guard: if typ == MetaReferenceShieldType.llama_guard: