Skip to content

Commit

Permalink
Fix up safety client for versioned API (#573)
Browse files Browse the repository at this point in the history
When running:
python -m llama_stack.apis.safety.client localhost 5000

The API server was logging:
INFO:    ::1:57176 - "POST /safety/run_shield HTTP/1.1" 404 Not Found

This patch uses the versioned API, uses the updated safety endpoint, and
updates the model name to what's being served. The above python command
now demonstrates a passing and failing example.
  • Loading branch information
stevegrubb authored Dec 5, 2024
1 parent 6eb5f2a commit a4daf4d
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions llama_stack/apis/safety/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
from pydantic import BaseModel
from termcolor import cprint

from llama_stack.apis.version import LLAMA_STACK_API_VERSION

from llama_stack.distribution.datatypes import RemoteProviderConfig

from llama_stack.apis.safety import * # noqa: F403
Expand Down Expand Up @@ -45,7 +47,7 @@ async def run_shield(
) -> RunShieldResponse:
async with httpx.AsyncClient() as client:
response = await client.post(
f"{self.base_url}/safety/run_shield",
f"{self.base_url}/{LLAMA_STACK_API_VERSION}/safety/run-shield",
json=dict(
shield_id=shield_id,
messages=[encodable_dict(m) for m in messages],
Expand Down Expand Up @@ -91,7 +93,7 @@ async def run_main(host: str, port: int, image_path: str = None):
]:
cprint(f"User>{message.content}", "green")
response = await client.run_shield(
shield_id="llama_guard",
shield_id="meta-llama/Llama-Guard-3-1B",
messages=[message],
)
print(response)
Expand Down

0 comments on commit a4daf4d

Please sign in to comment.